Gahow Wang
76fffb3b68
Megatron-style TP for Qwen3 on the 8x5090 (no-NVLink, PCIe) box: column/row split per layer, 2 AllReduces/layer, multi-thread one-rank-per-GPU model, NCCL, sharded weights, and the incremental implementation + verification plan. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
123 lines
7.1 KiB
Markdown
123 lines
7.1 KiB
Markdown
# Phase 17: Tensor Parallelism (TP)
|
||
|
||
> 目标：在单机多卡上做 **张量并行**，把 Qwen3-8B 的权重、计算和 KV cache 按
|
||
> head / 中间维切分到 TP 个 GPU 上，用 AllReduce 聚合，降低单卡显存压力并提升吞吐。
|
||
> 先做 **TP=2 / 4（组内）**，跳过投机解码（原 Phase 16）。
|
||
|
||
## 1. 硬件约束（dash5）
|
||
|
||
- 8× RTX 5090（32GB，SM120），**无 NVLink**，纯 PCIe Gen5。
|
||
- 拓扑：GPU 0–3 一组、4–7 一组，组内 `PHB`（同 host bridge，可 P2P），跨组 `NODE`。
|
||
- **TP 必须在组内**（0–3 或 4–7），否则 AllReduce 走跨组 PCIe，延迟更高。
|
||
- AllReduce 带宽受限于 PCIe（~单向 64GB/s），远低于 NVLink；通信会是 decode 的主要开销。
|
||
|
||
## 2. 切分方案（Megatron-style）
|
||
|
||
Qwen3-8B：`hidden=4096`、`num_heads=32`、`num_kv_heads=8`、`head_dim=128`、
|
||
`intermediate=12288`、`layers=36`、`vocab=151936`。TP=2/4/8 都能整除 32/8/12288。
|
||
|
||
每个 transformer block 的切分（设 world size = `T`，本 rank = `r`）：
|
||
|
||
### Attention（column → row）
|
||
| 权重 | 原 shape (已转置) | 切分 | 每 rank shape |
|
||
|------|-------------------|------|---------------|
|
||
| `q_proj_wt` | `[hidden, num_heads·head_dim]` | column（按 Q head） | `[hidden, (num_heads/T)·head_dim]` |
|
||
| `k_proj_wt` | `[hidden, num_kv_heads·head_dim]` | column（按 KV head） | `[hidden, (num_kv_heads/T)·head_dim]` |
|
||
| `v_proj_wt` | 同上 | column | 同上 |
|
||
| `o_proj_wt` | `[num_heads·head_dim, hidden]` | **row** | `[(num_heads/T)·head_dim, hidden]` |
|
||
|
||
- 每个 rank 只算自己的 `num_heads/T` 个 Q head 和对应的 `num_kv_heads/T` 个 KV head；
|
||
GQA 的 `n_rep = num_heads/num_kv_heads = 4` 在每个 rank 内保持不变。
|
||
- `q_norm`/`k_norm`（`[head_dim]`）逐 head 应用，**复制**到每个 rank。
|
||
- RoPE 逐 head、按 position 应用，每个 rank 独立做。
|
||
- **KV cache 也切分**：每个 rank 的 paged KV 只存自己的 `num_kv_heads/T` 个 head
|
||
→ 每卡 KV 显存降为 1/T（TP 的一大收益）。
|
||
- `attn = merge_heads(...) @ o_proj_wt` 得到**部分** `[T_tok, hidden]` → **AllReduce(sum)** → 完整。
|
||
|
||
### MLP / SwiGLU（column → row）
|
||
| 权重 | 原 shape | 切分 | 每 rank shape |
|
||
|------|----------|------|---------------|
|
||
| `gate_proj_wt` | `[hidden, intermediate]` | column | `[hidden, intermediate/T]` |
|
||
| `up_proj_wt` | 同上 | column | 同上 |
|
||
| `down_proj_wt` | `[intermediate, hidden]` | **row** | `[intermediate/T, hidden]` |
|
||
|
||
- `silu(gate)*up` 在切分后的 `[T_tok, intermediate/T]` 上逐元素做，无需通信。
|
||
- `down = (...) @ down_proj_wt` 得到部分 `[T_tok, hidden]` → **AllReduce(sum)** → 完整。
|
||
|
||
### 复制（不切分）
|
||
- 所有 RMSNorm 权重（`input_norm`/`post_norm`/最终 `norm`）：每个 rank 在 AllReduce 后
|
||
拿到完整 hidden，本地用复制的权重归一化。
|
||
- **第一版**：`embed_tokens` 和 `lm_head` 复制（各 ~1.2GB）。
|
||
后续优化：vocab-parallel embedding（local lookup + AllReduce）、column-parallel lm_head + AllGather。
|
||
|
||
### 通信点
|
||
每层 **2 次 AllReduce**（o_proj 后、down_proj 后）→ 每生成 1 token 共 `2·36 = 72` 次。
|
||
decode 时每次 AllReduce 张量是 `[batch, 4096]` BF16（单 token batch=1 时 8KB），**延迟主导**。
|
||
|
||
## 3. 进程 / 线程模型
|
||
|
||
**单进程、多线程**：每个 TP rank 一个 OS 线程，线程启动时 `cudaSetDevice(rank_device)` 并绑定。
|
||
|
||
选择理由：
|
||
- xserv 的 caching allocator 是 `thread_local`，每线程独立池 → 天然契合「一线程一卡一池」。
|
||
- CUDA context 隐式按 device/thread 管理；线程内只 set 一次 device、不再切换即可。
|
||
- HTTP server / 调度器仍在主线程（rank 0 协调），无需多进程 IPC，改动最小。
|
||
- 单机 8 卡足够；多进程（torchrun 式）留待真正跨节点时再说。
|
||
|
||
执行流：主线程调度器准备一个 step 的输入（tokens/positions/slots），广播给 `T` 个 rank 线程；
|
||
每个 rank 线程跑自己的分片 forward（含层内 AllReduce），rank 0 拿到完整 logits 后采样。
|
||
用 barrier / channel 同步每个 step。
|
||
|
||
## 4. 通信库：NCCL
|
||
|
||
用 **NCCL**（dash5 已装：`/usr/lib/x86_64-linux-gnu/libnccl.so.2`，`/usr/include/nccl.h`）。
|
||
|
||
- 新建 crate `xserv-distributed`：NCCL FFI（`ncclGetUniqueId`、`ncclCommInitRank`、
|
||
`ncclAllReduce`、`ncclGroupStart/End`）+ `TpContext`（rank/world/comm/stream）+
|
||
`all_reduce_sum(&mut GpuBuffer)` 原语。
|
||
- NCCL 多线程模式：主线程生成 `ncclUniqueId`，各 rank 线程用 `ncclCommInitRank(comm, world, id, rank)`
|
||
初始化（需 `ncclGroupStart/End` 包裹并发 init）。
|
||
- AllReduce 用 BF16（`ncclBfloat16`）+ `ncclSum`，在每个 rank 自己的 stream 上。
|
||
|
||
> **决策点**：collective 用 NCCL 还是自己手写 P2P ring/tree AllReduce？
|
||
> 本项目是「从零构建」，但 collective 属于基础设施（类比我们也用 cuBLAS 作为可切换后端）。
|
||
> 推荐先用 NCCL 把 TP 跑通、拿到正确性与加速比，后续可选做手写 AllReduce 作为学习项。
|
||
|
||
## 5. 权重分片加载
|
||
|
||
每个 rank 只加载/保留自己的分片，省显存：
|
||
- column-parallel 权重：按输出维切片取本 rank 的 `[*, dim/T]` 段。
|
||
- row-parallel 权重：按输入维切片取本 rank 的 `[dim/T, *]` 段。
|
||
- 复制权重（norm/embed/lm_head）：每个 rank 各留一份。
|
||
|
||
实现：`loader` 读 safetensors（mmap）时按 rank 只搬运需要的切片到该 rank 的 GPU；
|
||
`Qwen3::from_weights_tp(config, weights, rank, world)` 在转置/切分时按 rank 取段。
|
||
|
||
## 6. 实现步骤（逐步可验证）
|
||
|
||
1. **P17.1 — `xserv-distributed` 基础**：NCCL FFI + `TpContext` + `all_reduce_sum`。
|
||
验收：2 卡各放一个已知向量，AllReduce 后两卡结果都等于和。
|
||
2. **P17.2 — 分片权重加载**：`from_weights_tp`，每 rank 只持有自己的分片。
|
||
验收：各 rank 权重 shape 正确、显存占用约为 1/T（+ 复制项）。
|
||
3. **P17.3 — TP forward**：rank 内 attention/MLP + 层内 AllReduce。
|
||
验收：**TP=2 的 logits 与 TP=1 在 BF16 容差内一致**（top-1 一致，top-5 重合）。
|
||
4. **P17.4 — 接入 engine/server**：`--tp N`，多线程 rank workers + rank0 调度。
|
||
验收：`--tp 2` 端到端可服务；用现有 llama.cpp bench 跑正确性；
|
||
测 TP=2 vs TP=1 的吞吐 / TTFT / 每卡显存。
|
||
|
||
## 7. 预期与风险
|
||
|
||
- **显存**：每卡权重 + KV 降到约 1/T（embed/lm_head 暂复制）。TP=2 时单卡 ~8GB 权重 + 更大 KV/并发空间。
|
||
- **吞吐**：PCIe AllReduce 延迟会吃掉部分收益；decode 是延迟敏感的，72 次小 AllReduce/token
|
||
可能让 TP=2 的单流 tok/s **不一定线性提升**，但能跑更大 batch / 更长 context。先测实测数。
|
||
- **风险**：NCCL 多线程初始化的同步、每 rank stream 与现有 kernel stream 的协调、
|
||
KV cache 按 rank 切 head 后 paged kernel 的 head 维参数要用 per-rank 值。
|
||
- 正确性优先：先 TP=1 等价（logits 对齐），再谈性能。
|
||
|
||
## 8. 不在本阶段范围
|
||
|
||
- 跨组 TP=8、Pipeline Parallelism、多节点。
|
||
- vocab-parallel embedding / lm_head（先复制）。
|
||
- 手写 AllReduce（NCCL 跑通后可选）。
|
||
- 与 CUDA Graph decode 的结合（先走非 graph 路径）。
|