Holmes2011

ZeroPoint v0 research prototype: new from-scratch NumPy framework, memory-slot architecture, audited Transformer comparison

5c6c847 verified about 24 hours ago

Raw History Blame Contribute Delete

4.43 kB

"""各个实验脚本共用的部分：语料、批次、训练循环、计算量统计、结果存档。"""
import datetime
import glob
import os
import sys
import sysconfig
import time

import numpy as np

HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, os.path.join(HERE, ".."))
import zerocg as zc  # noqa: E402
from models import ZeroX  # noqa: E402
from zerocg import functional as F  # noqa: E402

ITERS = int(os.environ.get("ITERS", 2000))
T, B = 64, 32
CFG = dict(n_embd=64, n_head=4, n_layer=2)


def load_corpus(path=None):
    """默认语料：Python 标准库源码（纯 ASCII 的 .py 文件，约 150 万字符）。"""
    if path:
        return open(path, encoding="utf-8").read()
    parts, total = [], 0
    for p in sorted(glob.glob(os.path.join(sysconfig.get_paths()["stdlib"], "*.py"))):
        try:
            s = open(p, encoding="utf-8").read()
        except (UnicodeDecodeError, OSError):
            continue
        if s.isascii():
            parts.append(s)
            total += len(s)
        if total > 1_500_000:
            break
    return "\n".join(parts)


def prepare_data(path=None):
    text = load_corpus(path)
    chars = sorted(set(text))
    stoi = {c: i for i, c in enumerate(chars)}
    data = np.array([stoi[c] for c in text], dtype=np.int64)
    split = int(0.9 * len(data))
    return text, len(chars), data[:split], data[split:]


def batches(data, n, length, seed, batch=B):
    r = np.random.default_rng(seed)
    for _ in range(n):
        ix = r.integers(0, len(data) - length - 1, batch)
        yield np.stack([data[i:i + length] for i in ix]), np.stack([data[i + 1:i + length + 1] for i in ix])


def loss_of(logits, y):
    return F.cross_entropy(logits.reshape(-1, logits.shape[-1]), y.reshape(-1)).item()


def train(name, model, train_data, val_data, seed, lr=3e-3, iters=None, seq_len=None, batch=None):
    """返回验证 loss（固定在 64 字上评估，便于各实验比较）。
    学习率：100 步线性预热 + 余弦退火到 10%；AdamW，权重衰减 0.01，梯度裁剪 1.0。
    seq_len / batch 默认 T / B；改训练长度时请同时改 batch，保持每步的字数不变。"""
    iters = iters or ITERS
    opt = zc.optim.AdamW(model.parameters(), lr=lr, weight_decay=0.01)
    t0 = time.time()
    for it, (x, y) in enumerate(batches(train_data, iters, seq_len or T, seed, batch or B)):
        opt.lr = lr * min(1, (it + 1) / 100) * (0.1 + 0.9 * 0.5 * (1 + np.cos(np.pi * it / iters)))
        logits = model(x)
        loss = F.cross_entropy(logits.reshape(-1, logits.shape[-1]), y.reshape(-1))
        opt.zero_grad()
        loss.backward()
        zc.optim.clip_grad_norm(model.parameters(), 1.0)
        opt.step()
        if it % 500 == 0:
            print(f"  [{name} seed={seed} lr={lr:g}] step {it:5d}  train {loss.item():.3f}  "
                  f"{time.time() - t0:.0f}s", flush=True)
    return evaluate(model, val_data)


def evaluate(model, val_data, n=40):
    model.eval()
    with zc.no_grad():
        val = float(np.mean([loss_of(model(x), y) for x, y in batches(val_data, n, T, 999)]))
    model.train()
    return val


def forward_flops_per_char(model, V):
    """前向一次、平均每个字多少次浮点运算（只算矩阵乘）。训练一步约是它的 3 倍。"""
    x = np.zeros((1, T), dtype=np.int64)
    with zc.no_grad(), zc.count_flops() as c:
        model(x)
    return c.flops / T


def slot_usage(model):
    """每层"有效槽数"：最近一次 forward 的平均读/写地址分布的 exp(熵)，对头取平均。
    16 个槽都用得一样多 = 16；全挤在一个槽 = 1。"""
    out = []
    for b in model.blocks:
        read, write = b.mix.usage
        eff = lambda p: float(np.mean(np.exp(-(p * np.log(p + 1e-12)).sum(-1))))  # noqa: E731
        out.append((eff(read), eff(write)))
    return out


def save_result(name, lines):
    """结果写两份：experiments/results/<时间>-<name>.md（存档，不覆盖）和 experiments/<name>.md（最新）。"""
    text = "\n".join(lines) + "\n"
    stamp = datetime.datetime.now().strftime("%Y%m%d-%H%M")
    os.makedirs(os.path.join(HERE, "results"), exist_ok=True)
    for path in (os.path.join(HERE, "results", f"{stamp}-{name}.md"), os.path.join(HERE, f"{name}.md")):
        open(path, "w").write(text)
    print("\n" + text)


