如果你最近在研究 DiTDiffusion Transformer大概率会碰到一个很实际的困惑论文里那张 FID 对“计算规模”的曲线横坐标到底是怎么算出来的同样是 DiT-XL/2有人算出来一张图一次前向只要两百多 GFLOPs有人换一种口径却算出几十 TFLOPs数字能差几十倍。这其实不是模型变了而是“在哪个空间算”“按哪个口径算”没有先统一。这篇文章就从一个做项目落地的角度把 DiT 计算规模这事完整拆开讲一遍。我会把一个 DiT 的前向 FLOPs 公式手推出来拿 DiT-XL/2 的实际配置算一遍再补上训练、采样、显存这些工程上真正关心的口径最后给几个我在实际排查计算量时踩过的坑。不管你是要复现 DiT还是想估算自己改出来的模型在 A100/H100 上能不能跑这套方法都够用。1. 先分清DiT 计算规模到底在算哪几个量我最早被 DiT 计算规模搞晕就是因为把几个不同口径的量混在一起比了。这里先把口径钉死后面所有公式才不会乱套。口径含义常见数量级参数量Params模型里所有可学习权重的个数33M ~ 675M前向 FLOPs一张图做一次去噪推理需要的浮点运算量12G ~ 237G训练 FLOPs / 样本训练时处理一个样本的一次优化步大约为前向的 3 倍采样 FLOPs / 图生成一张图需要跑完整采样链路N个采样步 × 前向 FLOPs这几个概念里有两条最容易乱的分界线。第一条是“前向 FLOPs”和“训练 FLOPs”的区别。很多人以为扩散模型训练时要像采样那样跑完整个去噪链实际上训练时每个样本是随机抽一个时间步 t只跑一次加噪样本的前向和反向传播。所以训练成本大约等于“前向 × 3”这里的 3 指的是反向传播大约是前向的 2 倍。这个细节如果搞错你估出来的总训练算力会直接膨胀几十倍。第二条是“理论 FLOPs”和“实际耗时”的区别。FLOPs 只能告诉你计算量的数学上限不能直接等成 GPU 跑出来的时间。尤其是当序列长度 T 比较小的时候矩阵乘法形状也跟着变小A100 这种卡根本喂不满实际吞吐会比理论值低很多。所以我后面会专门说算完 FLOPs 之后怎么做工程层面的折算。还有一个前置条件必须确认DiT 是在 VAE 的 latent 空间里跑还是在 pixel 空间里跑。DiT 论文默认用的是 latent 空间256×256 的图经过 VAE 下采样 8 倍变成 32×32 的 latent 特征图再按 patch size 切块。这个选择对 T 的影响是数量级的因为序列长度 T 等于(latent 分辨率 / patch size)²而 FLOPs 里面大头是 T 的二次项和 d 的二次项。后续所有计算我都先明确告诉你是基于 latent 空间还是 pixel 空间。2. DiT 前向 FLOPs 的推导一个 DiT Block 里都有哪些乘法想快速估算 DiT 的计算规模不需要真的拿 profiler 去跑一遍先把一个 DiT Block 的结构拆开逐段数乘加次数就行。我按“乘加各算一次统称 1 FLOP”的常见口径来算。这里全部按前向计算来推不考虑反向。先约定几个符号LDiT Block 数量d模型宽度也就是 transformer 的 embedding 维度T输入 token 数也就是 patch 序列长度ppatch sizeC输入通道数。latent 空间通常 C4pixel 空间 C3h注意力头数多头的设置只影响线性层的分块方式不改变总 FLOPs 公式2.1 patch embedding 与输出解码层DiT 先要把二维的 latent 特征图切成 patch然后线性投影成 d 维向量。每个 patch 展平后长度是C × p²从C × p²映射到 d是一个全连接层patch embedding FLOPs 2 × T × d × C × p²对应的输出解码层也差不多从 d 再映射回C × p²成本同样是2 × T × d × C × p²。这部分通常只占总量的一小部分。拿 DiT-XL/2 来说T256, d1152, C4, p2嵌入层两端的 FLOPs 大约是 940 MFLOPs相比单 block 的 8.45 GFLOPs 完全可以忽略但参数别忘了它占显存。2.2 注意力模块QKV、attention map、加权求和DiT 用的是标准的非因果 multi-head self-attention没有 mask。这一步是很多估算公式最容易丢东西的地方。首先是把输入序列同时映射成 Q、K、V。这个操作通常实现成一个d → 3d的线性层QKV 投影 FLOPs 2 × T × d × 3d 6Td²然后计算 Q 和 K 的注意力分数。Q 的形状是T × dK 的转置是d × T矩阵乘法一次是T × d × T乘加各算一次attention score FLOPs 2 × T × d × T 2T²dsoftmax 本身是逐元素操作不做矩阵乘法通常可以忽略。接下来用 attention map 去加权 VQ 的加权注意矩阵T × T乘以 V 矩阵T × dattention 加权 FLOPs 2 × T × T × d 2T²d最后还有一个注意力输出的线性投影d → doutput projection FLOPs 2 × T × d × d 2Td²所以注意力部分汇总如下注意力模块 FLOPs 6Td² 2T²d 2T²d 2Td² 8Td² 4T²d这里的4T²d就是很多人会用错的地方。如果直接用视觉 Transformer 里流传的“每层 12Td²”这种速算公式会把注意力分数和注意力加权漏掉。对 ViT 这种 patch 数量只有一两百的模型漏掉问题不大但 DiT 一旦处理高分辨率 latent序列长度很容易涨到几千甚至上万T²d就会成为主导项。2.3 MLP 与 AdaNorm别把 4 倍宽度展开漏掉DiT Block 的 MLP 基本沿用 Transformer 的 MLP 结构隐藏层宽度通常取4d激活函数用 GELU。第一个线性层从 d 升到 4dMLP 第一层 FLOPs 2 × T × d × 4d 8Td²第二个线性层从 4d 压回 dMLP 第二层 FLOPs 2 × T × 4d × d 8Td²MLP 合计MLP FLOPs 16Td²这里我要提醒一下虽然激活函数 GELU 也有计算成本但在比较模型计算规模时行业里通常只统计矩阵乘法和卷积这类主导算子。逐元素的 GELU、LayerNorm、AdaNorm 的 scale/shift量级都是O(Td)跟Td²相比小两三个数量级日常估算直接忽略。2.4 合并成速算公式现在把一个 DiT Block 的前向 FLOPs 汇总FLOPs_block 注意力模块 MLP (8Td² 4T²d) 16Td² 24Td² 4T²d全模型前向FLOPs_forward L × (24Td² 4T²d) 2 × T × d × C × p²这个公式是我做计算规模估算时最常用的主力公式。有两个点需要说明这里用的是“乘加各算 1 FLOP”的口径。如果某些工具返回的是 MACs乘累加次数那转换成 FLOPs 需要再乘以 2否则数字会对不上。AdaNorm 和 modulation 的线性投影我没有放进主公示因为它们只在时间步 embedding 上做线性变换不跟在 batch、token 维度上展开总占比通常小于 1%。3. 拿 DiT-XL/2 手算一遍从配置到最终数字先看 DiT 论文里常用的四个模型配置这些数字在后面所有手算中都会用到。模型层数 L宽度 d多头数参数量DiT-S/212384633MDiT-B/21276812131MDiT-L/224102416458MDiT-XL/228115216675M如果你用的是 256×256 的输入图像并且走 VAE latent 空间下采样 8 倍后 latent 是 32×32通道数 C4。当 patch size2 时T (32 / 2)² 256对 DiT-XL/2 来说d1152, L28。先算单个 block24Td² 24 × 256 × 1152 × 1152 8.15e9 4T²d 4 × 256 × 256 × 1152 3.02e8 FLOPs_block ≈ 8.45e9全模型FLOPs_forward ≈ 28 × 8.45e9 ≈ 2.37e11也就是大约 236 GFLOPs。这样一个数才算出来。对应地DiT-S/2、DiT-B/2、DiT-L/2 的前向 FLOPs 分别大约是 12 GFLOPs、46 GFLOPs、161 GFLOPs。现在把口径变化一下你会看到非常夸张的差异。还是 DiT-XL/2口径结果单样本前向一次236 GFLOPs训练一个样本优化一步约 708 GFLOPs50 步采样生成一张图11.8 TFLOPs50 步采样并且开启 classifier-free guidance23.6 TFLOPs如果把模型搬到 pixel 空间跑同一张 256×256 图49.3 TFLOPs注意最后一行如果直接在像素空间跑T (256 / 2)² 16384序列长度是 latent 空间的 64 倍FLOPs 直接从 0.236 T 飙到 49.3 T这还是在没考虑 VAE 编解码成本的情况下。所以你看任何文章讨论 DiT 计算量第一反应一定是去确认它说的是 latent 空间还是 pixel 空间否则数字根本没可比性。4. 为什么计算规模对“分辨率”会二次爆炸复杂度敏感性分析DiT 的计算规模不是线性变化的它的脾气很鲜明T和d对计算量的影响完全不同而且这个差异会随着输入尺寸变化突然反转。看公式24Td² 4T²d这两项一个是线性于 T一个是平方于 T。什么时候 attention 的平方项开始主导让两项相等4T²d 24Td² T 6d也就是说当序列长度 T 超过大约 6 倍的模型宽度 d 时注意力矩阵计算的 FLOPs 就会反超 MLP 和 QKV 投影成为新的主要矛盾。拿 DiT-XL/2 来说d1152临界点大约是T6912。在 latent 空间里256×256 图像只有 256 个 token离临界点非常远这时候注意力平方项占比只有大约 3.6%计算规模主要由 d 决定。但是一旦跑到高分辨率场景比如 latent 变成 128×128T(128/2)²4096还不算太夸张再上去到 256×256T16384注意力平方项就会占到七成以上。这也是为什么高分辨率 DiT 一定要做窗口注意力或者改成 masked attention否则计算量会不可控。patch size 的影响就更直接了。latent 分辨率不变的情况下p 每翻一倍T 变成原来的四分之一attention 平方项理论上变成十六分之一。我拿 DiT-XL/2 在 32×32 latent 上算过patch sizeT前向 FLOPs2256236.8 GFLOPs46457.7 GFLOPs81614.3 GFLOPs从 p2 换到 p8计算量直接砍掉 16 倍以上。代价也很明确patch 太大等于信息在入口就被暴力压缩图像细节尤其是高频结构会丢失。DiT 论文里实验也验证了patch size8 时计算效率很好但 FID 会变差patch size2 的细节保留能力最强但最贵。实际项目里做这个权衡时我建议先看你的目标分辨率如果画面里有大量小物体或文字就不要为了省算力一上来就 p4 或者 p8。还有一个常见做法是“分辨率分层”先在低分辨率 latent 上用大步长 patch 走流程最后用一个小步长 patch 的高分辨率 DiT 做精修。这样总计算量不是简单相加因为高分辨率分支的T²会被限制在一个小范围内整体成本可控。5. 从单次前向到训练和采样总消耗5.1 训练总算力估算如果你要估整个训练任务可以用下面这条公式训练总 FLOPs ≈ 总训练样本数 × 迭代轮数 × 3 × FLOPs_forward这里的 3 就是“1 次前向 约 2 次反向”的训练开销系数。扩散模型训练时每个样本每次只随机抽一个时间步所以不需要乘扩散步数这一点我再强调一遍。举一个真实的项目估算例子。假设 DiT-XL/2 在 256×256 latent 空间训练FLOPs_forward236 GFLOPs训练集 128 万张图训练 300 个 epoch总样本数 1.28M × 300 3.84e8 训练总 FLOPs 3.84e8 × 3 × 236e9 2.72e20也就是大约 272 EFLOPs。如果单卡有效吞吐按 150 TFLOPs 算单卡需要大约 500 个小时。实际跑出来往往要再放宽 1.5 到 2 倍因为 DiT 在 latent 空间的矩阵形状偏小GPU 利用率很难跑到 50% 以上。5.2 采样总算力估算采样就完全是另一套算法了。生成一张图要迭代 N 个去噪步每步都要跑一次前向所以采样 FLOPs / 图 N × FLOPs_forward × guidance_倍数guidance_倍数只有在用 classifier-free guidance 时才需要乘。CFG 的一次采样里每个 step 都要同时跑条件模型和无条件模型所以相当于乘 2。50 步 DDIM 就已经很贵了如果用 Euler、Heun 这类更高阶的 solver实际 step 数还会叠加。这也是为什么 DiT 这类模型实际部署时主流方案都是蒸馏或者步数压缩。一个 DiT-XL/2 在 256×256 上50 步加 CFG 的生成成本大约是 23.6 TFLOPs对比 Stable Diffusion 那一类 U-Net 同样是 50 步加 CFG 的 512×512 图像成本已经不在一个量级了。部署前不做步数压缩推理费用会非常难看。5.3 显存估算公式之外的工程账本FLOPs 算完还要看显存否则你会出现“算力够但卡放不下”的窘境。显存这块有两条快速估算线。第一条线是参数量带来的权重和优化器显存。以 fp16 混合精度训练为例DiT-XL/2 的 675M 参数fp16 权重675M × 2B 1.35 GBfp16 梯度同样 1.35 GBfp32 的 master weight675M × 4B 2.7 GBAdam 的两个 fp32 状态675M × 8B 5.4 GB单模型一轮训练光权重和优化器就要吃掉大约 10.8 GB。所以 DiT-XL/2 想在 24 GB 的消费级卡上微调必须用 LoRA/QLoRA 一类的参数高效微调方案否则 activation 根本放不下。第二条线是 activation 显存。训练时要留存部分中间激活做反向传播DiT 的 activation 大头有两块一个是每层T × d的特征一个是某些实现里没做 flash attention 的batch × head × T × T注意力矩阵。后者在 T256 时不大但当 T 进入几千级别后一张草图就能把显存写满。所以我做高分辨率训练时一定会开 flash attention并且配合 activation checkpointing否则就算 FLOPs 再低显存也会先爆。6. 公式、工具和三个最容易踩的坑最终我一般会把公式直接写成一个 Python 函数方便在调整分辨率、patch size、模型宽度时快速出数。下面这段是我一直沿用的小工具def dit_forward_flops( d_model: int, layers: int, token_len: int, patch: int 2, latent_channels: int 4, ) - float: 估算 DiT 单图单次前向的 FLOPs乘加各算一次。 block_flops ( 24 * token_len * d_model * d_model 4 * token_len * token_len * d_model ) embedding_flops ( 2 * token_len * d_model * latent_channels * patch * patch ) return layers * block_flops 2 * embedding_flops if __name__ __main__: token_len (32 // 2) ** 2 # 32x32 latent, patch2 flops dit_forward_flops( d_model1152, layers28, token_lentoken_len, patch2, ) print(f{flops / 1e9:.2f} GFLOPs)如果想用工具做交叉验证常见选择有fvcore.nn.FlopCountAnalysis、thop、pytorch profiler。但我实测下来这类工具对 Transformer 的 attention 部分统计并不一致。有的工具把torch.matmul Q K.T正确算进去了有的直接当成普通乘加算有的还会漏掉 softmax 后面的attn V。所以工具只能做二次确认不能当唯一答案还是要以手推公式为准。我建议的验证办法很简单先拿一个很小配置比如 DiT-S/2用 profiler 跑一个真实输入把 profiler 给的总 FLOPs 和我公式算出来的做比较差异在 10% 以内就说明口径基本一致。如果差异很大多半是工具把某个重复计数了或者把非矩阵乘法的操作也算进去了。最后是我项目里总结的三个高频坑每个都曾经让我对不上数字第一个坑是把 MACs 当 FLOPs。很多工具返回的是 MACs一个乘加算 1换算 FLOPs 时要乘以 2。第二个坑是拿 ViT 的12Td²速算公式套 DiT。ViT 输入 patch 少attention 平方项占比低但 DiT 的 T 经常远大于临界点6d这时候套那个公式会严重低估。第三个坑是训练时乘了扩散步数。训练每样本随机抽一个 t只跑一次模型真正反复跑模型的是采样阶段。我个人现在复盘 DiT 计算规模的习惯是先把“latent 还是 pixel”“前向还是训练”“FLOPs 还是 MACs”这三个问题写在纸最上面再开始套公式。把口径锁死之后DiT 的计算规模其实很好算它比 U-Net 那种结构清晰的卷积网络更容易手推因为所有成本都集中在24Td² 4T²d这一个式子里。把这一个式子用熟以后不管是 Scaled DiT、DiT 变体还是视频 DiT核心估算都能快速给出靠谱答案。