docs: Phase T17 — process-per-GPU DDP design
torchrun-style: launcher spawns N worker processes, each with its own CUDA context; cross-process ncclUniqueId distributed via launcher-minted hex env injection (race-free, no shared FS / TCP); train_rank + grad all-reduce reused unchanged. Keeps thread-per-GPU path as regression baseline. ZeRO-1 dropped (user scope decision). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
217
docs/16-process-per-gpu.md
Normal file
217
docs/16-process-per-gpu.md
Normal file
@@ -0,0 +1,217 @@
|
||||
# Phase T17: Process-per-GPU DDP（torchrun 式独立 CUDA context）— Design Document
|
||||
|
||||
## Goal
|
||||
|
||||
T8 的 DDP 是**单进程 thread-per-GPU**：一个进程开 N 个 OS 线程，每线程 `cudaSetDevice` 绑一张卡、
|
||||
在**同一个 CUDA primary context** 里跑自己 rank 的训练。T11 修掉 per-op `cudaMalloc` 串行后，8 卡
|
||||
scaling 从 ~1.3× 恢复到 **~5×@8**，但**残留 5×@8 而非 ~8× 的非线性**——根因在 T11 doc / KI-5 已点明：
|
||||
**N 个 rank 线程共享同一个 CUDA context，driver 层很多调用（kernel launch、cuBLAS handle、stream
|
||||
排队）在单 context 内进程级串行**，pool allocator 只消掉了其中最大的一笔（malloc），剩下的 launch /
|
||||
cuBLAS 串行仍在。
|
||||
|
||||
T17 的目标 = **torchrun 式 process-per-GPU**：每个 rank 是一个**独立 OS 进程**，各自持有**独立的
|
||||
CUDA context**，彼此的 driver 调用不再在同一 context 排队 → 移除 thread-per-GPU 的残留串行，把 8 卡
|
||||
scaling 推向更接近线性。这是 Phase 2 里**改动最大**的一项（launcher 结构性重写 + 跨进程 NCCL
|
||||
bootstrap），所以本 doc 先行。
|
||||
|
||||
**Scope（用户已拍板）：process-per-GPU ONLY。ZeRO-1 / sharded optimizer 明确 drop**——本尺度
|
||||
optimizer state 小、收益薄。本任务只换**启动模型与 NCCL bootstrap**，**训练 step（grad all-reduce →
|
||||
本地 AdamW）原样复用、零改动**。
|
||||
|
||||
**保留 thread-per-GPU 路径**：T8 的 `launch()` + `train_ddp` bin 不删（回归保护 + 闸门 ①要求新旧路径
|
||||
loss 对齐）。process-per-GPU 作为**并列的新 launcher** 加上去。
|
||||
|
||||
验收（硬闸门全绿，诚实正确性，不放宽容差）：
|
||||
1. 多进程（world=2 / world=4）训练 loss **对单卡贴合**（进既有 DDP 容差 `<1e-3`），且对住旧 thread-per-GPU 路径；
|
||||
2. 跨 rank 参数一致（repo 既有 `<1e-6` 约定）；
|
||||
3. **8 卡线性度 before→after 实测**：thread-per-GPU baseline（~5×@8）vs process-per-GPU @ {1,2,4,8}，给数字；
|
||||
4. 全回归套绿（含 xserv 闭环 md5 / token-identical）；单卡与旧 thread-per-GPU 路径不回归。
|
||||
|
||||
## 什么变、什么不变
|
||||
|
||||
```
|
||||
thread-per-GPU (T8, 保留) process-per-GPU (T17, 新增)
|
||||
启动 1 进程 × N 线程 1 launcher 进程 → fork/exec N 个 worker 进程
|
||||
CUDA context N 线程共享 1 个 primary context 每 worker 进程 1 个独立 context
|
||||
rank/world/device 闭包捕获 + thread::scope env: RANK / WORLD_SIZE / LOCAL_RANK
|
||||
模型构建 每线程闭包内 build_model（!Send） 每进程 main 内 build_model（天然隔离）
|
||||
NCCL UniqueId 分发 move 一个 Copy struct 进线程闭包 launcher 生成 → hex 编码进子进程 env
|
||||
NCCL comm init DdpContext::init（不变） DdpContext::init（不变）
|
||||
─────────────────────────────────────────────────────────────────────────────────────
|
||||
grad all-reduce all_reduce_average_grads（不变） ← 同一份代码，零改动
|
||||
本地 AdamW step train_rank（不变） ← 同一份代码，零改动
|
||||
batch sharding i % world == rank（不变） ← 同一份代码，零改动
|
||||
参数一致性证明 同 init+同 grad+同 opt（不变） ← 同一论证
|
||||
```
|
||||
|
||||
**核心洞察**：T8 早把训练 step 写成「**per-rank**、接受 `&DdpContext`」的形状（`train_rank`）。
|
||||
thread-per-GPU 与 process-per-GPU **唯一的区别只在「怎么把 rank 跑起来 + 怎么把 UniqueId 递给每个
|
||||
rank」**——前者跨线程 move，后者跨进程 env。`train_rank` / `all_reduce_average_grads` / sharding /
|
||||
一致性论证**全部原样复用**。这正是把启动模型与训练逻辑解耦的回报。
|
||||
|
||||
## Module Layout
|
||||
|
||||
```
|
||||
crates/xtrain-distributed/src/
|
||||
├── lib.rs ← 加 pub mod proc; re-export hex_encode/decode_unique_id + run_worker entry
|
||||
├── proc.rs ← 新增：① launcher（spawn N worker 进程，env 注入 rank/world/local_rank/uid）
|
||||
│ ② worker entry（读 env → DdpContext::init → build_model → train_rank）
|
||||
│ ③ UniqueId hex 编解码（跨进程 env 传 128 字节）
|
||||
├── ddp.rs ← 不变（train_rank / build_model / DdpConfig 复用）
|
||||
├── lib.rs::DdpContext / all_reduce_average_grads / get_unique_id ← 不变
|
||||
└── bin/
|
||||
├── train_ddp.rs ← 不变（thread-per-GPU，保留）
|
||||
└── train_ddp_mp.rs ← 新增：multi-process launcher / worker 二合一入口
|
||||
crates/xtrain-distributed/tests/
|
||||
└── ddp_proc.rs ← 新增：spawn 多进程跑几步 → loss 对单卡 + 跨 rank 参数一致 +（顺手）before/after 吞吐
|
||||
docs/16-process-per-gpu.md ← 本文
|
||||
```
|
||||
|
||||
`proc.rs` 全程 `#[cfg(not(no_cuda))]` 门控（同 crate 既有约定）；本地无 nvcc 时 crate 编空，`cargo
|
||||
check` 过；dash5 上全量编译链 NCCL。
|
||||
|
||||
## Key Design Decisions
|
||||
|
||||
### ① Launch model：同一 binary 双模（launcher / worker）
|
||||
|
||||
`train_ddp_mp` 一个可执行文件，靠**环境变量是否存在**自判角色（torchrun 的 `LOCAL_RANK` 注入 worker
|
||||
是同一思路）：
|
||||
|
||||
- **launcher 模式**（直接被用户 / 测试调用，env 里没有 `XTRAIN_RANK`）：
|
||||
1. 看 `CUDA_VISIBLE_DEVICES` / `device_count()` 定 `world`；
|
||||
2. 调 `get_unique_id()` 生成一个 `ncclUniqueId`（128 字节），**hex 编码**成字符串；
|
||||
3. `for rank in 0..world`：`Command::new(current_exe())`，**复制自己全部 argv**（超参/路径透传），
|
||||
额外设 env `XTRAIN_RANK=rank`、`XTRAIN_WORLD=world`、`XTRAIN_LOCAL_RANK=rank`、
|
||||
`XTRAIN_NCCL_ID=<hex>`，spawn 为子进程；
|
||||
4. `wait()` 所有子进程，任一非零退出码 → launcher 以非零退出（CI / 闸门可感知）。
|
||||
- **worker 模式**（被 launcher spawn，env 里有 `XTRAIN_RANK`）：
|
||||
1. 从 env 读 `rank / world / local_rank / uid_hex`；
|
||||
2. `device::set_device(local_rank)` 绑卡（**每进程独立 primary context** 在此首次 CUDA 调用时建立）；
|
||||
3. hex 解码出 `NcclUniqueId`；`DdpContext::init(rank, world, id, local_rank)`（**复用 T8 的 init**）；
|
||||
4. `build_model(cfg, device)`（**复用 T8 的确定性 init** → 同种子 → 跨进程逐位同起点）；
|
||||
5. `train_rank(&ctx, &model, …, &cfg)`（**复用 T8 的训练 step，零改动**）；
|
||||
6. 退出码 0（成功）/ 非零（panic → 进程崩，launcher 感知）。
|
||||
|
||||
**单机 `CUDA_VISIBLE_DEVICES` 处理**：launcher 看到的 visible 设备集就是 `0..world`；每个 worker
|
||||
继承同一个 `CUDA_VISIBLE_DEVICES`（env 默认透传），`local_rank` 直接当作 visible 集内的 device
|
||||
ordinal → `set_device(local_rank)`。这与 thread-per-GPU 的 `devices = 0..count` 语义一致，单节点足够。
|
||||
（真·多节点要把 `LOCAL_RANK` 与全局 `RANK` 分离 + 每节点 `CUDA_VISIBLE_DEVICES` 切片，单节点不需要，
|
||||
记为 follow-up。）
|
||||
|
||||
### ② 跨进程 NCCL UniqueId 分发：launcher 生成 + hex-env 注入（**最简、无竞态**）
|
||||
|
||||
这是 T17 最该想清楚的一处。候选机制（任务列了文件 / TCP / 共享 FS）逐一权衡：
|
||||
|
||||
| 机制 | 怎么做 | 单节点取舍 |
|
||||
|---|---|---|
|
||||
| **共享文件** | rank0-worker 写 `/tmp/xtrain.id`，其余 worker **轮询读** | 要处理「文件还没写好」的 race（轮询 + 重试 + 超时），还要 worker 间约定谁是 rank0、何时清理 |
|
||||
| **TCP rendezvous** | 起一个 c10d-store 式小 server 派发 id | 最贴 torchrun，但要写 socket server/client、端口选择、握手协议——单节点 overkill |
|
||||
| **launcher 生成 + env 注入** ✅ | **launcher**（而非 rank0-worker）调 `ncclGetUniqueId`，hex 编码后**在 spawn 时就写进每个子进程的 env** | 无文件 race、无轮询、无 TCP server、无清理——env 在子进程出生前就备好。子进程读 env 即得 id |
|
||||
|
||||
**选 env 注入**，诚实理由：单节点下 launcher 是所有 worker 的**共同父进程**，env 是父→子最朴素的带外
|
||||
通道，且**在子进程创建那一刻就原子地确定**——天然没有「id 还没就绪」的竞态，比文件轮询 / TCP 握手都
|
||||
简单且更鲁棒。代价是 launcher 进程要链 NCCL 调 `ncclGetUniqueId`（它本就在 distributed crate 里，已链
|
||||
NCCL），可接受。
|
||||
|
||||
> **与「rank 0 生成」的关系**：torchrun 概念上是 rank 0 把 id 放进 c10d store、别的 rank 取。这里
|
||||
> **launcher 充当协调者**替 rank 0 生成——功能等价（id 只是个一次性握手 token，谁生成不影响正确性，
|
||||
> 只要全 rank 拿到同一个），但单节点下省掉了「worker 间再来一轮带外同步」。**128 字节 → hex = 256
|
||||
> 字符**，远低于环境变量长度上限，env 传输安全。
|
||||
|
||||
`hex_encode`/`decode_unique_id` 是 `proc.rs` 里两个纯函数（`[c_char;128] ↔ 256-char hex`），单测可在
|
||||
host 侧验 roundtrip（不需 GPU）。
|
||||
|
||||
### ③ 独立 CUDA context = 移除残留串行（这才是 T17 的 payoff）
|
||||
|
||||
thread-per-GPU 的残留非线性（KI-5 / T11 doc）来自：**N rank 线程共享同一 CUDA primary context**，driver
|
||||
对该 context 的很多操作（kernel launch 队列、cuBLAS handle、内部锁）是进程级 / context 级串行的——
|
||||
pool allocator 消掉了 malloc 这一最大笔，但 launch / cuBLAS 串行仍在，表现为 8 卡 ~5× 而非 ~8×。
|
||||
|
||||
process-per-GPU 下**每个 rank 是独立进程 → 独立 CUDA context → 独立 driver 状态**：各进程的 kernel
|
||||
launch / cuBLAS 调用**互不在同一 context 排队**，残留串行被结构性移除。这正是闸门 ③（before→after 线性度）
|
||||
要量出来的东西——若 process-per-GPU 把 8 卡从 ~5× 推到明显更高，即验证了根因诊断。**诚实原则**：若提升
|
||||
有限，如实报告（说明残留瓶颈转移到 NCCL all-reduce / PCIe 拓扑，那是另一层，非本任务 scope）。
|
||||
|
||||
### ④ 训练 step / 一致性论证：原样复用 T8，零改动
|
||||
|
||||
process-per-GPU 不碰任何训练数学：
|
||||
|
||||
- **grad all-reduce**：`all_reduce_average_grads(params)` 一字不改——NCCL collective 跨**进程**和跨
|
||||
**线程**对调用方完全一样（comm 是 rank 维度的，与进程/线程无关）。
|
||||
- **batch sharding**：`i % world == rank` 不变——每进程推进**同一个 seed 的 RNG**抽出整批 `B_global`
|
||||
序列、只算自己那片。各进程的并集 == 单卡同序批 → all-reduce 后的 grad 和与单卡逐序列一致。
|
||||
- **参数一致性**：同 ③个充分条件（T8 doc ④）——(a) 同确定性 `build_model`（同 LCG 种子，跨进程同样
|
||||
成立）；(b) NCCL all-reduce 跨 rank 返回逐位相同的归约（PCIe-only run-to-run 几 ULP 抖动，故闸门
|
||||
②用 `<1e-6` 而非 `==0`，与 T11 既有约定一致）；(c) 同 optimizer 超参/状态演化。
|
||||
- **对单卡**：与单卡只在 **fp 求和顺序**上差（单卡 tape SUM B 个；DDP 各 rank 先 SUM 分片再 NCCL SUM）
|
||||
→ `<1e-3` rel，不逐位。与 thread-per-GPU 路径则应**数值同量级**（同一 sharding + 同一 all-reduce）。
|
||||
|
||||
### ⑤ 进程生命周期 / 失败传播 / 资源清理
|
||||
|
||||
- **失败传播**：worker panic → 进程非零退出；launcher `wait()` 收集所有退出码，任一非零 → launcher
|
||||
非零退出并打印哪个 rank 挂了。NCCL comm 在进程退出时由 OS 回收 context（`DdpContext::Drop` 调
|
||||
`ncclCommDestroy`，正常退出路径走到；崩溃时 OS 兜底回收）。
|
||||
- **不需要跨进程 barrier**：每个 worker 独立跑完 `cfg.steps` 自然退出；NCCL collective 本身是同步点
|
||||
（所有 rank 必须到齐才返回），训练循环天然对齐。
|
||||
- **资源清理**：无临时文件（env 注入，无 `/tmp` id 文件）；ckpt 由 rank0-worker 写到 `--ckpt` 指定路径，
|
||||
与 thread-per-GPU 一致；测试用的 ckpt / 进程在测试结束清理。
|
||||
|
||||
## 验证方法（硬闸门全绿，dash5 实跑）
|
||||
|
||||
### 闸门 ①②：正确性 —— `tests/ddp_proc.rs`（`#[cfg(not(no_cuda))]`，<2 卡 skip）
|
||||
|
||||
测试本身是 launcher：用 `Command` spawn N 个 worker 进程（worker = 同测试 binary 的一个特殊模式，或复用
|
||||
`train_ddp_mp`），跑固定步数，worker 把最终 loss / 参数 dump 到各自的 stdout / 临时文件，测试父进程读回：
|
||||
|
||||
- **(a) loss 对单卡**：单卡 baseline（既有 `run_single_gpu`）vs 2-进程 / 4-进程 DDP，整条 loss 轨迹
|
||||
`max_rel < 1e-3`（与 thread-per-GPU 测试同容差）。
|
||||
- **(b) 跨 rank 参数一致**：`max|p_i - p_j| < 1e-6`（KI-5 既有约定）。
|
||||
- **(c) 对住 thread-per-GPU 路径**：同 config 同 seed，process-per-GPU 的 loss 轨迹 vs thread-per-GPU
|
||||
的 loss 轨迹应在 `<1e-3`（两者只差进程/线程，sharding+all-reduce 同）。
|
||||
|
||||
> **harness 注意**：分布式测试在共享 GPU 上并行会争卡 deadlock → 一律 `--test-threads=1`（已知 harness
|
||||
> 属性，capstone/known-issues 记过）。
|
||||
|
||||
### 闸门 ③：线性度 before→after —— `train_ddp`(thread) vs `train_ddp_mp`(process) @ {1,2,4,8}
|
||||
|
||||
固定**每卡 batch 32 / seq 256 / dim384**（与 T11 KI-5 表同口径，便于直接对比），各跑 steady-state tok/s：
|
||||
|
||||
```
|
||||
thread-per-GPU (T11 baseline) process-per-GPU (T17)
|
||||
world tok/s(global) speedup tok/s(global) speedup
|
||||
1 ~92K 1.00× ? 1.00×
|
||||
2 ~147K 1.59× ? ?
|
||||
4 ~270K 2.92× ? ?
|
||||
8 ~461K 4.99× ← 残留非线性 ? ? ← 目标更接近 8×
|
||||
```
|
||||
|
||||
8 卡跑时 `nvidia-smi` 抽样确认 8 卡 util。**资源纪律**：线性度 bench 合法地短用 8 卡，但**短跑**（每个
|
||||
world 几十~一两百步够测 steady-state），跑完清 ckpt / 中间物。
|
||||
|
||||
### 闸门 ④：全回归套（标准 `--test-threads=1`）
|
||||
|
||||
autograd / structural / batched / bf16 / recompute / overfit / AdamW / 既有 DDP loss-match + 跨 rank /
|
||||
flash / gqa / grad_accum / dropout，**+ xserv 闭环**（导出 → md5 对 registry → token-identical）。单卡与
|
||||
旧 thread-per-GPU 路径不得回归（process-per-GPU 是**新增**路径，旧路径代码未动 → 天然不回归，测试确认）。
|
||||
|
||||
### dash5 实跑
|
||||
|
||||
```bash
|
||||
export PATH=/usr/local/cuda/bin:/opt/wjh/.cargo/bin:$PATH
|
||||
# 正确性（多进程）：
|
||||
CUDA_VISIBLE_DEVICES=0,1 cargo test -p xtrain-distributed --release --test ddp_proc -- --nocapture --test-threads=1
|
||||
# 多进程训练 / 线性度 driver（process-per-GPU launcher）：
|
||||
CUDA_VISIBLE_DEVICES=0,1,2,3 cargo run -p xtrain-distributed --release --bin train_ddp_mp -- \
|
||||
/opt/wjh/models/gpt2/tokenizer.json data/tinystories-valid-3mb.txt \
|
||||
--dim 384 --heads 12 --head-dim 32 --layers 12 --ffn 1536 --steps 200 --batch 128 --seq 256
|
||||
```
|
||||
|
||||
实测数字回填见 xtrain.md T17 note / commit。
|
||||
|
||||
## 不做（本任务范围外，记 follow-up）
|
||||
|
||||
- **ZeRO-1 / sharded optimizer**：用户已 drop（本尺度 optimizer state 小、收益薄）。
|
||||
- **真·多节点 bootstrap**：本任务单节点（env 注入足够）；跨节点要 TCP rendezvous（c10d-store 式）+
|
||||
`LOCAL_RANK`/`RANK` 分离 + 每节点 `CUDA_VISIBLE_DEVICES` 切片 → 留 follow-up。
|
||||
- **NCCL 通信压缩 / overlap with backward**：与 T8/T11 同理由，all-reduce 当前非主瓶颈。
|
||||
- **删除 thread-per-GPU 路径**：保留（回归 baseline + 闸门 ①要求对齐）。
|
||||
Reference in New Issue
Block a user