[qescccczmr]

Problem

GLM-5.3 KDA training's BF16 chunk solve/state intermediates differ from serving's FP32 recurrence. On the same real layer-0 Q/K/V/gate/beta inputs (64 heads x 128 dimensions, 129 fixed tokens), the training core output has relative L2 difference 0.00459635426159536 versus actual LMDeploy recurrence. Gate activation was independently checked; it is not the source of that discrepancy in this fixture.

Change

Add the opt-in startup setting:

export XTUNER_GLM53_KDA_FP32=1

For the compile-friendly chunk_kda path, promote Q/K/V before L2 normalization, keep normalized vectors, chunk solve/state intermediates and backward recomputation in FP32, and cast only the visible output and input gradients back to the original activation dtypes. The fake implementations reflect the changed intermediate dtypes. The existing fixed-layout gate prefix, packed offsets, parameters and native FLA backward remain in place. No compression backward, new model class or second inference model is introduced.

Unset/0 retains the original path. Set before importing XTuner in all training workers; for Ray, propagate it in runtime_env.env_vars. XTUNER_HF_IMPL=1 selects FLA's own entry point and does not use this option. This change does not add initial/final state or context-parallel support to this wrapper.

FP32 intermediate storage is not a guarantee of IEEE FP32 GEMM: the installed FLA kernels may use TF32. This PR improves chunk precision; it does not claim bitwise equivalence to serving recurrence.

Validation

On 10.102.215.14 / H200, PyTorch 2.9.1+cu128, FLA 0.5.2, TileLang 0.1.11:

• The real-model same-input replay improves relative L2 difference from 0.0045963543 to 0.0026294753, a 42.79% reduction. Values are finite and nonzero gradients are produced for Q/K/V/gate/beta. This is core-output evidence, not an end-to-end logprob/KL measurement.
• 3 passed: independent FP64 recurrence and VJP checks, including packed lengths 37/28, a chunk boundary at 129, and the production 64-head x 128-dimension geometry at 65 tokens. Forward error must improve by at least 10% over the original; all five nonzero input gradients have relative VJP error below 2%. A planted 25% gradient scaling error must fail.
• Changed-file Ruff, formatting, syntax and diff checks pass.
• A broader dependency regression produced 11 passed / 14 failed. The first failing pre-existing rectangular-state case (K=16, V=24) raises TileLang's no available layout found; the same case and error were reproduced on the unmodified main KDA implementation. This is recorded as a dependency/compiler limitation, not a clean whole-suite or zero-added-failures claim. The real 128 x 128 geometry succeeds in the new independent VJP check.

CUDA-event timing of CUDA Graph forward replay, 50 samples after warmup (64 heads x 128 dimensions):

Tokens	Original p50 / max (us)	FP32 p50 / max (us)
129	71.712 / 94.144	112.096 / 116.960
2048	550.080 / 594.048	1039.648 / 1048.640
16384	4050.976 / 4117.728	7842.400 / 7882.624

This option increases memory/compute; the 16K core is about 1.94x slower. These are operator forward timings, not full training/serving throughput. The option remains off by default.

Scope / remaining work

This draft is intentionally limited to a measured precision option. The remaining recurrence difference is significant and still under diagnosis. A combined five-step reduced-25B MoE-only GRPO run completed with BF16 KPool alignment, FP8 activation rounding and the separate mHC/SwiGLU patches, prefix caching and MTP5. The end-to-end mismatch/mismatch_k3_kl <= 2e-4 target is not yet established and is not inferred from operator relative L2 statistics.

Five-step combined numerical check (2026-11-11, H200)

A fresh five-step check completed on 10.102.215.14 / 8 x H200, exit code 0, using the five-layer reduced GLM-5.3 Flash 25B checkpoint, GSM8K, EP4/SP1 training and TP1/EP4 rollout, prefix caching and MTP with num_speculative_tokens=5.

This is a combined recipe: BF16 KPool pooling/head alignment, FP8 MoE activation rounding, opt-in FP32 KDA chunks, and the separate mHC-post/SwiGLU alignment patches. The eight startup variables FP8, ROLLOUT_FP8, LMDEPLOY_FP8_MOE_ONLY, XTUNER_GLM53_KPOOL_BF16, LMDEPLOY_GLM53_KPOOL_BF16, XTUNER_GLM53_LMDEPLOY_SWIGLU, XTUNER_GLM53_KDA_FP32, and LMDEPLOY_MLA_DISABLE_DENSE_PREFILL were all 1. The actual environments of 8 training workers, 8 inference-compute workers and their engine/controller wrappers were checked; the prefix-cache replay probe also passed.

Step	mismatch/mismatch_k3_kl
1	2.6386266109e-04
2	2.6337933377e-04
3	2.6362264180e-04
4	2.6350014377e-04
5	2.6471854653e-04
Mean	2.6381666539e-04

The end-to-end <=2e-4 target was not reached (all five values remain above it). TensorBoard scalar values were cross-checked with the worker's full-precision log values. Local KPool score relative L2 at 2.2e-7 is a different metric and does not imply whole-model KL at that level.

Rewards, losses and gradient norms were zero; this is a numerical GRPO smoke test, not evidence of useful learning or nonzero parameter updates. No per-patch causal improvement is inferred from this combined run. Remaining differences, including the actual DeepEP dispatch/combine path and KDA recurrence, are still under diagnosis.