今天进行简单的训练框架的调研。和我们在日常的博客、知乎文章和其他零零碎碎的地方看到的开源资料不同,其实从Megatron-LM出世到现在,训练框架已经发生了长足的发展,乃至于接近发生了范式级别的变化,不过,就我个人的观察而言,目前的开源资料针对这些新变化的普及其实还不够充分。具体而言,至少在最近1~2年的时间跨度内,工业界的前沿热点已经逐渐从单纯的实现高性能、大规模的预训练,转向中等规模的微调、后训练等等面向不同具体领域和需求的后训练相关内容。
这个演化的方向和近1~2年中Agent的崛起密切相关,同时也和参数规模的膨胀相关。随着LLM的scaling law发力,在5年间,LLM的参数量已经从单卡可以推理的约100M增加到了令人发指的约1000B,而后者光是训练就需要以十万计数量的GPU,进行持续的为期数月甚至将近年为单位的后训练。许多更小的模型则随后从这个教师模型中提取知识进行蒸馏。这意味着一件事情:基模本身的训练越来越依赖于规模效应,而规模效应和硬件资源挂钩,硬件资源和财力上限挂钩。因此,越来越多的无法在这个维度上竞争的企业(包括企业内部的团队)转向后训练,即通过更小卡需求的方式和可控规模的训练数据量,实现基于基模的垂直领域性能表现提升。也因此,这一热点需求的改变反过来催生了训练框架范式的改变:从规整的、追求规模可扩展性、强调分布式设计和并行效率的大规模预训练框架,变成强调灵活协调性、追求复杂训练流程和支持不同训练算法的后训练和强化学习训练框架,同时可能还需要和推理框架有效整合。
值得注意的是,所有的框架几乎全部基于Python+Pytorch,底层通信依赖NCCL,通过PyTorch的torch.distributed封装调用。具体来说,PyTorch提供了 ProcessGroup、all_reduce、all_gather、reduce_scatter 等原语,各框架在此之上构建自己的并行策略。部分框架(Megatron-LM、DeepSpeed)还包含自定义的C++/CUDA扩展来加速特定算子。Triton kernel在较新的框架(TorchTitan、Nanotron)中也有使用。这再次提醒了Python语言(实际上是Pytorch平台)在LLM时代的统治力。即使它并不是性能最优化的语言,也依然在工业场景里没有被C/C++淘汰,很大程度上和早期的奠基者效应以及完善的生态依赖有关系。
什么是训练
准确的说,能够让模型性能提升的操作都叫训练。在实践中,它又分为两个领域,预训练和后训练。就像在前面提到的那样,前者追求极致性能和规模的scaling,而后者追求训练调度复杂性的处理。
预训练
预训练部分的技术重点是维度切分和集合通信。这一部分在先前的学习日志#1和#2当中已经比较详细的整理了,在此不做水字数的冗余讲解。就现在而言,在当下已经落地和可预见的大模型架构下,可切分的维度基本已经被完全发掘,数学的理论意义上已经几乎没有数量级层面的突破可能性。因此,这部分的热点目前在于“大规模”、“高性能”和“基于MoE的异构”。很明显,这些都远远超过了个人、小型团队甚至许多中型企业团队的能力范畴,这从近几年国内的有竞争力的基模团队和相应人才的演化趋势当中可见一斑:豆包(字节跳动-Seed)、Qwen(阿里巴巴)、混元(腾讯)、GLM(智谱华章)、Kimi(月之暗面)、Deepseek(幻方量化)、文心一言(百度),无一不是具有雄厚财力支撑下的结果(也许Kimi除外)。
后训练
监督微调(SFT)
监督微调虽然被归类为后训练,但其实其在训练的特性上和预训练没有过多的差别,甚至可以说很相似。它的通俗解释就是“对只会续写文字的预训练基模进行训练,使得其可以按照对话的形式,根据提问和上下文给出回答”。它的技术要点在于loss和梯度更新的mask:模型是不需要对固定的格式部分,例如<|assistant|>,只需要在assistant回复的内容部分上算loss。除此之外,没有特别需要注意的地方了,大部分工作都在于数据本身(system/user/assistant的多轮对话格式),而非训练的实现。
基于偏好的对齐(Preference Alignment)
在这个领域内,可以说产生的算法和算法家族是最多的,也是最复杂的。笔者本人在学习的时候,也很长时间没能整理出一个具体的分类关系和所以然来。经过反复的调研,个人认为以下的心智模型是最好的,记录在这里,供读者参考。
后训练中的对齐/强化学习方法,可以沿两个维度来分类:
第一个维度:用于评价好坏的数据来自哪里?是1、事先收集好的(offline),还是2、训练过程中模型实时生成的(online)。
第二个维度:reward 信号来自哪里?是1、人类标注的偏好,2、训练好的 reward model,还是说3、可验证的规则。
把这两个维度交叉起来,就可以对所有目前的主流方法进行无重复无遗漏的归类。
数据来源的两种模式
模式一:离线方法。这种模式意味着在训练开始之前,已经完整拥有了一个完整的偏好数据集,其中每条数据是一个prompt加上一个 chosen response和一个rejected response(正负样本对),这些数据是提前收集好的,其来源不一,可能由人类标注,也可能由更强的模型生成。这些数据接下来以其原始文本的形式直接被用于新模型的训练,而不是通过任何中间层学习(再重新生成)之后再用于训练,这是和在线方法的区别之处。
训练过程中,模型不需要生成任何东西。它直接读取这些固定的数据对,通过特定的loss函数(鼓励policy相对于reference更偏好 chosen)来更新参数。整个训练循环和 SFT 几乎一样,也和预训练更类似,其工作流大致是读batch、forward、算loss、backward、update。没有rollout,也就是模型不新生成这些固定数据对之外的数据以获得评价(在这个工作流下无法获得评价),不存在一个可以对模型本身生成内容的评估也就是没有reward model的在线打分。
实现简单且对数据量要求低,但这个模式存在distribution shift问题:偏好数据是由某个旧policy(或人类)生成的,但随着训练进行,模型被鼓励避免旧输出模式之后,当前policy的行为分布会偏离训练数据的分布。模型无法在自己新的输出模式上得到反馈。
模式二:在线方法。训练过程中,模型自己生成回复,然后在这些自己生成的回复上获得反馈并学习。一个online RLHF的step大概是这样的:
1. 从 prompt 池中采样一批 prompt
2. 当前 policy 对这些 prompt 生成回复(rollout)
3. 获取 reward 信号
4. 用 RL 算法更新 policy
5. 回到第 1 步
复杂很多,但模型始终在自己当前的分布上生成数据并获得反馈,不存在 distribution shift问题。不过,它本身有比较严重的reward hacking问题,因为reward model在泛化之后通常质量不如人类直接进行偏好标注。
Reward信号来源的三种模式
模式一:人类直接标注。 最原始的RLHF,人类看模型的输出,给偏好排序。好处是reward质量高(人类直接进行选择),坏处是成本极高、速度很慢,无法规模化。早期InstructGPT/ChatGPT用的是这种模式,但现在已经很少直接用了。
模式二:Reward Model打分。 先用人类标注的偏好数据训练一个reward model(本质上是一个回归模型,输入prompt+response,输出一个标量分数)。之后的RL训练中,用这个reward model代替人类做在线打分。PPO和GRPO传统上都用这种模式。好处是可以规模化,reward model速度比人类标注快几个数量级。坏处是reward model本身可能不准确,而且policy可能学会hack reward model,即找到 一个方法,可以使得reward model系统性的给高分,但实际回复的质量对人类来说并不高。
模式三:可验证的规则,基于有ground truth的reward信号进行后训练,这就是RLVR。这种reward信号有几个非常好的性质。第一,它是ground truth,有客观评价指标,没有模糊空间,因此可以避免reward hacking问题。第二,它是完全免费的,不需要人类标注和reward model,一个Python脚本就可以直接进行验证并且输出+1或者-1的reward。第三,它是无限可扩展的,通过自动方法可以生成无限多的数学题和编程题作为训练数据。不过,并不是所有问题都是可验证的,许多语言风格、道德倾向等等的训练目标完全无法通过自动形式验证,因此这个模式可以做的训练目标本身受限。
Offline方法
DPO
虽然写在前面,但其实这个方法的诞生晚于后面提到的PPO。其核心思路是绕开reward model,直接用偏好数据训练policy,从而简化训练流程和在小样本下提升训练的稳定性和可行性。每条数据(数据对)包含三个部分,一个prompt + 一个chosen response + 一个rejected response。训练时需要:
- 用一个冻结的reference model分别对chosen和rejected计算log probabilities
- 用当前policy model同样分别计算log probabilities
- Loss基于这四组log probabilities构造,鼓励policy相对于reference更偏好chosen
一个数据对的训练因此实际上要跑四次forward过程(ref×chosen、ref×rejected、policy×chosen、policy×rejected)。实际上,它很难真的和其他的后训练方法被放在一起,因为它实际上处理的是整个sequence level的优化,没有一个RL的过程,换句话说根本没有进入RL的框架。
此外,也存在SimPO和KTO这样的DPO变体,分别处理不需要reference model的情况和无数据对只有答案的好坏性质(不成对,单独存在)的情况,在此不逐一介绍,感兴趣的读者可以自行了解。它们的关系如下:
数据格式 reference model 长度处理
─────────────────────────────────────────────────────────
DPO 偏好对 需要 sum(有长度偏差)
(chosen, rejected)
SimPO 偏好对 不需要 average(无偏差)
(chosen, rejected) + margin
KTO 逐条标注 需要 sum
(desirable 或
undesirable)
Online RLHF
PPO
这是最经典和最复杂的后训练算法,其思想直接借鉴了早期的RL相关的理论和实践。其完整流程涉及四个模型:
- Rollout 阶段(推理):policy model对prompt做自回归的generation,生成回复
- Reward 阶段(推理):reward model对生成的回复打分
- Advantage 估计:value model估计每个token的 value,结合reward算GAE
- Policy update(训练):用PPO clip loss和相对reference model的KL散度更新policy,同时更新value model
区别在于多出了一个value model和一个reward model。它们是协作关系:reward model需要对每一个policy model的Prompt整体给出一个标量估计,而value model负责对于每一步状态给出一个估计,即从这个状态继续出发按照现有模型的策略和分布,最终期望能拿到什么程度的reward model,因此实际上是一个把reward从sequence的整体level分解到任意中间状态的辅助评估者。
实际上,从这个视角观察的话,reward model和value model在思路上有相似之处,都是基于模仿的扩展者。reward model从标注的数据对中学习,把少量的sequence泛化到对任意整体sequence都有一个标量可供评估,而value model从reward model中学习,把对任意整体sequence都有一个连续分布的标量泛化到不完整的任意中间状态sequence都有一个连续分布的标量可供评估。我认为这个心智模型是有助于理解和记忆的。
值得一提的是,技术上的关键难点之一在于rollout阶段是自回归推理,需要KV cache、高效 batching、高效通信,以及很多推理相关的技术。这就是为什么veRL/OpenRLHF要集成vLLM。
GRPO
DeepSeek 提出的简化版PPO,核心改进是只使用三个模型,也就是说去掉了value model。对每个prompt生成一组(比如8个)回复,用reward model打分后,在组内做归一化得到advantage。也就是说,它和DPO一样,是一个sequence level赋予reward而不是token level进行评估的方法,其消除噪声的方法在于通过一组的密集输出和reward的归一化来将sequence level的不精确性部分抵消。它的最终性能上限不一定有PPO好,但简单的多。这样就只需要policy + reference + reward三个模型,不需要value model去做critic。实现上比PPO简单不少,效果也被验证过,目前非常热门。
RLVR(RL with Verifiable Rewards)
这是最近几年被Deepseek带起来的时兴架构,也就是DeepSeek-R1、Qwen等模型采用的路线。它和传统RLHF的区别在于reward不来自训练好的reward model,而来自可验证的规则(数学题的答案是否正确、代码是否通过测试用例)。工程上它依然是RL训练,因此需要rollout、需要reward scoring、需要policy update,但reward的来源从一个神经网络或者一个人类的标注结果变成了一个确定性的验证器。
RLVR的产生是LLM本身能力目标发生范式转移的标志之一。随着DeepSeek-R1的成功,它展示了纯RL(不经过传统的SFT中间步骤)就能让模型学会复杂的推理行为:长链chain-of-thought、自我纠错、"wait let me reconsider"这种反思模式。从此,推理能力成为新战场。2024-2025年的LLM竞争焦点从"对话能力"转向"推理能力"(数学、代码、逻辑),而推理任务恰好是最容易做verifiable reward的领域。这让RLVR有了天然的应用场景。
训练会用到的其他技术特性
除了1F1B这类调度特性、TP-PP-DP-CP-SP-EP的并行切分之外,其实还有许多无法简单以谱系方法归类,但仍然值得一提的技术特性。
FP8精度训练
FP8精度训练是一种对计算进行优化的方法。在这个方法下,权重和激活以FP8格式直接喂给tensor core做矩阵乘法,GPU的计算单元直接使用8-bit精度做运算。H100/B200的tensor core原生支持FP8 GEMM,而且FP8的吞吐量是BF16的两倍(因为同样的硬件单元一个周期能处理两倍数量的8-bit元素)。因此,使用FP8训练既省显存(tensor变小了),又提升速度(FLOPS翻倍)。
具体来说,一个FP8训练的GEMM工作流如下:1、输入的X和W都是BF16格式,2、将X和W量化为X_fp8 = quantize(X / scale_X), W_fp8 = quantize(W / scale_W),均为FP8格式,3、计算:Y_fp32_acc = X_fp8 @ W_fp8,其中Y_fp32_acc是FP32格式,4、输出Y_bf16 = Y_fp32_acc * scale_X * scale_W,其中Y_bf16和输入一样,是BF16格式。
在更细节的层面,FP8有两种格式,针对不同场景:
E4M3(4-bit指数+3-bit尾数):动态范围较小但精度较高,用于forward pass的权重和激活。
E5M2(5-bit指数+2-bit尾数):动态范围较大但精度较低,用于backward pass的梯度。这是因为梯度的数值范围波动更大,需要更大的动态范围。
LoRA和Q-LoRA
LoRA的思想相对简单:用低维度矩阵及其乘法结果的地址矩阵作为可训练量参与计算,即:$y=Wx+ABx$,其中A、B是可训练的小矩阵,$A: [r, d_{in}]$,$B: [d_{out}, r]$,其中r是rank,通常取8、16、32、64这样的值,远小于d_in和d_out(例如,~4096)。这里不多展开。
QLoRA的出发点是:LoRA虽然省了可训练参数的显存,但冻结的基座权重W仍然占着14GB(7B×BF16,即2)。能不能把这个也压缩呢?对此,QLoRA的做法是:把冻结的基座权重量化到4-bit存储,forward时动态反量化回BF16做计算,LoRA部分仍然以BF16训练。
具体技术有三个关键点。
NF4(NormalFloat 4-bit)。这是QLoRA论文提出的一种4-bit量化格式。它的设计思路是:预训练模型的权重分布近似正态分布,那么 4-bit 的16个量化档位应该按正态分布的分位点来划分,而不是均匀划分。这样信息论意义上最优,即每个量化档位承载的信息量相等。和均匀量化相比,NF4在同样的bit数下能更好地保留权重信息。
Double Quantization。量化时每个block(比如64个权重)需要一个scale factor(FP32,4bytes)。如果block size=64,那么scale factor的额外开销是4/64=0.0625bytes per parameter。虽然看起来不多但7B参数下也有~0.4GB。Double quantization是对这些scale factor本身再做一次量化(量化到FP8),进一步压缩这个开销。
Paged Optimizers。QLoRA 还用了NVIDIA的unified memory(CUDA managed memory)来管理优化器状态,当GPU显存不够时自动page到CPU内存,避免OOM,这其实和后面ZeRO-Offload的思路类似。
方法 基座权重 优化器状态 总计
─────────────────────────────────────────────────
全量微调 BF16 14 GB 84 GB ~120 GB(需要多卡)
LoRA BF16 14 GB 0.5 GB ~30 GB(单卡 A100)
QLoRA 4-bit 3.5 GB 0.5 GB ~10 GB(单卡 RTX 4090)
QLoRA让单张24GB消费级GPU就能微调7B模型,这对学术界和个人开发者的意义巨大。但它也有其局限性。第一,LoRA的低秩约束意味着表达能力上界。第二,QLoRA的激进4bit量化会带来不可忽略的精度损失,而且反量化的计算引入了额外延迟,训练速度通常比纯BF16 LoRA慢 20-30%。第三,LoRA需求一个训练好的预训练基座, 因此不适合预训练。
Zero-1/2/3/Offload/Infinity
前三者不再赘述,可以参见之前的博客,这里讲后两个,他们都是关于存储层次利用的工作。先回顾ZeRO-3的状态:模型参数、梯度、优化器状态全部被shard到各张GPU上。每张GPU只持有1/N的完整状态,需要某个参数时通过all-gather临时拼回来。这已经把GPU显存的利用率压到很低了,但如果模型再大,即使shard之后每张卡分到的那1/N也放不下呢?
ZeRO-Offload的思路是:GPU显存不够,就往CPU内存搬。CPU内存通常比GPU显存大一个数量级(比如一台机器GPU显存总共640GB,但CPU内存可能有 1-2TB),而且便宜得多。
具体来说,ZeRO-Offload把训练过程中的不同计算和数据按照"计算密度"来划分——计算密集的操作留在GPU,内存密集但计算量不大的操作搬到CPU:
留在GPU上的:forward和backward的计算(矩阵乘法、attention等),因为这些是计算密集型,GPU的算力优势在这里是不可替代的。
搬到CPU上的:优化器状态(Adam的exp_avg和exp_avg_sq)和优化器更新计算(parameter update)。Adam的更新本质上是若干个逐元素运算,读四个 tensor、做加减乘除、写回去。计算量不大,但占显存的大头(fp32 master weights+两个矩估计=参数量的12倍字节)。把这些搬到CPU,用CPU做Adam update,然后把更新后的参数传回GPU。
关键的性能瓶颈在PCIe带宽。GPU和CPU之间的数据传输走PCIe,带宽大约32-64GB/s(PCIe Gen4/Gen5),比GPU内部的HBM带宽(几TB/s)和节点内NVLink(几百GB/s)差了一到两个数量级。所以ZeRO-Offload需要精心设计通信与计算的overlap,在GPU做当前层的backward的同时,把上一层的梯度通过PCIe传到CPU;在CPU做optimizer update的同时,GPU继续做下一个microbatch的forward。
ZeRO-Offload的收益场景是:GPU数量少但模型大。同时,ZeRO-Infinity是ZeRO-Offload的进一步延伸:CPU内存也不够的话,就往NVMe SSD搬。这个在工业上通常过慢,不具体展开。不过,其中的通算overlap和prefetch思想值得借鉴和考虑。
Activation Checkpointing
是一个很有意思的细粒度技术。它的主要思想是:保留部分中间计算值,丢弃部分中间计算值,以减少缓存压力。在没有保存中间值的部分,需要计算梯度的时候,则重新通过最近的激活值通过前向传播重新回到这个地方,进行计算,在计算完成后再丢弃。此外,和这个相关伴生的是反向传播情况下的Kernel设计和fusion相关的工作。几个著名的例子是Flash Attention的backward版本,以及optimizer fusion相关的Kernel,即通过合并Adam的四个需要更新的Tensor(param+grad+exp_avg+exp_avg_sq)的更新和合并逐元素操作来显著减少memory bandwidth压力;还有fused grad accumulation,即backward产生的梯度直接累加到梯度buffer里而不是先写出来再加,在gradient accumulation场景下省一次读写。
总的来说,推理的融合目标主要是减少kernel launch开销和memory bandwidth,训练的融合目标除了这些之外,还多了一个"减少中间 tensor的materialize以配合activation checkpointing的显存策略"。
Fault tolerance
这具体来说是指一类计数,即容错,也是典型的分布式场景的技术应用。这个概念其实远早于,但在LLM的训练场景下由于不同的workload特性被赋予了。在这一领域,目前经常采用的技术就是Checkpoint(存盘恢复),还有自动重启和弹性训练。此外,梯度和Loss的异常检测、静默故障检测、心跳检测和快速故障发现等技术也有所应用。这个领域目前的主要故障类型是硬件故障(ECC错误)、软件故障(NCCL通信超时和OOM)、静默错误,热点挑战之一是弹性的扩缩容和并行度的变化。例如,如果用TP=4、PP=2训练配置保存了checkpoint,故障恢复时某些机器挂了,随后想换成TP=2、PP=4重新开始,这时候就需要checkpoint格式支持不同并行拓扑之间的resharding,而这其实是一个non-trivial的课题。此外,网络层面的容错也需要进行相应的处理(慢节点、链路故障等等)。
Ray
准确的说这是一个分布式框架而不是一个计数特性,更准确的说它其实本来和LLM毫无关系:它本质上是一个分布式任务调度框架,其原始定位是通用的分布式Python程序执行框架。它的核心概念包括了两个基本类型:
Task:一个无状态的远程函数调用。写一个普通Python函数,加上 @ray.remote 装饰器,Ray会把它调度到集群中某台机器的某个进程上执行,返回一个future(异步句柄)。
Actor:一个有状态的远程对象。写一个普通Python类,加上 @ray.remote 装饰器,Ray会在某台机器上实例化它,之后就可以远程调用它的方法。每次调用都在同一个实例上执行,所以它有持久状态。
它最经典的应用就是在OpenRLHF,连接训练端的DeepSpeed ZeRO和推理端的vLLM。