核心抽象：DTensor / DeviceMesh / ParallelStyle

TorchTitan 不引入新概念，而是依赖 PyTorch 2.x 引入的三个原生抽象：DeviceMesh（拓扑）、DTensor（分布式张量）、ParallelStyle（切分策略）。理解这三个就懂了 TorchTitan 的全部分布式。

4.1DeviceMesh

DeviceMesh 是 PyTorch 对"多 GPU 拓扑"的抽象。可以是 1D / 2D / 3D 网格。

from torch.distributed.device_mesh import init_device_mesh

# 1D mesh：8 张 GPU 作 DP
mesh_1d = init_device_mesh("cuda", (8,))
mesh_1d["dp"]

# 2D mesh：8 GPU = 4 DP × 2 TP
mesh_2d = init_device_mesh("cuda", (4, 2), mesh_dim_names=("dp", "tp"))
mesh_2d["dp"]    # DP 通信组
mesh_2d["tp"]    # TP 通信组

# 3D mesh：DP × PP × TP
mesh_3d = init_device_mesh("cuda", (2, 2, 2), mesh_dim_names=("dp", "pp", "tp"))

TorchTitan 用 4D mesh：(dp_replicate, dp_shard, cp, tp, pp)。

4.2DTensor

DTensor = "切到 mesh 上的 tensor"。每个 DTensor 包含：

• local tensor：当前 rank 持有的那份；
• placement：在 mesh 各维上的切法（Shard / Replicate / Partial）。

from torch.distributed.tensor import DTensor, Replicate, Shard, Partial

# Replicate：每张卡都有完整副本
dt_rep = DTensor.from_local(tensor, mesh, [Replicate()])

# Shard(0)：沿第 0 维切，每张卡有 1/N
dt_shard = DTensor.from_local(tensor, mesh, [Shard(0)])

# Partial(sum)：每张卡持部分和，待 reduce
dt_partial = DTensor.from_local(tensor, mesh, [Partial("sum")])

4.3三种 placement 物理含义

4.4DTensor 算子自动通信

关键特性：DTensor 上做 op 会自动插通信。例：

# A: Shard(0) on TP mesh
# B: Replicate on TP mesh
y = A @ B
# DTensor 自动判断：
#   - A 已经按 row 切了
#   - 结果应该按 row 切（仍是 Shard(0)）
#   - 不需要通信
# 但如果 A: Shard(1), B: Replicate
# y = A @ B 会自动插入 all-reduce

看 torch/distributed/tensor/_ops/，几乎所有常用 op 都有 DTensor 实现。

4.5ParallelStyle：把 nn.Linear 切成 DTensor

PyTorch 提供几个 ParallelStyle，用 parallelize_module() 应用：

from torch.distributed.tensor.parallel import (
    parallelize_module,
    ColwiseParallel,        # 列切
    RowwiseParallel,        # 行切
    SequenceParallel,       # 序列切
    PrepareModuleInput,
    PrepareModuleOutput,
)

# 给 attention 加 TP
attn_plan = {
    "wq": ColwiseParallel(),
    "wk": ColwiseParallel(),
    "wv": ColwiseParallel(),
    "wo": RowwiseParallel(),
}
parallelize_module(model.attention, mesh["tp"], attn_plan)

这等价于 Megatron / ColossalAI 的"按 layer 替换 ParallelLinear"，但不用换 Linear 类——只是给原 nn.Linear 的 weight 重新 layout 成 DTensor。

4.6TorchTitan 怎么用三件套

看 torchtitan/distributed/parallelize_llama.py（简化）：

def parallelize_llama(model, mesh, parallel_dims, job_config):
    # 1) 应用 TP
    if parallel_dims.tp_enabled:
        apply_tp(model, mesh["tp"])

    # 2) 应用 PP
    if parallel_dims.pp_enabled:
        model = apply_pp(model, mesh["pp"])

    # 3) 应用 CP
    if parallel_dims.cp_enabled:
        apply_cp(model, mesh["cp"])

    # 4) 应用 FSDP2
    if parallel_dims.dp_shard_enabled:
        apply_fsdp(model, mesh["dp_shard"])

    return model

整个流程没有自定义并行类，全是 PyTorch 原生 API + DTensor。

4.7跟其他框架对比

4.8这章你需要带走的

• TorchTitan 全部分布式建立在 DeviceMesh + DTensor + ParallelStyle 三件套上；
• DeviceMesh 是 N 维拓扑，DTensor 是切到 mesh 上的张量，ParallelStyle 决定切法；
• DTensor op 自动插通信，这是 TorchTitan 不用造新轮子的根本原因；
• 三件套全在 torch.distributed.tensor，PyTorch 2.4+ 起稳定。