第 3 课：五大工具实战

DeepSpeed Flops Profiler、PyTorch FlopCounterMode、fvcore、ptflops、calflops 的代码用法与输出解读

1. 本课目标

前面两课搞定了“为什么算”和“怎么手算”。这节课把重点放到“怎么跑代码算”——分别用五种主流工具给同一个 Transformer 模型算 FLOPs，并学会对比它们的输出。

2. DeepSpeed Flops Profiler

DeepSpeed 提供两种使用方式：推理时一次性估算用 get_model_profile；训练流程中埋点用 FlopsProfiler 类。

方式 A：get_model_profile（推理/快速估算）

import torch
from transformers import BertForSequenceClassification
from deepspeed.profiling.flops_profiler import get_model_profile
from deepspeed.accelerator import get_accelerator

with get_accelerator().device(0):
    model = BertForSequenceClassification.from_pretrained("bert-base-uncased")
    model.eval()

    # 用 input_shape 自动生成张量
    flops, macs, params = get_model_profile(
        model=model,
        input_shape=(4, 128),          # (batch_size, seq_len)
        print_profile=True,
        detailed=True,
        as_string=True,                # 输出成 "1.5G" 这种可读格式
        warm_up=10,
    )

    print(f"FLOPs: {flops}, MACs: {macs}, Params: {params}")

方式 B：FlopsProfiler 类（训练流程中）

from deepspeed.profiling.flops_profiler import FlopsProfiler

prof = FlopsProfiler(model)
profile_step = 5

for step, batch in enumerate(dataloader):
    if step == profile_step:
        prof.start_profile()

    loss = model(**batch)
    loss.backward()
    optimizer.step()

    if step == profile_step:
        prof.stop_profile()
        flops = prof.get_total_flops()
        macs = prof.get_total_macs()
        params = prof.get_total_params()
        prof.print_model_profile(profile_step=profile_step)
        prof.end_profile()            # 注意：必须在 get_total_flops 之后调用

      关键细节：
      end_profile() 必须在 get_total_flops() / get_total_macs() / print_model_profile() 之后调用。
自定义算子、融合 kernel、共享权重容易统计不准。
get_model_profile 返回的 flops 已经是 FLOPs；如果你只看到 MACs，需要手动 ×2。

    

3. PyTorch FlopCounterMode

PyTorch 2.x 内置了 FlopCounterMode，不需要额外安装。它通过拦截 ATen 算子来估算前向 FLOPs，适合快速验证。

import torch
from transformers import AutoModel
from torch.utils.flop_counter import FlopCounterMode

model = AutoModel.from_pretrained("bert-base-uncased")
model.eval()

batch_size, seq_len = 4, 128
x = torch.randint(0, 1000, (batch_size, seq_len))

with FlopCounterMode(model, display=False) as fcm:
    out = model(x)
    total_flops = fcm.get_total_flops()
    print(f"Total FLOPs: {total_flops / 1e9:.2f} GFLOPs")

    # 也可以看每个算子层的贡献
    fcm.print_results()

FlopCounterMode 目前主要覆盖 aten::mm、aten::addmm、aten::convolution、aten::bmm 等常见算子。如果你的模型里用了自定义 CUDA kernel，它不会计入。

4. fvcore（Meta / Detectron2 生态）

fvcore 是 Meta 早期推出的计算复杂度工具，Detectron2 官方用它统计 MACs。它通过 FlopCountAnalysis 对模型做 hook 式前向遍历，并按算子类型查表累计。

import torch
from transformers import AutoModel
from fvcore.nn import FlopCountAnalysis, parameter_count

model = AutoModel.from_pretrained("bert-base-uncased")
model.eval()

batch_size, seq_len = 4, 128
x = torch.randint(0, 1000, (batch_size, seq_len))

# fvcore 默认返回的是 MACs，但变量名常写作 flops
analysis = FlopCountAnalysis(model, x)
macs = analysis.total()          # 注意：这是 MACs
params = parameter_count(model)

print(f"MACs: {macs / 1e9:.2f} GMACs")
print(f"FLOPs ≈ {2 * macs / 1e9:.2f} GFLOPs")
print(f"Params: {params / 1e6:.2f} M")

# 也可以按模块/算子展开
analysis.print_model_graph_analytics()

      关键细节：
      fvcore 的 FlopCountAnalysis.total() 返回的是 MACs，但文档里常把它叫做 FLOPs。和下游同事对齐时，务必确认口径。
需要手动构造 dummy 输入；对 HuggingFace 模型通常直接用 input_ids 张量即可。
对自定义算子、融合 kernel、共享权重同样可能统计不准。

    

5. ptflops（轻量 CNN 估算）

ptflops 是一个更轻量的 hook-based profiler，API 极简。它同样返回 MACs，适合快速估算标准 CNN。

import torch
from transformers import AutoModel
from ptflops import get_model_complexity_info

model = AutoModel.from_pretrained("bert-base-uncased")
model.eval()

# input_res 接受 tuple/list，表示单个样本的维度
macs, params = get_model_complexity_info(
    model,
    input_res=(128,),           # 单条序列长度；batch 维度由工具自动加
    as_strings=True,
    print_per_layer_stat=True,
)

print(f"MACs: {macs}, Params: {params}")

      关键细节：
      ptflops 返回的 "FLOPs" 其实是 MACs，需要 ×2 换算。
input_res 是单样本维度，不含 batch；对 Transformer 输入写 (128,) 即可。
对 LLM / 复杂 attention 结构支持较弱，更适合 ResNet、EfficientNet 等 CV 模型。

    

6. calflops（推荐用于 HuggingFace）

calflops 严格区分 FLOPs 与 MACs，并且支持 HuggingFace tokenizer 自动构造输入，对 LLM 最友好。

