1. 为什么非要把 DiT 的“算力账”算明白在扩散模型项目里泡久了你迟早会遇到一个绕不开的问题手上拿到一张图要训一个 DiT 模型到底该申请多少卡、租多久、用多大的 batch我见过太多人上来就按论文里的 FLOPs 数据拍脑袋定资源结果要么多花一倍的钱要么训练中途显存爆炸被迫改配置。DiTDiffusion Transformer这一类模型本质上就是把 U-Net 那套卷积结构换成 Transformer 块然后用 patch 化的输入图像来做扩散去噪。它的计算规模和传统 U-Net 扩散模型差异巨大核心变量不再是简单的通道数加分辨率而是序列长度、注意力维度和 transformer block 数量的组合关系。如果不会手动推算计算规模你连“为什么 DiT-XL/2 比 DiT-B/4 慢那么多”这种基础问题都解释不了更不用说后续的并行策略设计、激活值显存估算、甚至推理时的时延预算了。本篇文章会把 DiT 的计算规模从公式推导、代码实现、工程避坑三个层面完整拆开。适合正在训练或推理 DiT 的算法工程师也适合准备把 DiT 作为 baseline 做研究的同学——你会得到一个可以直接抄作业的 FLOPs 计算函数和一张常用配置的参数总表。2. DiT 计算规模的核心拆解公式2.1 你真的懂 DiT 的输入形态吗DiT 与经典 ViT 的最大区别在于输入不是纯 token 序列而是由图像 patch 化后得到的 latent token。假设输入图像经过 VAE 编码后尺寸为 f×fpatch 大小为 p×p那么送入 transformer 的序列长度 S 等于S (f / p)²DiT 论文的默认设定是 256×256 图像经 VAE 得到 32×32 的 latent默认 patch 为 2所以 S (32/2)² 256再经过 adaLN-zero 的条件拼接后序列不变。很多人把“序列长度 256”这份默认直接代进公式最后算出来的 FLOPs 和论文对不上原因就是没搞清楚 S 在不同 block 之间可能是变化的比如某些变体在浅层用大 patch深层用patch size。另外一点必须明确DiT block 里有两个核心计算单元一是多头自注意力MHSA二是 MLP 块。它们的 FLOPs 规模分别由序列长度和隐藏维度决定而且前者随图像分辨率呈平方级增长后者随 patch size 呈四次方级别衰减。用大白话说分辨率翻倍注意力计算翻四倍patch 从 2 变 4序列长度直接缩到原来的四分之一注意力成本缩到十六分之一。这一层关系是你估算任何扩散 Transformer 变体的地基。2.2 FLOPs 公式逐项拆解把公式先亮出来后面逐一解释。对于一个标准的 DiT block单次前向的乘加次数FLOPs可以分三个主要部分来算。第一部分是 patch embedder。输入从 f×f 的 latent 投影到 hidden dim D这一步的算子本质是一个卷积或者线性投影FLOPs 约等于 2 × S × D²。第二部分是 transformer block 内部注意力部分要算 Q、K、V 的投影3 次 2×S×D²然后计算注意力矩阵本身2×S²×D再乘 V2×S²×D最后输出投影2×S×D²。MLP 部分则是两个线性层通常膨胀比为 4所以是 2×S×D×4D再加一个 2×S×4D×D 的收缩层合起来约等于 16×S×D²。每个 block 再配 adaLN-zero 的调制层这部分很小但不算零头通常额外加 2×S×D² 量级。把上面的分项加起来单个 transformer block 的 FLOPs 约等于F_block ≈ 24 × S × D² 4 × S² × D到这里你会发现一个关键规律当 S 很大时S²×D 这一项会迅速占据主导这是 DiT 在低分辨率 latent 上计算效率高的根本原因——latent 空间远比像素空间小S 通常是 256 或者 1024不会像文本 transformer 动辄上万 token。最后把 block 数量 L 乘上再加上最后的 linear head通常是 2×S×D×D_out和输出解码层的开销就能得到整体前向 FLOPs。实际工程中DiT-S/B/L/XL 分别对应 D384/768/1152/1152L 对应 12/12/12/28。我, 通常会忽略 embedding 和 head因为占比低于 5%但在对比两个差异很小的模型时还是建议把这两项加上避免结论被 0.1 GFLOPs 的误差翻转。2.3 从单 block 到整模型的汇总逻辑把上面的 block FLOPs × L然后加上 VAE 和采样步数的乘数效应这就是完整的全图计算量了。扩散模型有一个特殊性训练时每一步 denoise 都要过一次 DiT所以总 FLOPs 单次前向 FLOPs × 训练步数 × 采样步数。如果你想估算训练一个 DiT 模型的总算力这一乘数最关键而且它不是固定值取决于你的 noise schedule 和采样器。DDPM 训练通常要 1000 步但 EDM 风格会用更少步数配合更好的调度器所以自定义 schedule 以前一定要把“总步数”作为一个显式超参写进计算脚本而不是靠直觉估。3. 手把手实现一个 DiT FLOPs 计算器3.1 先选定参数再写代码别倒着来最稳妥的流程是先把 DiT 配置字典定义好再写计算函数。下面这份代码可以直接复制使用替换参数即可适配任意 DiT 变种。我在这里把注意力的实现方式限定为 pytorch 标准多头注意力每个头维度为 D/h不单独计算 head split 带来的额外开销。import math def flops_dit(f, p, D, L, steps, batch1, mlp_ratio4): 估算 DiT 模型单次前向以及训练指定步数的总 FLOPs f: VAE 后 latent 尺寸 p: patch size D: transformer hidden dim L: transformer block 数量 steps: 扩散训练/采样步数 S (f // p) ** 2 # 序列长度 # patch embedder 开销近似 linear proj flops_embed 2 * S * D * D # 每一个 transformer block qkv_proj 3 * 2 * S * D * D # QKV 投影 attn_score 2 * S * S * D # QK^T 与缩放 attn_value 2 * S * S * D # 加权 V out_proj 2 * S * D * D # attention 输出投影 mlp 2 * S * D * mlp_ratio * D 2 * S * mlp_ratio * D * D adaln 2 * S * D # adaLN 调制参数近似 flops_block qkv_proj attn_score attn_value out_proj mlp adaln # 解码头简化影响小 flops_head 2 * S * D * D total_fwd flops_embed L * flops_block flops_head return total_fwd * steps * batch # 验证 DiT-XL/2 的配置 flops_per_step flops_dit(f32, p2, D1152, L28, steps1) print(f单步前向 FLOPs: {flops_per_step / 1e9:.2f} GFLOPs)这份代码跑出来的单步 FLOPs 大约在 118 GFLOPs 左右和 DiT 论文 openreview 里报告的数值基本一致差别主要来自 adaLN 的细节实现和 patch embedder 是否算卷积核开销。3.2 常见配置的 FLOPs 速查表把 DiT 官方 GitHub 里的几组配置代进上述函数整理成速查表。下面这张表按 patch 2/4/8 分列展示单位是 GFLOPs单次前向、batch1、256×256 输入。模型配置DBlock数patch2patch4patch8DiT-S384126.11.50.4DiT-B7681224.36.11.5DiT-L11521254.713.73.4DiT-XL115228118.429.67.4注意到 DiT-XL/2 和 DiT-L/2 只相差 block 数量FLOPs 几乎线性增长而 patch 从 2 改到 4FLOPs 直接掉到原来的四分之一。这说明如果你想在有限算力下硬跑大模型增加 patch size 是比减少 D 更高效的策略——代价是图像细节还原能力的下降这个 trade-off 必须在实验设计阶段就想清楚。另外有人会问我“为什么参数量差不多但 FLOPs 差好几倍”答案就在序列长度上。DiT-B/2 与 DiT-B/4 的参数量完全一样但前者的注意力计算量是后者的 16 倍因为 S 从 256 掉到 64而 S² 项缩得最狠。这就是为什么很多工程落地选 patch 4 而不是 2。3.3 单卡 A100 理论吞吐估算有 FLOPs 数据以后可以做一件非常实用的事估算单张 A100 在 fp16 下的理论训练吞吐上限。A100 80G 的 FP16 峰值算力约 312 TFLOPS带稀疏约 624这里用稠密值。考虑到 MFU模型浮点利用率通常做不到 1实际在 35%~50% 之间——如果你用 PyTorch 原生分布式训练40% 已经算不错的成绩了。拿 DiT-XL/2 为例单步 118.4 GFLOPs单卡每秒可以前向 312e12 / 118.4e9 ≈ 2635 次。但训练需要前向反向反向大约是前向的 2 倍也就是说一张卡每秒能处理约 878 个样本的步数。假设你训练 200k 步batch size 为 256那总样本数为 51.2M单卡步数为 200k×256 / 878 ≈ 58,300 秒约 16.2 小时。这个估算放在 40% MFU 下约 40 小时和经验值很吻合。注意这还没算 VAE 前向、数据加载和 mixed precision 的额外开销工程上建议在这个数字上再乘 1.3 的系数。4. 怎么用计算规模指导并行策略和显存规划4.1 一个公式推导出你的并行方案FLOPs 不只用来估算时间它直接决定了你的并行策略。DiT-L/2 的单步前向是 54.7 GFLOPs如果单卡吞吐跟不上你就需要考虑 tensor parallelTP还是 sequence parallelSP。最朴素的判断原则单卡计算量超过单卡算力的 60%就该考虑切模型。具体操作是用“单卡可承担的 FLOPs/s 峰值 × MFU”对比训练单步所需的总 FLOPs前向反向×优化器更新额外开销超过阈值就切。比如 A100 上 MFU 40% 意味着单卡约 125 TFLOPS 的实际算力DiT-XL/2 训练需要约 355 GFLOPs前向反向后 FWDBWD≈3 倍算下来单卡需 0.00284 秒一步勉强可以。但如果 batch 加到 64单步变成 22.7 GFLOPs×64 ≈ 1.45T单卡直接溢出必须上 TP。TP 的切分规则也很直白DiT 的 Linear 层按 D 维度切分注意力头按 head 维度切分。以 DiT-L 为例D1152要切成 4 卡 TP每卡 D2888 卡则 D144。注意力 head168 卡时每卡 2 个 head是整数——所以 DiT 系列选 head16 是刻意为之不必担心切分碎掉。如果不整除只能用 SP 或让框架自动 padding两种方案都会带来额外的 allreduce 通信开销。4.2 激活值显存估算让 OOM 远离你训练时最现实的坑是显存不够而 FLOPs 能帮你提前算出 activation 的规模。DiT 每个 block 需要保存的激活值约等于 S×D×(2mlp_ratio) 的规模乘以 block 数再加上注意力分数矩阵 S×S×D如果是显式物化。用 DiT-XL/2 算序列 256D1152每层激活约 256×1152×(24)×4 byte ≈ 6.3MB28 层约 176MBattention score 是 256×256×1152×2 byte ≈ 144MB。实际训练还需乘 batch、梯度、优化器状态——batch 为 16 时最原始的激活总占用就超过 5GB加上 Adam 的 fp32 状态和模型权重总显存直奔 70GB。这个数字没算 TMA 和通信缓存所以你会发现 A100 80G 其实很紧真正跑稳还得配合 activation checkpointing。激活重计算能省掉约 60%~80% 的激活显存代价是额外约 30% 的 FLOPs——这个 trade-off 在算力有余而显存不足时非常划算。5. 实操中的典型问题与排错速查表5.1 为什么你的 FLOPs 和理论值对不上我最常被问到的问题是“为什么我照着公式算出来 300 GFLOPs但 profiler 里显示 360”这类偏差一般有三个来源。第一个来源是显式 shape 的算子 padding。比如注意力头数不为 D 整除时框架会 padding 到下一个对齐单位带来 5%~10% 的额外计算。第二个是 softmax 的数值稳定实现有些后端会拆分两次 pass 来避免溢出导致注意力矩阵的 FLOPS 翻倍。第三个更隐蔽——如果你用了 flash attention它不会显式物化 S×S 矩阵但缩放因子、因果 mask 和分块策略会让理论上的 4×S²D 下降约 2 倍。所以实测比理论低不一定是 bug反而是优化生效的信号。对于 DiT 这种没有 mask 的标准注意力FlashAttention 的 FLOPs 和手动实现的一致真正的问题出在算子融合如果 MLP 的 GELU 被融合进前面的 GEMM你 profiler 里看不到单独的 GELU kernel理论计数不能叠加算两层否则重复计数。我自己习惯把 torch.profiler 的输出按 operator 聚合然后用表格和理论值对比超过 1.2 倍就需要检查是不是某处矩阵乘的 shape 搞错了。5.2 数据处理和精度对 FLOPs 的隐性影响这里再说一个很简单但总被人忽略的事实FLOPs 是理论乘加次数不代表硬件实际计算次数。当你用 fp16 混合精度训练时某些框架会把乘积累加到 fp32 里硬件上计算的实际上是一个 FP32 FMA而不是 FP16 FMA。很多开源 FLOPs 计算器报的是“算法 FLOPs”但你实测的 MFU 如果按这个数算出来超过 100%基本上是精度换算的问题。另外 DiT 训练常用 VAE 把图像编码后用 latent 做 diffusion如果你把 VAE 的 FLOPs 也算进去单次前向会给 DiT 本身加上约 15~20 GFLOPs 的开销。这一点论文里经常避而不谈工程落地对比不同方案的端到端延迟时必须带上否则采样性能会被严重高估。5.3 开源计算工具实测对比与选型建议市面上能直接算 DiT FLOPs 的工具不多我更推荐自己写函数但为了 sanity check以下工具实测可以辅助验证。工具名称原理对 DiT 的适配度备注torchinfo只统计参数量和 tensor shape低不涉及 FLOPs仅用来算参数量thop基于 hook 统计乘加次数中对 transformer 支持差常低估注意力部分fvcoreMeta 开源支持 transformer高注意 patch embedder 需要手动标记calflops较新的 PyPI 包高内置常见 DiT 模块支持实测下来 fvcore 的数值和手动公式算的误差在 3% 以内calflops 更省事但偶尔会把 adaLN 的归一化统计成 O(S) 而不是 O(SD)需要二次核对。真正到了 100G FLOPs 的大模型我都会保留一份手写公式函数的输出作为基准其余工具只做交叉验证。6. 关于计算规模我最后想多说两句DiT 的 FLOPs 计算本身不难难的是把“算出来的数字”和“真实的硬件表现”之间的差值解释清楚。我在项目里早就放弃了追求 100% 精确的 FLOPs 计数器反而更依赖一套相对稳定的估算流程先用公式算理论值再在单卡上跑 100 步看实测吞吐两者对比出来的 MFU 就是我判断模型、框架和硬件是否匹配的核心指标。如果你准备用一个全新的 DiT 配置跑大型实验我建议在正式训练前花半天时间把你手头所有可能的配置不同的 patch、D、L用本文的代码快速扫一遍记录每个配置下前向和训练的 FLOPs 表再对照你集群的单卡 MFU直接就能筛掉 80% 不可能在 deadline 前跑完的方案。后续只要改动任何一个超参把函数重新跑一遍你脑子里的“算力地图”就始终是新的不会出现训练到一半才发现资源缺口的问题。
1. 为什么非要把 DiT 的“算力账”算明白在扩散模型项目里泡久了你迟早会遇到一个绕不开的问题手上拿到一张图要训一个 DiT 模型到底该申请多少卡、租多久、用多大的 batch我见过太多人上来就按论文里的 FLOPs 数据拍脑袋定资源结果要么多花一倍的钱要么训练中途显存爆炸被迫改配置。DiTDiffusion Transformer这一类模型本质上就是把 U-Net 那套卷积结构换成 Transformer 块然后用 patch 化的输入图像来做扩散去噪。它的计算规模和传统 U-Net 扩散模型差异巨大核心变量不再是简单的通道数加分辨率而是序列长度、注意力维度和 transformer block 数量的组合关系。如果不会手动推算计算规模你连“为什么 DiT-XL/2 比 DiT-B/4 慢那么多”这种基础问题都解释不了更不用说后续的并行策略设计、激活值显存估算、甚至推理时的时延预算了。本篇文章会把 DiT 的计算规模从公式推导、代码实现、工程避坑三个层面完整拆开。适合正在训练或推理 DiT 的算法工程师也适合准备把 DiT 作为 baseline 做研究的同学——你会得到一个可以直接抄作业的 FLOPs 计算函数和一张常用配置的参数总表。2. DiT 计算规模的核心拆解公式2.1 你真的懂 DiT 的输入形态吗DiT 与经典 ViT 的最大区别在于输入不是纯 token 序列而是由图像 patch 化后得到的 latent token。假设输入图像经过 VAE 编码后尺寸为 f×fpatch 大小为 p×p那么送入 transformer 的序列长度 S 等于S (f / p)²DiT 论文的默认设定是 256×256 图像经 VAE 得到 32×32 的 latent默认 patch 为 2所以 S (32/2)² 256再经过 adaLN-zero 的条件拼接后序列不变。很多人把“序列长度 256”这份默认直接代进公式最后算出来的 FLOPs 和论文对不上原因就是没搞清楚 S 在不同 block 之间可能是变化的比如某些变体在浅层用大 patch深层用patch size。另外一点必须明确DiT block 里有两个核心计算单元一是多头自注意力MHSA二是 MLP 块。它们的 FLOPs 规模分别由序列长度和隐藏维度决定而且前者随图像分辨率呈平方级增长后者随 patch size 呈四次方级别衰减。用大白话说分辨率翻倍注意力计算翻四倍patch 从 2 变 4序列长度直接缩到原来的四分之一注意力成本缩到十六分之一。这一层关系是你估算任何扩散 Transformer 变体的地基。2.2 FLOPs 公式逐项拆解把公式先亮出来后面逐一解释。对于一个标准的 DiT block单次前向的乘加次数FLOPs可以分三个主要部分来算。第一部分是 patch embedder。输入从 f×f 的 latent 投影到 hidden dim D这一步的算子本质是一个卷积或者线性投影FLOPs 约等于 2 × S × D²。第二部分是 transformer block 内部注意力部分要算 Q、K、V 的投影3 次 2×S×D²然后计算注意力矩阵本身2×S²×D再乘 V2×S²×D最后输出投影2×S×D²。MLP 部分则是两个线性层通常膨胀比为 4所以是 2×S×D×4D再加一个 2×S×4D×D 的收缩层合起来约等于 16×S×D²。每个 block 再配 adaLN-zero 的调制层这部分很小但不算零头通常额外加 2×S×D² 量级。把上面的分项加起来单个 transformer block 的 FLOPs 约等于F_block ≈ 24 × S × D² 4 × S² × D到这里你会发现一个关键规律当 S 很大时S²×D 这一项会迅速占据主导这是 DiT 在低分辨率 latent 上计算效率高的根本原因——latent 空间远比像素空间小S 通常是 256 或者 1024不会像文本 transformer 动辄上万 token。最后把 block 数量 L 乘上再加上最后的 linear head通常是 2×S×D×D_out和输出解码层的开销就能得到整体前向 FLOPs。实际工程中DiT-S/B/L/XL 分别对应 D384/768/1152/1152L 对应 12/12/12/28。我, 通常会忽略 embedding 和 head因为占比低于 5%但在对比两个差异很小的模型时还是建议把这两项加上避免结论被 0.1 GFLOPs 的误差翻转。2.3 从单 block 到整模型的汇总逻辑把上面的 block FLOPs × L然后加上 VAE 和采样步数的乘数效应这就是完整的全图计算量了。扩散模型有一个特殊性训练时每一步 denoise 都要过一次 DiT所以总 FLOPs 单次前向 FLOPs × 训练步数 × 采样步数。如果你想估算训练一个 DiT 模型的总算力这一乘数最关键而且它不是固定值取决于你的 noise schedule 和采样器。DDPM 训练通常要 1000 步但 EDM 风格会用更少步数配合更好的调度器所以自定义 schedule 以前一定要把“总步数”作为一个显式超参写进计算脚本而不是靠直觉估。3. 手把手实现一个 DiT FLOPs 计算器3.1 先选定参数再写代码别倒着来最稳妥的流程是先把 DiT 配置字典定义好再写计算函数。下面这份代码可以直接复制使用替换参数即可适配任意 DiT 变种。我在这里把注意力的实现方式限定为 pytorch 标准多头注意力每个头维度为 D/h不单独计算 head split 带来的额外开销。import math def flops_dit(f, p, D, L, steps, batch1, mlp_ratio4): 估算 DiT 模型单次前向以及训练指定步数的总 FLOPs f: VAE 后 latent 尺寸 p: patch size D: transformer hidden dim L: transformer block 数量 steps: 扩散训练/采样步数 S (f // p) ** 2 # 序列长度 # patch embedder 开销近似 linear proj flops_embed 2 * S * D * D # 每一个 transformer block qkv_proj 3 * 2 * S * D * D # QKV 投影 attn_score 2 * S * S * D # QK^T 与缩放 attn_value 2 * S * S * D # 加权 V out_proj 2 * S * D * D # attention 输出投影 mlp 2 * S * D * mlp_ratio * D 2 * S * mlp_ratio * D * D adaln 2 * S * D # adaLN 调制参数近似 flops_block qkv_proj attn_score attn_value out_proj mlp adaln # 解码头简化影响小 flops_head 2 * S * D * D total_fwd flops_embed L * flops_block flops_head return total_fwd * steps * batch # 验证 DiT-XL/2 的配置 flops_per_step flops_dit(f32, p2, D1152, L28, steps1) print(f单步前向 FLOPs: {flops_per_step / 1e9:.2f} GFLOPs)这份代码跑出来的单步 FLOPs 大约在 118 GFLOPs 左右和 DiT 论文 openreview 里报告的数值基本一致差别主要来自 adaLN 的细节实现和 patch embedder 是否算卷积核开销。3.2 常见配置的 FLOPs 速查表把 DiT 官方 GitHub 里的几组配置代进上述函数整理成速查表。下面这张表按 patch 2/4/8 分列展示单位是 GFLOPs单次前向、batch1、256×256 输入。模型配置DBlock数patch2patch4patch8DiT-S384126.11.50.4DiT-B7681224.36.11.5DiT-L11521254.713.73.4DiT-XL115228118.429.67.4注意到 DiT-XL/2 和 DiT-L/2 只相差 block 数量FLOPs 几乎线性增长而 patch 从 2 改到 4FLOPs 直接掉到原来的四分之一。这说明如果你想在有限算力下硬跑大模型增加 patch size 是比减少 D 更高效的策略——代价是图像细节还原能力的下降这个 trade-off 必须在实验设计阶段就想清楚。另外有人会问我“为什么参数量差不多但 FLOPs 差好几倍”答案就在序列长度上。DiT-B/2 与 DiT-B/4 的参数量完全一样但前者的注意力计算量是后者的 16 倍因为 S 从 256 掉到 64而 S² 项缩得最狠。这就是为什么很多工程落地选 patch 4 而不是 2。3.3 单卡 A100 理论吞吐估算有 FLOPs 数据以后可以做一件非常实用的事估算单张 A100 在 fp16 下的理论训练吞吐上限。A100 80G 的 FP16 峰值算力约 312 TFLOPS带稀疏约 624这里用稠密值。考虑到 MFU模型浮点利用率通常做不到 1实际在 35%~50% 之间——如果你用 PyTorch 原生分布式训练40% 已经算不错的成绩了。拿 DiT-XL/2 为例单步 118.4 GFLOPs单卡每秒可以前向 312e12 / 118.4e9 ≈ 2635 次。但训练需要前向反向反向大约是前向的 2 倍也就是说一张卡每秒能处理约 878 个样本的步数。假设你训练 200k 步batch size 为 256那总样本数为 51.2M单卡步数为 200k×256 / 878 ≈ 58,300 秒约 16.2 小时。这个估算放在 40% MFU 下约 40 小时和经验值很吻合。注意这还没算 VAE 前向、数据加载和 mixed precision 的额外开销工程上建议在这个数字上再乘 1.3 的系数。4. 怎么用计算规模指导并行策略和显存规划4.1 一个公式推导出你的并行方案FLOPs 不只用来估算时间它直接决定了你的并行策略。DiT-L/2 的单步前向是 54.7 GFLOPs如果单卡吞吐跟不上你就需要考虑 tensor parallelTP还是 sequence parallelSP。最朴素的判断原则单卡计算量超过单卡算力的 60%就该考虑切模型。具体操作是用“单卡可承担的 FLOPs/s 峰值 × MFU”对比训练单步所需的总 FLOPs前向反向×优化器更新额外开销超过阈值就切。比如 A100 上 MFU 40% 意味着单卡约 125 TFLOPS 的实际算力DiT-XL/2 训练需要约 355 GFLOPs前向反向后 FWDBWD≈3 倍算下来单卡需 0.00284 秒一步勉强可以。但如果 batch 加到 64单步变成 22.7 GFLOPs×64 ≈ 1.45T单卡直接溢出必须上 TP。TP 的切分规则也很直白DiT 的 Linear 层按 D 维度切分注意力头按 head 维度切分。以 DiT-L 为例D1152要切成 4 卡 TP每卡 D2888 卡则 D144。注意力 head168 卡时每卡 2 个 head是整数——所以 DiT 系列选 head16 是刻意为之不必担心切分碎掉。如果不整除只能用 SP 或让框架自动 padding两种方案都会带来额外的 allreduce 通信开销。4.2 激活值显存估算让 OOM 远离你训练时最现实的坑是显存不够而 FLOPs 能帮你提前算出 activation 的规模。DiT 每个 block 需要保存的激活值约等于 S×D×(2mlp_ratio) 的规模乘以 block 数再加上注意力分数矩阵 S×S×D如果是显式物化。用 DiT-XL/2 算序列 256D1152每层激活约 256×1152×(24)×4 byte ≈ 6.3MB28 层约 176MBattention score 是 256×256×1152×2 byte ≈ 144MB。实际训练还需乘 batch、梯度、优化器状态——batch 为 16 时最原始的激活总占用就超过 5GB加上 Adam 的 fp32 状态和模型权重总显存直奔 70GB。这个数字没算 TMA 和通信缓存所以你会发现 A100 80G 其实很紧真正跑稳还得配合 activation checkpointing。激活重计算能省掉约 60%~80% 的激活显存代价是额外约 30% 的 FLOPs——这个 trade-off 在算力有余而显存不足时非常划算。5. 实操中的典型问题与排错速查表5.1 为什么你的 FLOPs 和理论值对不上我最常被问到的问题是“为什么我照着公式算出来 300 GFLOPs但 profiler 里显示 360”这类偏差一般有三个来源。第一个来源是显式 shape 的算子 padding。比如注意力头数不为 D 整除时框架会 padding 到下一个对齐单位带来 5%~10% 的额外计算。第二个是 softmax 的数值稳定实现有些后端会拆分两次 pass 来避免溢出导致注意力矩阵的 FLOPS 翻倍。第三个更隐蔽——如果你用了 flash attention它不会显式物化 S×S 矩阵但缩放因子、因果 mask 和分块策略会让理论上的 4×S²D 下降约 2 倍。所以实测比理论低不一定是 bug反而是优化生效的信号。对于 DiT 这种没有 mask 的标准注意力FlashAttention 的 FLOPs 和手动实现的一致真正的问题出在算子融合如果 MLP 的 GELU 被融合进前面的 GEMM你 profiler 里看不到单独的 GELU kernel理论计数不能叠加算两层否则重复计数。我自己习惯把 torch.profiler 的输出按 operator 聚合然后用表格和理论值对比超过 1.2 倍就需要检查是不是某处矩阵乘的 shape 搞错了。5.2 数据处理和精度对 FLOPs 的隐性影响这里再说一个很简单但总被人忽略的事实FLOPs 是理论乘加次数不代表硬件实际计算次数。当你用 fp16 混合精度训练时某些框架会把乘积累加到 fp32 里硬件上计算的实际上是一个 FP32 FMA而不是 FP16 FMA。很多开源 FLOPs 计算器报的是“算法 FLOPs”但你实测的 MFU 如果按这个数算出来超过 100%基本上是精度换算的问题。另外 DiT 训练常用 VAE 把图像编码后用 latent 做 diffusion如果你把 VAE 的 FLOPs 也算进去单次前向会给 DiT 本身加上约 15~20 GFLOPs 的开销。这一点论文里经常避而不谈工程落地对比不同方案的端到端延迟时必须带上否则采样性能会被严重高估。5.3 开源计算工具实测对比与选型建议市面上能直接算 DiT FLOPs 的工具不多我更推荐自己写函数但为了 sanity check以下工具实测可以辅助验证。工具名称原理对 DiT 的适配度备注torchinfo只统计参数量和 tensor shape低不涉及 FLOPs仅用来算参数量thop基于 hook 统计乘加次数中对 transformer 支持差常低估注意力部分fvcoreMeta 开源支持 transformer高注意 patch embedder 需要手动标记calflops较新的 PyPI 包高内置常见 DiT 模块支持实测下来 fvcore 的数值和手动公式算的误差在 3% 以内calflops 更省事但偶尔会把 adaLN 的归一化统计成 O(S) 而不是 O(SD)需要二次核对。真正到了 100G FLOPs 的大模型我都会保留一份手写公式函数的输出作为基准其余工具只做交叉验证。6. 关于计算规模我最后想多说两句DiT 的 FLOPs 计算本身不难难的是把“算出来的数字”和“真实的硬件表现”之间的差值解释清楚。我在项目里早就放弃了追求 100% 精确的 FLOPs 计数器反而更依赖一套相对稳定的估算流程先用公式算理论值再在单卡上跑 100 步看实测吞吐两者对比出来的 MFU 就是我判断模型、框架和硬件是否匹配的核心指标。如果你准备用一个全新的 DiT 配置跑大型实验我建议在正式训练前花半天时间把你手头所有可能的配置不同的 patch、D、L用本文的代码快速扫一遍记录每个配置下前向和训练的 FLOPs 表再对照你集群的单卡 MFU直接就能筛掉 80% 不可能在 deadline 前跑完的方案。后续只要改动任何一个超参把函数重新跑一遍你脑子里的“算力地图”就始终是新的不会出现训练到一半才发现资源缺口的问题。
