无 FSDP通常指传统的 DDP - 分布式数据并行假设有 8 张显卡每张显卡上都必须完整保存一份一模一样的模型副本。这就像让 8 个工人各自拿着一本完整的大字典在工作。如果这本字典模型太大任何一个工人的背包显存都装不下训练就直接崩溃了OOM显存溢出。FSDPFully Sharded Data Parallel完全分片数据并行。把这本大字典拆成了 8 份每个工人GPU只负责保管其中一份。只有当某个工人需要用到其他页的内容时他才会快速向别人“借”过来通过网络通信同步看完后立刻“销毁”借来的复印件以腾出空间。具体来说FSDP 将模型的参数Parameters、梯度Gradients和优化器状态Optimizer States全部分片并分散到了不同的 GPU 上。FSDPNoneGradient CheckpointingTrueFSDPfull_shard auto_wrapGradient CheckpointingTrueGradient Checkpointing梯度检查点有时也叫 Activation Checkpointing在标准的深度学习训练中分为前向传播和反向传播两步。1前向传播 数据从第一层一层层算到最后一层得出预测结果。在这个过程中每一层计算产生的中间结果我们称之为 Activations激活值都会被保存在显存里。2反向传播 模型根据预测误差从最后一层往回算更新参数。重点来了反向传播在计算梯度时必须用到前向传播时保存的那些“激活值”。问题在于对于 LLMs 来说层数极深如果把每一层的激活值全都存在显存里会占用极其庞大的空间。这往往是导致 CUDA Out of Memory显存溢出的罪魁祸首其占用的显存甚至比模型参数本身还要大得多。Example假设你在爬一座 100 层的塔前向传播。为了原路返回反向传播你原本需要在每一层都留一个记号保存激活值这会让你的背包显存塞满。开启 Gradient Checkpointing 后现在只在第 10 层、20 层、30 层……这些关键楼层留下记号也就是 Checkpoints。当你退回到第 19 层发现这里没有记号怎么办很简单你退到最近的第 10 层记号处重新往上爬一次重新计算一次前向传播 算出第 19 层的状态然后再继续往下退。Gradient Checkpointing 只在某些神经网络层的边界保存激活值。当反向传播需要用到被丢弃的激活值时模型会从最近的“检查点”重新进行一次局部的前向计算来临时恢复这些数据是一种“用计算时间换取显存空间”的技术。FSDPfull_shard auto_wrapGradient CheckpointingFalse. 【torch.OutOfMemoryError: CUDA out of memory. Tried to allocate 340.00 MiB. GPU 1 has a total capacity of 94.97 GiB of which 17.69 MiB is free. Including non-PyTorch memory, this process has 94.95 GiB memory in use. Of the allocated memory 86.44 GiB is allocated by PyTorch, and 7.53 GiB is reserved by PyTorch but unallocated. If reserved but unallocated memory is large try setting PYTORCH_CUDA_ALLOC_CONFexpandable_segments:True to avoid fragmentation. See documentation for Memory Management (https://docs.pytorch.org/docs/stable/notes/cuda.html#optimizing-memory-usage-with-pytorch-cuda-alloc-conf)】FSDPNoneGradient CheckpointingFalseFP16 (IEEE 754 Half-Precision)1位符号 5位指数 10位尾数优点精度较高因为有10位尾数动态范围极窄。它的最大值只能表示到 65504最小值是5.96×10−85.96×10^{−8}5.96×10−8BF16 (Brain Floating Point)1位符号 8位指数 7位尾数优点动态范围巨大。它的指数位和标准的 32位单精度浮点数FP32一模一样。这意味着它的最大值可以达到约3.4×10383.4×10^{38}3.4×1038缺点精度较低只有7位尾数。两个非常接近的小数在 BF16 下可能会被当作同一个数。但幸运的是神经网络对范围极其敏感对微小的精度误差却非常宽容甚至相当于一种正则化。BF16 FSDPFP32 FSDP显存消耗拆解在使用 Adam 优化器进行混合精度训练FP16/BF16时假设模型参数量为xxx单卡显存占用如下FP16 模型参数 (Parameters)2x2x2xbytesFP16 梯度 (Gradients)2x2x2xbytes优化器状态 (Optimizer States) Adam 需要保存 FP32 的参数备份4x4x4xbytes、Momentum4x4x4xbytes和 Variance4x4x4xbytes。共计12x12x12xbytes。因此显存占用总量是16x16x16xbytes.MB等单位以10为底数的指数MiB是以2为底数的指数如1KB10310001KB10^310001KB1031000,1MB10610000001000KB1MB10^610000001000KB1MB10610000001000KB,而1KiB21010241KiB2^{10}10241KiB2101024,1MiB22010485761024KiB1MiB2^{20}10485761024KiB1MiB22010485761024KiB数据并行Data Parallel用于模型参数量还能塞进单张显卡但数据量庞大时。PyTorch 中经历了从 DataParallel (DP) 到 DistributedDataParallel (DDP) 的演进。DP的问题单进程多线程Single-process, Multi-thread机制。受限于 Python 的全局解释器锁GIL多线程并不能实现真正的并行计算。通信瓶颈在于ParameterServer 架构每次 Forward 前主节点通常是 GPU 0需要把模型广播给所有 GPUBackward 后所有 GPU 的梯度都要汇聚到主节点进行求和与更新。GPU 0 成为了严重的通信瓶颈。DDP多进程Multi-process。每个 GPU 都有一个独立的进程独享属于自己的模型副本和优化器状态。DDP核心通信算法 Ring-AllReduce假设有 N 个 GPU每张卡上的梯度被切分成 N 个数据块。1Scatter-Reduce分散-规约 每张 GPU 将自己的第iii个数据块发送给下一个 GPU并接收前一个 GPU 发来的数据块进行累加。经过N−1N−1N−1步后每张 GPU 上都会拥有一块完整的、包含所有 GPU 累加结果的梯度块。2All-Gather全收集 每张 GPU 将自己已经计算完整的那个梯度块沿着环继续传递。再经过N−1N−1N−1步所有 GPU 都将拥有完整的、全局累加后的梯度。DDP Ring-AllReduce 数学分析不论有多少张 GPU 参与训练单张 GPU 的总通信数据量约为2(N−1)N×S\frac{2(N-1)}{N}\times SN2(N−1)​×S其中 S 是模型参数总大小。当NNN足够大时通信量近似为2S2S2S。它实现了通信开销与 GPU 数量NNN的解耦极其适合大规模集群。模型并行Model Parallelism当一个模型如 7B, 13B 甚至更大大到单张 GPU 的显存根本装不下时数据并行就失效了。必须切分模型本身。工业界主流采用的是 3D 并行TP PP DP。张量并行Tensor Parallelism, TP对模型中的矩阵乘法进行切分。主要参考 NVIDIA 的 Megatron-LM 架构。以 Transformer 的 MLP 层为例全连接层计算本质是YGeLU(XA)BY{\rm GeLU}(XA)BYGeLU(XA)B。1对权重矩阵 A 按列切分Column Parallelism每个 GPU 计算一部分列激活函数 GeLU 可以独立进行因为是 element-wise 的。2对权重矩阵 B 按行切分Row Parallelism各个 GPU 计算出的结果相加才能得到最终结果。【问题张量并行具体是如何进行切分的】TP 的代价TP 在每次切分层结束时需要进行一次 All-Reduce 通信操作。因为通信极其频繁TP 通常只在同一个物理节点Node内部的 GPU 之间进行依赖高带宽的 NVLink跨节点做 TP 会导致严重的网络延迟。流水线并行Pipeline Parallelism, PP核心思想是将模型的不同层按顺序分配到不同的 GPU 上。例如 GPU 0 计算 Layer 1-4GPU 1 计算 Layer 5-8。痛点在于如果是简单的顺序执行GPU 1 必须等待 GPU 0 计算完毕才能开始这会导致大量的 GPU 处于空闲等待状态即气泡。解决方案1Micro-batching微批次 将一个大的 Batch 切分成多个小 Micro-batch。GPU 0 计算完第一个 Micro-batch 的第一层后立刻传给 GPU 1 计算同时 GPU 0 开始计算第二个 Micro-batch。21F1B 调度One Forward, One Backward 是工业界常用策略指在稳态阶段GPU 交替执行一次前向传播和一次反向传播这能大幅降低显存中需要缓存的激活值数量。ZeRO 优化器Zero Redundancy Optimizer在传统的 DDP 中每个 GPU 都保存了完整的16x16x16x参数状态这造成了极大的冗余。ZeRO 通过切分这些状态来节省显存。1ZeRO-Stage 1切分优化器状态切分12x12x12xbytes 的优化器状态每张 GPU 只负责更新1/N1/N1/N的参数因此显存占用降低为4x12xN4x\frac{12x}{N}4xN12x​。2ZeRO-Stage 2切分梯度在 Stage 1 的基础上进一步切分梯度显存占用降低为2x14xN2x\frac{14x}{N}2xN14x​。3ZeRO-Stage 3切分参数将模型参数也完全切分。每张显卡上只保存1/N1/N1/N的参数。前向和反向传播时通过 All-Gather 动态获取需要的参数用完即弃。显存占用为16xN.\frac{16x}N.N16x​.梯度检查点Gradient Checkpointing时间换空间。在标准的前向传播Forward Pass中PyTorch 会将所有的中间激活值Activations保存在显存中以便在反向传播Backward Pass时计算梯度。对于层数极深的大模型这些激活值的显存占用通常远超模型参数本身。梯度检查点采用了一种“断点存档”的策略。在前向传播时它不再保存所有的中间激活值而是只保存部分关键层Checkpoint layers的激活值丢弃其余的。当反向传播进行到那些被丢弃了激活值的层时PyTorch 会利用最近保存的“检查点”重新做一次前向传播Recomputation来临时生成这些激活值。收益 可以将激活值的显存占用从O(N)O(N)O(N)降低到O(N)O(\sqrt N)O(N​)N 为层数极大地释放显存允许你跑更大的 Batch Size 或更长的序列。代价 引入了额外的计算开销。通常会导致整体训练时间增加 20% ~ 30% 左右。【问题梯度检查点中的“关键层”如何选取】KV Cache键值缓存是大语言模型如 GPT、Llama 等基于 Transformer 的模型在推理生成文本阶段用于加速计算的一项核心空间换时间的优化技术。。大语言模型生成文字的方式是自回归Autoregressive的也就是“逐字吐出”。为了准确生成下一个字模型必须联系上文所有已经生成过的内容。在 Transformer 的核心机制“自注意力Self-Attention”中输入的每一个 Token词元都会经过矩阵乘法被转化为 QKV 三个向量。如果不使用缓存模型每生成一个新词为了计算注意力都要把前面所有已经生成过的词重新计算一遍它们的 K 和 V。随着文章越来越长这种“每次都要从头算起”的冗余计算量是非常恐怖的会导致生成速度越来越慢。大模型的推理过程实际上分为两个截然不同的阶段Prefill预填充阶段 和 Decode解码阶段。假设我们给模型输入了提示词Prompt“大 型”我们希望它接着生成“语 言 模 型”。KV Cache 是如何工作的1缓存历史模型在处理之前的词时直接把计算出来的 K 矩阵和 V 矩阵保存在显存GPU VRAM中。2只算当下当需要生成第NNN个词时模型只计算当前这一个新词的 QKV。3提取记忆然后用当前新词的 Q去和显存里历史积累的所有 K 和 V 一起套入注意力公式进行计算。4更新缓存最后把这个新词的 KV 也追加到缓存池里供生成下一个词时使用。推理端常常面临“显存墙Memory Wall”的瓶颈。为了缓解这个问题在 KV Cache 之后业界后续又衍生出了 MQA (Multi-Query Attention)、GQA (Grouped-Query Attention) 以及 vLLM 框架中的 PagedAttention 等进阶技术。深度学习训练耗时训练耗时数据规模×\times×单步计算量///计算速率计算速率计算速率单设备计算速率×\times×设备数×\times×多设备并行效率加速比单设备计算速率混合精度、算子融合、梯度累加加速比数据并行、模型并行、流水并行设备数服务器架构、通信拓扑trainer.train() 会执行训练循环。1自动处理前向传播、计算 Loss、反向传播Backward、梯度裁剪Gradient Clipping以及调用优化器Optimizer和学习率调度器LR Scheduler更新权重。2自动处理数据在 GPU/TPU 上的放置并在需要时自动管理混合精度FP16/BF16以及复杂的分布式训练策略如 DDP、FSDP 或 DeepSpeed。3在训练的间隙它会根据设定的频率eval_steps 和 save_steps自动暂停训练在验证集上进行评估并将中途的模型和优化器状态保存为 Checkpoint。4如果你传入了 trainer.train(resume_from_checkpointTrue)它不会从头开始而是从指定的检查点恢复模型权重、优化器动量、随机数种子和当前步数继续训练。trainer.state.best_model_checkpoint 会追踪最佳模型路径记录在整个训练过程中表现最好的那个检查点Checkpoint所在的文件夹路径。1只有在 TrainingArguments 中设置了 load_best_model_at_endTrue 并且指定了 metric_for_best_model比如 ‘eval_loss’ 或 ‘accuracy’时这个属性才会被有效维护。2每次 Trainer 执行评估并保存 Checkpoint 时它都会对比当前的评估指标和历史最佳指标。如果当前模型更好它就会把这个 Checkpoint 的路径赋值给 trainer.state.best_model_checkpoint。3当 trainer.train() 执行到最后一步即将结束时如果 load_best_model_at_endTrueTrainer 会自动从这个路径读取权重覆盖掉内存中当前可能是过拟合的最后一步的模型权重。trainer.save_model(output_dir) 用于保存最终纯净模型。这是一个手动调用的方法用于将当前内存中的模型保存到硬盘上的指定目录中。它不仅会保存模型的权重pytorch_model.bin 或 model.safetensors还会保存配置文件config.json如果初始化 Trainer 时传入了分词器还会一并保存分词器文件tokenizer.json 等。与 Checkpoint 不同的是训练中途保存的 ckpt 文件通常体积非常大。因为它不仅包含模型权重还包含优化器状态Optimizer State、学习率调度器状态和随机数种子等。它的目的是为了恢复训练。Q在刚跑完的一组实验总steps数为535步save-steps为300最终一共保存了两个ckptcheckpoint-300 和 checkpoint-535。我的问题是为什么checkpoint-535会被保存下来A1触发 checkpoint-300常规的周期性保存。配置中指定了按步数保存save_strategy“steps”且保存频率为 300 步save_steps300。因此当训练顺利跑到第 300 步时Trainer 会按照你的指令雷打不动地保存下 checkpoint-300。2触发 checkpoint-535训练结束的强制保存。当训练继续进行跑到了第 535 步时达到了你设定的总步数上限可能因为 Epoch 跑完了或者触发了 max_steps训练循环正式宣告结束。此时虽然 535 不是 300 的整数倍但 Trainer 的底层源码中有一条硬性规则只要你开启了保存功能save_strategy 不是 “no”在整个训练循环完全结束的那一刻它必定会无条件地再保存最后一个 Checkpoint。remove_unused_columns (默认值: True)控制是否自动删除模型 forward() 方法签名中没有声明的数据列。默认情况下Hugging Face 的 Trainer 非常“智能”。当你把 Dataset 传给它时它会自动检查你底层模型比如 Qwen的 forward(input_ids, attention_mask, labels) 方法。如果它发现你的 Dataset 里包含模型不需要的额外字段比如 market_id, event_id它会在送入 GPU 之前静默地把这些字段删掉以节省显存和防止报错。ddp_find_unused_parameters (默认值: 视情况而定建议 False)控制在分布式训练 (DDP) 时是否花时间去寻找“没有参与反向传播的死参数”。在多卡训练DistributedDataParallel中每一层计算完梯度后所有卡必须互相通信All-Reduce同步梯度然后才能进行下一步。DDP 严格要求模型的所有参数都必须收到梯度如果有的参数在某次 forward 中没被用到DDP 就会死锁卡住。如果开启这个参数TruePyTorch 会在每次反向传播时从 Loss 往回极其吃力地遍历整个计算图把那些没用到的参数挑出来告诉其他卡“别等这些参数了它们是 0”。eval_accumulation_steps (默认值: None)控制在验证Evaluation阶段每隔多少步将预测结果从 GPU 显存转移到 CPU 内存。它解决的痛点 (Eval OOM)在验证阶段Trainer 需要把整个验证集的预测结果logits和真实标签labels全部攒在 GPU 显存里等跑完所有验证数据后再一次性计算诸如 MSE、Accuracy 等指标。对于大模型来说Logits 的张量通常极其庞大例如 [Batch_Size, Seq_Len, Vocab_Size]。如果你的验证集比较大验证跑到一半GPU 显存就会被这些积累的 Logits 撑爆报出经典的 CUDA Out of Memory。最让人崩溃的是训练没 OOM反而验证 OOM 了。如果你设置了比如 eval_accumulation_steps10Trainer 就会在每验证完 10 个 Batch 后把积压在 GPU 显存里的预测结果搬运到廉价且容量巨大的 CPU 内存中清空 GPU 显存以迎接下一批验证数据。eval_accumulation_steps 的核心本质是“用时间换空间”即牺牲一定的验证速度换取 GPU 显存的安全。如果设置得过小比如设为 1 或 2也就是每验证一两个 Batch 就把数据往外搬会带来非常明显的坏处主要体现在性能的急剧下降上1严重的 CPU-GPU 通信开销PCIe 瓶颈将张量Tensor从 GPU 显存转移到主机主板上的 CPU 内存需要经过 PCIe 总线。虽然现在的 PCIe 4.0/5.0 速度很快但跟 GPU 内部极其恐怖的显存带宽如 H100 的 3TB/s相比依然是“慢车道”。2GPU 算力严重闲置 (GPU Starvation) 在深度学习框架中跨设备的张量转移通常会引发同步阻塞Synchronization。GPU 处理前向传播Forward的速度极快如果每算完一小步就要等待数据搬运到 CPU就会导致 GPU 算一半停一半。load_best_model_at_end 控制在整个训练循环结束时自动将内存中的模型权重回滚替换为训练中途表现“最好”的那个版本而不是直接使用训练最后一步的模型。在深度学习中随着训练步数Epoch/Steps的增加模型在验证集上的表现通常是先变好后变差过拟合或者在后期呈现震荡。如果你不开启这个参数默认值 Falsetrainer.train() 结束后内存里保留的是最后一步Last Step的模型状态。这时候的模型大概率已经过拟合了并不是它的巅峰状态。load_best_model_at_end 必须配合以下参数一起使用才能生效。1save_strategy eval_strategy必须开启中途验证和保存比如设为 “steps” 或 “epoch”。如果你中途都不保存设为 “no”Trainer 去哪里给你找“最好的”那个模型呢2metric_for_best_model告诉裁判“最好”的标准是什么。默认是 “eval_loss”验证集损失越低越好。如果你在算分类准确率可以设为 “accuracy” 或 “f1”。3greater_is_better告诉裁判指标是越大越好准确率设为 True还是越小越好Loss设为 False。scatter_add_ 函数的核心作用是根据指定的索引index将源张量src中的元素累加到目标张量self的对应位置上。函数签名通常如下Tensor.scatter_add_(dim, index, src)。importtorch# 1. 初始化目标张量 (比如有 5 个节点初始特征全为 0)targettorch.zeros(5)# 2. 准备源张量 (要累加的值)srctorch.tensor([10.0,20.0,30.0,40.0,50.0])# 3. 准备索引张量 (告诉 src 里的数字分别要去 target 的哪个位置)# 比如 10 去位置 020 去位置 130 去位置 0...indextorch.tensor([0,1,0,2,2])# 4. 执行 scatter_add_ (沿着第 0 维)target.scatter_add_(dim0,indexindex,srcsrc)print(target)# 输出: tensor([40., 20., 90., 0., 0.])# 解释:# target[0] 0 10 (来自src[0]) 30 (来自src[2]) 40# target[1] 0 20 (来自src[1]) 20# target[2] 0 40 (来自src[3]) 50 (来自src[4]) 90