总目录›TorchTitan›第 9 章
Chapter 09

性能优化：compile / async TP / float8

📌 commit af33f76PyTorch 新特性的展览馆

TorchTitan 是 PyTorch 新特性最快落地的地方。开几个开关就能从 ~45% MFU 拉到 ~65%。本章把主要开关讲清楚。

9.1开关清单（按重要性）

开关典型收益条件
bf16 / fp16+50-100%所有 NVIDIA GPU
activation_checkpoint=selective+5% 显存—
torch.compile+10-20%torch 2.4+
Float8（fp8 训练）+25-40%H100/H200 + torchao
async TP+5-15%TP > 1
HSDP (replicate × shard)+10%（跨机）多机
Pipeline 1F1B显存大幅省大模型

9.2torch.compile

[compile]
enable = true
model_backend = "inductor"

TorchTitan 对每个 TransformerBlock 单独 compile（而非整个 model），收益更大、warm-up 更快。代码在 torchtitan/distributed/parallelize_llama.py::apply_compile。

注意：

9.3Float8（fp8 训练）

用 torchao 的 fp8 GEMM 替换 bf16 GEMM。仅 H100/H200/B200 上有 fp8 硬件。

[float8]
enable_float8_linear = true
recipe_name = "rowwise"           # rowwise / tensorwise / mxfp8
recipe说明
tensorwiseper-tensor 缩放，最稳，最慢
rowwiseper-row 缩放，常用
mxfp8MX FP8（H100 新格式），最快

实现：torchao.float8.convert_to_float8_training()。把所有 nn.Linear 换成 Float8Linear，weight / activation 在 forward 之前现 quant 到 fp8 算 GEMM，结果回 bf16。

9.4async TP

普通 TP 在 attention/MLP 出口要 all-reduce，等通信完才进下一层。Async TP 把 GEMM 切成微块，每算完一块就启动该块的 reduce-scatter，让通信和计算 overlap。

[experimental]
enable_async_tensor_parallel = true

需要 PyTorch 2.5+ 且 TP > 1。典型增益 5-15%。

9.5activation checkpoint

[activation_checkpoint]
mode = "selective"              # none / selective / full
selective_ac_option = "op"      # op / N（数字 N 表示每 N 层选一个）

三档：

TorchTitan 推荐 selective + op，是 SAC（Selective Activation Checkpointing）方案。

9.6HSDP

多机时让跨机做 DDP-style replicate（每 N 步同步梯度），机内做 FSDP shard。

[training]
data_parallel_replicate_degree = 2     # 节点间
data_parallel_shard_degree = 8         # 节点内

典型场景：2 节点 × 8 GPU 训 70B，效果比纯 FSDP 跨机快 ~10%。

9.7性能总览（70B 配方）

Llama-3 70B、32×H100、bf16、seqlen=8K：

配方MFUtok/s/GPU
FSDP2 + TP=4 baseline~48%~1,800
+ selective ac~49%~1,850
+ torch.compile~56%~2,100
+ async TP~60%~2,250
+ Float8 rowwise~68%~2,550

9.8调优顺序

  1. 先 bf16 + selective ac，跑通；
  2. 开 torch.compile，跑稳；
  3. 显存够 → 加大 batch 或减 ac；不够 → 加 TP；
  4. 多机：开 HSDP；
  5. H100：开 float8 + async TP。

9.9性能优化的源码位置

每个优化开关都对应一个文件，迷路时直接 grep：

优化源码位置核心 API
Async TP torchtitan/distributed/tensor_parallel.py:102-116 maybe_enable_async_tp() 设置 torch._inductor.config._micro_pipeline_tp = True
Context Parallel torchtitan/distributed/context_parallel.py:35-80 apply_cp_to_forward()，分 FlexAttention / SDPA 两条路
FSDP2 torchtitan/distributed/fsdp.py fully_shard() + DataParallelMeshDims
Pipeline Parallel torchtitan/distributed/pipeline_parallel.py 1F1B + looped schedules（torch.distributed.pipelining）
Float8 训练 torchtitan/components/quantization/float8.py:24-80Float8LinearConverter 替换 Linear；recipes: rowwise / rowwise_with_gw_hp
parallelize_llama torchtitan/models/llama3/parallelize.py:41-100 TP→CP→async_tp→AC→compile→FSDP 顺序施加

9.10ParallelDims：所有并行的统一抽象

看 torchtitan/distributed/parallel_dims.py:26-80：

@dataclass
class ParallelDims:
    dp_replicate: int   # DDP replication
    dp_shard: int       # FSDP shard
    cp: int             # Context parallel
    tp: int             # Tensor parallel
    pp: int             # Pipeline parallel
    ep: int             # Expert parallel (MoE)
    world_size: int
    full_dtensor: bool

    @classmethod
    def from_config(cls, cfg): ...
    def get_mesh(self): ...          # → torch.distributed.DeviceMesh
    @property
    def tp_enabled(self): ...

整个框架的 "5D + EP" 并行就是这 8 个字段。world_size = dp_replicate × dp_shard × cp × tp × pp（EP 与 TP/DP 复用 rank）。配置约束：seq_len % (tp × 2cp) == 0 在 parallelize.py 里硬校验。

9.11Float8 配方选择

recipe精度 vs bf16速度收益用场景
rowwise 近无损 +15-20% 默认
rowwise_with_gw_hp 无损 +10-15% 梯度敏感任务
tensorwise（不在 TorchTitan） 有时损 1pp +20-25% 不推荐

注意：Float8 跳过 dim 不能被 16 整除的 Linear（看 float8.py:Float8LinearConverter 的 skip 列表），所以 vocab_size 不是 16 倍数时 lm_head 会回退 bf16。

9.12这章你需要带走的