from transformers import AutoModel, AutoTokenizer
from calflops import calculate_flops

model = AutoModel.from_pretrained("bert-base-uncased")
tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")

flops, macs, params = calculate_flops(
    model=model,
    input_shape=(4, 128),
    transformer_tokenizer=tokenizer,
    print_results=True,
    output_as_string=True,
)

print(f"FLOPs: {flops}, MACs: {macs}, Params: {params}")

如果你不想下载模型权重，只想快速估算某个 HuggingFace 模型的 FLOPs，可以用 calculate_flops_hf：

from calflops import calculate_flops_hf

flops, macs, params = calculate_flops_hf(
    model_name="bert-base-uncased",
    input_shape=(4, 128),
    output_as_string=True,
)

7. 五种工具输出对比

维度	DeepSpeed	FlopCounterMode	fvcore	ptflops	calflops
依赖安装	deepspeed	PyTorch ≥ 2.0	fvcore	ptflops	calflops + transformers
计数层级	module hook + 公式表	ATen 算子拦截	module hook + 公式表	module hook + 公式表	module hook + 公式表
是否覆盖 backward	✅ 可以	❌ 仅 forward	❌ 仅 forward	❌ 仅 forward	⚠️ 部分支持 / 需估算
HuggingFace tokenizer	手动构造 kwargs	手动构造输入	手动构造输入	手动构造输入	✅ 自动
输出单位真相	可选 MACs / FLOPs	FLOPs	MACs（常标成 FLOPs）	MACs（常标成 FLOPs）	严格区分 FLOPs/MACs
输出详细度	模块级占比、latency	算子级占比	模块/算子级	模块级占比	模块级占比、百分比
常见坑	自定义 op / 共享权重不准	仅常见 ATen 算子；自定义 CUDA kernel 不计	MACs/FLOPs 标签混淆	LLM 支持弱；单样本输入	较新，某些特殊模块可能漏算

8. 动手练习：用五种工具算同一个 BERT

下面是完整的对比流程，建议你在自己的环境里跑一遍：

import torch
from transformers import AutoModel, AutoTokenizer
from deepspeed.profiling.flops_profiler import get_model_profile
from deepspeed.accelerator import get_accelerator
from torch.utils.flop_counter import FlopCounterMode
from fvcore.nn import FlopCountAnalysis, parameter_count
from ptflops import get_model_complexity_info
from calflops import calculate_flops

model_name = "bert-base-uncased"
batch_size, seq_len = 4, 128

# 1. DeepSpeed
with get_accelerator().device(0):
    model_ds = AutoModel.from_pretrained(model_name)
    flops_ds, macs_ds, params_ds = get_model_profile(
        model_ds, input_shape=(batch_size, seq_len),
        print_profile=False, detailed=False, as_string=False
    )

# 2. PyTorch FlopCounterMode
model_pt = AutoModel.from_pretrained(model_name)
model_pt.eval()
x = torch.randint(0, 1000, (batch_size, seq_len))
with FlopCounterMode(model_pt, display=False) as fcm:
    _ = model_pt(x)
    flops_pt = fcm.get_total_flops()

# 3. fvcore（注意：返回的是 MACs）
model_fv = AutoModel.from_pretrained(model_name)
model_fv.eval()
macs_fv = FlopCountAnalysis(model_fv, x).total()
params_fv = parameter_count(model_fv)

# 4. ptflops（注意：返回的是 MACs，input_res 不含 batch）
model_pf = AutoModel.from_pretrained(model_name)
model_pf.eval()
macs_pf, params_pf = get_model_complexity_info(
    model_pf, input_res=(seq_len,), as_strings=False, print_per_layer_stat=False
)

# 5. calflops
model_cf = AutoModel.from_pretrained(model_name)
tokenizer = AutoTokenizer.from_pretrained(model_name)
flops_cf, macs_cf, params_cf = calculate_flops(
    model_cf, input_shape=(batch_size, seq_len),
    transformer_tokenizer=tokenizer, print_results=False
)

print(f"DeepSpeed:      {flops_ds/1e9:.2f} GFLOPs")
print(f"FlopCounterMode:{flops_pt/1e9:.2f} GFLOPs")
print(f"fvcore:         {2*macs_fv/1e9:.2f} GFLOPs (MACs {macs_fv/1e9:.2f} GMACs)")
print(f"ptflops:        {2*macs_pf/1e9:.2f} GFLOPs (MACs {macs_pf/1e9:.2f} GMACs)")
print(f"calflops:       {flops_cf/1e9:.2f} GFLOPs")

      思考题： 如果五个结果有偏差，可能是哪些原因造成的？（提示：是否计入 embedding、是否计入 LayerNorm/Bias、是否把 MACs 当 FLOPs 显示、输入构造是否一致、是否自动加了 batch 维度。）
    

小测验

9. 推荐延伸阅读

DeepSpeed Flops Profiler 官方教程：包含完整的 API 说明与 print_profile 输出示例。
PyTorch FlopCounterMode 教程：官方对 ATen 算子拦截和结果解读的说明。
fvcore GitHub：Meta 官方仓库，README 里有 FlopCountAnalysis 与 Detectron2 集成示例。
ptflops GitHub：轻量 CNN FLOPs 估算工具，README 包含常见 CV 模型示例。
calflops GitHub：README 里有大量 LLM 和 CV 模型的使用示例。

10. 接下来可以问的问题

我的 DeepSpeed profiler 输出和 calflops 差了一倍，怎么排查？
训练时的 backward FLOPs 怎么用代码统计？
想给 LLaMA / GPT 类模型做一个 MFU 监控脚本，该怎么做？
FlopCounterMode 能统计 FlashAttention 的 FLOPs 吗？