分布式训练与FSDP2 | 摩尔线程 GPU 学习中心

一、为什么需要分布式训练？

当模型参数量突破百亿，单张 GPU 的显存和算力都无法承载完整训练流程时，分布式训练就成了必选项。主流分布式策略分两类：

数据并行（Data Parallel）：每张 GPU 上都有完整模型副本，各自计算不同 batch 的梯度后再同步。代表：DDP。
分片并行（Sharded / Model Parallel）：把模型参数、梯度、优化器状态切分到多张 GPU 上，显著降低单卡显存占用。代表：ZeRO、FSDP、FSDP2、Context Parallel。

在摩尔线程生态中，torch_musa v2.9.0（2026-03-17 发布）完成了对 FSDP2 和 Context Parallel（Ulysses） 的官方适配，让千亿参数级模型的国产 GPU 训练成为现实。本课聚焦于如何在 MTT S4000 / S5000 集群上跑通 FSDP2 分布式训练。

💡

FSDP2 vs FSDP（传统版）

FSDP2 是 PyTorch 2.4+ 引入的重构版本，采用 DTensor 为核心抽象，支持更细粒度的分片（per-parameter sharding）、更高效的通信调度，并原生兼容 torch.compile。FSDP2 的 API 以 fully_shard() 函数式接口为主，替代了旧版的 FullyShardedDataParallel 类。

二、环境准备

要跑通 FSDP2 on MUSA，至少需要：

MUSA SDK >= 4.3.2（torch_musa v2.9.0 的硬性要求）

torch_musa >= 2.9.0（对应 PyTorch 2.7.x）

MCCL（摩尔线程的 NCCL 等价物）已正确安装，多机场景下配置好主机互联

至少 2 张 MTT GPU（推荐 S4000 / S5000，显存 ≥ 48GB）

bash

# 确认版本
python -c "import torch, torch_musa; print(torch.__version__, torch_musa.__version__)"
# 期望输出类似：2.7.1 2.9.0

# 确认可见 GPU
python -c "import torch; print(torch.musa.device_count())"

三、先跑通 DDP：单机多卡的基线

在上 FSDP2 之前，先用 DDP 跑通最简单的数据并行。这让我们把 分布式启动 / 通信 / 训练循环 与 分片细节 解耦。

python

# ddp_train.py
import os, torch, torch.nn as nn, torch.distributed as dist
import torch_musa  # 激活 MUSA 后端
from torch.nn.parallel import DistributedDataParallel as DDP

def setup():
    dist.init_process_group(backend="mccl")  # 摩尔线程的 MCCL
    rank = dist.get_rank()
    torch.musa.set_device(rank)
    return rank

def cleanup():
    dist.destroy_process_group()

class ToyModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(2048, 4096), nn.GELU(),
            nn.Linear(4096, 4096), nn.GELU(),
            nn.Linear(4096, 1024),
        )
    def forward(self, x): return self.net(x)

def main():
    rank = setup()
    device = torch.device(f"musa:{rank}")
    model = ToyModel().to(device)
    ddp_model = DDP(model, device_ids=[rank])
    opt = torch.optim.AdamW(ddp_model.parameters(), lr=1e-4)

    for step in range(100):
        x = torch.randn(32, 2048, device=device)
        y = torch.randn(32, 1024, device=device)
        out = ddp_model(x)
        loss = ((out - y) ** 2).mean()
        opt.zero_grad(); loss.backward(); opt.step()
        if rank == 0 and step % 10 == 0:
            print(f"step {step}: loss={loss.item():.4f}")
    cleanup()

if __name__ == "__main__":
    main()

bash

# 用 torchrun 启动 2 卡 DDP
torchrun --standalone --nproc_per_node=2 ddp_train.py

四、升级到 FSDP2：函数式分片 API

FSDP2 的入口是 torch.distributed.fsdp.fully_shard。它不再需要把整个模型包一层，而是直接对任意 nn.Module 或其子模块调用：参数会被自动切分成 DTensor，训练时自动按需 all-gather、反向时自动 reduce-scatter。

python

# fsdp2_train.py
import os, torch, torch.nn as nn, torch.distributed as dist
import torch_musa
from torch.distributed.fsdp import fully_shard, MixedPrecisionPolicy
from torch.distributed.device_mesh import init_device_mesh

def setup():
    dist.init_process_group(backend="mccl")
    rank = dist.get_rank()
    torch.musa.set_device(rank)
    return rank

class Block(nn.Module):
    def __init__(self, dim=4096):
        super().__init__()
        self.attn = nn.Linear(dim, dim)
        self.mlp = nn.Sequential(nn.Linear(dim, 4*dim), nn.GELU(), nn.Linear(4*dim, dim))
        self.ln1, self.ln2 = nn.LayerNorm(dim), nn.LayerNorm(dim)
    def forward(self, x):
        x = x + self.attn(self.ln1(x))
        x = x + self.mlp(self.ln2(x))
        return x

class TinyLLM(nn.Module):
    def __init__(self, depth=8, dim=4096):
        super().__init__()
        self.blocks = nn.ModuleList([Block(dim) for _ in range(depth)])
        self.head = nn.Linear(dim, 32000)
    def forward(self, x):
        for b in self.blocks:
            x = b(x)
        return self.head(x)

