祝融说。

在前面的章节中，我们已经深入了解了大模型训练的各种模式与并行策略。现在，一个所有AI Infra工程师都必须面对的终极问题摆在了面前：“我需要多少资源？”

这个问题会以各种形式出现：

• 算法工程师：“我想微调一个70B的模型，需要多少张A800？”
• 你的老板：“我们计划从零预训练一个千亿模型，需要采购多少服务器？预算大概是多少？要训多久？”
• 你自己：“这个训练任务为什么OOM（Out of Memory）了？到底是哪里爆了显存？”

回答这些问题，不能再依靠经验和感觉。你需要的是一套科学、严谨的量化分析方法。本章，我们将一起化身为“算力精算师”，从第一性原理出发，推导和建立大模型训练的两大核心数学模型：显存占用模型和训练吞吐量模型。我们将用公式和计算器，将模糊的“资源需求”转化为精确的数字。最后，我们将通过一个选型实战（合成计算示例），将这些理论应用到集群规划的决策中。

掌握本章，你将获得透视AI训练资源消耗的“超能力”。当你下一次面对资源规划时，你将不再是一个被动的执行者，而是一个手握数据、运筹帷幄的决策者。

7.1 显存计算公式推导：参数、梯度、优化器、激活值=？

GPU显存（Memory）是AI训练中最宝贵、也最容易出现瓶颈的资源。一个训练任务能否跑起来，首先取决于它的显存峰值是否超过了单张GPU的物理显存容量。理解显存是如何被消耗的，是进行故障排查和性能优化的第一步。

一个Transformer大模型在训练过程中的显存占用，主要由四大部分构成：模型参数（Model Parameters）、梯度（Gradients）、优化器状态（Optimizer States）和激活值（Activations）。此外，还有一些零碎的开销，如临时缓冲区和CUDA内核本身占用的显存。

7.1.1 基础知识：数据类型与字节数

在计算前，我们必须明确不同数据精度所占用的字节数：

• FP32 (单精度浮点数): 4字节
• FP16 (半精度浮点数): 2字节
• BF16 (bfloat16): 2字节
• INT8 (8位整型): 1字节
• 1B (Billion) 参数 = 10亿个参数

7.1.2 显存占用的四大组成部分

假设我们有一个参数量为 P (单位：Billion, 十亿) 的模型。

模型参数 (Model Parameters)

这是最直观的一部分。模型本身需要被加载到显存中。

• 在标准的FP32训练中，模型参数占用 P * 10^9 * 4 字节。
• 在目前主流的混合精度训练 (Mixed Precision)中，通常会同时保留一份FP32的“主权重”（用于精确的梯度更新）和一份FP16的权重（用于高效的前向和反向计算）。
  • 模型参数占用：P * 10^9 * (4 + 2) = 6P GB (近似，10^9字节约等于1GB)
  • 但在很多现代框架的实现中（如PyTorch AMP），为了优化，可能只在显存中保留FP16的权重和FP32的梯度，参数更新在CPU或临时在GPU上完成。一个更常见的模型是，显存中有一份FP16的权重用于计算。
  • 为简化模型，我们通常关注主要的显存开销。在混合精度下，用于计算的权重是FP16的，所以至少有 P * 2 GB。

梯度 (Gradients)

反向传播算法会为模型中的每一个可训练参数都计算一个梯度。因此，梯度的数量与参数数量完全相同。

在混合精度训练中，为了保持更新的精度，梯度通常以FP32格式存储。

梯度占用：P * 10^9 * 4 = 4P GB

优化器状态 (Optimizer States)

这是最容易被忽略、也最容易导致OOM的“显存杀手”。优化器（如Adam或AdamW）在更新模型参数时，需要维护自身的“状态”。

• Adam/AdamW优化器： 它会为每一个模型参数维护两个状态：
  • 一阶动量 (Momentum): 存储过去梯度的指数移动平均，通常是FP32。
  • 二阶动量 (Variance): 存储过去梯度平方的指数移动平均，通常是FP32。
• 因此，Adam优化器会带来 2倍模型参数量的额外显存开销。
  • 优化器状态占用：P * 10^9 * 4 (动量) + P * 10^9 * 4 (方差) = 8P GB
