Records the key empirical finding: process-per-GPU is statistically identical to thread-per-GPU at this scale (thread 5.27x vs proc 5.31x @8, <1% noise; all 8 GPUs 95-99% util). The residual ~5.3x@8 non-linearity is the NCCL/PCIe communication wall, NOT single-CUDA-context launch/cuBLAS serialization as the old KI-5/T11 note speculated — measurement falsifies that hypothesis (same methodology as T11 falsifying "bucket the all-reduce"). Correctness all green: proc==thread loss 1.5e-7, cross-rank 1.2e-7, full regression + xserv md5 b04fc9f9 identical. Closes the process-per-GPU backlog item (measured no-op); default training path unchanged. evolution.md Infra row + README T17 row + known-issues entry. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
267 lines
19 KiB
Markdown
267 lines
19 KiB
Markdown
# 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)。
|
||
|
||
> ⚠️ **此假设被实测证伪**——见下方「实测结果 · 闸门 ③」:process-per-GPU 与 thread-per-GPU 吞吐统计上一致
|
||
> (~5.3×@8 都一样),且 8 卡全 95–99% util。残留非线性是通信/PCIe 墙,不是单 context 串行。结论钉死、留档。
|
||
|
||
### ④ 训练 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
|
||
```
|
||
|
||
实测数字见下方「实测结果」。
|
||
|
||
## 实测结果(dash5, 8× RTX 5090, sm_120)
|
||
|
||
### 正确性(闸门 ①②④ 全绿)
|
||
|
||
- **闸门 ① loss 对单卡 / 对 thread 路径**(`ddp_proc`, world=2,合成语料 20 步):
|
||
proc-per-GPU vs single-GPU `max_rel = 5.67e-7`;**proc-per-GPU vs thread-per-GPU `max_rel = 1.5e-7`**
|
||
(两条路径数值同量级,符合预期——只差进程/线程,sharding+all-reduce 同)。
|
||
- **闸门 ② 跨 rank 参数**:`max|p0−p1| = 1.19e-7`(< 1e-6,KI-5 既有 ULP 容差,PCIe NCCL run-to-run 抖动)。
|
||
- **闸门 ④ 全回归**:全 workspace `--test-threads=1` 全绿(autograd/structural/batched/bf16/recompute/
|
||
overfit/AdamW/既有 DDP/flash/gqa/grad_accum/dropout)+ **xserv 闭环**:v3 ckpt 用 T17 代码重导
|
||
safetensors 与 registry **md5 逐位一致 `b04fc9f9a0c9af04c47d9ca649aea12e`**(T17 不碰任何数值路径 → 必然一致)。
|
||
|
||
### 闸门 ③ 线性度 before→after —— **本任务的关键发现:process-per-GPU 在本尺度对吞吐中性**
|
||
|
||
固定每卡 batch 32 / dim384 / seq256 / 150 步(与 T11 KI-5 表同口径),steady-state tok/s:
|
||
|
||
| world | thread-per-GPU (`train_ddp`) | speedup | process-per-GPU (`train_ddp_mp`) | speedup |
|
||
|---|---|---|---|---|
|
||
| 1 | 93257 | 1.00× | 92952 | 1.00× |
|
||
| 2 | 149747 | 1.61× | 148809 | 1.60× |
|
||
| 4 | 278276 | 2.98× | 273308 | 2.94× |
|
||
| 8 | **491360** | **5.27×** | **493128** | **5.31×** |
|
||
|
||
(world=8 重复 2 次确认非噪声:thread 493671/493292,proc 491102/494123——**两路差异 < 1%,落在 run-to-run 噪声内**。)
|
||
|
||
→ **process-per-GPU 与 thread-per-GPU 吞吐统计上一致(~5.3×@8 都一样)**。本 doc 设计假设 ③
|
||
(「残留 5×@8 来自单 CUDA context 的 kernel-launch/cuBLAS 串行,process-per-GPU 给独立 context 即可移除」)
|
||
**被实测证伪**——这正是 ③ 里预留的「诚实原则」分支。
|
||
|
||
**根因重定位(实测佐证)**:proc-per-GPU world=8 跑时 `nvidia-smi` 抽样 **8 卡全部 95–99% util**
|
||
(每卡 ~23GB)——GPU **已 compute-bound 喂满、并非串行空转**(KI-5 当年「1–2/8 在忙」的串行病在 T11 的
|
||
caching allocator 就已治好)。8 卡已满载却仍只 5.3×,缺的 ~35% 吞吐只能去向**每步 grad all-reduce +
|
||
本机 PCIe-only 拓扑在 8 rank 下的通信开销**——即 T11 早已点明的「~7% all-reduce + 8 卡 PCIe 余量」那一层,
|
||
在 8 卡下被放大。换独立 context 不动这一层,故吞吐不变。
|
||
|
||
**这与 T11 自身的方法论一致**:T11 实测证伪了「分桶 all-reduce」;T17 实测证伪了「process-per-GPU 解残留
|
||
串行」。两次都靠 profile/measure 推翻假设而非硬上。**结论**:本尺度(dim384–1024、单机 8× PCIe RTX 5090)
|
||
残留非线性是**通信/拓扑墙**,不是 launch 模型;要再逼近线性得动 all-reduce overlap / 更快互联(NVLink),
|
||
那是另一条线,**非 T17 scope**。
|
||
|
||
**T17 的净价值(诚实记账)**:① 学到 / 落地了 torchrun 式 process-per-GPU 这条训练栈标准链路(独立进程 +
|
||
独立 CUDA context + 跨进程 NCCL bootstrap)——**项目本职「学训练全栈」的目标达成**;② **实测把「process-per-GPU
|
||
是残留非线性的解」这个长期挂在 KI-5/T11 doc 里的猜想钉死为「在本尺度无吞吐收益」**,移除一个误导性 backlog
|
||
项;③ 正确性零回归、与 thread 路径数值对齐。**吞吐上它与 thread-per-GPU 等价**——故默认训练路径**不变**
|
||
(thread-per-GPU 仍是 v1–v8 用的那条),process-per-GPU 作为并列可选路径 + 这条诊断结论留档。
|
||
|
||
## 不做(本任务范围外,记 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 + 闸门 ①要求对齐)。
|