AI Notebook - 技术栈

无 FSDP（通常指传统的 DDP - 分布式数据并行）：假设有 8 张显卡，每张显卡上都必须完整保存一份一模一样的模型副本。这就像让 8 个工人各自拿着一本完整的大字典在工作。如果这本字典（模型）太大，任何一个工人的背包（显存）都装不下，训练就直接崩溃了（OOM，显存溢出）。
FSDP（Fully Sharded Data Parallel）：完全分片数据并行。把这本大字典拆成了 8 份，每个工人（GPU）只负责保管其中一份。只有当某个工人需要用到其他页的内容时，他才会快速向别人"借"过来（通过网络通信同步），看完后立刻"销毁"借来的复印件以腾出空间。具体来说，FSDP 将模型的参数（Parameters）、梯度（Gradients）和优化器状态（Optimizer States）全部分片并分散到了不同的 GPU 上。
FSDP：None；Gradient Checkpointing：True
FSDP：full_shard auto_wrap；Gradient Checkpointing：True

Gradient Checkpointing（梯度检查点，有时也叫 Activation Checkpointing）：在标准的深度学习训练中，分为前向传播和反向传播两步。
（1）前向传播：数据从第一层一层层算到最后一层，得出预测结果。在这个过程中，每一层计算产生的中间结果（我们称之为 Activations，激活值）都会被保存在显存里。
（2）反向传播：模型根据预测误差，从最后一层往回算，更新参数。重点来了，反向传播在计算梯度时，必须用到前向传播时保存的那些"激活值"。
问题在于：对于 LLMs 来说，层数极深，如果把每一层的激活值全都存在显存里，会占用极其庞大的空间。这往往是导致 CUDA Out of Memory（显存溢出）的罪魁祸首，其占用的显存甚至比模型参数本身还要大得多。Example：假设你在爬一座 100 层的塔（前向传播）。为了原路返回（反向传播），你原本需要在每一层都留一个记号（保存激活值），这会让你的背包（显存）塞满。
开启 Gradient Checkpointing 后：现在只在第 10 层、20 层、30 层......这些关键楼层留下记号（也就是 Checkpoints）。当你退回到第 19 层，发现这里没有记号怎么办？很简单，你退到最近的第 10 层记号处，重新往上爬一次（重新计算一次前向传播）算出第 19 层的状态，然后再继续往下退。
Gradient Checkpointing 只在某些神经网络层的边界保存激活值。当反向传播需要用到被丢弃的激活值时，模型会从最近的"检查点"重新进行一次局部的前向计算来临时恢复这些数据，是一种"用计算时间换取显存空间"的技术。
FSDP：full_shard auto_wrap；Gradient Checkpointing：False. 【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_CONF=expandable_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)】
FSDP：None；Gradient Checkpointing：False

FP16 (IEEE 754 Half-Precision)：1位符号 + 5位指数 + 10位尾数
优点：精度较高（因为有10位尾数）
动态范围极窄。它的最大值只能表示到 65504，最小值是 5.96×10−85.96×10^{−8}5.96×10−8
BF16 (Brain Floating Point)：1位符号 + 8位指数 + 7位尾数
优点：动态范围巨大。它的指数位和标准的 32位单精度浮点数（FP32）一模一样。这意味着它的最大值可以达到约 3.4×10383.4×10^{38}3.4×1038
缺点：精度较低（只有7位尾数）。两个非常接近的小数在 BF16 下可能会被当作同一个数。但幸运的是，神经网络对范围极其敏感，对微小的精度误差却非常宽容（甚至相当于一种正则化）。
BF16 + FSDP
FP32 + FSDP