• 在一些框架中，为了节省显存，可能会使用FP16来存储优化器状态，但会牺牲一定的精度和稳定性。默认情况下，我们按FP32计算。

小结1：静态显存（与Batch Size无关）

我们把模型参数、梯度和优化器状态这三部分称为“静态显存”，因为它们的占用量只与模型参数量P有关，与我们用多大的Batch Size去训练无关。

在一次标准的、使用Adam优化器的全量微调（或预训练）中，静态显存占用为：

Memory_static = (P * 2) [FP16参数] + (P * 4) [FP32梯度] + (P * 8) [Adam状态]

Memory_static ≈ 14P GB

如果使用DeepSpeed ZeRO-2，它会将梯度和优化器状态分片到N个GPU上，则单卡的静态显存变为：

Memory_static_ZeRO2 = (P * 2) + (P * 4 / N) + (P * 8 / N) = 2P + 12P / N GB

如果使用DeepSpeed ZeRO-3，它会将模型参数也分片，则单卡的静态显存变为：

Memory_static_ZeRO3 = (P * 2 / N) + (P * 4 / N) + (P * 8 / N) = 14P / N GB

LoRA的显存优势： LoRA为什么省显存？因为它只训练极少数（假设为P'）的参数。因此，梯度和优化器状态也只为这P'个参数而存在。静态显存约为 P * 2 (原始模型FP16权重，冻结) + 14 * P' (P'远小于P)。这极大地降低了显存开销。

激活值 (Activations)

这是显存占用中最复杂、最动态的部分。

• 什么是激活值： 在前向传播过程中，每一层网络的输出结果，都需要被保存（缓存在显存）下来，因为在反向传播计算梯度时需要用到它们（根据链式法则）。这些被缓存的中间结果，就是激活值。
• 激活值的大小与什么有关？
  • 序列长度 (Sequence Length, s): 输入的文本越长，中间的激活值矩阵就越大。
  • 批次大小 (Batch Size, b): 一次处理的样本越多，需要缓存的激活值就越多。
  • 模型隐藏层维度 (h): 模型的“宽度”。
  • 注意力头数量 (a):
  • 模型层数 (L): 模型的“深度”。
• 估算公式（非常重要）：
  • 对于一个标准的Transformer模型，其激活值显存占用有一个近似的估算公式：Memory_activations ≈ s * b * h * L * (10 + 24/T) / T (在使用了序列并行时)
  • 一个更实用、被广泛引用的简化公式是（不考虑序列并行，并使用混合精度FP16）：Memory_activations_per_gpu = s * b * h * L * (1 + a/(h*T)) * 2 (近似)
  • 一个更粗略但好记的经验公式是（由BLOOM论文提供）：Memory_activations ≈ 2 * s * b * P / L GB (其中P是Billion为单位的参数)
  • 让我们采用一个业界广泛引用的、由Anyscale工程师推导的估算公式（假设使用FP16/BF16，即2字节）：Memory_activations_GB = (s * b * h * L * (34 + 5 * a * s / (h * T))) / 10^9，这里的 T 是张量并行的路数。该公式源自Megatron-LM论文中的激活显存分析，针对特定假设（无序列并行、无选择性重计算、标准Transformer结构）成立。

这个公式的推导非常复杂，它考虑了注意力机制中K,V缓存、Softmax输出、Dropout掩码等所有细节。对于AI Infra工程师，我们不需要从头推导，但必须会使用这个公式，并理解其变量的含义。

关键洞察：

• 激活值与 s 和 b 成正比。序列长度翻倍，激活值显存翻倍。
• 激活值与模型规模（h, L）成正比。
• 使用张量并行（T > 1）可以显著降低激活值显存，因为模型的一些隐藏状态也被切分了。
• 激活重计算（Activation Recomputation/Checkpointing）： 这是一种用计算换显存的技术，在前向传播时只保存少量“关键点”，反向传播需要时再从最近的关键点重新计算。它会增加额外的计算时间（视重计算粒度不同，开销通常在20%–33%的量级，具体取决于实现与策略），但可以将激活值显存占用降低一个数量级。DeepSpeed和Megatron-LM都支持该技术。

