NLP第二阶段学习 从零开始训练GPT模型代码# train_tiny_llm.pyimporttorchimporttorch.nnasnnimporttorch.nn.functionalasFfromtorch.utils.dataimportDataset,DataLoaderfromtokenizersimportTokenizer,models,trainers,pre_tokenizersimportjsonimportmathfromtqdmimporttqdm# 配置 classConfig:vocab_size6400# 词表大小dim512# 嵌入维度n_layers8# Transformer层数n_heads8# 注意力头数max_seq_len512# 最大序列长度dropout0.1# Dropout率# 训练配置batch_size4# CPU用小批次learning_rate5e-4epochs1warmup_steps100# 设备devicecpumixed_precisionFalse# CPU不支持FP16configConfig()# 1. 训练分词器 deftrain_tokenizer(texts,vocab_size6400,output_dirtokenizer):训练BPE分词器tokenizerTokenizer(models.BPE())tokenizer.pre_tokenizerpre_tokenizers.ByteLevel(add_prefix_spaceFalse)trainertrainers.BpeTrainer(vocab_sizevocab_size,special_tokens[,,,])tokenizer.train_from_iterator(texts,trainer)# tokenizer.save(f{output_dir}/tokenizer.json)tokenizer.save(ftokenizer.json)returntokenizer# 2. 模型组件 classRMSNorm(nn.Module):RMS归一化比LayerNorm更高效def__init__(self,dim,eps1e-6):super().__init__()self.epseps self.weightnn.Parameter(torch.ones(dim))defforward(self,x):returnx*torch.rsqrt(x.pow(2).mean(-1,keepdimTrue)self.eps)*self.weightclassRotaryPositionalEmbedding(nn.Module):旋转位置编码RoPEdef__init__(self,dim,max_seq_len512,base10000):super().__init__()inv_freq1.0/(base**(torch.arange(0,dim,2).float()/dim))self.register_buffer(inv_freq,inv_freq)defforward(self,seq_len):ttorch.arange(seq_len,deviceself.inv_freq.device)freqstorch.einsum(i,j-ij,t,self.inv_freq)embtorch.cat((freqs,freqs),dim-1)returnemb.cos(),emb.sin()defapply_rotary_pos_emb(q,k,cos,sin):应用旋转位置编码defrotate_half(x):x1,x2x.chunk(2,dim-1)returntorch.cat((-x2,x1),dim-1)q_embed(q*cos)(rotate_half(q)*sin)k_embed(k*cos)(rotate_half(k)*sin)returnq_embed,k_embedclassAttention(nn.Module):多头注意力机制def__init__(self,config):super().__init__()self.n_headsconfig.n_heads self.head_dimconfig.dim//config.n_heads self.scaleself.head_dim**-0.5self.qkvnn.Linear(config.dim,3*config.dim,biasFalse)self.projnn.Linear(config.dim,config.dim,biasFalse)self.dropoutnn.Dropout(config.dropout)self.ropeRotaryPositionalEmbedding(self.head_dim,config.max_seq_len)defforward(self,x):B,T,Cx.shape# QKV投影qkvself.qkv(x).reshape(B,T,3,self.n_heads,self.head_dim).permute(2,0,3,1,4)q,k,vqkv[0],qkv[1],qkv[2]# 应用旋转位置编码cos,sinself.rope(T)q,kapply_rotary_pos_emb(q,k,cos,sin)# 注意力计算attn(q k.transpose(-2,-1))*self.scale attnattn.masked_fill(torch.triu(torch.ones(T,T),diagonal1).bool(),float(-inf))attnF.softmax(attn,dim-1)attnself.dropout(attn)out(attn v).transpose(1,2).reshape(B,T,C)outself.proj(out)returnoutclassFeedForward(nn.Module):前馈网络使用SwiGLU激活def__init__(self,config):super().__init__()hidden_dim4*config.dim self.w1nn.Linear(config.dim,hidden_dim,biasFalse)self.w2nn.Linear(hidden_dim,config.dim,biasFalse)self.w3nn.Linear(config.dim,hidden_dim,biasFalse)defforward(self,x):returnself.w2(F.silu(self.w1(x))*self.w3(x))classTransformerBlock(nn.Module):Transformer块def__init__(self,config):super().__init__()self.norm1RMSNorm(config.dim)self.attnAttention(config)self.norm2RMSNorm(config.dim)self.ffFeedForward(config)defforward(self,x):xxself.attn(self.norm1(x))xxself.ff(self.norm2(x))returnxclassTinyLLM(nn.Module):完整的语言模型def__init__(self,config):super().__init__()self.configconfig self.token_embnn.Embedding(config.vocab_size,config.dim)self.layersnn.ModuleList([TransformerBlock(config)for_inrange(config.n_layers)])self.normRMSNorm(config.dim)self.headnn.Linear(config.dim,config.vocab_size,biasFalse)# 权重初始化self.apply(self._init_weights)def_init_weights(self,module):ifisinstance(module,nn.Linear):torch.nn.init.normal_(module.weight,mean0.0,std0.02)ifmodule.biasisnotNone:torch.nn.init.zeros_(module.bias)elifisinstance(module,nn.Embedding):torch.nn.init.normal_(module.weight,mean0.0,std0.02)defforward(self,idx,targetsNone):B,Tidx.shape xself.token_emb(idx)forlayerinself.layers:xlayer(x)xself.norm(x)logitsself.head(x)lossNoneiftargetsisnotNone:lossF.cross_entropy(logits.view(-1,logits.size(-1)),targets.view(-1))returnlogits,losstorch.no_grad()defgenerate(self,idx,max_new_tokens,temperature1.0,top_kNone):生成文本for_inrange(max_new_tokens):idx_condidx[:,-self.config.max_seq_len:]logits,_self(idx_cond)logitslogits[:,-1,:]/temperatureiftop_kisnotNone:v,_torch.topk(logits,min(top_k,logits.size(-1)))logits[logitsv[:,[-1]]]-float(Inf)probsF.softmax(logits,dim-1)idx_nexttorch.multinomial(probs,num_samples1)idxtorch.cat((idx,idx_next),dim1)returnidx# 3. 数据集 classTextDataset(Dataset):def__init__(self,texts,tokenizer,max_length512):self.tokenizertokenizer self.max_lengthmax_length self.data[]fortextintqdm(texts,descTokenizing):encodingtokenizer.encode(text)tokensencoding.ids self.data.extend(tokens)def__len__(self):returnmax(0,len(self.data)-self.max_length)def__getitem__(self,idx):chunkself.data[idx:idxself.max_length1]xtorch.tensor(chunk[:-1],dtypetorch.long)ytorch.tensor(chunk[1:],dtypetorch.long)returnx,y# 4. 训练循环 deftrain_model(model,train_loader,config):model.to(config.device)optimizertorch.optim.AdamW(model.parameters(),lrconfig.learning_rate)total_stepslen(train_loader)*config.epochs schedulertorch.optim.lr_scheduler.CosineAnnealingLR(optimizer,T_maxtotal_steps)model.train()step0forepochinrange(config.epochs):pbartqdm(train_loader,descfEpoch{epoch1}/{config.epochs})forbatch_idx,(x,y)inenumerate(pbar):x,yx.to(config.device),y.to(config.device)logits,lossmodel(x,y)optimizer.zero_grad()loss.backward()torch.nn.utils.clip_grad_norm_(model.parameters(),1.0)optimizer.step()scheduler.step()ifstep%100:pbar.set_postfix({loss:f{loss.item():.4f},lr:f{scheduler.get_last_lr()[0]:.2e}})step1returnmodel# 5. 主程序 if__name____main__:# 示例文本数据texts[人工智能是计算机科学的一个分支致力于创造能够模拟人类智能的系统。,机器学习是人工智能的核心技术之一通过数据训练模型来做出预测。,深度学习使用多层神经网络能够自动学习数据的层次化表示。,自然语言处理让计算机能够理解、解释和生成人类语言。,Transformer架构 revolutionized 了自然语言处理领域。,大语言模型通过在海量文本上预训练获得了强大的语言理解能力。,训练大模型需要大量的计算资源和高质量的数据。,微调技术可以在预训练模型的基础上适配特定的下游任务。,]*100# 重复以扩大数据集print(*50)print(步骤1: 训练分词器)print(*50)tokenizertrain_tokenizer(texts,vocab_sizeconfig.vocab_size)print(\n*50)print(步骤2: 准备数据集)print(*50)datasetTextDataset(texts,tokenizer,config.max_seq_len)train_loaderDataLoader(dataset,batch_sizeconfig.batch_size,shuffleTrue)print(f数据集大小:{len(dataset)}个序列)print(\n*50)print(步骤3: 创建模型)print(*50)modelTinyLLM(config)total_paramssum(p.numel()forpinmodel.parameters())print(f模型参数量:{total_params:,}({total_params/1e6:.2f}M))print(\n*50)print(步骤4: 开始训练)print(*50)modeltrain_model(model,train_loader,config)print(\n*50)print(步骤5: 保存模型)print(*50)torch.save(model.state_dict(),tiny_llm.pt)print(模型已保存到 tiny_llm.pt)print(\n*50)print(步骤6: 测试生成)print(*50)model.eval()prompt人工智能prompt_idstorch.tensor([tokenizer.encode(prompt).ids],dtypetorch.long)generatedmodel.generate(prompt_ids,max_new_tokens20,temperature0.8)# 解码简化版实际需要完整解码逻辑print(f提示:{prompt})print(f生成: [token IDs:{generated[0].tolist()}])print(\n训练完成)运行效果最终生成的模型