大语言模型训练面临显存与通信瓶颈，本文详解了 CPU Offload、Checkpointing、量化压缩等显存优化技术，以及 Ring AllReduce、混合精度训练等通信与精度策略。重点阐述了 Zero 优化器的三个阶段及其显存节省原理，对比了数据并行、流水线并行、张量并行及 3D 并行的优缺点。此外，文章提供了 Transformer 架构的 FLOPs 计算公式，帮助开发者评估计算成本。通过合理组合这些技术，可有效解决大规模模型训练中的资源限制问题。
随着大语言模型（LLM）参数规模的爆炸式增长，训练过程面临着显存容量不足、通信带宽瓶颈以及计算效率低下等严峻挑战。为了在有限的硬件资源下成功训练大规模模型，业界发展出了一系列关键的优化技术。本文详细解析了从显存管理、精度控制到并行策略的核心训练技巧。
原理：用额外的通讯开销换取显存空间。对于模型计算的中间结果（如 Activation、优化器状态等），暂时将其从 GPU 显存迁移到系统内存（CPU RAM）中。当计算需要这些数据时，再通过 PCIe 总线传输回 GPU。
适用场景：适用于单卡显存不足以容纳整个 Batch 或模型状态的情况。虽然能显著降低显存峰值占用，但频繁的 CPU-GPU 数据传输会引入显著的延迟，可能降低训练吞吐量。
原理：用额外的计算时间换取显存空间。在前向传播过程中，不保存所有中间激活值（Activations），而是只保存部分关键节点或丢弃它们。在反向传播计算梯度时，根据需要的输入重新执行前向计算来恢复这些激活值。
优势：可以将显存占用减少约一半，特别适合深层网络。代价是增加了反向传播的计算量，通常增加 30%-50% 的训练时间。
原理：通过减少参数表示的位数来减小模型存储量和计算量。例如将 FP32 转换为 FP16、INT8 甚至 INT4。
影响：通常会带来一定的模型精度损失，但在大模型训练中，这种损失往往是可以接受的。量化不仅减少了显存占用，还能利用低精度指令集加速计算。常见的量化方案包括 Post-Training Quantization (PTQ) 和 Quantization-Aware Training (QAT)。
Ring AllReduce 是一种高效的分布式集合通信算法，常用于数据并行中的梯度同步。
工作流程：
效果：相比传统的 AllReduce 实现，Ring AllReduce 能够充分利用网络带宽，降低通信延迟，适合多机多卡环境。
背景：模型通常使用 float32 精度进行训练，但随着模型越来越大，训练的硬件成本和时间成本急剧增加。采用 float16 精度可以解决这一问题。
问题：直接使用 float16 可能导致梯度值太小，超出 float16 表示范围（下溢），导致权重不再更新，模型难以收敛。
解决方案：
显存分析：假设参数量为 X，参数和梯度使用 float16（各占 2X），优化器存储 float32 副本及状态（共 8X），总显存约为 12X。相比纯 float32 的 32X 显存需求，节省显著。
零冗余优化器（Zero Redundancy Optimizer, ZeRO）是一种高效的数据并行策略，旨在克服标准数据并行中每个 GPU 都保存完整模型状态的缺点。ZeRO 通过对模型状态（优化器状态、梯度、权重）进行划分后存储在单个 GPU 上，然后需要的时候通过动态通信调度来降低单卡显存占用。
将优化器状态划分成 Nd 份，每一份存到不同的 GPU 上。每个 GPU 只需要存储和更新总优化器状态的 1/Nd。
在优化器状态划分的基础上，将梯度划分成 Nd 份，每一份存到不同的 GPU 上。
在前两者的基础上，将参数划分成 Nd 份，每一份存到不同的 GPU 上。在前向和反向传播时，通过广播（Broadcast）获取完整参数。
不同设备执行相同的模型，处理不同的数据批次。这是最基础的并行方式，但受限于单卡显存大小。
当一个模型大到单个 GPU 无法训练时，最直接的想法是对模型层进行划分，将划分后的部分放置在不同的 GPU 上。
GPipe 将 minibatch 划分为更小且相等尺寸的 microbatch 来提高效率。前一个计算设备计算出该 microbatch 对应的结果会马上传给下一个计算设备，同时接着对下一个 microbatch 进行计算。这样能同时进行通信和计算。
张量并行的核心是将矩阵乘法进行拆分，分配到多个 GPU 上计算，降低对单个 GPU 的计算需求。TP 需要大量通讯，因此不建议跨多个节点进行张量并行。实际中，若一个节点有 4 个 GPU，最高的张量并行度通常为 4。
基于流水线并行将模型按 stage 划分到不同的管道，每个管道处理一个模型的 stage；基于张量并行将模型的每个 stage 按张量切分，划分成不同块；最后数据并行可以将这种 2D 组合扩展到任意数量的 GPU 上。
示例配置：
显存估算： 已知 Transformer encoder 的参数为：embedding(E)，sequence(s)，attention head(ah)，vocabulary(v)，hidden size(h)，layer(n)。
设 DP=8，TP=8，PP=16，使用基于 ZeRO 的 3D 并行，单张 GPU 的模型参数量将大幅降低，具体取决于 ZeRO Stage 的设置。
FLOPs（Floating Point Operations）意指浮点运算数，用来衡量算法/模型的复杂度。基于标准 Transformer decoder 结构的模型的 FLOPs 计算方法如下：
2 × seq_len × vocab_size × d_model2 × 3 × seq_len × d_model × (key_size × num_heads)2 × seq_len × seq_len × (key_size × num_heads)3 × num_heads × seq_len × seq_len2 × seq_len × seq_len × (key_size × num_heads)2 × seq_len × (key_size × num_heads) × d_model2 × seq_len × (d_model × ffw_size + d_model × ffw_size)2 × seq_len × d_model × vocab_sizeTotal forward pass FLOPs = embeddings + num_layers × (total_attention + dense_block) + logits Total backward pass FLOPs = 2 × Total forward pass FLOPs Total FLOPs = Total forward pass FLOPs + Total backward pass FLOPs
Total FLOPs ≈ 6DN，其中 D 是总的训练 tokens 数，N 是模型的参数量。这个公式提供了快速评估训练计算成本的依据。
大语言模型的训练是一个系统工程，需要在显存、通信和计算之间寻找平衡。选择合适的优化策略组合至关重要：
通过上述技术的综合应用，可以在现有的硬件条件下实现更大规模模型的训练与微调，推动人工智能技术的发展。
微信公众号「极客日志」，在微信中扫描左侧二维码关注。展示文案：极客日志 zeeklog
使用加密算法（如AES、TripleDES、Rabbit或RC4）加密和解密文本明文。 在线工具，加密/解密文本在线工具，online
生成新的随机RSA私钥和公钥pem证书。 在线工具，RSA密钥对生成器在线工具，online
基于 Mermaid.js 实时预览流程图、时序图等图表，支持源码编辑与即时渲染。 在线工具，Mermaid 预览与可视化编辑在线工具，online
解析常见 curl 参数并生成 fetch、axios、PHP curl 或 Python requests 示例代码。 在线工具，curl 转代码在线工具，online
将字符串编码和解码为其 Base64 格式表示形式即可。 在线工具，Base64 字符串编码/解码在线工具，online
将字符串、文件或图像转换为其 Base64 表示形式。 在线工具，Base64 文件转换器在线工具，online