7.1.3 终极公式与工具箱

单卡总显存占用 (Full Fine-Tuning, FP16, Adam):Memory_Total = (2P) [参数] + (4P) [梯度] + (8P) [优化器] + Memory_activations、Memory_Total ≈ 14P + (s * b * h * L * ...)

这个公式是理论峰值。在实际中，通过ZeRO等并行策略，可以将梯度和优化器部分分散到N张卡上。

单卡总显存占用 (ZeRO-3 + 激活重计算):Memory_Total_optimized ≈ 14P / N + Memory_activations_recompute

构建你的Excel/Python算力计算器：这是每一位AI Infra工程师都应该拥有的“神器”。

• 输入：
  • 模型参数量 P (Billion)
  • 模型层数 L
  • 隐藏层维度 h
  • 注意力头数量 a
  • 序列长度 s
  • 单卡批次大小 b (Micro Batch Size)
  • 并行策略：DP路数, TP路数, PP路数
  • 优化技术：是否开启ZeRO (哪个Stage), 是否开启激活重计算
• 输出：
  • 单卡静态显存 (GB)
  • 单卡激活值显存 (GB)
  • 单卡总显存峰值 (GB)

这个计算器，将是你进行OOM问题分析、资源申请评估时的最强武器。

7.2 算力估算模型：基于Token量与Flops估算训练天数

解决了显存问题（能不能跑起来），我们现在要解决时间问题（要跑多久）。这直接关系到项目排期和成本预算。

估算训练时间的核心，在于计算两样东西：总共需要多少计算量（Total FLOPs），以及我们的集群每秒能提供多少有效计算量（Effective TFLOPS）。

训练总时间 = 总计算量 / 集群有效算力

7.2.1 总计算量估算 (Total FLOPs)

对于Transformer模型，其训练过程的计算量（FLOPs - Floating Point Operations，浮点运算次数）有一个广为接受的经验公式：

Total FLOPs ≈ 6 * P * D

• P: 模型参数量（注意，这里是实际的参数个数，不是Billion）。例如，70B模型就是 70 * 10^9。
• D: 训练数据集的总Token数量。
• 6 这个系数是怎么来的？
  • 前向传播： 对于一个参数量为P的模型，处理1个Token，其计算量约为 2 * P FLOPs。这个2是因为每个参数在矩阵乘法中都参与了一次乘法和一次加法。
  • 反向传播： 根据链式法则，反向传播的计算量大约是前向传播的2倍，即 4 * P FLOPs。
  • 总计： 2P + 4P = 6P FLOPs。

示例： 预训练一个70B模型，使用2T (2万亿) Tokens的数据集。

• P = 70 * 10^9
• D = 2 * 10^12
• Total FLOPs = 6 * (70 * 10^9) * (2 * 10^12) = 8.4 * 10^23 FLOPs。这是一个天文数字！

7.2.2 集群有效算力 (Effective TFLOPS)

这是估算中最考验经验的部分。它不是简单地把所有GPU的理论峰值算力加起来。 集群有效算力 = 单卡理论峰值算力 * GPU数量 * MFU

• 单卡理论峰值算力 (Peak TFLOPS):

  • 从硬件规格书中查到。关键：要使用与训练精度匹配的算力值。
  • 例如，NVIDIA A100在FP16/BF16下的理论峰值是 312 TFLOPS。
  • 昇腾910B在FP16下的理论峰值是 320 TFLOPS。

• GPU数量 (N): 集群中用于该任务的GPU总数。

• MFU (Model FLOPs Utilization, 模型算力利用率):

  • 这是最重要的、也是最不确定的一个变量。它代表了我们实际压榨出了硬件理论性能的多少百分比。
  • MFU受到无数因素的影响：网络通信开销（特别是All-Reduce）、数据加载和预处理的瓶颈、没有优化的CUDA Kernel、内存墙等等。
  • MFU的取值范围（经验值）：
    • 糟糕 (10-20%): 网络瓶颈严重，或代码优化差。
    • 一般 (20-30%): 常见的水平，仍有优化空间。
    • 良好 (30-50%): 经过良好优化的集群和代码，例如使用了InfiniBand网络、通信计算有重叠。
    • 优秀 (50%以上): 业界顶尖水平，通常是像NVIDIA、Google这种软硬件垂直整合的公司才能达到。
  • 在做事前规划时，使用一个相对保守的MFU值（如30%）是比较稳妥的。

