docs: KI-5 re-diagnosis — all-reduce is NOT the DDP bottleneck (T11)

T11 set out to coalesce/overlap the gradient all-reduce per the original
KI-5 hypothesis. Profiling on dash5 (8× RTX 5090, dim384, per-rank batch
32, seq 256) falsifies that hypothesis:

  - grad all-reduce is only ~6-7% of each step;
  - per-rank fwd+bwd inflates ~linearly with world (136→780 ms for the
    SAME per-rank workload) and dominates;
  - coalescing the ~150 per-tensor all-reduces into one grouped/flat
    launch gives ~0 scaling gain AND breaks cross-rank bit-identity
    (max|p0-p1| 0.0 → 1.49e-8), violating the T8 correctness gate — so
    the coalescing commit (b8b5821) was reverted.

Real bottleneck (NOCOMM=1 still inflates; util shows 1-2 of 8 GPUs busy
at a time; CPU not starved; per-thread default stream doesn't help):
single-process thread-per-GPU ranks serialize on the single CUDA
context's per-op cudaMalloc / driver calls. Fix direction (out of T11
scope): a caching/pool allocator, or process-per-GPU. Recorded in
docs/known-issues.md with the measured table; KI-5 stays Open.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

crates/xtrain-cuda/src/device.rs

@@ -14,15 +14,3 @@ pub fn set_device(device: u32) -> Result<()> {
 pub fn synchronize() -> Result<()> {
     error::check(unsafe { ffi::cudaDeviceSynchronize() })
 }
-
-/// Device-to-device copy of `count` bytes (`dst <- src`) on the same GPU. Issued
-/// on the null stream (like every other xtrain kernel), so it orders with the
-/// surrounding work. Used by the DDP bucketed all-reduce to pack/unpack grads
-/// into a flat scratch buffer.
-///
-/// # Safety
-/// `dst`/`src` must point to at least `count` valid bytes of device memory on the
-/// current device, with no overlap.
-pub unsafe fn copy_d2d(dst: *mut u8, src: *const u8, count: usize) -> Result<()> {
-    error::check(unsafe { ffi::cudaMemcpy(dst, src, count, ffi::CUDA_MEMCPY_D2D) })
-}

crates/xtrain-cuda/src/ffi.rs

@@ -5,7 +5,6 @@ pub type CudaStream = *mut c_void;
 
 pub const CUDA_MEMCPY_H2D: i32 = 1;
 pub const CUDA_MEMCPY_D2H: i32 = 2;
-pub const CUDA_MEMCPY_D2D: i32 = 3;
 
 pub const CUDA_SUCCESS: i32 = 0;
 pub const CUDA_ERROR_OUT_OF_MEMORY: i32 = 2;

crates/xtrain-distributed/src/lib.rs

@@ -4,11 +4,10 @@
 //! rank thread binds its device, builds its own model (xtrain's `Var` graph is
 //! `Rc`-based and not `Send`, so it must be constructed thread-locally — only the
 //! `UniqueId` and scalar config cross the thread boundary), processes a disjoint
-//! shard of the global batch, then **coalesces every parameter's `.grad()` into a
-//! few large buckets and all-reduces each bucket once** (Phase T11 — see
-//! `all_reduce_average_grads`), averages by world size, and runs its own
-//! `GpuAdamW.step`. Identical init + identical optimizer state across ranks keeps
-//! the parameters consistent without ever re-syncing the weights.
+//! shard of the global batch, then AllReduces every parameter's `.grad()` device
+//! buffer in place, averages by world size, and runs its own `GpuAdamW.step`.
+//! Identical init + identical optimizer state across ranks keeps the parameters
+//! consistent without ever re-syncing the weights.
 //!
 //! NCCL is issued on the legacy null stream — every xtrain kernel launches on the
 //! null stream (`std::ptr::null_mut()`), so the AllReduce stays correctly ordered
@@ -27,7 +26,6 @@ use std::ffi::c_void;
 use ffi::{NcclComm, NcclUniqueId};
 use xtrain_autodiff::tape::Var;
 use xtrain_cuda::device;
-use xtrain_tensor::{Device, Tensor};
 
 pub use ffi::NcclUniqueId as UniqueId;
 
@@ -103,7 +101,7 @@ impl DdpContext {
     }
 
     /// AllReduce every parameter's `.grad()` across ranks and divide by `world`,
-    /// the one collective DDP needs per step — **coalesced (bucketed)**.
+    /// the one collective DDP needs per step.
     ///
     /// Each rank ran forward+backward on its own shard of `b` sequences, so
     /// `.grad()` holds the SUM over that shard (the tape's fan-out rule). After
@@ -114,99 +112,38 @@ impl DdpContext {
     /// mean gradient the single-GPU loop computes from a batch of `B_global`.
     /// Params without a grad are skipped.
     ///
-    /// **Coalescing (KI-5 fix, Phase T11)**: instead of one tiny `ncclAllReduce`
-    /// per parameter tensor (~150 serial launches for dim512 → DDP's dominant cost
-    /// once T10's batched forward made compute fast), pack the grads into a few
-    /// large contiguous scratch buckets and all-reduce each bucket ONCE. The packed
-    /// buffer is just the concatenation of the grad tensors, so NCCL's element-wise
-    /// sum over a bucket equals the per-tensor sums — the result is **bit-identical**
-    /// to the un-bucketed path; only the launch/latency overhead is removed. The
-    /// `1/world` average folds into one per-bucket scale. The per-bucket all-reduces
-    /// are wrapped in `ncclGroupStart/End` so NCCL fuses them into one operation.
+    /// A single-process group barrier is unnecessary: the all-reduces serialize
+    /// on the comm, and the in-place scale runs on the same null stream after.
     pub fn all_reduce_average_grads(&self, params: &[Var]) {
         if self.world == 1 {
             return;
         }
-        // Collect this step's grads (in `params()` order) and plan buckets.
-        let grads: Vec<Tensor> = params.iter().filter_map(|p| p.grad()).collect();
-        if grads.is_empty() {
-            return;
+        // 1. Sum every grad across ranks (in place, on the null stream).
+        for p in params {
+            if let Some(g) = p.grad() {
+                let n = g.numel();
+                self.all_reduce_sum_f32_ptr(g.data_ptr() as *mut c_void, n);
+            }
         }
-        let buckets = plan_buckets(&grads, BUCKET_CAP_ELEMS);
-
+        // 2. Average: scale each summed grad by 1/world (null-stream kernel,
+        //    ordered after the AllReduce that produced it).
         let inv_world = 1.0 / self.world as f32;
-        let device = Device::Cuda(self.device);
-        for bucket in &buckets {
-            let total: usize = bucket.iter().map(|g| g.numel()).sum();
-            // Flat scratch buffer for this bucket (fully overwritten by the pack
-            // below; `cudaFree` on drop synchronizes, so it outlives its copies).
-            let flat = Tensor::zeros(&[total], xtrain_tensor::DType::F32, device);
-            let flat_ptr = flat.data_ptr() as *mut u8;
-            // Pack: D2D-copy each grad into the bucket at its running offset.
-            let mut off = 0usize;
-            for g in bucket {
-                let bytes = g.numel() * 4;
+        for p in params {
+            if let Some(g) = p.grad() {
                 unsafe {
-                    device::copy_d2d(flat_ptr.add(off), g.data_ptr(), bytes)
-                        .expect("pack grad bucket");
+                    xtrain_cuda::ffi::launch_scale_inplace_f32(
+                        g.data_ptr() as *mut f32,
+                        inv_world,
+                        g.numel() as i32,
+                        std::ptr::null_mut(),
+                    );
                 }
-                off += bytes;
-            }
-            // One AllReduce(sum) over the whole bucket (fused via the group), then
-            // one scale by 1/world — same math as per-tensor, far fewer launches.
-            ffi::check(unsafe { ffi::ncclGroupStart() }, "ncclGroupStart(bucket)");
-            self.all_reduce_sum_f32_ptr(flat_ptr as *mut c_void, total);
-            ffi::check(unsafe { ffi::ncclGroupEnd() }, "ncclGroupEnd(bucket)");
-            unsafe {
-                xtrain_cuda::ffi::launch_scale_inplace_f32(
-                    flat_ptr as *mut f32,
-                    inv_world,
-                    total as i32,
-                    std::ptr::null_mut(),
-                );
-            }
-            // Unpack: D2D-copy each averaged slice back into its grad tensor.
-            let mut off = 0usize;
-            for g in bucket {
-                let bytes = g.numel() * 4;
-                unsafe {
-                    device::copy_d2d(g.data_ptr() as *mut u8, flat_ptr.add(off), bytes)
-                        .expect("unpack grad bucket");
-                }
-                off += bytes;
             }
         }
         device::synchronize().expect("grad all-reduce sync failed");
     }
 }
 
-/// Target bucket size in F32 elements (~25 MB). Big enough to amortize NCCL
-/// launch latency across many params, small enough that the scratch allocation
-/// stays modest. The exact value is not load-bearing for correctness.
-const BUCKET_CAP_ELEMS: usize = 25 * 1024 * 1024 / 4;
-
-/// Greedily group `grads` (in order) into buckets whose total element count stays
-/// under `cap` — except a single grad larger than `cap`, which gets its own
-/// bucket. Order is preserved so packing offsets are deterministic across ranks.
-fn plan_buckets(grads: &[Tensor], cap: usize) -> Vec<Vec<Tensor>> {
-    let mut buckets: Vec<Vec<Tensor>> = Vec::new();
-    let mut cur: Vec<Tensor> = Vec::new();
-    let mut cur_n = 0usize;
-    for g in grads {
-        let n = g.numel();
-        if cur_n > 0 && cur_n + n > cap {
-            buckets.push(std::mem::take(&mut cur));
-            cur_n = 0;
-        }
-        cur.push(g.clone());
-        cur_n += n;
-    }
-    if !cur.is_empty() {
-        buckets.push(cur);
-    }
-    buckets
-}
-
 impl Drop for DdpContext {
     fn drop(&mut self) {
         if !self.comm.is_null() {

docs/known-issues.md

@@ -7,13 +7,26 @@
 
 ## Open
 
-_(none — KI-1 fixed in T10; see Fixed below. KI-5 is the DDP-scaling follow-up batching newly exposed.)_
+_(KI-1 fixed in T10. KI-5 仍 Open，但 T11 实测把根因从「all-reduce 未分桶」**改诊断**为「单进程多 rank 的逐 rank compute 互相串行」——见下。原拟修复（分桶 all-reduce）经实测证伪。)_
 
-### KI-5 · DDP 弱扩展性（all-reduce 每步全参数，未分桶 / 未与 backward overlap）— `P2` · 由 T10 暴露
+### KI-5 · DDP 弱扩展性 — `P2` · 由 T10 暴露，T11 重新诊断（all-reduce **不是**瓶颈）
 - **现象**：batched forward 修掉单卡 launch-bound 后，dim384/per-rank batch 32：1 卡 40.3K → 4 卡 47.2K tok/s（global），仅 ~1.17×。
-- **根因**：单卡 compute 快了 15–24× 后，每步对全部 ~67M 参数的 **eager all-reduce + host 侧 optimizer/clip 同步**成了 DDP 主导开销，不随卡数缩小。注意**单卡 batch 32 = 40K tok/s 已是 KI-1 时代 4 卡(3163)的 ~12×**——根因已修，这是新的、更上层的瓶颈。
-- **拟修复**：梯度 **bucketed all-reduce + 与 backward 计算 overlap**（即 KI-1 修复项 2，此前 all-reduce 非瓶颈做了无收益，batched 之后才有意义）。可选：optimizer/clip 进一步去 host 同步。
-- **重启条件**：v3 训练若被 DDP 扩展性卡住再做；单卡吞吐已足够，v3 可先单卡 / 小 world 跑。
+- **T11 实测（dash5, 8× RTX 5090, dim384/12L, per-rank batch 32, seq 256, 原 ungrouped all-reduce, 50 步均, ms/step）**：
+
+  | world | fwd+bwd | grad all-reduce | clip+opt+zero | TOTAL | tok/s(global) | speedup |
+  |---|---|---|---|---|---|---|
+  | 1 | 136 | 0 | 8.6 | 145 | 36582 | 1.00× |
+  | 2 | 202 | 21 | 15 | 238 | 47267 | 1.29× |
+  | 4 | 342 | 29 | 21 | 392 | 51466 | 1.41× |
+  | 8 | 780 | 54 | 47 | 882 | 47719 | 1.30× |
+
+  → grad all-reduce 每步只占 **~6–7%**；真正爆炸的是**逐 rank 的 fwd+bwd 时间随 world 线性膨胀**（同一 per-rank workload，136→780ms，~6×）。
+- **「分桶 all-reduce」拟修复经 T11 实测证伪**：
+  - ① 把 ~150 个 per-tensor `ncclAllReduce` 用 `ncclGroupStart/End` 融成一发 → 1/2/4/8 卡 = 1.00/1.30/1.42/1.34×，**与不分桶几乎无差**（因为 all-reduce 本就只占 7%）。
+  - ② 分桶/grouped/flat 还会**破坏跨 rank bit-identical**：correctness 测试里 `max|p0−p1|` 从 `0.0` 变 `1.49e-8`（1 ULP，逐步 AdamW 累积）——NCCL 只对**单个 ungrouped collective** 保证跨 rank 逐位一致，grouped/大 message 会换 algorithm/chunking 扰动结果。**违反 T8 硬闸门**，故保留原 ungrouped 路径。
+- **重新定位的根因**：**单进程 thread-per-GPU 模型下，N 个 rank 线程各自跑独立训练却互相串行**——`NOCOMM=1`（完全不做任何跨 rank 通信/barrier）时 fwd+bwd 仍 136→378→800ms 膨胀；`nvidia-smi` 抽样显示 8 卡同一时刻只有 1–2 张在忙、轮流跑。排除项：CPU 不缺（187 核, load 2.5）；`nvcc --default-stream per-thread` 不解决。**剩余怀疑：每个 op 输出走 `Tensor::zeros`→`cudaMalloc`+`cudaMemset`，而 `cudaMalloc` 是同步、进程级串行的 driver 调用；单 CUDA context 下 N rank 每步几百次 alloc 互相排队**——即 DDP 真瓶颈是 **per-op 显存分配 / driver 调用在单进程内串行**，不是梯度通信。
+- **真正的修复方向（待定，非 T11 范围）**：① **caching/pool allocator**（op 输出复用显存，消掉每步几百次 `cudaMalloc`，单卡也受益）；或 ② **process-per-GPU**（每 rank 独立 CUDA context，torchrun 式，彻底解串行，但要改 launcher + 跨进程 UniqueId 分发）。先做 ① 再实测是否解 DDP 串行。
+- **重启条件**：多卡 v4 需要扩展性时做。**单卡 batched 已 40K tok/s（v3 即单卡训完）**，多卡当前只有 ~1.4× 上限，v4 若要多卡须先修上面的真瓶颈。
 
 ---