显存消耗拆解 ：在使用 Adam 优化器进行混合精度训练（FP16/BF16）时，假设模型参数量为 xxx，单卡显存占用如下：FP16 模型参数 (Parameters)： 2x2x2x bytes；FP16 梯度 (Gradients)： 2x2x2x bytes；优化器状态 (Optimizer States)： Adam 需要保存 FP32 的参数备份（4x4x4x bytes）、Momentum（4x4x4x bytes）和 Variance（4x4x4x bytes）。共计 12x12x12x bytes。因此显存占用总量是 16x16x16x bytes.
（MB等单位以10为底数的指数，MiB是以2为底数的指数，如：1KB=103=10001KB=10^3=10001KB=103=1000, 1MB=106=1000000=1000KB1MB=10^6=1000000=1000KB1MB=106=1000000=1000KB,而 1KiB=210=10241KiB=2^{10}=10241KiB=210=1024,1MiB=220=1048576=1024KiB1MiB=2^{20}=1048576=1024KiB1MiB=220=1048576=1024KiB）
数据并行（Data Parallel）：用于模型参数量还能塞进单张显卡，但数据量庞大时。PyTorch 中经历了从 DataParallel (DP) 到 DistributedDataParallel (DDP) 的演进。
DP的问题：单进程多线程（Single-process, Multi-thread）机制。受限于 Python 的全局解释器锁（GIL），多线程并不能实现真正的并行计算。 通信瓶颈在于 P arameter Server 架构，每次 Forward 前，主节点（通常是 GPU 0）需要把模型广播给所有 GPU；Backward 后，所有 GPU 的梯度都要汇聚到主节点进行求和与更新。GPU 0 成为了严重的通信瓶颈。
DDP：多进程（Multi-process）。每个 GPU 都有一个独立的进程，独享属于自己的模型副本和优化器状态。
DDP核心通信算法 Ring-AllReduce：假设有 N 个 GPU，每张卡上的梯度被切分成 N 个数据块。（1）Scatter-Reduce （分散-规约）：每张 GPU 将自己的第 iii 个数据块发送给下一个 GPU，并接收前一个 GPU 发来的数据块进行累加。经过 N−1N−1N−1 步后，每张 GPU 上都会拥有一块完整的、包含所有 GPU 累加结果的梯度块。（2）All-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 层为例：全连接层计算本质是 Y=GeLU(XA)BY={\rm GeLU}(XA)BY=GeLU(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-4，GPU 1 计算 Layer 5-8。痛点在于如果是简单的顺序执行，GPU 1 必须等待 GPU 0 计算完毕才能开始，这会导致大量的 GPU 处于空闲等待状态（即气泡）。
解决方案：（1）Micro-batching（微批次）：将一个大的 Batch 切分成多个小 Micro-batch。GPU 0 计算完第一个 Micro-batch 的第一层后，立刻传给 GPU 1 计算，同时 GPU 0 开始计算第二个 Micro-batch。（2）1F1B 调度（One Forward, One Backward）：是工业界常用策略，指在稳态阶段，GPU 交替执行一次前向传播和一次反向传播，这能大幅降低显存中需要缓存的激活值数量。
ZeRO 优化器（Zero Redundancy Optimizer）：在传统的 DDP 中，每个 GPU 都保存了完整的 16x16x16x 参数状态，这造成了极大的冗余。ZeRO 通过切分这些状态来节省显存。（1）ZeRO-Stage 1（切分优化器状态 ）：切分 12x12x12x bytes 的优化器状态，每张 GPU 只负责更新 1/N1/N1/N 的参数，因此显存占用降低为 4x+12xN4x+\frac{12x}{N}4x+N12x。（2）ZeRO-Stage 2（切分梯度）：在 Stage 1 的基础上，进一步切分梯度，显存占用降低为 2x+14xN2x+\frac{14x}{N}2x+N14x。（3）ZeRO-Stage 3（切分参数）：将模型参数也完全切分。每张显卡上只保存 1/N1/N1/N 的参数。前向和反向传播时，通过 All-Gather 动态获取需要的参数，用完即弃。显存占用为 16xN.\frac{16x}N.N16x.

深度学习训练耗时：训练耗时===数据规模×\times×单步计算量///计算速率
计算速率：计算速率===单设备计算速率×\times×设备数×\times×多设备并行效率（加速比）
单设备计算速率：混合精度、算子融合、梯度累加
加速比：数据并行、模型并行、流水并行
设备数：服务器架构、通信拓扑