首页 / 资讯中心 / 文章详情

大模型训练底层逻辑:递推、递归与三明治三重框架解析

大模型训练底层逻辑:递推、递归与三明治三重框架解析 ★ FEATURED ARTICLE
1. 从“调包”到“造轮子”为什么你必须理解训练底层逻辑很多人做大模型相关的工作日常就是from transformers import AutoModelForCausalLM然后model.train()一把梭。跑通了指标涨了就觉得自己“会了”。但一旦遇到loss不收敛、梯度爆炸、显存溢出、训练速度慢得离谱就彻底抓瞎——因为你根本不知道model.train()那一行代码背后数据是怎么流的、梯度是怎么传的、参数是怎么更新的。我见过太多这样的情况某开发者用现成的训练框架跑一个文本分类任务换了三套超参都不收敛最后发现是数据预处理阶段把padding token也算进了loss。这种问题你调一万次包也找不到原因但只要你理解训练底层的逻辑链条五分钟就能定位。这篇文章要讲的“递推、递归、三明治”三重逻辑不是什么学术黑话而是我在实际训练大模型过程中总结出来的一套理解框架。递推解决的是“时间维度上信息如何一步步往前传”的问题递归解决的是“网络结构上梯度如何一层层往回传”的问题三明治解决的是“整个训练流程中数据、模型、优化器三者如何配合”的问题。这三个逻辑串起来就是大模型训练的完整底层图景。适合谁看如果你已经会用PyTorch或类似框架搭模型、跑训练但对反向传播的细节、优化器的更新规则、数据管道的设计还停留在“大概知道”的水平那这篇文章就是写给你的。如果你是完全的新手也没关系我会用生活化的类比把每个概念讲清楚保证你能跟上。接下来我会先拆解这三个逻辑各自的含义和它们之间的关系然后逐个深入最后给出可以直接复现的实操方案和踩坑经验。2. 三重逻辑的整体设计与思路拆解2.1 为什么是“递推、递归、三明治”这三个词先说清楚这三个词是我自己总结的不是教科书上的标准术语。但它们精准地对应了大模型训练中三个最核心的机制。递推对应的是序列数据在时间步上的展开。你喂给模型一句话“今天天气真好”模型不是一眼看全的而是一个字一个字或者一个token一个token地处理。每处理一个token就把当前的隐藏状态传递给下一个时间步。这个过程就是递推——用前一步的结果作为后一步的输入之一。RNN是这样Transformer的自回归生成也是这样只不过Transformer用注意力机制替代了简单的状态传递但递推的本质没变。递归对应的是计算图的回溯。你定义了一个多层网络前向传播的时候数据从输入层一路走到输出层算出loss。然后反向传播的时候梯度从loss出发沿着计算图一层一层往回传每一层都用链式法则算出自己该更新多少。这个过程是递归的——每一层的梯度计算都依赖于后一层的梯度结果层层嵌套直到传到最底层。三明治对应的是整个训练循环的结构。一个完整的训练step本质上是三件事夹在一起前向传播数据进、loss出、反向传播梯度出、参数更新优化器干活。这三件事像三明治一样叠在一起每一层都依赖上一层的结果。而且这个三明治不是做一次就完了你要反复做几万次、几十万次每次的数据不一样但结构完全一样。2.2 三个逻辑之间的关系不是孤立的是咬合的很多人学的时候会把这三个东西分开学先学RNN的递推再学反向传播的递归最后学训练循环。但实际上它们是咬合在一起的。你想想一个完整的训练过程是什么样的你从数据管道里取一个batch的数据这个数据可能是序列比如文本需要按时间步递推地送进模型模型前向传播算出loss然后反向传播递归地算梯度最后优化器更新参数。更新完参数之后下一个batch的数据又来了又走一遍同样的流程。所以递推是“数据在模型内部怎么流”递归是“梯度在模型内部怎么回传”三明治是“整个训练循环怎么转”。三者缺一不可而且互相影响。比如你的序列长度变了递推的步数就变了计算图的大小就变了反向传播的递归深度和显存占用也跟着变。再比如你换了优化器三明治的第三层就变了但前两层不变。2.3 为什么理解这些能让你少踩坑我举个真实的例子。之前有个做文本生成的项目模型训练的时候loss一直震荡调学习率、换优化器、加warmup都不管用。后来我让他把训练循环拆开看发现他在每个batch之后都调用了optimizer.zero_grad()但他在计算loss的时候用了reductionsum而不是reductionmean。这意味着每个batch的loss大小取决于batch内token的数量而不同batch的token数量差异很大导致梯度大小剧烈波动。这个问题你如果只理解“三明治”的结构知道前向传播算loss、反向传播算梯度、优化器更新参数你就能定位到是loss计算这一层出了问题。如果你不理解就只能瞎调超参。再比如你做长序列训练的时候显存爆了。如果你理解递推的逻辑你就知道显存占用和序列长度是平方关系Transformer的注意力机制你就会想到用梯度累积、梯度检查点、或者滑动窗口注意力来缓解。如果你不理解就只能减小batch size然后发现训练速度慢得没法接受。所以这三个逻辑不是理论摆设是你排查问题、优化性能、设计方案的底层工具箱。3. 递推逻辑序列数据在时间步上的展开3.1 递推的本质用前一步的状态参与当前步的计算递推这个词听起来数学味很重但说白了就是当前时刻的输出依赖于当前时刻的输入和上一时刻的状态。用生活化的例子来说。你读一本书读到第100页的时候你对书的理解不仅取决于第100页的内容还取决于你读前99页积累下来的理解。这个“积累下来的理解”就是状态它从第1页开始一页一页地传递和更新这就是递推。在RNN里这个状态叫hidden state公式是h_t f(W_h * h_{t-1} W_x * x_t b)其中h_t是当前时刻的隐藏状态h_{t-1}是上一时刻的隐藏状态x_t是当前时刻的输入。你看h_t的计算用到了h_{t-1}这就是递推。在Transformer里递推的表现形式不一样。Transformer没有显式的hidden state传递它用自注意力机制让每个位置都能看到之前所有位置的信息。但在自回归生成的时候你生成第t个token需要把前t-1个token都作为输入这本质上还是一种递推——只不过Transformer用KV Cache来缓存之前计算过的key和value避免重复计算。3.2 递推在训练中的具体表现时间步展开与计算图训练的时候递推是怎么体现的关键在计算图的展开。假设你有一个序列长度为5的输入RNN会把它展开成5个时间步。每个时间步都有自己的计算节点但这些节点之间是有依赖关系的第2步依赖第1步的hidden state第3步依赖第2步的hidden state以此类推。这就意味着你的计算图不是5个独立的子图而是一条链。反向传播的时候梯度要沿着这条链一路往回传。链越长梯度传得越远越容易出问题——要么梯度消失传着传着变成0了要么梯度爆炸传着传着变成无穷大了。Transformer虽然用注意力机制缓解了长距离依赖的问题但计算图依然很大。序列长度为N的时候注意力矩阵是N×N的计算复杂度是O(N²)。这就是为什么长序列训练这么吃显存——你的计算图大小和序列长度是平方关系。注意很多人以为Transformer没有递推其实在自回归生成和训练中递推逻辑依然存在只是实现方式从显式的状态传递变成了隐式的注意力计算。3.3 实操要点如何处理变长序列和padding实际训练中你的序列长度是不固定的。有的句子10个token有的句子100个token。你不可能为每个句子单独跑一次前向传播那样效率太低。所以你会把同一个batch内的句子padding到相同长度。但padding会带来一个问题padding token也会参与计算也会产生梯度。如果你不处理模型就会学到“padding token也是有效信息”这个错误的模式。处理方法很简单在计算loss的时候把padding位置的loss mask掉。具体来说你可以创建一个attention mask把padding位置的attention score设为负无穷这样softmax之后这些位置的权重就是0。然后在计算loss的时候用mask把padding位置的loss置零最后只对有效位置求平均。# 假设logits是模型输出shape为[batch_size, seq_len, vocab_size] # labels是真实标签shape为[batch_size, seq_len] # mask是有效位置掩码shape为[batch_size, seq_len]有效位置为1padding位置为0 loss_fn nn.CrossEntropyLoss(reductionnone) loss loss_fn(logits.view(-1, vocab_size), labels.view(-1)) loss loss.view(batch_size, seq_len) loss (loss * mask).sum() / mask.sum()这个操作看起来简单但很多人会忘。忘了的后果就是loss曲线看起来在下降但模型在实际生成的时候表现很差因为它把大量算力浪费在学习padding上了。3.4 递推逻辑带来的性能瓶颈与优化思路递推逻辑最大的性能瓶颈在于串行计算。RNN的每个时间步必须等上一个时间步算完才能开始没法并行。这就是为什么RNN训练慢——你的GPU有几千个核心但大部分时间都在等。Transformer通过注意力机制实现了并行计算所有位置可以同时算。但代价是O(N²)的计算复杂度和显存占用。序列长度翻倍计算量和显存占用变成4倍。优化思路有几个方向梯度累积把一个大batch拆成几个小batch分别算梯度然后累加最后一起更新。这样可以在显存有限的情况下模拟大batch的效果。梯度检查点前向传播的时候不保存所有中间激活值只保存部分关键节点。反向传播的时候重新计算需要的激活值。用时间换空间。滑动窗口注意力每个位置只关注前W个位置而不是所有位置。把O(N²)降到O(N×W)。序列并行把长序列切分成多段分配到不同的GPU上通过通信来交换边界信息。这些优化手段的选择取决于你的具体场景。如果你的序列长度在512以内标准的注意力机制完全够用。如果序列长度到了4096甚至更长就必须考虑优化了。4. 递归逻辑梯度在计算图上的反向传播4.1 递归的本质链式法则的层层嵌套递归这个词在计算机科学里指的是“函数调用自身”但在训练语境下我用它来指代梯度计算沿着计算图层层回溯的过程。你定义了一个多层网络前向传播的时候数据从第一层流到最后一层算出loss。反向传播的时候梯度从loss出发沿着相反的方向流回去。每一层都要计算“如果我这一层的参数变化一点点loss会变化多少”这个计算依赖于后一层传回来的梯度。用数学公式表达就是链式法则∂L/∂W_1 ∂L/∂y_n * ∂y_n/∂y_{n-1} * ... * ∂y_2/∂y_1 * ∂y_1/∂W_1你看这个公式本身就是递归的——每一层的梯度都依赖于后一层的梯度结果。你没法跳过中间层直接算第一层的梯度必须一层一层往回传。4.2 计算图的构建与释放PyTorch的autograd机制PyTorch的autograd机制就是递归逻辑的具体实现。当你对一个requires_gradTrue的tensor做任何操作PyTorch都会在后台构建一个计算图。每个操作都是一个节点节点之间用边连接表示数据的依赖关系。前向传播的时候PyTorch一边算结果一边记录计算图。反向传播的时候调用loss.backward()PyTorch从loss节点出发沿着计算图反向遍历用链式法则算出每个参数的梯度。这里有个关键细节计算图默认是动态的每次前向传播都会重新构建。这意味着你可以在训练循环里改变网络结构比如条件分支计算图会自动适应。这是PyTorch比早期TensorFlow静态图灵活的地方。但动态图也有代价构建和释放计算图有开销。如果你训练的是小模型这个开销可能比实际计算还大。所以PyTorch提供了torch.jit和torch.compile来优化。注意每次调用loss.backward()之后计算图会被释放除非你设置了retain_graphTrue。如果你试图对同一个loss调用两次backward()第二次会报错。这是新手常踩的坑。4.3 梯度消失与梯度爆炸递归逻辑的固有风险递归逻辑最大的风险就是梯度在回传过程中出问题。梯度消失如果每一层的梯度都小于1那么经过多层连乘之后梯度会指数级衰减传到最底层的时候已经接近0了。底层参数几乎不更新模型学不到东西。这在深层网络和RNN中特别常见。梯度爆炸反过来如果每一层的梯度都大于1连乘之后梯度会指数级增长变成无穷大。参数更新幅度过大loss直接飞掉。解决方案残差连接让梯度可以跳过某些层直接回传缓解梯度消失。ResNet的核心思想就是这个。层归一化把每层的激活值归一化到均值为0、方差为1稳定梯度分布。梯度裁剪当梯度的范数超过某个阈值时把它缩放到阈值以内。防止梯度爆炸。合适的初始化权重初始化不能太大也不能太小。Xavier初始化和He初始化是常用的方案。# 梯度裁剪示例 torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm1.0)这行代码放在loss.backward()之后、optimizer.step()之前。max_norm1.0意味着如果所有参数梯度的总范数超过1.0就等比例缩小。4.4 实操要点如何调试梯度问题梯度出问题的时候loss曲线会有明显的异常。梯度消失表现为loss几乎不下降或者下降极慢。梯度爆炸表现为loss突然变成NaN或者剧烈震荡。调试方法打印梯度范数在每个训练step之后计算所有参数梯度的范数看看是不是在合理范围内。正常情况应该在0.1到10之间。逐层检查如果发现梯度异常逐层打印梯度范数找到问题出在哪一层。可视化用TensorBoard或类似工具把梯度分布画出来直观看梯度是不是在衰减或爆炸。小规模实验用少量数据、小模型先跑通确认梯度正常之后再扩大规模。# 打印每层梯度范数 for name, param in model.named_parameters(): if param.grad is not None: print(f{name}: grad_norm{param.grad.norm().item():.6f})这个操作在调试阶段非常有用。我一般会在训练的前100个step里每隔10步打印一次确认梯度稳定之后再关掉。5. 三明治逻辑训练循环的三层结构5.1 三明治的第一层前向传播与loss计算前向传播是数据从输入到输出的过程。你喂给模型一个batch的数据模型经过层层计算输出预测结果。然后你用损失函数比较预测结果和真实标签算出loss。这一步看起来简单但细节很多。数据准备你的原始数据可能是文本、图像、音频需要经过预处理变成模型能吃的tensor。文本要tokenize图像要归一化音频要提取特征。这一步的坑最多——数据格式不对、维度不匹配、类型错误都会导致前向传播失败。模型计算数据进入模型之后经过embedding层、注意力层、前馈层、输出层最后得到logits。每一层的计算都要注意数值稳定性。比如softmax的时候如果输入值太大exp之后会溢出。标准做法是先减去最大值再算exp。Loss计算根据任务类型选择损失函数。分类任务用交叉熵回归任务用MSE生成任务用负对数似然。注意reduction参数的选择——mean、sum、none各有用途。大多数情况下用mean但如果你需要手动mask就用none然后自己加权平均。# 标准的前向传播loss计算 model.train() optimizer.zero_grad() outputs model(input_ids, attention_maskattention_mask) logits outputs.logits shift_logits logits[..., :-1, :].contiguous() shift_labels labels[..., 1:].contiguous() loss nn.CrossEntropyLoss()(shift_logits.view(-1, vocab_size), shift_labels.view(-1))这段代码是语言模型训练的标准写法。注意shift操作——语言模型预测的是下一个token所以输入序列和标签序列要错开一位。5.2 三明治的第二层反向传播与梯度计算反向传播是loss对每个参数求偏导的过程。PyTorch的loss.backward()一行代码就搞定了但背后发生的事情很复杂。计算图遍历PyTorch从loss节点出发沿着计算图反向遍历。每个节点都有一个backward函数定义了如何计算该节点的梯度。梯度累加PyTorch的梯度是累加的不是覆盖的。这意味着如果你不调用optimizer.zero_grad()上一个batch的梯度会保留下来和当前batch的梯度加在一起。这有时候是故意的比如梯度累积但大多数时候是bug。梯度类型不同操作的梯度计算方式不同。加法操作的梯度是直接传递乘法操作的梯度是交叉相乘矩阵乘法操作的梯度涉及转置。理解这些有助于你排查梯度异常。注意loss.backward()只能对标量调用。如果你的loss是向量需要先reduce成标量比如求平均。这是新手常见的错误。5.3 三明治的第三层优化器更新参数优化器根据梯度更新参数。最基础的是SGDW W - lr * grad但实际训练大模型的时候几乎没人用纯SGD。常用的是Adam和AdamW。Adam的核心思想是自适应学习率每个参数的学习率不是固定的而是根据该参数梯度的历史统计量动态调整。梯度变化大的参数学习率调小梯度变化小的参数学习率调大。AdamW是Adam的改进版主要区别在于权重衰减的实现方式。Adam把权重衰减加在梯度里AdamW把权重衰减直接加在参数上。看起来差别不大但实际效果差很多——AdamW的训练更稳定泛化性能更好。# AdamW优化器的标准配置 optimizer torch.optim.AdamW( model.parameters(), lr1e-4, betas(0.9, 0.999), eps1e-8, weight_decay0.01 )参数说明lr学习率大模型常用1e-4到1e-5betas一阶和二阶矩估计的衰减率默认0.9和0.999eps数值稳定项防止除零weight_decay权重衰减系数常用0.015.4 三明治的完整循环从数据到参数更新的全流程把三层夹在一起一个完整的训练step是这样的从数据管道取一个batch数据预处理tokenize、padding、转tensor前向传播算lossoptimizer.zero_grad()清空梯度loss.backward()算梯度梯度裁剪可选optimizer.step()更新参数记录日志loss、学习率、梯度范数重复1-8直到跑完所有数据这个循环看起来简单但每个环节都有优化空间。比如数据管道可以用多进程预取减少GPU等待时间。梯度累积可以模拟大batch。混合精度训练可以减少显存占用、加速计算。# 完整的训练循环示例 scaler torch.cuda.amp.GradScaler() for epoch in range(num_epochs): for step, batch in enumerate(dataloader): input_ids batch[input_ids].to(device) attention_mask batch[attention_mask].to(device) labels batch[labels].to(device) optimizer.zero_grad() with torch.cuda.amp.autocast(): outputs model(input_ids, attention_maskattention_mask) loss compute_loss(outputs.logits, labels) scaler.scale(loss).backward() scaler.unscale_(optimizer) torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm1.0) scaler.step(optimizer) scaler.update() if step % log_interval 0: print(fEpoch {epoch}, Step {step}, Loss {loss.item():.4f})这段代码包含了混合精度训练和梯度裁剪是实际项目中常用的配置。6. 三重逻辑的协同从理论到落地的完整链路6.1 数据管道设计如何高效喂数据数据管道的效率直接决定训练速度。如果你的GPU利用率只有30%大概率是数据管道拖了后腿。PyTorch DataLoader的关键参数num_workers用几个子进程加载数据。设置成CPU核心数的一半到全部。设太大反而会因为进程切换开销降低效率。pin_memory把数据固定在内存中加速CPU到GPU的传输。设为True。prefetch_factor每个worker预取多少个batch。默认是2可以适当调大。persistent_workers保持worker进程存活避免每个epoch重新创建。数据预处理的位置能在CPU上做的预处理就不要放到GPU上做。比如tokenize、padding、mask生成这些都可以在DataLoader的collate_fn里完成。数据格式尽量用连续内存的tensor避免用Python list。tensor的传输和计算效率高得多。6.2 混合精度训练省显存、提速度的利器混合精度训练的核心思想是前向传播和反向传播用FP16半精度参数更新用FP32单精度。这样既能享受FP16的速度和显存优势又能保持FP32的数值稳定性。PyTorch的实现是torch.cuda.amp。用法很简单scaler torch.cuda.amp.GradScaler() with torch.cuda.amp.autocast(): outputs model(input_ids) loss compute_loss(outputs, labels) scaler.scale(loss).backward() scaler.step(optimizer) scaler.update()GradScaler的作用是动态调整loss的缩放因子。因为FP16的表示范围比FP32小梯度太小的时候会变成0。GradScaler把loss放大梯度也跟着放大避免下溢。更新参数之前再缩回去。注意混合精度训练不是万能的。有些操作在FP16下会不稳定比如softmax、layer norm。autocast会自动把这些操作保持在FP32但你自定义的操作可能需要手动处理。6.3 分布式训练从单卡到多卡的扩展单卡显存不够、训练太慢的时候就需要上分布式。数据并行每个GPU有一份完整的模型副本数据分成多份每个GPU算自己那份的梯度然后所有GPU的梯度求平均再各自更新参数。PyTorch的DistributedDataParallel是这个方案的标准实现。模型并行模型太大一张卡放不下就把模型切分成多份每张卡放一部分。前向传播的时候数据在卡之间传递。实现起来比数据并行复杂得多。流水线并行把模型按层切分成多个阶段每个阶段放在不同的GPU上。数据像流水线一样从第一个阶段流到最后一个阶段。微批处理可以进一步提高利用率。对于大多数场景数据并行就够了。只有模型大到单卡放不下的时候才需要考虑模型并行和流水线并行。6.4 实操中的性能调优清单训练速度慢的时候按这个清单逐项排查排查项检查方法优化手段GPU利用率nvidia-smi看GPU-Util低于80%说明数据管道是瓶颈数据加载时间在DataLoader里计时增加num_workers用pin_memory显存占用torch.cuda.max_memory_allocated()混合精度、梯度检查点、减小batch计算图大小打印模型参数量和序列长度梯度累积、滑动窗口注意力通信开销分布式训练时看NCCL耗时梯度压缩、通信重叠学习率看loss曲线warmup、余弦退火、ReduceLROnPlateau这张表是我在实际项目中总结的每次遇到性能问题就按这个顺序查一遍基本能定位到瓶颈。7. 常见问题与排查技巧实录7.1 Loss不收敛的排查思路Loss不收敛是最常见的问题原因可能出在递推、递归、三明治的任何一层。先查数据把batch里的数据打印出来看看输入和标签是不是对应的。我遇到过标签错位的情况——数据预处理的时候把输入和标签搞反了模型当然学不会。再查Loss计算确认loss函数的reduction参数、ignore_index设置、mask逻辑是否正确。特别是padding位置有没有被mask掉。再查梯度打印梯度范数看是不是消失了或者爆炸了。如果梯度正常但loss不降可能是学习率太小。最后查模型确认模型结构对不对有没有层被意外冻结初始化是不是合理。7.2 显存溢出的应急方案显存溢出的时候按这个顺序尝试减小batch size最直接但可能影响训练稳定性开启混合精度训练省显存效果明显开启梯度检查点用时间换空间减小序列长度如果任务允许清理不必要的缓存torch.cuda.empty_cache()用梯度累积模拟大batch注意torch.cuda.empty_cache()只能释放PyTorch缓存中未使用的显存不能释放正在使用的显存。频繁调用反而会降低性能。7.3 训练速度慢的定位方法用PyTorch Profiler可以精确定位瓶颈with torch.profiler.profile( activities[torch.profiler.ProfilerActivity.CPU, torch.profiler.ProfilerActivity.CUDA], scheduletorch.profiler.schedule(wait1, warmup1, active3), on_trace_readytorch.profiler.tensorboard_trace_handler(./log) ) as prof: for step, batch in enumerate(dataloader): train_step(batch) prof.step()跑完之后在TensorBoard里看能清楚看到每个操作占用了多少时间。常见瓶颈包括数据加载慢、GPU利用率低、通信开销大、某些操作没有GPU加速。7.4 独家避坑技巧汇总随机种子训练前固定随机种子保证结果可复现。torch.manual_seed(42)、np.random.seed(42)、random.seed(42)三件套。学习率warmup大模型训练一定要用warmup前几百步学习率从0线性增加到设定值。避免一开始梯度太大把参数带偏。梯度累积显存不够的时候用梯度累积模拟大batch。注意loss要除以累积步数否则梯度会放大。检查点保存定期保存模型检查点包括优化器状态。训练中断的时候可以从检查点恢复不用从头开始。日志记录loss、学习率、梯度范数、显存占用这些都要记录。出问题的时候日志是唯一的线索。小规模验证正式训练之前用少量数据跑通全流程确认没有bug之后再上全量数据。8. 从理解到落地一套可直接复用的训练模板8.1 项目结构设计一个规范的训练项目应该有这样的结构project/ ├── configs/ # 配置文件 │ └── default.yaml ├── data/ # 数据目录 │ ├── raw/ │ └── processed/ ├── src/ │ ├── data/ # 数据加载和预处理 │ │ ├── dataset.py │ │ └── collator.py │ ├── models/ # 模型定义 │ │ └── model.py │ ├── train/ # 训练逻辑 │ │ ├── trainer.py │ │ └── optimizer.py │ └── utils/ # 工具函数 │ ├── logger.py │ └── checkpoint.py ├── scripts/ # 启动脚本 │ └── train.sh └── requirements.txt这种结构的好处是职责清晰数据、模型、训练逻辑分开方便调试和复用。8.2 核心训练代码模板import torch import torch.nn as nn from torch.utils.data import DataLoader from torch.cuda.amp import autocast, GradScaler from tqdm import tqdm class Trainer: def __init__(self, model, train_dataset, config): self.model model self.config config self.device torch.device(cuda if torch.cuda.is_available() else cpu) self.model.to(self.device) self.dataloader DataLoader( train_dataset, batch_sizeconfig.batch_size, shuffleTrue, num_workersconfig.num_workers, pin_memoryTrue, collate_fncollate_fn ) self.optimizer torch.optim.AdamW( model.parameters(), lrconfig.lr, weight_decayconfig.weight_decay ) self.scheduler torch.optim.lr_scheduler.CosineAnnealingLR( self.optimizer, T_maxconfig.num_epochs * len(self.dataloader) ) self.scaler GradScaler() self.global_step 0 def train(self): self.model.train() for epoch in range(self.config.num_epochs): total_loss 0 progress_bar tqdm(self.dataloader, descfEpoch {epoch}) for batch in progress_bar: loss self.train_step(batch) total_loss loss progress_bar.set_postfix({loss: f{loss:.4f}}) avg_loss total_loss / len(self.dataloader) print(fEpoch {epoch} finished, avg_loss{avg_loss:.4f}) self.save_checkpoint(epoch) def train_step(self, batch): input_ids batch[input_ids].to(self.device) attention_mask batch[attention_mask].to(self.device) labels batch[labels].to(self.device) self.optimizer.zero_grad() with autocast(): outputs self.model(input_ids, attention_maskattention_mask) loss self.compute_loss(outputs.logits, labels) self.scaler.scale(loss).backward() self.scaler.unscale_(self.optimizer) torch.nn.utils.clip_grad_norm_(self.model.parameters(), max_norm1.0) self.scaler.step(self.optimizer) self.scaler.update() self.scheduler.step() self.global_step 1 return loss.item() def compute_loss(self, logits, labels): shift_logits logits[..., :-1, :].contiguous() shift_labels labels[..., 1:].contiguous() loss_fn nn.CrossEntropyLoss() return loss_fn(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1)) def save_checkpoint(self, epoch): checkpoint { epoch: epoch, model_state_dict: self.model.state_dict(), optimizer_state_dict: self.optimizer.state_dict(), scheduler_state_dict: self.scheduler.state_dict(), global_step: self.global_step } torch.save(checkpoint, fcheckpoint_epoch_{epoch}.pt)这个模板包含了混合精度、梯度裁剪、学习率调度、检查点保存可以直接拿来改改用。8.3 超参数配置建议参数推荐值说明batch_size16-64根据显存调整learning_rate1e-4到5e-5大模型用小的weight_decay0.01-0.1AdamW常用0.01warmup_steps500-2000总步数的5%-10%max_grad_norm1.0梯度裁剪阈值num_epochs3-10根据数据量调整dropout0.1防止过拟合这些值不是绝对的需要根据具体任务和数据调整。但作为起点它们能让你少走很多弯路。8.4 训练监控与日志训练过程中要监控的指标Loss训练loss和验证loss都要看。训练loss降但验证loss升说明过拟合了。学习率确认调度器按预期工作。梯度范数确认没有消失或爆炸。显存占用确认没有泄漏。吞吐量每秒处理多少token用来评估训练效率。用TensorBoard或WandB记录这些指标训练完之后可以回看分析。from torch.utils.tensorboard import SummaryWriter writer SummaryWriter(runs/experiment_1) # 在训练循环里记录 writer.add_scalar(Loss/train, loss.item(), global_step) writer.add_scalar(Learning_rate, scheduler.get_last_lr()[0], global_step) writer.add_scalar(Grad_norm, total_norm, global_step)这些日志在排查问题的时候非常有用。我一般会保留最近几次实验的日志方便对比。9. 我个人的一些实操体会训练大模型这件事调包能跑通只是起点。真正让你和别人拉开差距的是你对底层逻辑的理解深度。我见过太多人模型跑不通就换框架、换库、换版本折腾半天发现是数据里有个NaN。如果你理解三明治的每一层在干什么你就能快速定位到问题出在哪一层。递推逻辑教会我的是序列数据是有时间结构的你不能把它当成一堆独立的样本处理。padding、mask、位置编码这些细节决定了模型能不能学到正确的模式。递归逻辑教会我的是梯度是模型的反馈信号信号太弱或太强都不行。梯度裁剪、残差连接、层归一化这些技术都是为了保持信号的稳定。三明治逻辑教会我的是训练是一个系统工程数据、模型、优化器三者必须配合好任何一个环节出问题都会拖累整体。最后分享一个小技巧每次开始一个新项目先用极小的数据比如100条跑通全流程确认loss能降到接近0。这一步能帮你排除掉90%的bug。然后再上全量数据调超参做优化。这个习惯让我省下了大量 debugging 的时间。
阅读完成 · 觉得有帮助?
咨询建站