def is_zeropoint(model):
    return isinstance(model, ZeroX)

	"""各个实验脚本共用的部分：语料、批次、训练循环、计算量统计、结果存档。"""
	import datetime
	import glob
	import os
	import sys
	import sysconfig
	import time

	import numpy as np

	HERE = os.path.dirname(os.path.abspath(__file__))
	sys.path.insert(0, os.path.join(HERE, ".."))
	import zerocg as zc # noqa: E402
	from models import ZeroX # noqa: E402
	from zerocg import functional as F # noqa: E402

	ITERS = int(os.environ.get("ITERS", 2000))
	T, B = 64, 32
	CFG = dict(n_embd=64, n_head=4, n_layer=2)


	def load_corpus(path=None):
	"""默认语料：Python 标准库源码（纯 ASCII 的 .py 文件，约 150 万字符）。"""
	if path:
	return open(path, encoding="utf-8").read()
	parts, total = [], 0
	for p in sorted(glob.glob(os.path.join(sysconfig.get_paths()["stdlib"], "*.py"))):
	try:
	s = open(p, encoding="utf-8").read()
	except (UnicodeDecodeError, OSError):
	continue
	if s.isascii():
	parts.append(s)
	total += len(s)
	if total > 1_500_000:
	break
	return "\n".join(parts)


	def prepare_data(path=None):
	text = load_corpus(path)
	chars = sorted(set(text))
	stoi = {c: i for i, c in enumerate(chars)}
	data = np.array([stoi[c] for c in text], dtype=np.int64)
	split = int(0.9 * len(data))
	return text, len(chars), data[:split], data[split:]


	def batches(data, n, length, seed, batch=B):
	r = np.random.default_rng(seed)
	for _ in range(n):
	ix = r.integers(0, len(data) - length - 1, batch)
	yield np.stack([data[i:i + length] for i in ix]), np.stack([data[i + 1:i + length + 1] for i in ix])


	def loss_of(logits, y):
	return F.cross_entropy(logits.reshape(-1, logits.shape[-1]), y.reshape(-1)).item()


	def train(name, model, train_data, val_data, seed, lr=3e-3, iters=None, seq_len=None, batch=None):
	"""返回验证 loss（固定在 64 字上评估，便于各实验比较）。
	学习率：100 步线性预热 + 余弦退火到 10%；AdamW，权重衰减 0.01，梯度裁剪 1.0。
	seq_len / batch 默认 T / B；改训练长度时请同时改 batch，保持每步的字数不变。"""
	iters = iters or ITERS
	opt = zc.optim.AdamW(model.parameters(), lr=lr, weight_decay=0.01)
	t0 = time.time()
	for it, (x, y) in enumerate(batches(train_data, iters, seq_len or T, seed, batch or B)):
	opt.lr = lr * min(1, (it + 1) / 100) * (0.1 + 0.9 * 0.5 * (1 + np.cos(np.pi * it / iters)))
	logits = model(x)
	loss = F.cross_entropy(logits.reshape(-1, logits.shape[-1]), y.reshape(-1))
	opt.zero_grad()
	loss.backward()
	zc.optim.clip_grad_norm(model.parameters(), 1.0)
	opt.step()
	if it % 500 == 0:
	print(f" [{name} seed={seed} lr={lr:g}] step {it:5d} train {loss.item():.3f} "
	f"{time.time() - t0:.0f}s", flush=True)
	return evaluate(model, val_data)


	def evaluate(model, val_data, n=40):
	model.eval()
	with zc.no_grad():
	val = float(np.mean([loss_of(model(x), y) for x, y in batches(val_data, n, T, 999)]))
	model.train()
	return val


	def forward_flops_per_char(model, V):
	"""前向一次、平均每个字多少次浮点运算（只算矩阵乘）。训练一步约是它的 3 倍。"""
	x = np.zeros((1, T), dtype=np.int64)
	with zc.no_grad(), zc.count_flops() as c:
	model(x)
	return c.flops / T


	def slot_usage(model):
	"""每层"有效槽数"：最近一次 forward 的平均读/写地址分布的 exp(熵)，对头取平均。
	16 个槽都用得一样多 = 16；全挤在一个槽 = 1。"""
	out = []
	for b in model.blocks:
	read, write = b.mix.usage
	eff = lambda p: float(np.mean(np.exp(-(p * np.log(p + 1e-12)).sum(-1)))) # noqa: E731
	out.append((eff(read), eff(write)))
	return out


	def save_result(name, lines):
	"""结果写两份：experiments/results/<时间>-<name>.md（存档，不覆盖）和 experiments/<name>.md（最新）。"""
	text = "\n".join(lines) + "\n"
	stamp = datetime.datetime.now().strftime("%Y%m%d-%H%M")
	os.makedirs(os.path.join(HERE, "results"), exist_ok=True)
	for path in (os.path.join(HERE, "results", f"{stamp}-{name}.md"), os.path.join(HERE, f"{name}.md")):
	open(path, "w").write(text)
	print("\n" + text)


	def is_zeropoint(model):
	return isinstance(model, ZeroX)