从零写大模型，我花了三天把考研数学变成了代码

大家好我是阿浪说实话几个月前我还觉得“大模型”这三个字离我很远。什么Transformer、注意力机制、千亿参数听起来就像是天才才能碰的东西。直到有一天我硬着头皮把《Attention Is All You Need》那篇论文打印出来关掉所有浏览器标签页拿着笔一点点推公式。然后我发现了一个让我既震惊又兴奋的事实——大模型的底层全是我考研时背烂了的数学。线性代数、高等数学、概率统计这三门课几乎覆盖了Transformer的全部核心计算。矩阵乘法、复合函数、链式求导、Softmax归一化……这些东西我们当年刷了成百上千道题结果它们就是大模型的骨架。于是我决定做一件事从零写一个大模型一行预训练代码都不调全部手写。三天之后我跑通了一个最简版本。今天就把这个过程分享出来代码全附上不藏私。大模型到底是什么先别急着看代码咱们先把事情说透。如果你把大模型抽象到最顶层它其实就是一个超级复合函数输入文字 → 一系列数学变换 → 输出下一个文字这个函数有多层嵌套每一层都是一个数学变换。你给它一段话它预测下一个词最可能是什么。就这么简单。那这个“超级复合函数”由什么组成拆开看就是三样东西注意力机制Attention ——让模型知道一句话里谁跟谁关系最密切前馈网络Feed-Forward ——对每个位置独立做非线性变换位置编码Position Encoding ——告诉模型词语的顺序这三样东西堆叠起来就是Transformer。堆得越多模型越大。第一步输入层——把文字变成数字线性代数计算机不认识“你好世界”这四个字它只认识数字。所以第一步就是把文字变成矩阵。这个过程分两步词嵌入 位置编码。词嵌入离散→连续的线性映射词嵌入的本质就是一个矩阵乘法。我们有一个词表比如1000个词每个词对应一个128维的向量。这个向量就是模型对这个词的“理解”。importnumpyasnp# 词表大小和向量维度vocab_size1000hidden_dim128# 随机初始化词嵌入矩阵可学习参数embedding_matrixnp.random.randn(vocab_size,hidden_dim)*0.01# 模拟输入假设输入了4个token索引分别是10, 20, 30, 40input_tokensnp.array([10,20,30,40])# 查表得到每个token的向量表示input_embedsembedding_matrix[input_tokens]print(输入特征矩阵形状,input_embeds.shape)# (4, 128)这里我想多说一句很多人觉得词嵌入很神秘其实它就是一张大表格。每个词占一行每行是一个向量。模型训练的过程就是不断调整这张表格里的数字。位置编码用三角函数告诉模型顺序高等数学词嵌入解决了“每个词是什么”但没解决“每个词在什么位置”。Transformer是并行计算的它不像RNN那样一个一个读所以需要主动告诉它位置信息。位置编码用的是正弦和余弦函数defget_positional_encoding(seq_len,d_model):生成位置编码矩阵pos_encodingnp.zeros((seq_len,d_model))forposinrange(seq_len):foriinrange(0,d_model,2):# 偶数维度用sinpos_encoding[pos,i]np.sin(pos/(10000**(i/d_model)))# 奇数维度用cosifi1d_model:pos_encoding[pos,i1]np.cos(pos/(10000**(i/d_model)))returnpos_encoding# 假设序列长度是4维度是128pos_encodingget_positional_encoding(4,128)# 最终输入 词嵌入 位置编码final_inputinput_embedspos_encoding[:4,:]说句心里话我第一次看到这个公式的时候也是一脸懵为什么要用sin和cos后来我想明白了——因为sin和cos能产生不同频率的波形不同的位置就有不同的“指纹”模型就能区分谁在前谁在后。这就是高数里周期函数的实际应用。第二步自注意力——模型怎么知道谁跟谁有关系概率统计这是Transformer最核心的部分也是很多人觉得最难的部分。但其实它的逻辑特别朴素模型在处理一个词的时候会去“看”句子里的其他词判断谁跟当前词关系最紧密。举个例子“小明喜欢吃苹果因为它很甜。”这里的“它”指谁模型通过注意力机制会发现“苹果”和“它”的关系最密切。数学公式别怕我拆开讲自注意力的公式长这样Attention(Q,K,V)softmax(Q×K^T/√d_k)×V拆成三步算相关性Q和K做矩阵乘法得到一个“相关性分数”矩阵Softmax归一化把分数变成概率加起来等于1加权求和用这些概率去加权V得到最终输出手写代码importtorchimporttorch.nnasnnimporttorch.nn.functionalasFclassSelfAttention(nn.Module):def__init__(self,embed_size):super().__init__()self.embed_sizeembed_size# 三个线性层分别生成Q、K、Vself.Wqnn.Linear(embed_size,embed_size,biasFalse)self.Wknn.Linear(embed_size,embed_size,biasFalse)self.Wvnn.Linear(embed_size,embed_size,biasFalse)defforward(self,x):# x的形状: (batch_size, seq_len, embed_size)# 1. 生成Q、K、VQself.Wq(x)# (batch, seq_len, embed_size)Kself.Wk(x)# (batch, seq_len, embed_size)Vself.Wv(x)# (batch, seq_len, embed_size)# 2. 计算注意力分数Q × K^T# K.transpose(-2, -1) 把最后两维转置scorestorch.matmul(Q,K.transpose(-2,-1))# (batch, seq_len, seq_len)# 3. 缩放除以√d_k防止softmax进入梯度极小区域scoresscores/(self.embed_size**0.5)# 4. Softmax归一化成概率attention_weightsF.softmax(scores,dim-1)# (batch, seq_len, seq_len)# 5. 加权求和outtorch.matmul(attention_weights,V)# (batch, seq_len, embed_size)returnout这里有个细节我想特别说一下为什么要除以√d_k因为如果向量维度很大Q和K的点积结果会很大导致Softmax之后梯度极小模型学不动。除以√d_k就是把数值拉回到一个合适的范围。这就是为什么考研数学里的“量纲分析”在实际工程中这么重要。第三步多头注意力——让模型从多个角度理解一个注意力头只能从一个角度“看”句子。Transformer的做法是同时用多个头每个头关注不同的东西。比如头1可能关注语法结构头2可能关注语义关联头3可能关注指代关系classMultiHeadAttention(nn.Module):def__init__(self,embed_size,num_heads):super().__init__()self.num_headsnum_heads self.head_dimembed_size//num_heads# 用一个线性层同时生成所有头的Q、K、Vself.qkv_projnn.Linear(embed_size,embed_size*3,biasFalse)self.out_projnn.Linear(embed_size,embed_size,biasFalse)defforward(self,x):batch_size,seq_len,embed_sizex.shape# 1. 生成Q、K、Vqkvself.qkv_proj(x)# (batch, seq_len, 3 * embed_size)# 2. 拆分成多个头qkvqkv.reshape(batch_size,seq_len,3,self.num_heads,self.head_dim)qkvqkv.permute(2,0,3,1,4)# (3, batch, num_heads, seq_len, head_dim)Q,K,Vqkv[0],qkv[1],qkv[2]# 3. 每个头独立计算注意力scorestorch.matmul(Q,K.transpose(-2,-1))/(self.head_dim**0.5)attention_weightsF.softmax(scores,dim-1)outtorch.matmul(attention_weights,V)# (batch, num_heads, seq_len, head_dim)# 4. 合并所有头outout.permute(0,2,1,3).reshape(batch_size,seq_len,embed_size)# 5. 输出投影outself.out_proj(out)returnout第四步组装Transformer层有了注意力机制再加上前馈网络和残差连接就是一个完整的Transformer层了。classTransformerBlock(nn.Module):def__init__(self,embed_size,num_heads,ff_hidden_size):super().__init__()self.attentionMultiHeadAttention(embed_size,num_heads)self.feed_forwardnn.Sequential(nn.Linear(embed_size,ff_hidden_size),nn.ReLU(),nn.Linear(ff_hidden_size,embed_size))self.norm1nn.LayerNorm(embed_size)self.norm2nn.LayerNorm(embed_size)defforward(self,x):# 残差连接 层归一化attn_outself.attention(x)xself.norm1(xattn_out)# 残差连接输入输出ff_outself.feed_forward(x)xself.norm2(xff_out)# 残差连接输入输出returnx残差连接就是代码里的 x attn_out是我觉得特别巧妙的设计。它的作用是让梯度能直接传回浅层避免深层网络梯度消失。通俗说就是给信息开了一条“高速公路”不管中间经过多少层变换原始信息都能直接传到后面。第五步堆叠多层形成完整模型把多个Transformer层堆叠起来再加上输入层和输出层就是一个完整的大模型了。classMiniGPT(nn.Module):def__init__(self,vocab_size,embed_size,num_layers,num_heads,ff_hidden_size,max_seq_len):super().__init__()self.embed_sizeembed_size# 词嵌入self.token_embeddingnn.Embedding(vocab_size,embed_size)# 位置编码用可学习的位置编码简单起见self.pos_embeddingnn.Embedding(max_seq_len,embed_size)# Transformer层self.layersnn.ModuleList([TransformerBlock(embed_size,num_heads,ff_hidden_size)for_inrange(num_layers)])# 输出层self.ln_finalnn.LayerNorm(embed_size)self.lm_headnn.Linear(embed_size,vocab_size)defforward(self,input_ids):seq_leninput_ids.shape[1]# 词嵌入xself.token_embedding(input_ids)# (batch, seq_len, embed_size)# 位置编码positionstorch.arange(seq_len,deviceinput_ids.device).unsqueeze(0)xxself.pos_embedding(positions)# 通过所有Transformer层forlayerinself.layers:xlayer(x)# 输出层xself.ln_final(x)logitsself.lm_head(x)# (batch, seq_len, vocab_size)returnlogits训练让模型学会说话链式求导 梯度下降模型搭好了但它现在是个“傻子”——参数全是随机的输出全是乱码。怎么让它学会说话训练。训练的本质就是链式求导 梯度下降给模型输入一段文字它预测下一个词计算预测和真实答案之间的差距损失函数用链式求导算出每个参数对差距的贡献反向传播用梯度下降调整参数缩小差距importtorch.optimasoptim# 初始化模型modelMiniGPT(vocab_size1000,embed_size128,num_layers4,num_heads4,ff_hidden_size512,max_seq_len128)# 损失函数和优化器criterionnn.CrossEntropyLoss()optimizeroptim.Adam(model.parameters(),lr3e-4)# 模拟一次训练deftrain_step(input_ids,target_ids):model.train()optimizer.zero_grad()# 前向传播logitsmodel(input_ids)# (batch, seq_len, vocab_size)# 计算损失预测下一个词losscriterion(logits.view(-1,logits.shape[-1]),target_ids.view(-1))# 反向传播链式求导loss.backward()# 梯度下降更新参数optimizer.step()returnloss.item()# 模拟数据batch_size4seq_len32input_idstorch.randint(0,1000,(batch_size,seq_len))target_idstorch.randint(0,1000,(batch_size,seq_len))losstrain_step(input_ids,target_ids)print(f训练损失:{loss:.4f})你想过没有我们考研时背得滚瓜烂熟的链式法则在这里就是让模型变聪明的核心引擎。每一次 loss.backward()PyTorch都在做你当年手算过的那些求导。完整代码我把上面所有的代码整合成了一个完整的文件importnumpyasnpimporttorchimporttorch.nnasnnimporttorch.nn.functionalasFimporttorch.optimasoptim# 位置编码 defget_positional_encoding(seq_len,d_model):pos_encodingnp.zeros((seq_len,d_model))forposinrange(seq_len):foriinrange(0,d_model,2):pos_encoding[pos,i]np.sin(pos/(10000**(i/d_model)))ifi1d_model:pos_encoding[pos,i1]np.cos(pos/(10000**(i/d_model)))returntorch.tensor(pos_encoding,dtypetorch.float32)# 多头注意力 classMultiHeadAttention(nn.Module):def__init__(self,embed_size,num_heads):super().__init__()assertembed_size%num_heads0self.num_headsnum_heads self.head_dimembed_size//num_heads self.qkv_projnn.Linear(embed_size,embed_size*3,biasFalse)self.out_projnn.Linear(embed_size,embed_size,biasFalse)defforward(self,x):batch_size,seq_len,embed_sizex.shape qkvself.qkv_proj(x)qkvqkv.reshape(batch_size,seq_len,3,self.num_heads,self.head_dim)qkvqkv.permute(2,0,3,1,4)Q,K,Vqkv[0],qkv[1],qkv[2]scorestorch.matmul(Q,K.transpose(-2,-1))/(self.head_dim**0.5)attention_weightsF.softmax(scores,dim-1)outtorch.matmul(attention_weights,V)outout.permute(0,2,1,3).reshape(batch_size,seq_len,embed_size)outself.out_proj(out)returnout# Transformer层 classTransformerBlock(nn.Module):def__init__(self,embed_size,num_heads,ff_hidden_size):super().__init__()self.attentionMultiHeadAttention(embed_size,num_heads)self.feed_forwardnn.Sequential(nn.Linear(embed_size,ff_hidden_size),nn.ReLU(),nn.Linear(ff_hidden_size,embed_size))self.norm1nn.LayerNorm(embed_size)self.norm2nn.LayerNorm(embed_size)defforward(self,x):attn_outself.attention(x)xself.norm1(xattn_out)ff_outself.feed_forward(x)xself.norm2(xff_out)returnx# 完整模型 classMiniGPT(nn.Module):def__init__(self,vocab_size,embed_size,num_layers,num_heads,ff_hidden_size,max_seq_len):super().__init__()self.token_embeddingnn.Embedding(vocab_size,embed_size)self.pos_embeddingnn.Embedding(max_seq_len,embed_size)self.layersnn.ModuleList([TransformerBlock(embed_size,num_heads,ff_hidden_size)for_inrange(num_layers)])self.ln_finalnn.LayerNorm(embed_size)self.lm_headnn.Linear(embed_size,vocab_size)defforward(self,input_ids):seq_leninput_ids.shape[1]xself.token_embedding(input_ids)positionstorch.arange(seq_len,deviceinput_ids.device).unsqueeze(0)xxself.pos_embedding(positions)forlayerinself.layers:xlayer(x)xself.ln_final(x)logitsself.lm_head(x)returnlogits# 训练 deftrain_step(model,input_ids,target_ids,optimizer,criterion):model.train()optimizer.zero_grad()logitsmodel(input_ids)losscriterion(logits.view(-1,logits.shape[-1]),target_ids.view(-1))loss.backward()optimizer.step()returnloss.item()# 主程序 if__name____main__:# 超参数vocab_size1000embed_size128num_layers4num_heads4ff_hidden_size512max_seq_len128batch_size4seq_len32epochs100# 初始化modelMiniGPT(vocab_size,embed_size,num_layers,num_heads,ff_hidden_size,max_seq_len)criterionnn.CrossEntropyLoss()optimizeroptim.Adam(model.parameters(),lr3e-4)# 模拟训练print(开始训练...)forepochinrange(epochs):input_idstorch.randint(0,vocab_size,(batch_size,seq_len))target_idstorch.randint(0,vocab_size,(batch_size,seq_len))losstrain_step(model,input_ids,target_ids,optimizer,criterion)if(epoch1)%100:print(fEpoch{epoch1}/{epochs}, Loss:{loss:.4f})print(训练完成)把上面的代码复制下来跑一遍。改改参数看看效果。把每一行代码都搞清楚它在做什么。你会发现大模型真的没有那么神秘。它就是一个超级复合函数。而你早就学过怎么拆解它。