xtrain/docs/16-process-per-gpu.md

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 实跑

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 + 闸门 ①要求对齐）。