先回答一个我经常被问到的问题DiT的计算规模到底怎么算。DiT也就是Diffusion Transformer这两年从论文一步步成了图像生成和视频生成模型的主流底座像我们熟知的Sora、Stable Diffusion 3这些内部都跑着类似DiT的结构。我最近刚好帮团队做了一次DiT-XL/2的训练成本评估从参数量、单次前向的FLOPs一路算到训练完整个模型需要多少张A100整套推导过程不复杂但要是不把口径对齐数字很容易差出一倍甚至几个数量级。这篇文章把我实际用的方法、公式和实测验证全部整理出来适合要训DiT、要做推理部署预算、或者想从论文反推别人训练配置的朋友参考。1. 先把口径对齐你要算的是参数量、单次FLOPs还是卡时很多人在聊“计算规模”的时候其实说的是三个完全不同的东西模型参数量、单次前向FLOPs、训练总计算量。这三个数字对应完全不同的决策问题混着聊就会闹笑话。我整理成一个表格方便你对号入座口径常用单位解决什么问题怎么得到参数量M / B模型文件多大、显存够不够、参数量对比网络结构逐层累加单次前向FLOPsGFLOPs / TFLOPs跑一次推理消耗多少算力逐层计算乘加次数训练总计算量PFLOPs / ZFLOPs训练一个模型需要多少GPU资源FLOPs × 样本数 × 前后向系数举个例子有人问我“DiT-XL/2有多少算力”我第一反应是你到底想问这个模型多大还是想训一次需要多少卡如果是前者675M参数量就够了如果是后者你需要的是一张完整的卡时预算表。参数量和FLOPs之间的关系也不是固定的序列长度一变同样参数量的模型FLOPs能差出两三倍。我自己的习惯是在任何计算规模相关的文档开头先写清楚口径比如“本文FLOPs均按一次乘加各计一次且不含偏置项”。否则数字传到别人那里可能被按MACs口径再除一次2最后对不上账还得回头查。2. 从模型配置反推参数量一个DiT Block就够了2.1 DiT Block的参数构成DiT的block结构和ViT非常接近一个条件注入层一个self-attention一个MLP前面各带一个LayerNorm。但有几个细节直接决定了参数量的公式。先看条件注入。DiT用的是adaLN也就是把时间步和类别标签的embedding相加之后过一个线性层输出6倍的hidden_size用来生成attention前和MLP前的scale、shift、gate。关键是DiT官方代码里这层adaLN_modulation是“一个SiLU激活加一个Linear(D, 6D)”不是有些人以为的两层MLP。这一点必须记住否则参数量会多算一大截。再看attention部分QKV是一个线性层把D维映射到3D维参数是3D²输出投影把D维映射回D维参数是D²。MLP部分DiT的mlp_ratio默认是4也就是先扩到4D再压缩回D两个线性层参数分别是4D²和4D²合计8D²。两个LayerNorm在DiT里是elementwise_affineFalse不含可学习参数所以不贡献参数量。把这几项相加一个DiT Block的参数就是模块参数量公式说明adaLN调制6D²只有一层线性QKV投影3D²带bias量级可忽略输出投影D²MLP两层8D²4倍扩维再压回合计18D²单block近似公式D就是hidden_sizeL是block层数。整个DiT的参数量约等于18D² × L再加上patch embed、timestep embed、label embed、final layer这些小头。2.2 三个常见配置的算例拿DiT论文里最常用的几个配置来验算。DiT-S/2D384L1218×384²×12 ≈ 31.8M官方公布约33M。DiT-B/2D768L1218×768²×12 ≈ 127.4M官方公布约130M。DiT-L/2D1024L2418×1024²×24 ≈ 452.9M官方公布约458M。DiT-XL/2D1152L2818×1152²×28 ≈ 668.9M官方公布约675M。你会发现用这个近似公式算出来的数字和官方只差区区几个M误差来自patch embed、timestep embed、label embed和final layer。这几个模块加一起通常不到总参数的1-2%所以我在做预算时直接用18D²×L估算就够了误差对卡时预算的影响可以忽略。2.3 容易被忽略的小参数块既然上面提到了小头我顺便把它们的参数量公式列出来万一你需要精确值patch embed输入维度是patch_size²×3比如patch_size2时输入是12维Linear(12, D)参数约12D。timestep embedMLP把256维频率嵌入映射到D再映射到D约2D²。label embedImageNet有1000类Embedding(1000, D)参数1000D。final layer把D维映射回patch_size²×3约12D。以DiT-XL/2来说这些小模块加起来也不超过10M参数相比675M的总量确实可以忽略。但如果你的应用场景是自定义数据集、类别数很多label embed这一项会变大那时候别把它省掉。2.4 参数公式的一个使用技巧当你手头没有模型代码、只有论文里的配置表时这个18D²×L公式特别有用。比如看到一篇视频生成论文说“我们使用了deep28、hidden1152的DiT骨干”你30秒就能算出模型大约6.7亿参数不用等代码开源。我经常拿这个公式在会议现场快速估算判断对方模型规模和已知模型处于什么量级。3. 单次前向的FLOPs拆账Attention平方项才是主角3.1 线性层FLOPs的通用公式先说通用规则一个Linear(d_in, d_out)处理s个tokenFLOPs 2 × s × d_in × d_out。因为每个输出元素要做d_in次乘法和d_in次加法FLOPs统计时乘和加各算一次。注意这里有个口径问题很多工具返回的是“MACs”也就是乘加次数一个乘加在数值上等于一次乘法和一次加法但有的文档把它当成1次FLOPs。同样一个模型FLOPs和MACs直接差2倍。后面我讲的都是“乘加各算一次”的FLOPs口径这是大多数硬件厂商和论文使用的口径。凡是看到数字对不上优先检查是不是把MACs当成FLOPs用了。3.2 Attention要拆成三笔账设序列长度s (H/patch_size) × (W/patch_size)hidden_size D。对每个DiT BlockAttention相关的FLOPs可以分为四笔QKV投影3个线性层总计 6sD²。QK^T计算注意力矩阵输出是s×s的矩阵每个元素D维点积总计 2s²D。PV加权求和同样s×s×D的规模总计 2s²D。输出投影一个Linear(D, D)总计 2sD²。很多讲Transformer算力的文章只统计QKV投影和输出投影忽略QK^T和加权求和这在短文本序列上问题不大但在DiT这种长序列图像模型上就完全不行。256×256的图用patch_size2序列长度是16384token数量比普通文本任务高了一个量级Attention的平方项会迅速盖过线性投影项。3.3 FFN、adaLN和patch embed的贡献FFN两层第一层D→4D第二层4D→D每层都是2sD²×4合计16sD²。adaLN调制作用于condition向量而不是序列每block约2×D×6D12D²和s无关在s很大时可以忽略。patch embed和final layer各一个线性层合计约0.9GFLOPs跟主体比是两个数量级的差距。所以单个DiT Block的FLOPs可以近似写成FLOPs_block ≈ 6sD² 2s²D 2s²D 2sD² 16sD² 24sD² 4s²D注意24sD²里已经包含了QKV、输出投影和FFN4s²D则纯粹是Attention矩阵的计算量。3.4 DiT-XL/2 256×256完整算表下面用DiT-XL/2、输入256×256、patch_size2来完整算一遍。s16384D1152。组件公式FLOPsQKV投影6sD²130.5GQK^T注意力矩阵2s²D618.5GAV加权求和2s²D618.5G输出投影2sD²43.5GMLP两层16sD²347.9G单block小计约1758.8G28个block合计约49.2Tpatch embed final layer约0.9G总计forward约49.3T这个49.3T是单次前向传播的量。如果你看到有人直接用“2 × 参数量 × token数”估算得到约22T那就是没有算Attention矩阵的平方项偏差超过一倍。DiT的计算规模之所以比同等参数量的文本模型高核心就在这个4s²D上。我再给你一个直观对比把patch_size从2换成4同样一张256×256的图序列长度从16384降到4096Attention平方项直接降到原来的1/16整体FLOPs会从49.3T降到大约16-18T。这就是为什么很多实际部署的DiT推理模型倾向于用更大的patch_size算力削掉一大半图像质量只损失一点。4. 扩散模型特有的计算放大器训练步数、反向传播与CFG4.1 训练时每张图只采一个timestep到这里单次前向的FLOPs已经有了但计算规模远没有完。扩散模型的训练和普通分类模型最大的区别在于它不是对每张图跑一次前向就完事而是要为每个样本随机采样一个diffusion timestep然后预测这个timestep下的噪声。新手最容易踩的坑是看到DDPM采样要跑几百步就以为训练也要跑几百步。实际上训练时每张图只跑一次前向和一次反向唯一的额外成本是随机采一个t、把t和类别标签送进embedding这部分几乎可以忽略。换句话说训练的总FLOPs是“样本数 × 单次forward FLOPs × 前后向系数”而不是“样本数 × 采样步数 × 单次forward FLOPs”。前后向系数方面业界常用的经验是取3也就是backward的FLOPs大约是forward的2倍forwardbackward合计约3倍。这个系数在不同框架、不同混合精度下会有波动但拿来估算卡时足够用。4.2 采样时多步denoising叠加上CFG推理阶段就是另一套算法了。生成一张图要跑N步denoising每一步都是一次完整的前向传播。DiT论文里常用的采样配置是DDPM 250步或DDIM 50步。更关键的是classifier-free guidance也就是CFG。为了提升生成质量推理时会同时跑条件模型和无条件模型两个前向所以CFG开启后每步的FLOPs直接翻倍。我之前评估过一个视频生成模型的推理成本一开始只按单模型、50步去算结果实际线上配置是CFG加双模型推理算力差了整整4倍预算直接重做。以DiT-XL/2 256×256为例单次forward是49.3T250步DDPM无CFG就是12.3PFLOPs开CFG就是24.6PFLOPs。一张图消耗几十PFLOPs听起来很大但现代GPU峰值也在几百TFLOPs每秒所以延迟问题主要不在总FLOPs而在于每一步的kernel launch和访存开销。4.3 训练和推理的计算规模为什么差那么多训练侧是“每个样本一次前后向”采样侧是“每个样本几十次前向”。在模型刚训完、用户量不大的时候训练成本占绝对大头一旦模型上线、调用量上来推理总量会迅速反超。还有一个容易忽视的点训练是纯算力密集利用率可以做到40%-60%推理尤其是图像生成这种短序列场景访存占比高FLOPs利用率经常只有个位数。所以拿FLOPs直接除以GPU峰值来估算推理延迟会得到一个极度乐观的下限实际要比理论值慢得多。这一点在后面做预算和部署时非常重要。5. 从FLOPs到卡时预算DiT-XL/2训练一次要多少张A1005.1 卡时计算公式把FLOPs转换成GPU卡时公式很简单GPU卡时(天) 总FLOPs / (GPU峰值FLOPs × 实际利用率 × 86400)其中GPU峰值要按你实际使用的精度去查。A100 80G的FP16稠密算力是312 TFLOPsH100的BF16能达到989 TFLOPs左右。实际利用率在训练场景通常取40%-55%取决于框架优化、数据加载、通信开销。做预算时我习惯取50%作为基准再用40%做保守上界。以单张A100、50%利用率为例每秒有效算力就是156 TFLOPs约等于1.56×10¹⁴ FLOPs每秒。这个数字不需要背每次现算即可。5.2 完整算例DiT-XL/2训练400个epoch假设在ImageNet上训练DiT-XL/2分辨率256×256、patch_size2。ImageNet训练集约128万张图训400个epoch。单样本forward FLOPs49.3T即4.93×10¹³。训练前后向系数取3单样本训练FLOPs 147.9T。总样本数 128万 × 400 5.12亿。总FLOPs 5.12亿 × 147.9T ≈ 7.57×10²²约76 ZFLOPs。单卡A100有效算力按50%利用率 1.56×10¹⁴ FLOPs/s。GPU秒数 7.57×10²² / 1.56×10¹⁴ ≈ 4.85×10⁸秒。换算成GPU天约5600卡天。如果你手头有256张A100大概需要22天512张就是11天。这个量级和DiT同规模模型在真实训练场景下的工期基本吻合可以用来验证你的预算表是否靠谱。需要说明的是如果训练分辨率提升到512×512序列长度变成原来的4倍FLOPs会显著上涨如果改用更大的patch_size又会明显下降所以这套模板要根据实际配置重算。5.3 反向推导别人的训练配置同一个公式反过来用也很有价值。看到论文里写“我们用了64张A100训了15天”你可以快速判断它的总FLOPs量级总FLOPs 卡数 × 单卡有效算力 × 训练天数 × 86400然后用“总FLOPs ÷ (3 × 数据集大小 × epochs)”还能反推出单次forward FLOPs再对照是不是符合论文描述的网络规模。这个方法帮我识破过一次明显虚标的训练配置。那个论文号称参数量只有300M但按训练卡时反推的单样本FLOPs比同等参数量的标准DiT高出一大截明显是训练配方里加了什么没写清楚的东西或者参数标错了。5.4 算预算时还要把显存和通信一起看卡时只是算力维度显卡租用成本还受显存限制。DiT-XL/2有675M参数混合精度训练下AdamW优化器状态、梯度、激活值都会吃显存单卡显存不够就要梯度 checkpointing或者张量并行这会进一步降低利用率。所以算预算时我会先把MLP的激活值大小粗估一下判断是否需要重计算再决定MFU按50%还是按40%算。激活值估算比较复杂但至少要知道序列越长、batch越大激活值增长越快梯度checkpointing的收益越大代价是FLOPs额外增加约30%以上。6. 用代码实测验证以及我踩过的六个坑6.1 用工具实测参数量和FLOPs公式算完一定要拿工具实测一轮。我的做法是用fvcore的FlopCountAnalysis输入随机张量走一遍forward直接读统计结果。DiT的forward签名一般是(x, t, y)t是时间步相关的embeddingy是类别标签注意别只喂一个x进去。import torch from fvcore.nn import FlopCountAnalysis, parameter_count_table model DiT_XL_2(input_size256, patch_size2) model.eval() x torch.randn(1, 3, 256, 256) t torch.randn(1, model.num_classes) # 时间步embedding输入 y torch.randint(0, 1000, (1,), dtypetorch.long) flops FlopCountAnalysis(model, (x, t, y)) print(flops.total()) # FLOPs print(parameter_count_table(model)) # 参数量实测出来的数字和我手算的49.3T有差异是正常的因为每个profile工具对LayerNorm、SiLU、bias的计数口径不完全一样差异通常在10%以内。如果实测值差了一倍大概率是工具返回的是MACs而不是FLOPs或者模型定义里悄悄多了某个大模块。6.2 六个容易踩的坑第一个坑MACs和FLOPs混淆。这个前面反复强调过同一个模型两种口径差2倍跨工具对比时必踩。我个人的防护习惯是所有手算表格里都把公式列出来这样别人可以按公式反推口径。第二个坑用采样步数去算训练成本。扩散模型训练只采一个timestep跟推理跑几百步完全是两回事。我第一次给一个扩散项目做预算时就是把250步乘进去了结果算出来一个天文数字排查半天才发现是这一步错了。第三个坑忽略CFG。推理预算里CFG意味着每步跑两个模型FLOPs翻倍而CFG在扩散模型采样里几乎是标配。如果你只按单模型算部署成本会低估一半。第四个坑patch_size对Attention平方项的影响。我见过有人从paper里抄了一个FLOPs数字然后改了patch_size继续用旧数字。实际上patch从2改成4token数变成1/4Attention平方项变成1/16整体FLOPs大变千万不能沿用旧值。第五个坑拿FLOPs直接推导延迟。FLOPs给的是算力下限但图像生成推理在短序列下严重受限于内存带宽和kernel调度实际延迟往往远高于理论值。我在A100上实测DiT-XXL这种模型单步生成FLOPs只占理论算力的很小一部分大部分时间花在数据搬运上。第六个坑MFU假设太乐观。训练预算时很多人直接按GPU峰值算结果把卡时低估了一半以上。公共云上的实际训练利用率达不到你想象的峰值做预算时宁可按40%算多出的卡数就当安全余量。7. 最后分享一套我自己固定的估算顺序接到“DiT计算规模”这类问题我现在基本30分钟出全套结果先用第2节的18D²×L公式30秒粗算参数量再按第3节的表格手算单次forward FLOPs重点算清Attention平方项然后按第5节的卡时公式把训练预算、推理预算分别列出来最后用fvcore跑一轮实测校准手算误差。这套流程里最值钱的经验其实是第四节的“放大器”概念。DiT的FLOPs只是起点训练前后向的3倍系数、采样步数、CFG翻倍、GPU利用率这几个乘数往往比模型本身的FLOPs更能决定你的真实账单。任何一次预算讨论我都会要求把这几项系数白纸黑字列出来因为这正是所有对不上的数字背后最常见的来源。
DiT计算规模估算指南：参数量、FLOPs与A100训练卡时全解析
先回答一个我经常被问到的问题DiT的计算规模到底怎么算。DiT也就是Diffusion Transformer这两年从论文一步步成了图像生成和视频生成模型的主流底座像我们熟知的Sora、Stable Diffusion 3这些内部都跑着类似DiT的结构。我最近刚好帮团队做了一次DiT-XL/2的训练成本评估从参数量、单次前向的FLOPs一路算到训练完整个模型需要多少张A100整套推导过程不复杂但要是不把口径对齐数字很容易差出一倍甚至几个数量级。这篇文章把我实际用的方法、公式和实测验证全部整理出来适合要训DiT、要做推理部署预算、或者想从论文反推别人训练配置的朋友参考。1. 先把口径对齐你要算的是参数量、单次FLOPs还是卡时很多人在聊“计算规模”的时候其实说的是三个完全不同的东西模型参数量、单次前向FLOPs、训练总计算量。这三个数字对应完全不同的决策问题混着聊就会闹笑话。我整理成一个表格方便你对号入座口径常用单位解决什么问题怎么得到参数量M / B模型文件多大、显存够不够、参数量对比网络结构逐层累加单次前向FLOPsGFLOPs / TFLOPs跑一次推理消耗多少算力逐层计算乘加次数训练总计算量PFLOPs / ZFLOPs训练一个模型需要多少GPU资源FLOPs × 样本数 × 前后向系数举个例子有人问我“DiT-XL/2有多少算力”我第一反应是你到底想问这个模型多大还是想训一次需要多少卡如果是前者675M参数量就够了如果是后者你需要的是一张完整的卡时预算表。参数量和FLOPs之间的关系也不是固定的序列长度一变同样参数量的模型FLOPs能差出两三倍。我自己的习惯是在任何计算规模相关的文档开头先写清楚口径比如“本文FLOPs均按一次乘加各计一次且不含偏置项”。否则数字传到别人那里可能被按MACs口径再除一次2最后对不上账还得回头查。2. 从模型配置反推参数量一个DiT Block就够了2.1 DiT Block的参数构成DiT的block结构和ViT非常接近一个条件注入层一个self-attention一个MLP前面各带一个LayerNorm。但有几个细节直接决定了参数量的公式。先看条件注入。DiT用的是adaLN也就是把时间步和类别标签的embedding相加之后过一个线性层输出6倍的hidden_size用来生成attention前和MLP前的scale、shift、gate。关键是DiT官方代码里这层adaLN_modulation是“一个SiLU激活加一个Linear(D, 6D)”不是有些人以为的两层MLP。这一点必须记住否则参数量会多算一大截。再看attention部分QKV是一个线性层把D维映射到3D维参数是3D²输出投影把D维映射回D维参数是D²。MLP部分DiT的mlp_ratio默认是4也就是先扩到4D再压缩回D两个线性层参数分别是4D²和4D²合计8D²。两个LayerNorm在DiT里是elementwise_affineFalse不含可学习参数所以不贡献参数量。把这几项相加一个DiT Block的参数就是模块参数量公式说明adaLN调制6D²只有一层线性QKV投影3D²带bias量级可忽略输出投影D²MLP两层8D²4倍扩维再压回合计18D²单block近似公式D就是hidden_sizeL是block层数。整个DiT的参数量约等于18D² × L再加上patch embed、timestep embed、label embed、final layer这些小头。2.2 三个常见配置的算例拿DiT论文里最常用的几个配置来验算。DiT-S/2D384L1218×384²×12 ≈ 31.8M官方公布约33M。DiT-B/2D768L1218×768²×12 ≈ 127.4M官方公布约130M。DiT-L/2D1024L2418×1024²×24 ≈ 452.9M官方公布约458M。DiT-XL/2D1152L2818×1152²×28 ≈ 668.9M官方公布约675M。你会发现用这个近似公式算出来的数字和官方只差区区几个M误差来自patch embed、timestep embed、label embed和final layer。这几个模块加一起通常不到总参数的1-2%所以我在做预算时直接用18D²×L估算就够了误差对卡时预算的影响可以忽略。2.3 容易被忽略的小参数块既然上面提到了小头我顺便把它们的参数量公式列出来万一你需要精确值patch embed输入维度是patch_size²×3比如patch_size2时输入是12维Linear(12, D)参数约12D。timestep embedMLP把256维频率嵌入映射到D再映射到D约2D²。label embedImageNet有1000类Embedding(1000, D)参数1000D。final layer把D维映射回patch_size²×3约12D。以DiT-XL/2来说这些小模块加起来也不超过10M参数相比675M的总量确实可以忽略。但如果你的应用场景是自定义数据集、类别数很多label embed这一项会变大那时候别把它省掉。2.4 参数公式的一个使用技巧当你手头没有模型代码、只有论文里的配置表时这个18D²×L公式特别有用。比如看到一篇视频生成论文说“我们使用了deep28、hidden1152的DiT骨干”你30秒就能算出模型大约6.7亿参数不用等代码开源。我经常拿这个公式在会议现场快速估算判断对方模型规模和已知模型处于什么量级。3. 单次前向的FLOPs拆账Attention平方项才是主角3.1 线性层FLOPs的通用公式先说通用规则一个Linear(d_in, d_out)处理s个tokenFLOPs 2 × s × d_in × d_out。因为每个输出元素要做d_in次乘法和d_in次加法FLOPs统计时乘和加各算一次。注意这里有个口径问题很多工具返回的是“MACs”也就是乘加次数一个乘加在数值上等于一次乘法和一次加法但有的文档把它当成1次FLOPs。同样一个模型FLOPs和MACs直接差2倍。后面我讲的都是“乘加各算一次”的FLOPs口径这是大多数硬件厂商和论文使用的口径。凡是看到数字对不上优先检查是不是把MACs当成FLOPs用了。3.2 Attention要拆成三笔账设序列长度s (H/patch_size) × (W/patch_size)hidden_size D。对每个DiT BlockAttention相关的FLOPs可以分为四笔QKV投影3个线性层总计 6sD²。QK^T计算注意力矩阵输出是s×s的矩阵每个元素D维点积总计 2s²D。PV加权求和同样s×s×D的规模总计 2s²D。输出投影一个Linear(D, D)总计 2sD²。很多讲Transformer算力的文章只统计QKV投影和输出投影忽略QK^T和加权求和这在短文本序列上问题不大但在DiT这种长序列图像模型上就完全不行。256×256的图用patch_size2序列长度是16384token数量比普通文本任务高了一个量级Attention的平方项会迅速盖过线性投影项。3.3 FFN、adaLN和patch embed的贡献FFN两层第一层D→4D第二层4D→D每层都是2sD²×4合计16sD²。adaLN调制作用于condition向量而不是序列每block约2×D×6D12D²和s无关在s很大时可以忽略。patch embed和final layer各一个线性层合计约0.9GFLOPs跟主体比是两个数量级的差距。所以单个DiT Block的FLOPs可以近似写成FLOPs_block ≈ 6sD² 2s²D 2s²D 2sD² 16sD² 24sD² 4s²D注意24sD²里已经包含了QKV、输出投影和FFN4s²D则纯粹是Attention矩阵的计算量。3.4 DiT-XL/2 256×256完整算表下面用DiT-XL/2、输入256×256、patch_size2来完整算一遍。s16384D1152。组件公式FLOPsQKV投影6sD²130.5GQK^T注意力矩阵2s²D618.5GAV加权求和2s²D618.5G输出投影2sD²43.5GMLP两层16sD²347.9G单block小计约1758.8G28个block合计约49.2Tpatch embed final layer约0.9G总计forward约49.3T这个49.3T是单次前向传播的量。如果你看到有人直接用“2 × 参数量 × token数”估算得到约22T那就是没有算Attention矩阵的平方项偏差超过一倍。DiT的计算规模之所以比同等参数量的文本模型高核心就在这个4s²D上。我再给你一个直观对比把patch_size从2换成4同样一张256×256的图序列长度从16384降到4096Attention平方项直接降到原来的1/16整体FLOPs会从49.3T降到大约16-18T。这就是为什么很多实际部署的DiT推理模型倾向于用更大的patch_size算力削掉一大半图像质量只损失一点。4. 扩散模型特有的计算放大器训练步数、反向传播与CFG4.1 训练时每张图只采一个timestep到这里单次前向的FLOPs已经有了但计算规模远没有完。扩散模型的训练和普通分类模型最大的区别在于它不是对每张图跑一次前向就完事而是要为每个样本随机采样一个diffusion timestep然后预测这个timestep下的噪声。新手最容易踩的坑是看到DDPM采样要跑几百步就以为训练也要跑几百步。实际上训练时每张图只跑一次前向和一次反向唯一的额外成本是随机采一个t、把t和类别标签送进embedding这部分几乎可以忽略。换句话说训练的总FLOPs是“样本数 × 单次forward FLOPs × 前后向系数”而不是“样本数 × 采样步数 × 单次forward FLOPs”。前后向系数方面业界常用的经验是取3也就是backward的FLOPs大约是forward的2倍forwardbackward合计约3倍。这个系数在不同框架、不同混合精度下会有波动但拿来估算卡时足够用。4.2 采样时多步denoising叠加上CFG推理阶段就是另一套算法了。生成一张图要跑N步denoising每一步都是一次完整的前向传播。DiT论文里常用的采样配置是DDPM 250步或DDIM 50步。更关键的是classifier-free guidance也就是CFG。为了提升生成质量推理时会同时跑条件模型和无条件模型两个前向所以CFG开启后每步的FLOPs直接翻倍。我之前评估过一个视频生成模型的推理成本一开始只按单模型、50步去算结果实际线上配置是CFG加双模型推理算力差了整整4倍预算直接重做。以DiT-XL/2 256×256为例单次forward是49.3T250步DDPM无CFG就是12.3PFLOPs开CFG就是24.6PFLOPs。一张图消耗几十PFLOPs听起来很大但现代GPU峰值也在几百TFLOPs每秒所以延迟问题主要不在总FLOPs而在于每一步的kernel launch和访存开销。4.3 训练和推理的计算规模为什么差那么多训练侧是“每个样本一次前后向”采样侧是“每个样本几十次前向”。在模型刚训完、用户量不大的时候训练成本占绝对大头一旦模型上线、调用量上来推理总量会迅速反超。还有一个容易忽视的点训练是纯算力密集利用率可以做到40%-60%推理尤其是图像生成这种短序列场景访存占比高FLOPs利用率经常只有个位数。所以拿FLOPs直接除以GPU峰值来估算推理延迟会得到一个极度乐观的下限实际要比理论值慢得多。这一点在后面做预算和部署时非常重要。5. 从FLOPs到卡时预算DiT-XL/2训练一次要多少张A1005.1 卡时计算公式把FLOPs转换成GPU卡时公式很简单GPU卡时(天) 总FLOPs / (GPU峰值FLOPs × 实际利用率 × 86400)其中GPU峰值要按你实际使用的精度去查。A100 80G的FP16稠密算力是312 TFLOPsH100的BF16能达到989 TFLOPs左右。实际利用率在训练场景通常取40%-55%取决于框架优化、数据加载、通信开销。做预算时我习惯取50%作为基准再用40%做保守上界。以单张A100、50%利用率为例每秒有效算力就是156 TFLOPs约等于1.56×10¹⁴ FLOPs每秒。这个数字不需要背每次现算即可。5.2 完整算例DiT-XL/2训练400个epoch假设在ImageNet上训练DiT-XL/2分辨率256×256、patch_size2。ImageNet训练集约128万张图训400个epoch。单样本forward FLOPs49.3T即4.93×10¹³。训练前后向系数取3单样本训练FLOPs 147.9T。总样本数 128万 × 400 5.12亿。总FLOPs 5.12亿 × 147.9T ≈ 7.57×10²²约76 ZFLOPs。单卡A100有效算力按50%利用率 1.56×10¹⁴ FLOPs/s。GPU秒数 7.57×10²² / 1.56×10¹⁴ ≈ 4.85×10⁸秒。换算成GPU天约5600卡天。如果你手头有256张A100大概需要22天512张就是11天。这个量级和DiT同规模模型在真实训练场景下的工期基本吻合可以用来验证你的预算表是否靠谱。需要说明的是如果训练分辨率提升到512×512序列长度变成原来的4倍FLOPs会显著上涨如果改用更大的patch_size又会明显下降所以这套模板要根据实际配置重算。5.3 反向推导别人的训练配置同一个公式反过来用也很有价值。看到论文里写“我们用了64张A100训了15天”你可以快速判断它的总FLOPs量级总FLOPs 卡数 × 单卡有效算力 × 训练天数 × 86400然后用“总FLOPs ÷ (3 × 数据集大小 × epochs)”还能反推出单次forward FLOPs再对照是不是符合论文描述的网络规模。这个方法帮我识破过一次明显虚标的训练配置。那个论文号称参数量只有300M但按训练卡时反推的单样本FLOPs比同等参数量的标准DiT高出一大截明显是训练配方里加了什么没写清楚的东西或者参数标错了。5.4 算预算时还要把显存和通信一起看卡时只是算力维度显卡租用成本还受显存限制。DiT-XL/2有675M参数混合精度训练下AdamW优化器状态、梯度、激活值都会吃显存单卡显存不够就要梯度 checkpointing或者张量并行这会进一步降低利用率。所以算预算时我会先把MLP的激活值大小粗估一下判断是否需要重计算再决定MFU按50%还是按40%算。激活值估算比较复杂但至少要知道序列越长、batch越大激活值增长越快梯度checkpointing的收益越大代价是FLOPs额外增加约30%以上。6. 用代码实测验证以及我踩过的六个坑6.1 用工具实测参数量和FLOPs公式算完一定要拿工具实测一轮。我的做法是用fvcore的FlopCountAnalysis输入随机张量走一遍forward直接读统计结果。DiT的forward签名一般是(x, t, y)t是时间步相关的embeddingy是类别标签注意别只喂一个x进去。import torch from fvcore.nn import FlopCountAnalysis, parameter_count_table model DiT_XL_2(input_size256, patch_size2) model.eval() x torch.randn(1, 3, 256, 256) t torch.randn(1, model.num_classes) # 时间步embedding输入 y torch.randint(0, 1000, (1,), dtypetorch.long) flops FlopCountAnalysis(model, (x, t, y)) print(flops.total()) # FLOPs print(parameter_count_table(model)) # 参数量实测出来的数字和我手算的49.3T有差异是正常的因为每个profile工具对LayerNorm、SiLU、bias的计数口径不完全一样差异通常在10%以内。如果实测值差了一倍大概率是工具返回的是MACs而不是FLOPs或者模型定义里悄悄多了某个大模块。6.2 六个容易踩的坑第一个坑MACs和FLOPs混淆。这个前面反复强调过同一个模型两种口径差2倍跨工具对比时必踩。我个人的防护习惯是所有手算表格里都把公式列出来这样别人可以按公式反推口径。第二个坑用采样步数去算训练成本。扩散模型训练只采一个timestep跟推理跑几百步完全是两回事。我第一次给一个扩散项目做预算时就是把250步乘进去了结果算出来一个天文数字排查半天才发现是这一步错了。第三个坑忽略CFG。推理预算里CFG意味着每步跑两个模型FLOPs翻倍而CFG在扩散模型采样里几乎是标配。如果你只按单模型算部署成本会低估一半。第四个坑patch_size对Attention平方项的影响。我见过有人从paper里抄了一个FLOPs数字然后改了patch_size继续用旧数字。实际上patch从2改成4token数变成1/4Attention平方项变成1/16整体FLOPs大变千万不能沿用旧值。第五个坑拿FLOPs直接推导延迟。FLOPs给的是算力下限但图像生成推理在短序列下严重受限于内存带宽和kernel调度实际延迟往往远高于理论值。我在A100上实测DiT-XXL这种模型单步生成FLOPs只占理论算力的很小一部分大部分时间花在数据搬运上。第六个坑MFU假设太乐观。训练预算时很多人直接按GPU峰值算结果把卡时低估了一半以上。公共云上的实际训练利用率达不到你想象的峰值做预算时宁可按40%算多出的卡数就当安全余量。7. 最后分享一套我自己固定的估算顺序接到“DiT计算规模”这类问题我现在基本30分钟出全套结果先用第2节的18D²×L公式30秒粗算参数量再按第3节的表格手算单次forward FLOPs重点算清Attention平方项然后按第5节的卡时公式把训练预算、推理预算分别列出来最后用fvcore跑一轮实测校准手算误差。这套流程里最值钱的经验其实是第四节的“放大器”概念。DiT的FLOPs只是起点训练前后向的3倍系数、采样步数、CFG翻倍、GPU利用率这几个乘数往往比模型本身的FLOPs更能决定你的真实账单。任何一次预算讨论我都会要求把这几项系数白纸黑字列出来因为这正是所有对不上的数字背后最常见的来源。
