总目录 › ColossalAI › 第 10 章
Chapter 10

性能优化：fused kernel / chunked communication / grad ckpt

📌 commit 4f9953be335e 把吞吐从 30% MFU 推到 55%+ 的工程开关一览

跑通了不代表跑得快。本章把 ColossalAI 里影响吞吐的 8 个核心开关讲清楚， 让你能按需选用、知道每一项的代价。

10.1性能基线：MFU 怎么算

MFU（Model FLOPs Utilization）：

$$\text{MFU} = \frac{\text{实际 FLOPs/s}}{\text{硬件 peak FLOPs/s}} = \frac{6 \cdot N \cdot B \cdot T}{\text{step\_time} \cdot \text{TFLOPs\_peak}}$$

其中 $N$ 是模型参数量、$B$ 是 batch tokens（bs × seqlen）、$T$ = 1 step。 A100 80GB 的 bf16 peak ≈ 312 TFLOPS；H100 ≈ 989 TFLOPS。 LLaMA-7B 训练在 A100 上 MFU > 50% 算合格、> 55% 优秀。

10.2开关 1：Flash Attention

原版 attention 是 O(n²) 显存 + O(n²) 计算；Flash Attention 用 tiling 把显存降到 O(n)，速度也快 2-4 倍。

plugin = HybridParallelPlugin(
    ...,
    enable_flash_attention=True,   # ColossalAI 自动用 flash-attn-2 或 SDPA
)

实现：colossalai/shardformer/modeling/llama.py::llama_attention_forward 会优先调 flash_attn_func。 长序列收益最大（>= 4K 时 1.8× 加速）。

10.3开关 2：Fused LayerNorm / RMSNorm

原版 LayerNorm 在 PyTorch eager 模式下要发多个 kernel；apex 的 fused 版只发一个，速度快 2×。

plugin = HybridParallelPlugin(
    ...,
    enable_fused_normalization=True,
)

实现：colossalai/shardformer/layer/normalization.py::FusedRMSNorm 直接调 apex C++ kernel。 LLaMA 有 64 个 RMSNorm，开了这个开关 step 时间能减 5%-8%。

10.4开关 3：Fused QKV 投影

原版 Attention 的 q_proj / k_proj / v_proj 是三个独立 Linear，发三个 GEMM kernel。Fused QKV 把它们合并成一次 GEMM。

plugin = HybridParallelPlugin(
    ...,
    enable_jit_fused=True,
)

实现：colossalai/shardformer/layer/qkv_fused_linear.py。 注意有些模型（特别是 GQA 的 K/V head 数和 Q 不一样）不能直接 fuse，需要手动指定。

10.5开关 4：Gradient Checkpointing

用计算换显存：每个 transformer block 不保存 activation，反向时重算一遍。显存能减 60-80%，速度损失 20-30%。

from colossalai.shardformer.shard import ShardConfig, GradientCheckpointConfig

gc_config = GradientCheckpointConfig(
    gradient_checkpointing_ratio=0.5,    # 一半的 layer 开 ckpt
    # 或 num_ckpt_layers_per_stage=[2, 4, 4, 2]  显式指定
)

plugin = HybridParallelPlugin(
    ...,
    gradient_checkpoint_config=gc_config,
)

关键参数 gradient_checkpointing_ratio：

10.6开关 5：Chunked Sequence All-Reduce

开 TP 后每层 attention/MLP 都有 all-reduce，short seq 时通信比 = 计算/通信 太低（小通信 latency 主导）。 解决办法：把 sequence 切 chunk，依次执行 + overlap 通信和计算。

# 在 enable_sequence_parallelism + split_gather 模式下自动启用
plugin = HybridParallelPlugin(
    ...,
    enable_sequence_parallelism=True,
    sequence_parallelism_mode="split_gather",
)

10.7开关 6：BF16 / FP8

bf16 在 A100/H100 上 throughput 是 fp32 的 2×，比 fp16 数值范围更稳定（不会溢出）。 ColossalAI 全面支持 bf16；fp8 在 ShardFormer 里部分支持（FP8Linear），需要 H100。

# bf16（推荐）
plugin = HybridParallelPlugin(precision="bf16")

# fp8（实验性，H100 only）
plugin = HybridParallelPlugin(precision="fp8")

10.8开关 7：Zero-Bubble Pipeline

1F1B 调度下的气泡比 $(p-1)/m$，pp=4、micro_batch=8 时 ~37.5% 时间是空泡。 Zero-Bubble pipeline（2024 论文）把 backward 拆成 B_x 和 B_w 两步，能把气泡降到接近 0。

plugin = HybridParallelPlugin(
    pp_size=4,
    pp_style="zbv",        # zero-bubble v-shape
    num_model_chunks=2,
    num_microbatches=8,
)

实现在 colossalai/pipeline/schedule/zero_bubble_pp.py，需要配 colossalai/pipeline/weight_grad_store.py 缓存 weight gradient。

10.9开关 8：Async Gradient All-Reduce

DP 之间的梯度 all-reduce 默认在 backward 完后做。打开 async 后，backward 进行中就开始 all-reduce 已完成的层，通信和计算 overlap。

plugin = HybridParallelPlugin(
    ...,
    overlap_p2p=True,           # PP 的 p2p 也 overlap
)

# Gemini 路径
plugin = GeminiPlugin(
    ...,
    enable_async_reduce=True,
)

10.10性能调优清单

按这个顺序检查，能命中 90% 的性能问题：

#检查项影响
1bf16 / fp16 是否启用 +50–100% 吞吐
2flash_attention 是否启用 +30–80%（长序列）
3fused norm 是否启用 +5–10%
4TP size 是否过大（> 8） 通信瓶颈
5是否在显存够时还开了 Gemini 用 HybridParallel 更快
6micro_batch 数 vs PP stage 数比例 气泡比
7gradient checkpoint ratio 是否过高 显存够就关掉
8data loader 是否成为瓶颈 看 GPU SM 利用率

10.11profile：用 PyTorch Profiler

from torch.profiler import profile, ProfilerActivity, schedule, tensorboard_trace_handler

with profile(
    activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
    schedule=schedule(wait=1, warmup=2, active=3, repeat=1),
    on_trace_ready=tensorboard_trace_handler("./tb_log"),
) as prof:
    for i, batch in enumerate(dl):
        loss = model(**batch).loss
        booster.backward(loss, optimizer)
        optimizer.step()
        prof.step()
        if i >= 6: break

# 在另一终端 tensorboard --logdir tb_log

看 timeline，重点观察：

10.12这章你需要带走的