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>
7.9 KiB
xtrain — Known Issues & Perf Backlog
已知问题(性能 / 正确性 / 建模)与延后项的活文档:记录现象、复现、根因、拟修复、优先级、状态。
发现即记,修复即标 FIXED(附 commit)。
Open
(KI-1 fixed in T10. KI-5 仍 Open,但 T11 实测把根因从「all-reduce 未分桶」改诊断为「单进程多 rank 的逐 rank compute 互相串行」——见下。原拟修复(分桶 all-reduce)经实测证伪。)
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×。
-
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 路径。
- ① 把 ~150 个 per-tensor
-
重新定位的根因:单进程 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 若要多卡须先修上面的真瓶颈。
Fixed
KI-1 · 单序列 launch-bound("DDP 弱扩展性"的根因)— FIXED (T10, batched forward)
-
修复:T10 给 model + autograd 加 batch 维——linears 摊平成
[B*S, dim]一个大 GEMM 填满 GPU;attention 走 fused 批量 SDPA(cublasSgemmStridedBatched×2 + 一个 causal-softmax kernel),RoPE 位置 per-sequence 复位(row % S);训练 loop 用真 batch 一次 forward/backward 替代 "loop B 次 + SUM"。详见 docs/09-batched-forward.md。 -
before → after(dim384/12L/12h, batch 16, seq 256, 1 卡, back-to-back A/B):
tok/s GPU util 显存 before(单序列 launch-bound) ~1653 0–15% ~3 GB after(batched) 25627(batch16)/ 40263(batch32) 37% 均值 / 54% 峰 ~10 GB → 单卡 ~15.5×(batch16)/ ~24×(batch32),util 0–15% → 37–54%。
-
正确性(全绿,无回归):15 算子 grad-check(新增 batched-rope / transpose_4d12 / batched-attention dQ/dK/dV)、batched==looped 单序列(logits 0.0、grad 6.4e-4)、PyTorch 对拍 B>1(loss 5e-8 / logits 6.9e-6 / 全参数 grad 在 rtol 2e-2)、overfit 27/27、checkpoint 逐位、AdamW 对 torch、DDP loss 对单卡 5.7e-7 + 跨 rank 参数 bit-identical(0.0)、xserv 加载导出权重对 xtrain 贪心仍逐 token 一致(top token 同序、BF16 漂移 ~0.03)。
-
commit:见 T10 提交链(
perf: KI-1 fixed — GPU util / tok/s那条带 before/after)。 -
DDP 残留弱扩展性 → KI-5(这是 batching 后新暴露的 all-reduce 瓶颈,不是 KI-1 的单序列根因)。
-
历史诊断保留如下(v2 暴露 → v3 重诊断的过程,证明根因不是 all-reduce):
KI-1 历史诊断 · DDP 弱扩展性(吞吐受单序列 launch-bound 限制)— v2 暴露,v3 重新诊断
- 现象:4 卡 DDP 仅 ~3.2K tok/s,几乎不快于单卡(≈2× over 单卡,远低于近线性;T8 在 tiny micro-bench 为 3.0×@4)。
- 复现:
dim384/12L, world=4, seq 256。 - v3 实测(dash5, 4× RTX 5090, dim384, 隔离 back-to-back A/B):
global_batch 每卡 tok/s(4卡) GPU util 显存 32 8 3163 5–69%(spiky) ~2–3 GB / 32 GB 256 64 3200 0–15% ~2–3 GB / 32 GB → 加大 8× batch 仅 +1.2% 吞吐(噪声内)。1 卡 dim384 ≈ 1653 tok/s,4 卡 3163 ≈ 2.1×。 - 原"拟修复"(加大 global batch)经 v3 实测 falsified:gbatch256 时每 token 的 all-reduce 次数只有 gbatch32 的 1/8,若瓶颈是 all-reduce 应大幅提速——实际没有 → all-reduce / 通信不是瓶颈。
- 重新诊断的根因:瓶颈是单序列模型设计(T5:每个 sequence 各跑一次独立 forward/backward,逐 op kernel-launch 开销,见 docs/06 延迟瓶颈)。GPU util 仅 0–15%、显存仅占 ~8% → 严重 launch-bound / under-utilized;GEMM 太小喂不饱 GPU。加大 batch 只是按比例增加串行 launch 次数,无法摊薄。4 卡相对单卡 ~2× 的固定天花板来自跨 rank 同步税,但不是靠调 batch 能修的。
- 真正的修复(需实作,非调参):
- batched(多序列)forward——把一个 step 的多条序列在 batch 维一次性过模型,让 GEMM 大到能填满 GPU(这是 launch-bound 的根本解,但要改 T4/T5 的 single-sequence autograd/model,工作量大、有正确性风险);
- 在 (1) 之后,梯度 all-reduce 分桶 + 与 backward 重叠(bucketed / overlapped all-reduce)才会有意义(当前 all-reduce 已非瓶颈,做了也无收益)。
- 参考:docs/07-distributed.md、docs/06-performance.md。
Deferred(来自 T7,放大后重启)
KI-2 · bf16 混合精度(fp32 master)— deferred
- T7 延后理由:tiny 规模延迟瓶颈、bf16 改变数值会威胁 fp32 正确性闸门。
- 重启条件:模型放大(v2+
dim≥384)后 GEMM 渐成 compute-bound,tensor-core 收益显现。需 fp32 master weights + 单独 looser-tol 测试 + 收敛对比。
KI-3 · 激活重计算(gradient checkpointing)— deferred
- T7 延后理由:单序列、显存不紧。
- 重启条件:更大模型 / 更长 seq / 更大 batch 后显存成约束。
Modeling notes
KI-4 · 大词表 embedding 占比过高
- gpt2
vocab=50257在 dim 小时让 embed+lm_head 主导参数:v1 25.7M/34M、v2 38.6M/66.9M;core transformer 才是学习主体。 - 后续可考虑更贴合 TinyStories 的小 vocab(会牺牲 xserv gpt2-tokenizer 复用);或在更大 dim 下让 core 自然成为主体(继续 scaling 即可缓解占比)。