[hallerite]

Summary

Adds the text model of zai-org/GLM-5.3-Flash (glm5_next, 313B text params / ~18B active), building on the closed #3439 by @samsja. Based on main with #4037.

• KDA linear attention (3 of every 4 layers) on FLA's chunk_kda with GLM's safe, lower-bounded forget gate (lower_bound * sigmoid(exp(A_log) * (f + dt_bias)), computed in-kernel). The checkpoint's three depthwise convs run as one conv over concatenated q/k/v.
• NoPE sparse MLA in the remaining layers, absorbed form (512-wide latent as key and value), on the existing sparse_mla kernels. Both DSA backends now accept a q/k without a RoPE tail: TileLang skips the tail tile when it is zero-width (from Add GLM-5.3 training and stack configs #3439), FlashMLA/cuDNN take the 512 shape natively.
• The DSA indexer with k-pool compression: the piece Add GLM-5.3 training and stack configs #3439 dropped, without which attention is wrong past 2048 tokens. Keys are pooled per document in groups of 4 starting at each document's first token, as a per-channel softmax(gate + ape)-weighted mean; a query scores the complete pools ending at or before it, keeps the top 512, and also attends to its own incomplete pool. Following vLLM, queries and pooled keys are Hadamard-rotated (same fp32 butterfly order) before FP8 scoring, which reuses the existing fp8_indexer over pools. Index rows are padded to a multiple of 128 for FlashMLA. The indexer is frozen through freeze_sparse_indexer, like the other sparse indexers. Its per-document pool layout (PoolLayout) is built once per forward, outside the compiled decoder layers.
• mHC via the existing DeepseekV4HyperConnection (identical math; fused Sinkhorn and write-back kernels), final stream mean before the norm. MoE with ClampedSilu(10) (GLM's clamped SwiGLU) on routed and shared experts and the dense MLPs.
• Compile- and AC-safe cu_seqlens: maybe_mark_dynamic (a one-document row gets its own graph instead of failing the dynamic-size constraint), FLA's CP context built eagerly (as fix(trainer): keep Qwen3.5 CP context setup eager #4011 does for Qwen3.5), and a per-layer cu_seqlens copy in KDA so FLA's identity-keyed metadata cache cannot desynchronize activation-checkpoint replay (see below).
• FP8 trainer layout that matches vLLM: vLLM serves all of this model's attention in BF16 (KDA, sparse MLA and its indexer) and quantizes only the MLPs and experts, for the FP8 checkpoint and online FP8 alike. Models can now add FP8 ignore patterns through PreTrainedModelPrimeRL.quantization_ignore_patterns; GLM-5.3 Flash keeps every self_attn linear in BF16.
• fp32 hyper-connections and KDA decay: vLLM computes the mHC tensors (fn, base, scale) and KDA's A_log/dt_bias in fp32, and the checkpoint stores all but fn in fp32, but FSDP cast them to bf16 for the trainer's forward. Models can now list such modules in PreTrainedModelPrimeRL.fp32_modules: each becomes its own fp32 FSDP unit (fp32 compute dtype under full offload) that does not cast its inputs and stays unsharded until the backward. KDA's decay parameters live on a self_attn.decay submodule for this (converted from the checkpoint's self_attn.A_log), and hc fn goes to vLLM in fp32 with the rest of the hyper-connection.
• Composite glm5_next / text glm5_next_text configs registered like Qwen3.8 Flash Next (model.language_model.* layout, VLM layer prefix). Checkpoint keys the trainer does not model pass through, so an exported config.json still loads in vLLM. The conversion stacks experts, moves the flat hc_attn_fn-style tensors onto the hyper-connection modules, and drops the vision tower and MTP layer.
• deepseek_sparse_attention is not a layer type the pinned transformers accepts, which broke both our config and vLLM's Glm5NextConfig; the transformers compat shim now takes a generic list, applied in the trainer and the vLLM plugin.
• Two shared fixes found on the way: the grouped per-block FP8 weight cast used int32 offsets and crashed once a rank holds all 288 experts (ep=1); tools/convert_dcp_to_bf16.py crashed on models without a generation config (also Qwen3.8 Flash Next, Qwen3.5, gpt-oss, Nemotron-H).

vLLM 0.31.0 already serves the model; no inference bump.

Validation

Scripts and logs: /home/konstantin/glm53-flash-validation (H200).

• Checkpoint mapping: every text tensor of the real checkpoint maps HF→prime→HF with matching shapes (1372 tensors, 313.3B params), only vision/MTP dropped.

• Against the transformers 5.19 reference (tiny random 8-layer model, bf16 vs fp32 reference, routing replayed): dense-regime logits KL 9.0e-4 vs the reference's own bf16 error of 1.35e-3; per-layer error grows smoothly to ~1%; gradients at a median 0.89x the reference's own bf16 error. On the reference's exact layer inputs with an exact fp32 scorer, the indexer reproduces transformers' selection on 100% of queries (pooling, tail, packing, causal ranges); the FP8 scorer agrees on 78–83% of exact sets (Jaccard 0.97–0.98) in this random model. Packed vs separate documents, compiled vs eager (including one-document rows), TileLang vs FlashMLA/cuDNN: all within bf16 noise.

• Expert and context parallelism on real weights (8 real layers of GLM-5.3-Flash-BF16, the trainer's own setup_model on 8 GPUs, 8 packed rows of 8192 real tokens, routing replayed, full per-parameter gradients against an FSDP-only reference):

  Config	logits rel	mismatch KL	grad median rel
  reference rerun (noise floor)	0.46%	8.1e-5	1.33%
  ep=8, ep=4, ep=8 + compile + full AC + activation offloading	0.46%	8.2e-5 – 8.3e-5	1.34% – 1.35%
  cp=2, cp=4, ep=8 + cp=2, ep=8 + cp=2 + compile + selective AC	0.73%	1.6e-4 – 1.7e-4	1.88%

  EP, compile, AC and offloading sit exactly on the noise floor. CP sits ~2x above it because of the smaller per-rank shape, not CP logic: the same rows cut to 4096 tokens without CP deviate exactly as much (1.54e-4 vs 1.52e-4) and match cp=2 at the floor.

• Context parallelism, tiny model (real 2-rank NCCL, three row shapes including a one-document row, each twice): eager, compile, compile + full AC, eager + full AC and compile + selective AC all match an eager cp=1 reference (dense logits ≤0.8% rel, gradients ≤1.3% median rel).

• Trainer vs vLLM on real weights (first 8 layers, 1k–30k-token prompts, prefill and decode, fp32 router and lm_head on both sides): mismatch KL ≈ 5e-4 below 2048 tokens, 6.5e-4 above (k-pool sparse path), ≈ 9e-4 for decoded tokens; identical on both DSA backends and packed. No query in ~52k chose among zero-score pools, so trainer/vLLM tie-breaking differences cannot contribute there.

• SFT on real weights (4 layers, EP=8, compile, fp32 router, packed 8k rows, 20 steps on main with revert(trainer): compile around the checkpoint again for MoE blocks with an fp32 router #4037): full AC and selective AC, both with activation offloading (on by default with AC), finite and matching step for step (step 20 loss 14.8000 vs 14.7998), peak 30.6 vs 38.2 GiB per GPU, ~1.2–2.8 s per warm step.

• Checkpoint save/resume (same setup, lr 5e-5 so 20 steps move the weights): forking a run from its own step-10 checkpoint reproduces it (step 11 loss differs by 3.8e-6, i.e. forward-pass noise; later steps stay within run-to-run noise), and the data stream resumes at the same samples.

• HF export (tools/convert_dcp_to_bf16.py): an lr=0 export is bit-identical to the original checkpoint on all 972 text tensors, including the fp32 hc_*, A_log, dt_bias and selection-bias tensors. vLLM loads a trained export (text-only, see below) at trainer-vs-vLLM KL 1.0e-3 prefill / 3.4e-4 decode, versus 2.6e-4 on the original weights.

• RL smoke run (8-layer cut, 4 trainer GPUs with EP=4 + 4 vLLM GPUs, NCCL weight broadcast, glm-5.3 renderer, reverse-text with the debug algorithm since the cut only writes gibberish): vLLM serves the model, the startup weight broadcast reloads every language-model layer (only the unmodelled vision tower is skipped), rollouts flow, and 3 full RL steps complete, each followed by a weight broadcast into vLLM. Live mismatch KL 7e-4 – 8e-4 per step.

• FP8 (same RL smoke run with the FP8 trainer, DeepGEMM dense linears + deepgemm_fp8 MoE, against vLLM quantizing the BF16 weights online with fp8_per_block as in the GLM-5.3 recipe): 5 steps with weight broadcasts into the FP8 server, mismatch KL flat at 3.7e-3 – 4.1e-3 with vLLM's defaults, and 2.3e-3 – 2.4e-3 with VLLM_USE_DEEP_GEMM_E8M0=0 (exact DeepGEMM scales, as the trainer uses on SM90, instead of power-of-two scales). Offline on the 8-layer cut (52k tokens): trainer BF16 vs vLLM FP8 2.0e-3 (default) / 1.4e-3 (exact scales); trainer FP8 vs vLLM FP8 with exact scales 1.8e-3.

• Full model in vLLM (all 45 layers, BF16, one node at TP=8): loads in ~20 min, answers chat prompts correctly with reasoning, and prompt logprobs on 1k–32k tokens of real text average −1.3 to −1.8.

• Full model in the RL loop (all 313B params; 3 trainer nodes with EP=8, compile and full AC, plus one vLLM node at TP=8; reverse-text with real rewards, before the fp32 change above): 50 steps with SignSGD and offload = "full" complete in 66 min, with an NCCL weight broadcast into vLLM after every step, 63–73 GiB peak per trainer GPU and 33–45 s per typical step. Reward stays at 0.98–0.997 (the task starts near its ceiling), but mismatch KL drifts from 2.1e-3 (mean of steps 1–5) to 4.7e-3 (steps 46–50) at flat entropy (0.09–0.13). AdamW does not fit on 3 trainer nodes: offload = "optimizer" moves its state back to the GPU for the step (~208 GB per GPU), and offload = "full" needs 16 bytes per param of CPU RAM (~1.52 TiB per node, 1.48 TiB available), so AdamW on this model needs at least 4 trainer nodes.

• fp32 hyper-connections and KDA decay (8 real layers, the parallelism harness above): EP=8 + compile + full AC + offloading and EP=8 + CP=2 + compile + selective AC agree with an FSDP-only run of the same code at the earlier floors (mismatch KL 8.3e-5 and 1.6e-4), and all three move away from the bf16-parameter reference identically (KL 2.6e-4), so the fp32 units are correct under EP, CP, compile, AC and offloading. 20 SFT steps track the bf16 run (step 20 loss 14.7903 vs 14.7998) at the same step time and memory. In the 8-layer RL smoke run (full offload) mismatch KL goes from 7.4e-4 to 7.0e-4.

• Renderer: the existing glm-5.3 renderer reproduces GLM-5.3 Flash's own chat template token for token (plain, multi-turn with reasoning, tool call).

Coverage

Area	Status
Forward/backward vs transformers, packing, both DSA backends	✅ numerically verified (tiny random model)
k-pool indexer selection	✅ exact on the reference's own inputs
Trainer vs vLLM on real weights	✅ first 8 layers, up to 30k tokens
torch.compile, incl. one-document rows	✅ numerically verified
Expert parallelism	✅ ep=4/8 vs ep=1 on real weights, on the noise floor
Context parallelism (ring) + compile + full / selective AC	✅ cp=2/4 on real weights and tiny model
Activation checkpointing (full, selective) + activation offloading	✅ numerically verified, 20 SFT steps each on real weights
Checkpoint save/resume	✅ exact
HF export	✅ weights bit-identical; ⚠️ vision files and generation config are not carried over (text-only serving)
RL loop (rollouts → trainer step → NCCL weight broadcast → vLLM)	✅ 3 steps on the 8-layer cut
FP8 trainer (fp8 dense linears + DeepGEMM FP8 MoE)	✅ runs, layout matches vLLM; ⚠️ mismatch KL 2.4e-3 vs 7e-4 in BF16
Zero-score pool ties	✅ none on real weights up to 30k tokens
Full 45-layer model	✅ vLLM serving (BF16, TP=8); ✅ 50 RL steps with SignSGD on 3 trainer nodes + 1 vLLM node; ⚠️ mismatch KL drifts 2.1e-3 → 4.7e-3 over 50 steps; ❌ AdamW (needs ≥4 trainer nodes)
fp32 hyper-connections and KDA decay	✅ numerically verified under EP, CP, compile, AC and offloading
Context parallelism, Ulysses style	⚠️ not run; both styles shard identically and this model does its own gathers
Router replay, sampling-mask capture	❌ untested on this model

🚧 Still missing

1. Mismatch drift on the full model: the 50-step SignSGD run doubled the mismatch KL. It ran before the fp32 hyper-connections and KDA decay, which target exactly the parameters that drift (keeping them in fp32 halved the prefill KL of a trained export), and SignSGD moves those tiny tensors by the full lr every step. A rerun with the fp32 change, ideally with AdamW on at least 4 trainer nodes (offload = "full", ~1.14 TiB CPU RAM per node), should show whether the drift remains. Trainer-vs-vLLM KL at 64k–128k context on the full model is still unmeasured.
2. FP8 at full scale: online FP8 vLLM at TP=8 dies in CUDA graph capture (CUDA error: invalid argument) on the full model; not investigated yet. Trainer FP8 and vLLM FP8 still round independently (trainer FP8 raises the mismatch against FP8 vLLM instead of lowering it), so FP8 buys speed and memory, not accuracy.
3. Router replay and sampling masks, if production runs use them (GLM-5.3 Flash's generation config sets top_p = 0.95, which turns sampling masks on).
4. Ulysses CP has not been run.

🩹 Patches and workarounds

Needed but not in this PR:

• vLLM 0.31 KDA decode bug: its kernel launches num_seqs × heads on CUDA grid axis z (max 65535), so at TP=1 / data-parallel serving the server dies in warmup unless max_num_seqs < 1024. Workaround: set inference.vllm.max_num_seqs (e.g. 512).
• vLLM 0.31 k-pool indexer block size at TP>1 on Hopper: vLLM aligns hybrid KV blocks to index_kpool × 32 tokens, but DeepGEMM's paged-MQA kernel takes 32-pool pages only on SM100 and needs 64 on SM90. At TP=8 the block lands on 640 tokens and engine startup fails with block_kv == 64 or (arch_major == 10 and block_kv == 32). Workaround: inference.vllm.block_size = 768 (a multiple of 256). TP=1 happens to get 4352 and works.
• FP8 serving: set VLLM_USE_DEEP_GEMM_E8M0=0 in inference.env_vars on Hopper; it cut FP8 mismatch KL from ~4e-3 to ~2.4e-3.
• Stale conversion cache: the shared HF snapshot of GLM-5.3-Flash-BF16 on our cluster carries a prime/ cache from Add GLM-5.3 training and stack configs #3439's run in the old key layout. Point model.conversion_dir elsewhere until it is removed.

No longer needed on main:

• fix(trainer): recompile AC-recomputed compiled blocks with the forward's dynamo settings #3951 / fix(trainer): preserve checkpoint replay across autotuning and compile fallback #4029: before revert(trainer): compile around the checkpoint again for MoE blocks with an fp32 router #4037, a 20-step SFT run on the real weights crashed at step 5 with aten.empty_like ... not found in storage; with both cherry-picked it passed. On main with revert(trainer): compile around the checkpoint again for MoE blocks with an fp32 router #4037 the same run passes without them, and CP + AC passes too. With revert(trainer): compile around the checkpoint again for MoE blocks with an fp32 router #4037, checkpointed blocks run eagerly under compile; this model's blocks would anyway, since FLA's compile-disabled KDA calls break the graph inside the compiled checkpoint, so compile gives no speedup while AC is on.

Included in this PR:

• Layer-type shim: the pinned transformers rejects the deepseek_sparse_attention layer type in both our config and vLLM's Glm5NextConfig. The compat shim now takes a generic list and is applied in the trainer and the vLLM plugin.
• No-rope sparse MLA: the TileLang kernels accept a q/k without a RoPE tail; index rows are padded to FlashMLA's multiple of 128.
• One-document rows under compile: maybe_mark_dynamic on cu_seqlens (a 2-entry cu_seqlens otherwise fails the hard mark_dynamic constraint with ConstraintViolationError), plus the indexer layout built outside the compiled layers.
• CP + activation checkpointing: FLA's CP context is built eagerly, and KDA takes a per-layer cu_seqlens copy. FLA memoizes cu_seqlens metadata by tensor identity; with one tensor shared by every layer, a layer's forward and its recompute ran different op counts, torch's checkpoint replay (matched by op invocation order) handed a cached int CPU copy to the wrong op, and every CP + AC backward failed with Autograd not support dtype: Int.
• FP8 ignore hook, int64 FP8 weight cast, export tool fix, config pass-through: see the summary.

Found along the way, not addressed here:

• The FLA tensor-cache / checkpoint-replay mismatch is the bug fix(qwen3_5, qwen3_8_flash_next, nemotron_h): per-layer cu_seqlens so FLA's tensor cache matches across AC recompute #3880 described and is still on main for the other FLA models. It is confirmed on Qwen3.8 Flash Next: CP=2 with full, selective or compiled AC crashes with the same error, and either a per-layer cu_seqlens.clone() or FLA_DISABLE_TENSOR_CACHE=1 fixes it (feat(models): add native Qwen3.8 Flash Next #3434's validation run set the latter, which likely masked it). Qwen3.5 and Nemotron-H are untested.
• The one-document compile crash reproduces in other models that hard-mark_dynamic cu_seqlens (qwen3, qwen3_5, gpt_oss, nemotron_h; fix(qwen3.8): allow sequence-boundary shape specialization #4028 fixes Qwen3.8 only).
• HF exports of composite VLM checkpoints carry no vision tower, processor_config.json or generation_config.json, so vLLM's default multimodal setup cannot load them; only text-only serving (limit_mm_per_prompt = {image: 0, video: 0}) works. Exports of model.debug.num_layers runs keep the full layer_types list.
• After resuming SFT from an interval checkpoint, progress/num_tokens stays one step behind: sft/train.py saves interval checkpoints before updating progress.total_tokens.
• The SFT fake-data loader breaks with length = "variable" and input_ids = "random": it always draws seq_len + 1 ids.

Notes

• Keeping the hyper-connections and KDA decay in fp32 lowered the trainer-vs-vLLM KL by ~18% on the original weights and ~2x (prefill) on a trained export. DeepSeek V4's hyper-connections have the same bf16 cast and could opt into fp32_modules. PrimeRL conversion caches written by earlier versions of this branch use the old KDA key layout and must be regenerated.
• The RL smoke runs used the debug algorithm (constant advantage), because the 8-layer cut only writes gibberish and GRPO filters every zero-variance group, and full CPU offload, because AdamW state did not fit on 4 GPUs.
• Vision and the MTP draft layer are not modelled.

🤖 Generated with Claude Code