7.2.3 训练时间计算公式

训练总时间 (秒) = (6 * P * D) / (单卡峰值TFLOPS * N * MFU * 10^12)

为了方便换算，我们通常将时间单位转为天：训练总时间 (天) = 训练总时间(秒) / (3600 * 24)

构建你的Excel/Python训练时间计算器：

• 输入：
  • 模型参数量 P (Billion)
  • 训练数据量 D (Trillion Tokens)
  • 单卡峰值算力 (TFLOPS)
  • GPU数量 N
  • 预估的MFU (%)
• 输出：
  • 预估训练总天数
  • 预估总成本 (如果输入了单卡小时成本)

这个工具将成为你向管理层汇报项目周期和预算的有力支撑。

7.3 选型实战：70B模型需要多少张卡？如何规划集群规模？

现在，我们把前两节的公式用到一个选型决策场景中（以下为合成计算示例，用于演示方法，具体数字应按你的模型结构、硬件型号与实测MFU重新计算）。

场景： 某公司决定跟随开源潮流，基于一个高质量的中文数据集，从零开始预训练一个与Llama 2 70B规模相当的基础模型。CEO给你下达了任务：规划所需的计算集群，并给出预算和周期。

已知条件与约束：

• 目标模型： P = 70B (700亿参数)。假设其结构为 L=80层, h=8192, a=64。
• 训练数据： D = 2T (2万亿) Tokens。
• 硬件选项：
  • 选项A: NVIDIA A800 (80GB显存, 312 TFLOPS@FP16)
  • 选项B: 华为昇腾910B (64GB显存, 320 TFLOPS@FP16)
• 时间要求： 希望在3个月内（约90天）完成训练。

第一步：显存分析（决定单机内并行策略和最小GPU数）

我们需要找到一种并行策略，使得70B模型的训练任务能够放进单张GPU。我们将使用混合精度训练(BF16)和AdamW优化器。序列长度s通常设为4096。

• 计算激活值显存（粗略估算）： Memory_activations ≈ 2 * s * b * P / L = 2 * 4096 * b * 70 / 80 ≈ 7168 * b MB 为了把激活值显存控制在GB级别，micro batch size b 必须非常小，通常设为1或2。 假设 b=1，激活值显存约为 7.2 GB。这看起来不大，但这是在没有激活重计算的情况下。这是一个巨大的负担。

  如果我们开启激活重计算，可以将这部分显存降低到1GB以下。 因此，在训练大模型时，激活重计算是必选项。

• 计算静态显存（Full-Tuning）： Memory_static ≈ 14 * P = 14 * 70 = 980 GB。这是一个惊人的数字，没有任何单张GPU能装下。

• 引入并行策略来降低显存： 我们必须使用ZeRO、TP、PP等技术。一个典型的、高效的配置是 8路TP + ZeRO-1（因为TP已经将模型参数和优化器状态大幅切分）。

  • 张量并行(TP=8): 将模型参数和优化器状态切分到8张卡上。这通常在一个8卡节点内部完成。
  • 单卡静态显存（TP=8）： 14 * P / 8 = 980 / 8 = 122.5 GB。
  • 这个数字依然超过了A800的80GB和910B的64GB！

• 组合拳：TP + PP + ZeRO

  • 我们需要引入流水线并行(PP)来进一步切分模型。
  • 假设我们使用PP=4。
  • 单卡模型参数： P / (TP * PP) = 70 / (8 * 4) = 2.18B。
  • 单卡静态显存（粗略，不严谨）： 14 * P / (TP * PP) 吗？不完全是。ZeRO的切分和PP的切分作用机制不同。
  • 一个更工程化的方法是使用现有的开源配置作为参考。例如，Megatron-LM在训练类似规模模型时，通常会使用 TP=8, PP=16 甚至更高。
  • 我们来尝试一个配置：TP=8, PP=8, DP=N'
    • 每个流水线阶段负责 80/8=10 层。
    • 单卡上的模型参数和优化器状态主要由TP切分，约为 122.5 GB。依然太大。

