LLM 通信原语详解（一）：从 Send/Recv 到 All-to-All

act = compute(x)                              # 本 stage 的输出激活
if rank > 0:
    dist.send(act, dst=rank - 1, tag=10)      # 激活发下游
g_act = torch.empty_like(act)
if rank < world_size - 1:
    dist.recv(g_act, src=rank + 1, tag=20)    # 梯度从下游回来

# Broadcast：root(src=0) 的值覆盖所有卡 —— 权重/超参下发
buf = load_checkpoint() if rank == 0 else torch.empty_like(shape)
dist.broadcast(buf, src=0)                    # 之后所有卡的 buf 完全一致

# Scatter：root 把完整张量按 rank 切片分发（scatter_list 只在 root 上有效）
mine = torch.empty(1, dtype=torch.float64)
scatter_list = [torch.tensor([float(i)], dtype=torch.float64)
                for i in range(world_size)] if rank == 0 else None
dist.scatter(mine, scatter_list=scatter_list, src=0)      # rank i 拿到第 i 片

# Gather：Scatter 的逆，各卡的份收拢到 root 拼接
gather_list = [torch.empty(1, dtype=torch.float64)
               for _ in range(world_size)] if rank == 0 else None
dist.gather(mine, gather_list=gather_list, dst=0)

# Reduce：归约（默认 sum）后只落 root —— loss 归集
total = mine.clone()
dist.reduce(total, dst=0)                     # 只有 rank 0 拿到总和

shard = param_shard()                              # 本卡持有 1/M
full = torch.empty(world_size * shard.numel(), dtype=shard.dtype)
dist.all_gather_into_tensor(full, shard)           # 拼回完整参数