def main():
    rank = setup()
    world = dist.get_world_size()
    device = torch.device(f"musa:{rank}")

    # 1D mesh：纯数据分片
    mesh = init_device_mesh("musa", (world,))

    model = TinyLLM(depth=8, dim=4096).to(device)

    # 混合精度策略：参数 bf16、梯度 bf16、reduce 用 fp32
    mp = MixedPrecisionPolicy(
        param_dtype=torch.bfloat16,
        reduce_dtype=torch.float32,
    )

    # 每个 Block 单独 shard → 细粒度通信调度
    for block in model.blocks:
        fully_shard(block, mesh=mesh, mp_policy=mp)
    fully_shard(model, mesh=mesh, mp_policy=mp)  # 根模块也 shard

    opt = torch.optim.AdamW(model.parameters(), lr=1e-4)

    for step in range(50):
        x = torch.randn(4, 1024, 4096, device=device)
        out = model(x)
        loss = out.float().mean()
        opt.zero_grad(); loss.backward(); opt.step()
        if rank == 0 and step % 5 == 0:
            print(f"step {step}: loss={loss.item():.4f}")

    dist.destroy_process_group()

if __name__ == "__main__":
    main()

✅

为什么逐 Block 调用 fully_shard？

这样每个 Block 的 all-gather / reduce-scatter 可以与下一个 Block 的前向计算重叠，通信效率远高于一次性 shard 整棵模型。这是 FSDP2 最重要的性能开关之一。

五、Context Parallel（Ulysses）：序列维度切分

当序列长度 seq_len 达到 64K、128K 甚至更长时，即便参数已经分片，激活值也会把显存撑爆。此时需要把序列维度切分到多卡上，这就是 Context Parallel（CP），torch_musa v2.9.0 官方适配了 Ulysses 风格 CP，它结合 all-to-all 通信让每张卡只处理一段序列。

python

# 启用 2D mesh：dp × cp
from torch.distributed.device_mesh import init_device_mesh

# 例如 8 卡：4 路数据并行 × 2 路 context parallel
mesh = init_device_mesh("musa", (4, 2), mesh_dim_names=("dp", "cp"))
dp_mesh = mesh["dp"]
cp_mesh = mesh["cp"]

# 参数沿 dp 维分片
for block in model.blocks:
    fully_shard(block, mesh=dp_mesh, mp_policy=mp)

# 注意力层在 cp_mesh 上做 Ulysses all-to-all
# 具体实现参见 torch_musa 示例仓库的 ulysses_attention.py

Ulysses 的本质是：Q/K/V 在 head 维度切分，attention 内部做 all-to-all 把 head 维切换成 seq 维。计算结束后再 all-to-all 切回去。对用户代码侵入极小，但需要按 torch_musa 的示例模板改造 attention 实现。

六、torch.compile 的 reduce-overhead 模式

torch_musa v2.9.0 还启用了 torch.compile 的 reduce-overhead 模式，它通过 CUDA Graph 捕获训练 step，把 Python 端调度开销降到几乎为 0，对小 batch / 短 step 训练效果尤其明显。

python

compiled_model = torch.compile(
    model,
    mode="reduce-overhead",
    fullgraph=False,  # FSDP2 场景常用 False
)

for step in range(50):
    x = torch.randn(4, 1024, 4096, device=device)
    out = compiled_model(x)    # 首次会触发编译 / graph 捕获
    loss = out.float().mean()
    opt.zero_grad(); loss.backward(); opt.step()

⚠️

首次 step 变慢属于正常

reduce-overhead 会在第一次前向时捕获图，耗时显著高于后续 step。测评吞吐时请丢掉 warmup 的前 10 步再取平均。

七、多机启动脚本

以 2 机 × 8 卡 = 16 卡为例：

bash

# 主节点（node rank 0，假设 IP = 10.0.0.1）
torchrun \
  --nnodes=2 --node_rank=0 --nproc_per_node=8 \
  --rdzv_backend=c10d --rdzv_endpoint=10.0.0.1:29500 \
  fsdp2_train.py

# 从节点（node rank 1）
torchrun \
  --nnodes=2 --node_rank=1 --nproc_per_node=8 \
  --rdzv_backend=c10d --rdzv_endpoint=10.0.0.1:29500 \
  fsdp2_train.py

八、性能调优清单

逐层 shard：对每个 transformer block 单独调用 fully_shard，让通信与计算流水重叠。
混合精度：参数 / 梯度用 bf16，但 reduce 保留 fp32，兼顾吞吐与数值稳定。
activation checkpointing：对长 seq 开启，节省最多 50% 激活显存。
关注 MCCL 带宽：用 pymtml 查询 NIC / 总线利用率，避免通信瓶颈。
dp × cp 维度拆分：长 seq 场景，把总 GPU 切成 dp×cp 网格。

✅

本课小结

torch_musa v2.9.0 把 PyTorch 最新的分布式能力带到了摩尔线程 GPU 上：FSDP2 的函数式 fully_shard 替代了旧 API；Context Parallel（Ulysses） 让长序列训练可行；torch.compile reduce-overhead 进一步压榨吞吐。结合 MCCL 后端与合理的 mesh 划分，你可以在国产 GPU 集群上训练百亿参数级模型。参考：torch_musa v2.9.0 Release。