train_tiny_lm.py

小型Transformer真实训练。对应第03章。依赖：PyTorch。

原始Python文件 · 58行

"""可选PyTorch CPU真实训练实验。仅用合成循环序列，无商业LLM/API。训练下降证明优化链路能工作，不证明自然语言泛化或真实代码能力。"""from __future__ import annotationsimport json,time,mathfrom pathlib import Path  def run(steps: int=100):    import torch    from torch import nn    import torch.nn.functional as F    torch.set_num_threads(1);torch.manual_seed(17)    vocab,width,heads,length=16,32,4,12    class Block(nn.Module):        def __init__(self):            super().__init__();self.ln1=nn.LayerNorm(width);self.ln2=nn.LayerNorm(width)            self.qkv=nn.Linear(width,width*3,bias=False);self.proj=nn.Linear(width,width,bias=False)            self.mlp=nn.Sequential(nn.Linear(width,width*2),nn.GELU(),nn.Linear(width*2,width))        def forward(self,x):            b,t,d=x.shape;q,k,v=self.qkv(self.ln1(x)).chunk(3,dim=-1)            q,k,v=[z.view(b,t,heads,d//heads).transpose(1,2) for z in (q,k,v)]            [email protected](-1,-2)/math.sqrt(d//heads)            mask=torch.ones(t,t,dtype=torch.bool,device=x.device).tril()            a=score.masked_fill(~mask,float('-inf')).softmax(-1)@v            x=x+self.proj(a.transpose(1,2).reshape(b,t,d));return x+self.mlp(self.ln2(x))    class TinyLM(nn.Module):        def __init__(self):            super().__init__();self.emb=nn.Embedding(vocab,width);self.pos=nn.Embedding(length,width)            self.blocks=nn.Sequential(Block(),Block());self.norm=nn.LayerNorm(width);self.head=nn.Linear(width,vocab)        def forward(self,ids):            x=self.emb(ids)+self.pos(torch.arange(ids.size(1),device=ids.device))            return self.head(self.norm(self.blocks(x)))    # 同一规律的合成数据：下一个Token=(当前+1) mod 16。    seq=(torch.arange(length+1)[None,:]+torch.arange(vocab)[:,None])%vocab    x,y=seq[:,:-1],seq[:,1:];model=TinyLM();opt=torch.optim.AdamW(model.parameters(),lr=.006)    def loss():return F.cross_entropy(model(x).reshape(-1,vocab),y.reshape(-1))    before=float(loss().detach());start=time.perf_counter()    history=[]    for i in range(steps):        opt.zero_grad(set_to_none=True);l=loss();l.backward();torch.nn.utils.clip_grad_norm_(model.parameters(),1.)        opt.step()        if i%20==0:history.append({'step':i,'loss':float(l.detach())})    model.eval()    with torch.no_grad():        after=float(loss());acc=float((model(x).argmax(-1)==y).float().mean())        # 改变未来Token，不得改变位置0的logits。        changed=x.clone();changed[:,5:]=(changed[:,5:]+3)%vocab        causal_error=float((model(x)[:,:5]-model(changed)[:,:5]).abs().max())    assert after<before and causal_error<1e-5    return {'kind':'真实小型模型CPU训练；合成数据；未运行商业模型或GPU吞吐基准','torch_version':torch.__version__,'parameter_count':sum(p.numel() for p in model.parameters()),'steps':steps,'initial_train_loss':before,'final_train_loss':after,'train_token_accuracy':acc,'causal_prefix_error':causal_error,'elapsed_seconds':time.perf_counter()-start,'history':history,'limitation':'测试集泛化尚未评价；高训练准确率不是通用语言理解'} if __name__=='__main__':    import argparse    p=argparse.ArgumentParser();p.add_argument('--output',type=Path);p.add_argument('--steps',type=int,default=100);a=p.parse_args()    if not 1<=a.steps<=2000:p.error('steps必须在1..2000')    result=run(a.steps);text=json.dumps(result,ensure_ascii=False,indent=2);print(text)    if a.output:a.output.write_text(text,encoding='utf-8')