Chapter 04
核心抽象：ParallelContext 与 Process Group
nanotron 把 dp / pp / tp 三维 process group 装进一个 ParallelContext 对象。本章拆它的源码（不到 100 行）。
4.1ParallelContext 是什么
看 src/nanotron/parallel/context.py：
class ParallelContext:
def __init__(self, data_parallel_size, pipeline_parallel_size, tensor_parallel_size):
# 1) 全局初始化
if not dist.is_initialized():
dist.init_process_group(backend="nccl")
self.world_size = dist.get_world_size()
assert dp * pp * tp == self.world_size
# 2) 构造三维 mesh
self.dp_pg = self._build_pg("dp", ...)
self.pp_pg = self._build_pg("pp", ...)
self.tp_pg = self._build_pg("tp", ...)
# 3) 当前 rank 在每维的位置
self.dp_rank = ...
self.pp_rank = ...
self.tp_rank = ...4.2三维 process group 怎么排
nanotron 默认按 (pp, dp, tp) 顺序排：
rank = pp_rank * (dp_size * tp_size) + dp_rank * tp_size + tp_rank这意味着同一个 tp group 里的卡 rank 相邻，让 TP 的 all-reduce 在 NVLink 域内最快。
4.316 卡示例
假设 16 GPU，dp=2, pp=2, tp=4：
flowchart TB
subgraph PP0["pp_rank=0 (rank 0-7)"]
subgraph DP0a["dp_rank=0"]
G0(0) G1(1) G2(2) G3(3)
end
subgraph DP1a["dp_rank=1"]
G4(4) G5(5) G6(6) G7(7)
end
end
subgraph PP1["pp_rank=1 (rank 8-15)"]
subgraph DP0b["dp_rank=0"]
G8(8) G9(9) G10(10) G11(11)
end
subgraph DP1b["dp_rank=1"]
G12(12) G13(13) G14(14) G15(15)
end
end
group 划分：
- TP group：每行 4 卡（如 {0,1,2,3}）；
- DP group：相同 PP / TP 的两组 {0,4}, {1,5}, ...；
- PP group：相同 DP / TP 的两组 {0,8}, {1,9}, ...。
4.4从 ParallelContext 拿 group
ctx = ParallelContext(dp=2, pp=2, tp=4)
# TP all-reduce 用 tp_pg
dist.all_reduce(tensor, group=ctx.tp_pg)
# DP all-reduce 用 dp_pg
dist.all_reduce(grad, group=ctx.dp_pg)
# PP send/recv 用 pp_pg
dist.send(tensor, dst=next_pp_rank, group=ctx.pp_pg)4.5NanotronParameter
普通 nn.Parameter 没有"我在哪个 group 上切了"信息。nanotron 自己包了 NanotronParameter（看 parallel/parameters.py）：
class NanotronParameter(nn.Parameter):
def __new__(cls, tensor, ...):
...
# 关键 metadata
self.is_tied = False # tied weight（如 embed/lm_head）
self.is_sharded = True/False
self.shard_metadata = ShardMetadata(group, dim, ...)
self.is_pipeline_metadata = ...有了这层包装，optimizer / checkpoint 都能知道每个参数切到哪个 group 的哪个 rank，从而正确处理 reduce / save。
4.6跟 PyTorch DeviceMesh 比
| nanotron ParallelContext | PyTorch DeviceMesh | |
|---|---|---|
| 实现 | 自家 ~100 行 | PyTorch 内核 |
| 跟 DTensor 配合 | 不 | 原生 |
| 清晰度 | ★★★★★（小巧易读） | ★★★★（更通用） |
| 跟 PT 主版本绑定 | 不 | 是 |
4.7ParallelContext 源码定位
| 符号 | 位置 | 用途 |
|---|---|---|
ParallelContext.__init__ | src/nanotron/parallel/context.py:12-58 | 校验 tp×cp×dp×pp == world_size，建所有 process group |
_init_parallel_groups() | parallel/context.py:59 | 真正调 torch.distributed.new_group |
NanotronParameter | parallel/parameters.py | 带 sharding metadata 的 Parameter |
TensorPointer | parallel/pipeline_parallel/ | PP 中跨 stage 引用的占位符 |
| 分布式辅助 | distributed.py:1-50 | new_group / reduce_scatter_tensor 等 torch.distributed 薄封装 |
5 个并行维度（tp/pp/dp/cp/ep）都在 ParallelContext 一个对象上。后面所有 layer / optimizer / dataloader 拿到这一个对象就够了。
4.8这章你需要带走的
- ParallelContext = 3 维 process group 容器（dp / pp / tp）；
- rank 排列按 (pp, dp, tp)，让 TP 在 NVLink 域内最快；
- NanotronParameter 给参数加 shard / pipeline metadata；
- 整个抽象只 ~100 行代码，是学习分布式的好材料。