diff --git a/docs/16-process-per-gpu.md b/docs/16-process-per-gpu.md new file mode 100644 index 0000000..3946f18 --- /dev/null +++ b/docs/16-process-per-gpu.md @@ -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=`,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 + 闸门 ①要求对齐)。