Repository navigation
Conversation
…NER_DETERMINISTIC CI glm5-3-flash-sft-25B-vl double runs at ebf73a2 (same image, same code) diverge in grad_norm from step 1 with bitwise-identical loss. Op-level probe shows the cuDNN sparse-attention backward accumulates dKV with unordered atomics: forward and dq are bitwise stable across runs, dkv is not. The tiny dkv noise amplifies chaotically over steps (loss diff reaches 3.03 by step 14 in the VL task). When XTUNER_DETERMINISTIC=true, route the backward to an ordered FP32 reference implementation: fixed-order query chunks, gather + einsum recomputation of the softmax backward, and dkv accumulation via index_add_ (deterministic under torch.use_deterministic_algorithms). Verified bitwise repeatable across processes and within 3e-3 relL2 of the cuDNN backward; the default path is unchanged.
… segment reduction in sparse-MLA backward PyTorch's deterministic index_add_ pays an internal sort + serial reduction (~1.5s, ~64GB workspace at GLM-5.3-Flash's DSA shape: total_q=16384, topk=2048, dim=512). Stage per-chunk dKV contributions in a bf16 buffer, sort by kv index once per 2048-query flush, and reduce each kv's segment serially in sorted order with a Triton kernel into an FP32 accumulator. The fixed kv-major + in-segment order keeps the backward bitwise repeatable while the reduce itself drops to ~20ms/layer (total backward 1900ms -> 540ms). -1 padding rows are dropped before the reduction instead of being clamped into segment 0: they are exact-zero contributions, and clamping them formed a single ~2.1M-row monster segment that one program had to sum serially (~436ms on the first flush). Falls back to deterministic index_add_ when triton is unavailable or head_dim is not a power of two.
… split-slice accumulation The cuDNN SM90 DSA backward accumulates dKV with cross-CTA unordered atomics. This adds cudnn_dsa_det.py: a multi-launch driver (128-query launches over a widened FP32 accumulator) plus a subclass of the stock CuTe-DSL kernel that overrides only scatter_dkv_atomic to offset each CTA into its own slice; slices are reduced in fixed order afterwards. Bitwise repeatable across processes at 57ms per DSA-layer backward (stock 36ms), dq bitwise identical to stock, dkv relL2 5.1e-5. XTUNER_DETERMINISTIC=true prefers this path; the einsum+Triton reference backward (deterministic_bwd.py) remains as fallback. No vendored kernel copy: the subclass reuses the installed cudnn package's kernel.
… selector Removes the deterministic index_add_ fallback path: the reference backward now always uses the sort + Triton ordered segment reduction and asserts when triton or a power-of-two head_dim is unavailable. The planned XTUNER_GLM53_DSA_DET_BWD selector knob is reverted; dispatch stays a simple capability chain: XTUNER_DETERMINISTIC -> fork (SM90+dim512) -> Triton reference -> assert.
…eterministic path Removes deterministic_bwd.py (Triton ordered segment reduction) and the einsum reference backward: XTUNER_DETERMINISTIC=true now always uses the split-slice cuDNN fork and raises if its runtime is unavailable (SM90 + CuTe-DSL), instead of falling back to a slower alternative implementation.
…_backward_op The XTUNER_DETERMINISTIC branch now lives inside the shared backward op (cudnn_dsa.py) instead of flash_mla_cudnn.py's autograd entry: flash_mla_cudnn.py returns to main unchanged, and glm5.2's cudnn_dsa backend gets the deterministic path for free since it routes through the same op.
…c driver The fork driver always runs sink-free (GLM DSA passes attn_sink=-inf), so the sink validation, state variables, compile-key entries and return variants are dead code. No behavior change: tests still pass.
Review summary感谢这个实现。整体方向我认为是合理的：复用 cuDNN DSA backward 的主体，只通过一个很窄的子类覆写 我也认可目标 workload 上的结果：提供的 41 步对比中，两次 deterministic 训练的 loss、grad norm、LR、token、TGS 和 memory 曲线重合。我随后在同一 PR head、同一 CI 配置和 8×H200 上独立复跑，也确认两次 deterministic 的非计时训练指标及显存记录逐项一致，说明方案确实解决了目标确定性问题。 不过，基于当前 head P1: slice ID 只包含
|
| Arm | Peak allocated | 相对调用前的峰值增量 |
|---|---|---|
| stock | 4.0938 GiB | 1.0626 GiB |
| deterministic | 8.0938 GiB | 5.0626 GiB |
deterministic 的额外峰值恰好是 4 GiB，与上述 128-slice 公式一致。
按同一公式外推到 S_k=524288, D=576：
- 主 workspace 为 154,618,822,656 bytes，即 144 GiB；
- reduced accumulator 还需要额外 1.125 GiB；
- 如果
S_q也是 512K，当前实现还会执行 4,096 次 main launch。
提供的 41 步 CI 图显示，目标约 16K-token workload 的 deterministic memory 大约为 95–98 GB，stock 大约为 87–90 GB，说明当前目标形状可以容纳这份开销。但通用 dispatch 接受的长序列形状可能无法运行，因此建议不要只以目标 CI 形状作为显存契约。
建议至少选择一种处理方式：
- 实现 memory-bounded 的 slice/chunk 策略，并让 slice 数与实际并发 CTA/grid 一致；
- 按 KV tile 做分块累加和归约，避免保存“每 CTA × full KV”的 workspace；
- 使用另一种不保留全部 per-CTA full-KV accumulator 的确定性两阶段归约；
- 如果暂时不准备支持长序列，明确记录 deterministic 路径的显存契约和支持范围，并对必然超过资源的形状给出清晰错误，而不是静默 OOM。
需要特别注意：不能只把 n_splits 调小、同时保留原来的并发 CTA 数量，否则多个 CTA 会再次共享 slice，与第一个 P1 冲突。workspace 策略和 CTA-to-slice 映射最好一起修改、一起测试。
P2: deterministic 分支缺少对应的 runtime capability check
相关位置：
xtuner/v1/ops/sparse_mla/cudnn_dsa.py:112-130xtuner/v1/ops/sparse_mla/cudnn_dsa.py:172-188xtuner/v1/ops/sparse_mla/cudnn_dsa_det.py:10-23xtuner/v1/ops/sparse_mla/cudnn_dsa_det.py:75-76xtuner/v1/ops/sparse_mla/_cudnn_dsa_bwd_det_subclass.py:11-13
当前 availability probe 主要覆盖 stock wrapper 和 SM90+，但 deterministic 分支还依赖 cuDNN private class、CuTe helper、特定的内部调用签名，并且实现本身针对 SM90。子类注释提到了 import/capability fallback，但实际 dispatch 中没有看到对应的 guard。
建议增加 deterministic-specific capability probe，至少覆盖：
- private imports 和所需方法是否可用；
- 已验证的 cuDNN API/version；
major == 9，而不只是major >= 9；- dispatch 与 test skip 共用同一判断，避免测试环境和实际运行路径不一致。
不支持时应该给出清晰错误，或者走另一个仍然确定性的实现；不建议静默回退到 stock backward，因为调用方设置 XTUNER_DETERMINISTIC=true 后，通常会把“启用成功”理解为确定性契约已经生效。
P2: 建议把已经验证的关键边界合入 CI
当前新增测试使用 64 queries、64 heads，因此只覆盖一次 main launch，也没有触发 qhead_tiles > 1。建议补充：
S_q=64/128/129/257，覆盖 single-launch、完整 chunk 和 partial final chunk；- 64 和 128 heads，后者用于防止只按
blockIdx.x分配 slice 的回归； - 512/576 两种 head dim，576 目前只有数值回归，缺少跨 launch bitwise 覆盖；
- interspersed
-1、duplicate valid indices，并同时比较 stock/det 的数值语义； - workspace/slice 映射的小型可执行测试，以及至少一个代表性长上下文内存检查。
此外，tests/conftest.py:9 默认设置了 XTUNER_DETERMINISTIC=true。这意味着既有通用测试默认不再保护 stock backward 分支。建议共享测试显式参数化 deterministic=false/true，避免 deterministic 成为默认后，stock 路径失去原有回归覆盖。
P3: docformatter 检查尚未通过
我在当前 head 上执行的结果是：
ruff check: passed
ruff format --check: passed
py_compile: passed
git diff --check: passed
docformatter --check --wrap-descriptions 119: failed
docformatter 涉及：
xtuner/v1/ops/sparse_mla/_cudnn_dsa_bwd_det_subclass.pyxtuner/v1/ops/sparse_mla/cudnn_dsa_det.pytests/ops/test_flash_mla_cudnn_sparse_mla.py
按项目 CI 规则格式化后重跑即可。
我这边的验证范围
验证环境为 H200、PyTorch 2.9.1+cu128、cuDNN 1.28.0，代码为上述 PR head。
- PR 新增 deterministic 测试：2 passed；
- 既有 576-wide cuDNN 数值回归：1 passed；
- 额外 9 组 bitwise 探针，每组重复 30 次：全部逐位一致；
- 覆盖 64/128 heads、single/multi-launch、129/257/513 queries、duplicate/holes 和不同 head-tile 梯度量级；
- 固定同一份 FlashMLA forward 输出和 LSE，只切换 stock/deterministic backward：unique、duplicate、duplicate+holes 均通过既有数值容差；
- 16K 微基准中，两组交替 A/B 的算子时间分别为 8.33→11.43 ms、10.72→11.37 ms。这个数字只代表该算子、该形状，不外推为整步性能结论。
41 步完整训练独立复跑
我使用原始 autotest/config/glm53_25b_vl.py、相同模型和 ci_vl 数据，在同一源码快照上依次运行 deterministic A、deterministic B 和 stock A。三组均自然完成 41/41 步，退出码均为 0；除 XTUNER_DETERMINISTIC 外，配置和环境开关保持一致。
确定性结果通过：
- 8 个 rank × 41 步，共 6,232 个非计时训练指标值，deterministic A/B 全部完全相等，0 项差异；包括 loss、grad norm、LR、token 和负载统计；
- 两次 deterministic 的 allocated/reserved memory 记录也逐步、逐 rank 完全相等；
- 这里排除了
time/*和 TGS 等墙钟/吞吐字段，不把性能数值相同冒称训练语义相同。
性能和显存采用第 11–41 步作为共同稳态窗口。step_time 先对同一步的 8 个 rank 取最大值，再取窗口中位数；全局 TGS 为各 rank 唯一 text token 之和除以同一个 max-rank step-time。显存为第 41 步时各 rank 历史峰值的最大值，单位按源码实际口径记为 GiB：
| Arm | max-rank step-time median | 全局 TGS median | peak allocated | peak reserved |
|---|---|---|---|---|
| deterministic A | 3.5729 s | 33,222.9 | 94.694 GiB | 133.623 GiB |
| deterministic B | 3.7609 s | 33,031.4 | 94.694 GiB | 133.623 GiB |
| stock A | 3.2479 s | 38,555.0 | 94.275 GiB | 132.469 GiB |
在这组同源码、同环境对照中：
- deterministic A/B 的 step-time 分别比 stock 高 10.0% / 15.8%；
- 全局 TGS 分别低 13.8% / 14.3%；
- deterministic 的最大 allocated/reserved 分别高 0.419 / 1.154 GiB，没有复现描述中的 +7 GB peak memory。
目前只有一组 stock，运行间波动也客观存在，所以我不把这些数字直接定性为普遍性能回归；但它们足以说明“TGS 基本不降”不能作为当前实现的通用结论。建议作者核对自己的 stock 基座、源码 commit、环境、统计窗口和 TGS 口径，最好补充同源码开关切换的重复或交错 A/B。这个性能测量差异不改变上面的确定性复现结论，也不改变两个 P1 的代码问题。
总体而言，我认为这个方案值得继续推进。优先把完整 CTA grid 的 slice ownership 和 workspace 策略处理好，再补 capability guard 与对应回归，就能让实现从“目标形状实测稳定”变成“确定性不变量和资源边界都可检查”。
Motivation
CI task
glm5-3-flash-sft-25B-vl(added in #2215) fails its grad_norm assertion: two runs at thesame commit (
ebf73a232), same image, same data,XTUNER_DETERMINISTIC=truediverge. Tracker-levelevidence (
qa-llm-cicd/test_output/{40617516054,39953013246}):loss/reduced_llm_lossis bitwise identical (diff = 0.00e+00), butgrad_normalready differs at step 1 (23.309881 vs 23.308867) → forward and data are deterministic; the
first backward injects the noise.
Op-level bisection (GLM-5.3-Flash-25B DSA geometry: NoPE, heads=64, head_dim=512, topk=2048, two
separate processes): forward output and dq are bitwise stable; dkv is not (sha256 differs).
Source: cuDNN's
dsa_bwd_sm90kernel scatters per-CTA dKV contributions into one shared FP32accumulator with
atomic_add_fp32x4— the cross-CTA completion order differs every run. dq has noshared accumulation and is deterministic.
This is the same factor known from GLM-5.2 (whose CI config selects the
tilelangbackend, whichhas the same class of issue in another kernel). GLM-5.3's NoPE 512-wide geometry is only supported
by
flash_mla_cudnn, so there was no escape hatch. Full analysis, alternatives measured, anditeration history: see the design doc linked below.
Change (final form: subclassed cuDNN kernel + split-slice ordered reduction)
When
XTUNER_DETERMINISTIC=true, the backward takes a deterministic path; the default path isunchanged. No vendored kernel copy — everything subclasses/reuses the installed cudnn package:
xtuner/v1/ops/sparse_mla/_cudnn_dsa_bwd_det_subclass.py(96 lines): subclass of the stockCuTe-DSL kernel overriding only
scatter_dkv_atomic(+2 lines: offset each CTA's scatter targetby
blockIdx.x % n_splits * slice_sizeinto a widened FP32 accumulator, so CTAs never shareaddresses). FA2's deterministic backward uses the same split-buffer + ordered-reduction pattern.
xtuner/v1/ops/sparse_mla/cudnn_dsa_det.py(526 lines): phase-controllable interface +multi-launch driver (
flash_attn_bwd_sm90_multilaunch): preprocess once, 128-query launcheswhose slices persist, then a fixed-order slice reduction and one postprocess.
cudnn_dsa.py(+27/-2): theXTUNER_DETERMINISTICdispatch lives inside the shared backwardop
_cudnn_dsa_sparse_mla_backward_op, so every backend routing through it gets thedeterministic path (
flash_mla_cudnnand glm5.2'scudnn_dsa);flash_mla_cudnn.pyisbyte-identical to main.
tests/ops/test_flash_mla_cudnn_sparse_mla.py: bitwise-repeatability + torch-referenceequivalence tests.
Iteration history (kept in this branch for the record): v1 einsum reference backward with
deterministic
index_add_(worked, but 1.9s/layer), v2 Triton ordered segment reduction(540ms/layer), both removed in favor of the fork; the final diff only adds the fork.
Validation
(reduction-order level), bitwise repeatable across processes, 57ms/layer (stock 36ms).
XTUNER_DETERMINISTIC=true):all training metrics bitwise identical across runs (only wall-clock fields differ);
median step time 3.5s vs ~2s stock (tgs nearly unchanged), +7GB peak memory for the slices.
Impact / notes
XTUNER_DETERMINISTIC=trueruns change behavior; production default is untouched.scatter_dkv_atomic; a cudnn upgrade that changes thatmethod's internals must be handled by guarding/falling back (noted in the file docstring).
cudnn_dsa/flash_mla(tilelang bwd) backends share the same unordered-atomic dKVissue and are intentionally out of scope.