grad_shard = torch.empty(grad_flat.numel() // world_size, dtype=grad_flat.dtype)
dist.reduce_scatter_tensor(grad_shard, grad_flat)  # 归约+分片一步到位
grad_shard /= world_size                           # 平均 —— FSDP 在这里就地更新

x = torch.arange(2 * world_size, dtype=torch.float64) + 100 * rank  # 各卡 partial

shard = torch.empty(2, dtype=torch.float64)
dist.reduce_scatter_tensor(shard, x)                # 半步 1：RS
via_two_step = torch.empty(2 * world_size, dtype=torch.float64)
dist.all_gather_into_tensor(via_two_step, shard)    # 半步 2：AG

direct = x.clone()
dist.all_reduce(direct)                             # 直接 AR
assert torch.equal(via_two_step, direct)            # AG(RS(x)) == AR(x) ✓

send = stack([块_for_rank_0, 块_for_rank_1, ...])   # 第 j 块发给 rank j
recv = torch.empty_like(send)
dist.all_to_all_single(recv, send)                   # 转置后: recv[i] 来自 rank i

# ① 元数据交换：发往各卡的 token 数（各卡不同 → 先同步计数）
counts = torch.bincount(router, minlength=world_size)   # 本卡发往各卡的量
in_counts = torch.empty(world_size, dtype=torch.long)
dist.all_to_all_single(in_counts, counts)               # 收到: 各卡发给我的量

# ② padding 到统一容量 cap = max(counts)，做等长 A2A（dispatch）
cap = int(counts.max() 全局同步)
send = torch.zeros(world_size, cap, d)                  # 按目标分组 + 补零
for j in range(world_size):
    send[j, :counts[j]] = tokens[router == j]
recv = torch.empty_like(send)
dist.all_to_all_single(recv, send)                      # 现在 recv[j] = 来自卡 j

# ③ 专家计算后原路转置回去（combine），按 in_counts 剥掉 padding 还原 token 顺序

tp_size = 2
tp_group = dist.new_group(ranks_for_this_tp_column)   # 如 [0,1] 和 [2,3]
dp_group = dist.new_group(ranks_for_this_dp_row)      # 如 [0,2] 和 [1,3]

x = torch.full((2,), float(rank))
t, d = x.clone(), x.clone()
dist.all_reduce(t, group=tp_group)   # 只和 TP 组内的邻居归约
dist.all_reduce(d, group=dp_group)   # 同一张量，DP 组内另一个结果

"""comm_primitives_demo.py — torch.distributed 通信原语可运行演示
CPU + gloo 即可运行(无需 GPU):
    torchrun --nproc-per-node=4 comm_primitives_demo.py
"""
import os
import torch
import torch.distributed as dist

R = int(os.environ["RANK"])
W = int(os.environ["WORLD_SIZE"])
dist.init_process_group("gloo")
torch.manual_seed(0)  # 所有 rank 同种子 → 同初始化

def log(msg):
    if R == 0:
        print(f"[R0] {msg}", flush=True)

def head(s):
    if R == 0:
        print(f"\n=== {s} ===", flush=True)

# ---------- 1. Send/Recv: 模拟 PP 相邻 stage ----------
head("1. Send/Recv — 流水线: 激活向右, 梯度向左")
act = torch.tensor([float(R + 1), 0.0])
recv_l, recv_r = torch.zeros(2), torch.zeros(2)
if R > 0:
    dist.recv(recv_l, src=R - 1, tag=10)                    # 收上游激活
if R < W - 1:
    dist.send(act, dst=R + 1, tag=10)                       # 激活发下游
if R < W - 1:
    dist.recv(recv_r, src=R + 1, tag=20)                    # 收下游梯度
if R > 0:
    dist.send(torch.tensor([1.0, 1.0]), dst=R - 1, tag=20)  # 梯度发上游
log(f"stage{R}: 上游激活={recv_l.tolist()} 下游梯度={recv_r.tolist()}")

# ---------- 2. root 族 ----------
head("2. Broadcast / Scatter / Gather / Reduce")
b = torch.arange(4, dtype=torch.float64) + R * 100
dist.broadcast(b, src=0)
assert torch.equal(b, torch.arange(4, dtype=torch.float64))

mine = torch.empty(1, dtype=torch.float64)
src = [torch.tensor([float(i)], dtype=torch.float64) for i in range(W)] if R == 0 else None
dist.scatter(mine, scatter_list=src, src=0)                 # rank i 拿到第 i 片
assert mine.item() == R

gathered = [torch.empty(1, dtype=torch.float64) for _ in range(W)] if R == 0 else None
dist.gather(mine, gather_list=gathered, dst=0)
if R == 0:
    assert [g.item() for g in gathered] == list(range(W))

red = mine.clone()
dist.reduce(red, dst=0)
if R == 0:
    assert red.item() == sum(range(W))
log("broadcast/scatter/gather/reduce 断言全部通过")

# ---------- 3. 全员族 + AR = RS + AG ----------
head("3. All-Gather / Reduce-Scatter / All-Reduce")
inp = torch.arange(2 * W, dtype=torch.float64) + 100 * R    # 每卡一份 partial
shard = torch.empty(2, dtype=torch.float64)
dist.reduce_scatter_tensor(shard, inp)                      # RS: 各持 1/W 归约分片
full = torch.empty(2 * W, dtype=torch.float64)
dist.all_gather_into_tensor(full, shard)                    # AG: 分片拼回完整
ref = inp.clone()
dist.all_reduce(ref)                                        # 直接 AR 做对照
assert torch.equal(full, ref)                               # AG(RS(x)) == AR(x)
log(f"RS 分片={shard.tolist()}, 恒等式 AG(RS(x)) == AR(x) 成立 ✓")

# ---------- 4. All-to-All: mini MoE dispatch/combine ----------
head("4. All-to-All — MoE dispatch/combine (ragged 变长)")
g = torch.Generator().manual_seed(100 + R)
n_tok = 2
tok = torch.randn(n_tok, 2, generator=g)                    # 本卡 token
router = torch.randint(0, W, (n_tok,), generator=g)         # token → 专家 → 卡
counts = torch.bincount(router, minlength=W)                # 发往各卡的量(不均!)

in_counts = torch.empty(W, dtype=torch.long)
dist.all_to_all_single(in_counts, counts)                   # ① 元数据交换
cap_t = torch.tensor([int(counts.max())], dtype=torch.long)
dist.all_reduce(cap_t, op=dist.ReduceOp.MAX)                # 全局最大容量
cap = int(cap_t)

send = torch.zeros(W, cap, 2)                               # ② 按目标分组 + padding
for e in range(W):
    rows = tok[router == e]
    send[e, : rows.shape[0]] = rows
recv = torch.zeros(W, cap, 2)
dist.all_to_all_single(recv, send)                          # 等长 A2A

send2 = torch.zeros(W, cap, 2)                              # ③ 本地"专家"= *10, 转置回去
for j in range(W):
    send2[j, : int(in_counts[j])] = recv[j, : int(in_counts[j])] * 10
recv2 = torch.zeros(W, cap, 2)
dist.all_to_all_single(recv2, send2)
final = torch.empty(n_tok, 2)                               # ④ 按 router 还原顺序
for e in range(W):
    idx = (router == e).nonzero().squeeze(1)
    final[idx] = recv2[e, : len(idx)]
assert torch.allclose(final, tok * 10)
log(f"router={router.tolist()} 发送计数={counts.tolist()} (ragged), "
    f"dispatch→expert→combine 往返后 == tok*10 ✓")

# ---------- 5. new_group: TP × DP 子组 ----------
head("5. new_group — 同一集群, TP/DP 各自的通信组")
tp_size = 2
tp_g = [dist.new_group([0, 1]), dist.new_group([2, 3])]
dp_g = [dist.new_group([0, 2]), dist.new_group([1, 3])]
tp, dp = tp_g[R // tp_size], dp_g[R % tp_size]
a = torch.full((2,), float(R))
t, d = a.clone(), a.clone()
dist.all_reduce(t, group=tp)
dist.all_reduce(d, group=dp)
log(f"TP 组内和={t.tolist()} DP 组内和={d.tolist()} (同一张量, 两个组两个结果)")

# ---------- 6. DDP vs Zero-2: 同一笔 2V 通信账的两种花法 ----------
head("6. DDP vs Zero-2 — AR 的两种花法")
lin = torch.nn.Linear(8, 8, bias=False)                     # 同初始化
gd = torch.Generator().manual_seed(1000 + R)
x = torch.randn(2, 8, generator=gd)                         # 各卡不同数据(DP)
lin(x).pow(2).sum().backward()
lr = 0.1

g = lin.weight.grad.clone()                                 # DDP: 梯度 AR + /W
dist.all_reduce(g)
g /= W
p_ddp = lin.weight.data - lr * g

flat = lin.weight.grad.reshape(-1)                          # Zero-2: RS 分片
n = flat.numel()
shard = torch.empty(n // W)
dist.reduce_scatter_tensor(shard, flat)
shard /= W
p_local = lin.weight.data.reshape(-1)[R * n // W : (R + 1) * n // W] - lr * shard
p_full = torch.empty(n)
dist.all_gather_into_tensor(p_full, p_local)                # 参数 AG 重建
assert torch.allclose(p_full, p_ddp.reshape(-1))
log(f"64 参数: DDP 每卡常驻 64 梯度+64 优化器态; Zero-2 每卡只持 {n // W}; "
    f"两条路径更新结果一致 ✓ (通信量同为 2V)")

dist.destroy_process_group()

=== 1. Send/Recv — 流水线: 激活向右, 梯度向左 ===
[R0] stage0: 上游激活=[0.0, 0.0] 下游梯度=[1.0, 1.0]

=== 2. Broadcast / Scatter / Gather / Reduce ===
[R0] broadcast/scatter/gather/reduce 断言全部通过

=== 3. All-Gather / Reduce-Scatter / All-Reduce ===
[R0] RS 分片=[600.0, 604.0], 恒等式 AG(RS(x)) == AR(x) 成立 ✓

=== 4. All-to-All — MoE dispatch/combine (ragged 变长) ===
[R0] router=[2, 0] 发送计数=[1, 0, 1, 0] (ragged), dispatch→expert→combine 往返后 == tok*10 ✓

=== 5. new_group — 同一集群, TP/DP 各自的通信组 ===
[R0] TP 组内和=[1.0, 1.0] DP 组内和=[2.0, 2.0] (同一张量, 两个组两个结果)

=== 6. DDP vs Zero-2 — AR 的两种花法 ===
[R0] 64 参数: DDP 每卡常驻 64 梯度+64 优化器态; Zero-2 每卡只持 16; 两条路径更新结果一致 ✓ (通信量同为 2V)