• 终极方案：参考成熟框架的显存计算器

  • DeepSpeed和Megatron-LM的文档中通常会提供更精确的显存计算器或经验公式。根据NVIDIA公开资料的估算，一个70B模型，在TP=8, PP=1的设置下，使用BF16和激活重计算，峰值显存约为105GB。这仍然超过了A800。
  • 要使其适配80GB的A800，需要引入PP。一个可行的配置是：
    • TP = 4, PP = 4
    • 在这种配置下，每个GPU大约负责1/16的模型，显存占用可以被控制在80GB以内。

结论1： 训练一个70B模型，至少需要一个16卡的集群（TP=4, PP=4）才能启动一个副本。对于8卡/节点的服务器，这意味着至少需要2台服务器。

第二步：时间分析（决定集群总规模）

我们希望在90天内完成训练。现在来计算需要多少张卡。

• 总计算量： Total FLOPs = 6 * (70 * 10^9) * (2 * 10^12) = 8.4 * 10^23 FLOPs
• 目标有效算力：Required Effective TFLOPS = Total FLOPs / (90 * 24 * 3600 * 10^12) ≈ 107.5 * 10^3 TFLOPS = 107,500 TFLOPS
• 计算所需GPU数量 (N):
  • N = Required Effective TFLOPS / (单卡峰值TFLOPS * MFU)
  • 我们使用一个相对乐观但可达到的 MFU = 40% (0.4) 来进行规划。
  • 对于A800 (312 TFLOPS):N = 107500 / (312 * 0.4) ≈ 861 张卡。
  • 对于昇腾910B (320 TFLOPS):N = 107500 / (320 * 0.4) ≈ 840 张卡。
• 向上取整并考虑冗余：
  • 计算结果是861张卡。考虑到硬件故障、维护等因素，我们需要一定的冗余。同时，为了方便组成标准的并行单元（例如，一个DP副本是16卡），集群规模最好是并行单元的整数倍。
  • 我们可以规划一个 861 / 16 ≈ 54 个DP副本的集群。
  • 16 * 54 = 864 张卡。再增加一些冗余，例如5%， 864 * 1.05 ≈ 907 张卡。
  • 为了凑整和方便网络拓扑设计，最终我们可以规划一个由 1024张卡 构成的集群（例如，128台8卡服务器）。

第三步：最终方案与预算

• 集群规模： 1024张A800（或910B）的集群，由128台8卡服务器组成。
• 并行策略： TP=4, PP=4, DP=64。
• 网络要求： 必须使用高性能RDMA网络。对于128台服务器的规模，一个胖树（Fat-Tree）架构的InfiniBand或RoCE v2网络是必需的。
• 存储要求： 一个能够提供至少100GB/s聚合带宽的并行文件系统，容量在PB级别。
• 周期复核：实际训练天数 = (8.4e23) / (312 * 1024 * 0.4 * 1e12 * 3600 * 24) ≈ 76 天。这个结果在90天的目标之内，留出了一定的缓冲时间。
• 预算估算（量级参考，随市场行情波动）：
  • 硬件成本： 128台8卡服务器（含CPU、内存、硬盘）+ 1024张A800/910B + 高性能交换机 + 机柜、PDU等。这是一个数亿元人民币量级的投资。
  • 电力成本： 假设单机柜功率40KW，1024卡约需32-40个机柜，总功率超过1兆瓦。三年的电费也是千万量级。
  • 人力成本： 需要一个专业的AI Infra团队来建设和运维。

最终汇报给CEO的结论（基于本示例假设）：

“为了在3个月内完成70B模型的预训练，我们建议构建一个由1024张A800（或昇腾910B）组成的专用计算集群。该方案在假设MFU为40%的前提下，预计训练周期为76天，硬件投资在数亿元量级，后续运营成本巨大。作为替代方案，我们也可以考虑租用公有云的算力，虽然单价高，但可以免去前期巨大的固定资产投入和运维负担。”

这个基于定量分析得出的结论，远比一句“我们需要很多很多卡”要有力得多。这就是数学赋予AI Infra工程师的力量。