04-从零训练语言模型小模型跑通全流程再说大的

从零训练语言模型：小模型跑通全流程再说大的

很多人上来就想训7B、13B模型，结果环境配不好、OOM、loss不收敛，到处踩坑。我的建议是：先用小模型跑通整个训练流程，确认每一步都理解了、没问题了，再放大。

这篇从零开始，手写一个完整的语言模型训练代码。不依赖HuggingFace Trainer，不用任何封装，每一步都看得见摸得着。

模型：一个极简的语言模型

先写一个最小的语言模型：Embedding + 4层Transformer + LM Head。这里不展开Transformer的实现细节（下篇专门讲），先用PyTorch内置的nn.TransformerEncoder。

importtorchimporttorch.nnasnnimportmathclassSmallLM(nn.Module):"""一个小型因果语言模型"""def__init__(self,vocab_size:int=32000,d_model:int=512,n_heads:int=8,n_layers:int=4,max_seq_len:int=2048,dropout:float=0.1,):super().__init__()self.d_model=d_model# Token embedding + 位置编码self.token_embedding=nn.Embedding(vocab_size,d_model)self.pos_embedding=nn.Embedding(max_seq_len,d_model)self.dropout=nn.Dropout(dropout)# Transformer编码器encoder_layer=nn.TransformerEncoderLayer(d_model=d_model,nhead=n_heads,dim_feedforward=d_model*4,dropout=dropout,activation="gelu",batch_first=True,# 输入格式 (batch, seq, dim)norm_first=True,# Pre-LN，大模型标配)self.encoder=nn.TransformerEncoder(encoder_layer,num_layers=n_layers,norm=nn.LayerNorm(d_model),# 最后加一个LayerNorm)# 语言模型头：映射回词表空间self.lm_head=nn.Linear(d_model,vocab_size,bias=False)# 权重共享：embedding和lm_head用同一套权重self.lm_head.weight=self.token_embedding.weight# 因果掩码：下三角矩阵，防止看到未来的tokenself.register_buffer("causal_mask",torch.tril(torch.ones(max_seq_len,max_seq_len)).bool())# 初始化权重self.apply(self._init_weights)def_init_weights(self,module:nn.Module):"""GPT风格的权重初始化"""ifisinstance(module,nn.Linear):torch.nn.init.normal_(module.weight,mean=0.0,std=0.02)ifmodule.biasisnotNone:torch.nn.init.zeros_(module.bias)elifisinstance(module,nn.Embedding):torch.nn.init.normal_(module.weight,mean=0.0,std=0.02)defforward(self,input_ids:torch.Tensor,labels:torch.Tensor|None=None,)->dict:B,T=input_ids.shape# Embeddingpositions=torch.arange(T,device=input_ids.device).unsqueeze(0)# (1, T)x=self.token_embedding(input_ids)+self.pos_embedding(positions)x=self.dropout(x)# Transformer + 因果掩码