LLM 分布式推理基础（一）🎯
MoE 分布式推理基础（基于 vLLM）
基于
allenai/OLMoE-1B-7B-0924（\(L=16\)、\(E=64\)、\(K=8\)、\(H=2048\)、\(I=1024\)、\(H_q=H_{kv}=16\)、\(d=128\)）。本文档用 algorithm 风格伪代码 + 实际源码对照，描述单卡与分布式 MoE / Attention 的执行流程。
记号约定：
- \(x \in \mathbb{R}^{n \times H}\)：一层 decoder 的输入 hidden states（\(n\) 为 token 数）
- \(W_{qkv} \in \mathbb{R}^{H \times (H_q + 2H_{kv})d}\)、\(W_o \in \mathbb{R}^{H_q d \times H}\)：attention 权重
- \(W_{gate} \in \mathbb{R}^{H \times E}\)：MoE router 权重
- \(W_{13}^{(e)} \in \mathbb{R}^{H \times 2I}\)、\(W_2^{(e)} \in \mathbb{R}^{I \times H}\)：expert \(e\) 的权重
- \(h_q = H_q / tp\)、\(h_{kv} = H_{kv} / tp\)：TP 切后每 rank 的 head 数
- \(E_\ell = E / ep\)：EP 切后每 rank 持有的 expert 数
- \(\mathrm{AG}\)、\(\mathrm{AR}\)、\(\mathrm{A2A}\)：
all_gather、all_reduce、all_to_all通信原语
一、单卡实现
1.1 Attention Forward
源码（transformers modeling_olmoe.py）
# OlmoeAttention.forward (单卡)
qkv = self.qkv_proj(hidden_states) # (n, H) @ (H, (H_q+2·H_kv)·d) → (n, 6144)
q, k, v = qkv.split([H_q·d, H_kv·d, H_kv·d], dim=-1) # 各 (n, 2048)
q, k = self.q_norm(q), self.k_norm(k) # RMSNorm 在 (H_q·d,) 维
q, k = self.rotary_emb(positions, q, k)
attn_out = self.attn(q, k, v) # FlashAttention → (n, H_q·d)
out = self.o_proj(attn_out) # (n, H_q·d) @ (H_q·d, H) → (n, H)
Algorithm 1：单卡 Attention Forward
ALGORITHM 1: Single-GPU Attention Forward
Input: x ∈ R^{n×H}, positions pos ∈ N^n, KV cache (K_c, V_c)
Output: y_attn ∈ R^{n×H}
1: [q; k; v] ← x · W_{qkv} // (n, (H_q + 2H_kv)·d)
2: q ← reshape(q, (n, H_q, d)); k, v 同理
3: q ← RMSNorm(q); k ← RMSNorm(k) // norm 在全 head 维
4: q, k ← RoPE(q, pos), RoPE(k, pos)
5: K_c ← concat(K_c, k); V_c ← concat(V_c, v)
6: a ← FlashAttention(q, K_c, V_c) // (n, H_q, d)
7: y_attn ← reshape(a, (n, H_q·d)) · W_o // (n, H)
8: return y_attn
1.2 MoE Forward
源码（transformers modeling_olmoe.py）
# OlmoeTopKRouter
router_logits = F.linear(hidden_states, self.weight) # (n, H) @ (H, E) → (n, E)
router_probs = softmax(router_logits, dim=-1) # (n, E)
top_val, top_idx = torch.topk(router_probs, k=K) # 各 (n, K)
if norm_topk_prob:
top_val /= top_val.sum(-1, keepdim=True)
return router_logits, top_val, top_idx
# OlmoeExperts.forward
expert_mask = F.one_hot(top_k_index, E).permute(2, 1, 0) # (E, K, n)
expert_hit = nonzero(expert_mask.sum(...))
for e in expert_hit:
top_k_pos, token_idx = where(expert_mask[e])
X_e = hidden_states[token_idx] # gather, (n_e, H)
G_e = silu(X_e @ W1[e]) * (X_e @ W3[e]) # (n_e, I)
Y_e = G_e @ W2[e] # (n_e, H)
Y_e *= top_k_weights[token_idx, top_k_pos, None]
final_hidden_states.index_add_(0, token_idx, Y_e) # scatter-add
Algorithm 2：单卡 MoE Forward
ALGORITHM 2: Single-GPU MoE Forward (OlmoeExperts)
Input: x ∈ R^{n×H}, W_gate ∈ R^{H×E},
{W_13^{(e)}, W_2^{(e)}}_{e=0}^{E-1} // W_13: (H, 2I), W_2: (I, H)
Output: y ∈ R^{n×H}
// Router
1: ℓ ← x · W_gate // (n, E)
2: p ← softmax(ℓ) // (n, E)
3: (w_top, e_top) ← TopK(p, K) // (n, K) each
// 命中统计：n_e = Σ_{i,k} 1[e_top[i,k] = e]
4: for e ∈ {0..E-1} do
5: I_e ← {(i, k) : e_top[i,k] = e} // expert e 的 (token, slot) 集合
6: if |I_e| = 0: continue
7: X_e ← x[{i : (i,k) ∈ I_e}] // gather, (n_e, H)
8: G_e ← SiLU(X_e · W_13^{(e)}[:, :I]) * (X_e · W_13^{(e)}[:, I:])
9: Y_e ← G_e · W_2^{(e)} // (n_e, H)
10: for (i, k) ∈ I_e:
11: y[i] += Y_e[idx(i)] · w_top[i, k] // scatter-add
12: end for
13: return y
单卡耗时模型：\(T_{single} = \sum_{e: n_e > 0} \big( T_{GEMM}(n_e) + T_{gather/scatter}(n_e) \big)\)，热 expert 主导延迟。
SwiGLU expert 的两段 GEMM
W13[e]: (H, 2I) # gate_proj + up_proj 拼接
W2[e]: (I, H) # down_proj
gate_up = X_e @ W13[e] # (n_e, H) @ (H, 2I) → (n_e, 2I)
gate, up = gate_up.chunk(2, -1) # 各 (n_e, I)
h = silu(gate) * up # (n_e, I)
Y_e = h @ W2[e] # (n_e, I) @ (I, H) → (n_e, H)
每 token 对单 expert FLOPs ≈ \(6 \cdot n_e \cdot H \cdot I\)。
单卡的"不均衡"
总命中 \(\sum_e n_e = n \cdot K\)，平均 \(nK/E\)。router 输出长尾分布，\(n_e\) 高度不均：
\[T_{single} = \sum_{e:\, n_e > 0} \big( T_{GEMM}(n_e) + T_{gather/scatter}(n_e) \big)
\]
T_GEMM(n_e) 对 \(n_e\) 亚线性（GEMM 利用率随 \(n_e\) 增大），热 expert 主导延迟。
量化：counts[l, e] 为第 l 层 expert e 的命中数，\(\mathrm{CV} = \sigma_e(n_e) / \mu_e(n_e)\)。
二、分布式实现（vllm 0.10.2）
核心事实：vllm
enable_expert_parallel在 \(dp=1\) 时不触发 all2all dispatch/combine。它的 EP 走「权重分片 + expert map mask + all_reduce」路径。vllm 有完整的 dispatch/combine
all_to_all实现（get_ep_group().dispatch()/.combine()，见 layer.py:1783-1824），但 \(do\_naive\_dispatch\_combine\) 的判定是 \(dp\_size > 1\)——\(dp=1\) 时被跳过。此外还支持 DeepEP/pplx kernel 级 all2all（
use_deepep_ht_kernels、use_pplx_kernels），也是同样的判定入口。
2.1 TP / DP / EP 三个维度分别切什么
vllm 里 ep_size = tp_size · dp_size。EP 不是独立维度，是 "TP×DP 范围内 MoE 层选 expert 分片" 的子集。
| 维度 | 切的对象 | 对 hidden/token | 对 expert | 通信原语 |
|---|---|---|---|---|
| TP | 权重矩阵按列/行切 | \(x\) 不切，\(W\) 维度切 \(1/tp\) | EP 开启时 MoE 内部 tp 强制为 1 | all_reduce |
| DP | batch / 请求切 | 每 rank 拿自己 \(n_i\) 个 token | 不切（或用 EP） | dispatch/combine all_to_all |
| EP | 只切 expert id | \(x\) 不切 | 每 rank \(E/ep\) 个完整 expert | all_reduce（\(dp=1\)）或 all_to_all（\(dp>1\)） |
DP 是请求级切分，不是 sequence 切片：\(dp=m\) 起 \(m\) 个独立 engine，rank \(i\) 处理自己 batch 的 \(n_i\) 个 token。
2.2 TP Attention（rank \(i\) 视角）
源码（vllm olmoe.py）
# OlmoeAttention.__init__ (tp=n 时)
self.num_heads = self.total_num_heads // tp_size # H_q / n
self.num_kv_heads = self.total_num_kv_heads // tp_size # H_kv / n
self.qkv_proj = QKVParallelLinear(hidden_size, head_dim,
total_num_heads, total_num_kv_heads) # 列切
self.o_proj = RowParallelLinear(total_num_heads * head_dim,
hidden_size) # 行切
# OlmoeAttention.forward (tp>1)
def _apply_qk_norm(self, q, k):
if self.tp_size > 1:
q = tensor_model_parallel_all_gather(q.contiguous()) # (n, h_q·d) → (n, H_q·d)
k = tensor_model_parallel_all_gather(k.contiguous())
q = self.q_norm(q); k = self.k_norm(k)
if self.tp_size > 1:
splitter = partial(split_tensor_along_last_dim, num_partitions=self.tp_size)
q = splitter(q)[self.tp_rank]
k = splitter(k)[self.tp_rank]
return q, k
def forward(self, positions, hidden_states):
qkv, _ = self.qkv_proj(hidden_states) # (n, (h_q + 2·h_kv)·d)
q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
q, k = self._apply_qk_norm(q, k)
q, k = self.rotary_emb(positions, q, k)
attn_output = self.attn(q, k, v) # (n, h_q·d)
output, _ = self.o_proj(attn_output) # all_reduce inside
return output
Algorithm 3：TP Attention Forward on rank \(i\)（tp=n, dp=1）
ALGORITHM 3: TP Attention Forward on rank i (tp = n, dp = 1)
Input: x ∈ R^{n×H}（所有 tp rank 共享同一份）, W_qkv^{(i)}, W_o^{(i)}
Output: y_attn ∈ R^{n×H}
// Step 1: qkv_proj 列切（无通信）
1: [q_i; k_i; v_i] ← x · W_{qkv}^{(i)} // q_i: (n, h_q·d)
// W_{qkv} = [W_{qkv}^{(0)} | ... | W_{qkv}^{(n-1)}]，rank i 持第 i 份
// Step 2: qk_norm 需要全 head 维 → all_gather
2: q^full ← AG(q_0, q_1, ..., q_{n-1}) // (n, H_q·d)
3: k^full ← AG(k_0, k_1, ..., k_{n-1}) // (n, H_kv·d)
4: q^full ← RMSNorm(q^full); k^full ← RMSNorm(k^full)
5: q_i ← split_along_last_dim(q^full, n)[i] // 拿回 (n, h_q·d)
6: k_i ← split_along_last_dim(k^full, n)[i]
// Step 3: RoPE + attn kernel（无通信）
7: q_i, k_i ← RoPE(q_i), RoPE(k_i)
8: a_i ← Attn(q_i, K_c^{(i)}, V_c^{(i)}) // (n, h_q·d)
// Step 4: o_proj 行切 + all_reduce
9: z_i ← a_i · W_o^{(i)} // (n, H)，部分和
// W_o = [W_o^{(0)}; W_o^{(1)}; ...; W_o^{(n-1)}]，按输入维切
10: y_attn ← AR(z_0, z_1, ..., z_{n-1}) = Σ_i z_i // (n, H)
11: return y_attn
通信账（每层 attn）：
all_gather×2：量 = \(n \cdot H_q d \cdot (tp-1)/tp\)（q 和 k 各一次）all_reduce×1：量 = \(n \cdot H\)
为啥 q_norm 用 all_gather 不用 all_reduce：RMSNorm 需要每个元素的真实值算 rms，求和不能代替。
为啥 o_proj 用 all_reduce 不用 all_gather：分块矩阵乘 \(y_{attn} = a W_o = \sum_i a_i W_o^{(i)}\)，各 rank 输出是部分和，加起来是完整值。
2.3 vllm EP MoE（tp=n, ep=n, dp=1）
关键澄清：vllm enable_ep 下 MoE 内部 tp 被重置为 1
很多人会误以为 tp=n, ep=n 是 "\(E\) 个 expert ÷ \(n\) + 每个 expert GEMM 再 ÷ \(n\)"（要 \(n^2\) 卡）。看 fused_moe/config.py:294-302：
# In EP, each device owns a set of experts fully. There is no tensor parallel
ep_size = tp_size # 外层 tp_size 被借去当 ep_size
ep_rank = tp_rank
return FusedMoEParallelConfig(tp_size=1, tp_rank=0, ...) # ← MoE 内部 tp=1!
EP 和 MoE 内部 TP 互斥。外层 tensor_parallel_size 只作用于 attention/router；启用 EP 后 MoE 层把这个值当 \(ep\_size\) 用，内部 \(tp=1\)。
| 模块 | tp=n, ep=n 实际并行 |
每 rank 持有 |
|---|---|---|
| Attention | TP=\(n\) | \(H_q / n\) head |
Router (gate) |
replicated | \(W_{gate} \in \mathbb{R}^{H \times E}\) 全量 |
| Expert FFN | EP=\(n\)（完整 expert） | \(E/n\) 个 × \(W_{13}, W_2\) 全尺寸 |
总卡数 = tp_size × dp_size = n（不是 n²）。E/n 个完整 expert / rank。
expert_map 的工作机制
# determine_expert_map (layer.py:677) 生成 (E,) 张量：
# expert_map[global_e] = local_id 若 e 在 rank_i 的 E_ℓ 个 expert 里
# expert_map[global_e] = -1 否则
# 例: E=8, ep_size=2, ep_rank=1
expert_map = torch.tensor([-1, -1, -1, -1, 0, 1, 2, 3], dtype=torch.int32)
# rank 1 持 expert 4..7，本地 id 0..3；expert 0..3 不归它管 → -1
topk_ids 中每个全局 id 经 expert_map 改写后送进 moe_align_block_size，-1 的 block 在 kernel 里写 0。
monkey-patch determine_expert_map 就能在不动权重、不动 router 的前提下，完全改写"哪个 global expert 在哪张卡"——这是做 best/worst expert 重排布的注入点。
源码（vllm layer.py + fused_moe.py）
# OlmoeMoE.forward
router_logits, _ = self.gate(hidden_states) # ReplicatedLinear, (N, E)
final_hidden_states = self.experts(hidden_states=hidden_states,
router_logits=router_logits)
# FusedMoE.forward_impl
topk_weights, topk_ids = FusedMoE.select_experts(
hidden_states=x, router_logits=router_logits, top_k=top_k, ...)
# ↓ UnquantizedFusedMoEMethod.forward_cuda 内 fused_experts:
# fused_experts_impl (fused_moe.py:1498)
sorted_token_ids, expert_ids, num_tokens_post_padded = (
moe_align_block_size(curr_topk_ids, BLOCK_SIZE_M,
global_num_experts, expert_map))
invoke_fused_moe_kernel(qcurr_hidden_states, w1, intermediate_cache1,
..., sorted_token_ids, expert_ids, ...)
moe_align_block_size（moe_align_block_size.py:12）把 topk_ids 展平成 N·K 个 (token, expert) 对，按 expert id 排序、按 BLOCK_SIZE_M 对齐 padding：
# 输入: topk_ids (N, K) → flatten 成 N·K 个 (token_idx, expert_id) 对
# 输出:
sorted_token_ids: (max_num_tokens_padded,) # 按 expert 排序后的 token 下标
expert_ids: (max_num_m_blocks,) # 每个 block 对应哪个 expert
num_tokens_post_padded: int # 对齐后总 token 数
# expert_map 在返回前把 expert_ids 改成 local_id 或 -1
if expert_map is not None:
expert_ids = expert_map[expert_ids]
举例（topk_ids = [[2,3,4],[1,2,4],[1,3,4],[1,2,3]], BLOCK_SIZE_M=4, E=4）：
原始 12 个 (token, expert) 对，每个 expert 命中 3 个 token
按 expert 排序 + 补齐到 BLOCK_SIZE_M 的倍数 → 16 个槽位
sorted_token_ids = [3, 6, 9, PAD, 0, 4,10, PAD, 1, 7,11, PAD, 2, 5, 8, PAD]
expert_ids = [ expert 0 , expert 1 , expert 2 , expert 3 ]
kernel 内对每个 block（fused_moe.py:162-170）：
off_experts = tl.load(expert_ids_ptr + pid_m)
if off_experts == -1:
write_zeros_to_output(...) # 不归我管的 expert：写 0 立即返回
return
# 否则正常走 grouped GEMM
Algorithm 4：vllm EP MoE Forward on rank \(i\)（tp=n, ep=n, dp=1）
ALGORITHM 4: vllm EP MoE Forward on rank i (dp=1, ep_size = n)
Input: x ∈ R^{N×H}（所有 rank 拿到全量）,
W_gate ∈ R^{H×E}（replicated）,
{W_13^{(e)}, W_2^{(e)}}_{e ∈ my_experts}, // 本 rank E_ℓ = E/n 个完整 expert
expert_map: (E,) → {-1} ∪ [0, E_ℓ)
Output: y ∈ R^{N×H}
// Step 1: Router（replicated 计算）
1: ℓ ← x · W_gate // (N, E)，所有 rank 相同
2: p ← softmax(ℓ)
3: (w_top, e_top) ← TopK(p, K) // (N, K)，所有 rank 相同
// Step 2: moe_align_block_size 排序 + padding
4: flat_ids ← flatten(e_top) // (N·K,)
5: flat_tok ← repeat_interleave(arange(N), K) // (N·K,)
6: (sorted_tok, expert_ids, num_padded)
← moe_align_block_size(flat_ids, BLOCK_M, E)
// sorted_tok: (max_padded,) 按 expert 排序后的 token 下标
// expert_ids: (max_num_blocks,) 每个 block 对应哪个 expert
// Step 3: expert_map 标 -1
7: expert_ids ← expert_map[expert_ids] // 不在本 rank → -1
// Step 4: invoke_fused_moe_kernel（-1 block 写零）
8: y_partial ← zeros(N, H)
9: for each block b in grid:
10: e_local ← expert_ids[b]
11: if e_local = -1: write_zeros(y_partial[block_tokens(b)]); continue
12: T_b ← sorted_tok[b·BLOCK_M : (b+1)·BLOCK_M]
13: G_b ← SiLU(x[T_b] · W_13^{(e_local)}[:, :I]) * (x[T_b] · W_13^{(e_local)}[:, I:])
14: Y_b ← G_b · W_2^{(e_local)} // (BLOCK_M, H)
15: for (t, k_pos) ∈ T_b:
16: y_partial[t] += Y_b[idx(t)] · w_top[t, k_pos]
17: end for
// Step 5: all_reduce 合并
18: y ← AR(y_partial_0, y_partial_1, ..., y_partial_{n-1}) // Σ_i，(N, H)
19: return y
本 rank 工作量：\(\mathrm{local\_hits}_i = \sum_{t,k} \mathbb{1}[\mathrm{expert\_map}[e_{top}[t,k]] \neq -1]\)。
关键性质：
- 通信量 = \(N \cdot H \cdot \mathrm{dtype}\)（
all_reduce，固定，与路由无关） - 不均衡只体现在
local_hits_i差异 →all_reduce同步点等最忙 rank - Router/TopK 被重复算 \(n\) 次（所有 rank 各算一遍同样的）
2.4 vllm EP MoE（tp=1, ep=n, dp=n）真 dispatch/combine
dp=1 时所有 rank 看到同一份 x，topk_ids 相同 → 用 expert_map mask + all_reduce 就够。
dp>1 时这套失效：
# dp=2 时，两 rank 是不同的请求
rank_0: x_A ∈ (n_A, H) → topk_ids_A ∈ (n_A, K)
rank_1: x_B ∈ (n_B, H) → topk_ids_B ∈ (n_B, K) # 完全不同的 token!
all_reduce 没法把 x_A 的输出和 x_B 的输出加起来——它们是不同 token，求和无意义。
源码（vllm layer.py:1782）
do_naive_dispatch_combine = (self.dp_size > 1
and not use_deepep_ht
and not use_flashinfer_cutlass)
if do_naive_dispatch_combine:
hidden_states, router_logits = get_ep_group().dispatch(
hidden_states, router_logits)
# ... 本地 expert GEMM ...
if do_naive_dispatch_combine:
states = get_ep_group().combine(states)
Algorithm 5：真 all2all EP MoE on rank \(i\)（tp=1, ep=n, dp=n）
ALGORITHM 5: True all2all EP MoE on rank i (dp=n, ep=n)
Input: x_i ∈ R^{n_i×H}（本 rank 自己 batch 的 token）,
W_gate, {W_13^{(e)}, W_2^{(e)}}_{e ∈ my_experts}
Output: y_i ∈ R^{n_i×H}
// Step 1: 本地 Router + TopK（只算自己 n_i 个 token）
1: ℓ_i ← x_i · W_gate // (n_i, E)
2: (w_top_i, e_top_i) ← TopK(softmax(ℓ_i), K) // (n_i, K)
// Step 2: dispatch —— all_to_all 把 token 发到对应 expert 所在 rank
3: send_sizes[i→j] ← Σ_{t,k} 1[⌊e_top_i[t,k] / E_ℓ⌋ = j] // 要发给 rank j 的 (t,k) 对数
4: all_to_all_single(recv_sizes, send_sizes) // 交换 sizes
5: send_buf ← pack(x_i[t], w_top_i[t,k], t) for all (t,k)
6: recv_buf ← A2A(send_buf,
output_split_sizes=recv_sizes,
input_split_sizes=send_sizes)
// recv_buf 现含所有 rank 发给"我的 expert"的 (token, expert) 对
// Step 3: 本地 expert GEMM
7: y_partial_i ← fused_experts(recv_buf, my_w13, my_w2)
// Step 4: combine —— all_to_all 反方向
8: send_back_i ← pack(y_partial_i)
9: final_buf ← A2A(send_back_i,
output_split_sizes=send_sizes,
input_split_sizes=recv_sizes)
10: y_i ← scatter_add(final_buf, w_top_i) // (n_i, H)
11: return y_i
关键差异 vs Algorithm 4：
| 项 | Algorithm 4（dp=1） |
Algorithm 5（dp>1） |
|---|---|---|
topk_ids 在所有 rank |
相同（replicated） | 不同（各自 batch） |
| Router GEMM | \(N \cdot H \cdot E\)（重复 n 次） |
\(n_i \cdot H \cdot E\)（只算自己） |
| Token 通信 | 无（只 all_reduce） |
all_to_all dispatch + combine |
| 通信量 | \(N \cdot H\) 固定 | \(\sum_j \mathrm{send}_{i \to j} \cdot H\) ∝ 路由 |
expert_map 用途 |
标 -1 mask 不归我算的 | 本地分派收到的 (token, expert) 对 |
| 输出 | (N, H) 全量 partial |
(n_i, H) 自己 token |
两个同步点：dispatch 后等齐 + combine 后等齐 —— 都被最忙 rank 拖慢。
2.5 tp=n, ep=n·m, dp=m —— 多机扩展
ALGORITHM 6: TP+EP+DP 混合 MoE (ep_size = n·m)
Setup: rank 二维坐标 (tp_rank ∈ [0,n), dp_rank ∈ [0,m))
同 dp_rank 内 n 个 tp rank 看到同一份 x_{dp_rank}
MoE 内部 tp=1（vllm 设计）
Per-rank 持有:
- Attention 权重: 1/n 份（按 head 切）
- Expert 权重: E/(n·m) 个完整 expert
- Router: replicated（同 Algorithm 4）
通信组:
- TP 组 (n 个 rank): attn 的 AG + AR
- EP 组 (n·m 个 rank): MoE 的 dispatch/combine A2A
流程（单个 dp 组内）:
1. x_{dp_rank} 进入 attention（Algorithm 3，tp=n 组内通信）
2. 进入 MoE：
a. router/topk 在 tp=n 组内 replicated 算（同 Algorithm 4 Step 1）
b. dispatch/combine 在 ep=n·m 组内做 A2A（同 Algorithm 5 Step 2/4）
c. 本地 expert GEMM 算收到的 (token, expert) 对
通信组合：attn 的 AG/AR 只在 tp=n 组内；dispatch/combine 在 ep=n·m 组内跨 n·m 卡。
三、关键对比
不均衡到底影响谁
| 配置 | 计算不均 | 通信不均 | 同步等待点 |
|---|---|---|---|
tp=n, ep=n, dp=1 |
local_hits_i 差 |
❌ all_reduce 量固定 |
all_reduce 前等最忙 rank |
tp=1, ep=n, dp=n |
recv_i 差 |
✅ send_{i→j} 量差 |
dispatch 后等齐 + combine 后等齐 |
tp=n, ep=n·m, dp=m |
✅ | ✅ | ep 组内 A2A + tp 组内 AR |
要测通信不均必须 \(dp > 1\)；\(tp\) 只决定"每个 GEMM 多大"和 all_reduce 量，不影响 all_to_all 是否发生。
通信原语对比
# all_gather (AG): rank_i 有 x_i → 所有 rank 拿 [x_0; x_1; ...; x_{p-1}]
# 发送 = n·H/p, 接收 = n·H·(p-1)/p
# 出现在 _apply_qk_norm 拼齐 head 维
# all_reduce (AR): rank_i 有 x_i → 所有 rank 拿 Σ_i x_i
# 通信 = n·H
# 出现在 o_proj 输出、MoE 末尾（dp=1 路径）
# all_to_all (A2A): rank_i 有 [x_{i→0}, ..., x_{i→p-1}] → 按目标 rank 分发
# rank_i 拿到 [x_{0→i}, ..., x_{p-1→i}]
# 通信 = n·H，但分布随路由不均
# 出现在 dispatch/combine（dp>1）
关键区别：
all_gather/all_reduce：每 rank 收发 buffer 固定大小，不均衡只造成同步等待all_to_all：每 rank 收发的量由split_sizes决定，直接 ∝ 路由分布——不均衡真正放大的通信原语
浙公网安备 33010602011771号