feat(glm5-next): add GLM-5.3-Flash training support

[HuiyingLi]

What does this PR do ?

Adds native AutoModel training support for zai-org/GLM-5.3-Flash, including its hybrid Kimi Delta Attention (KDA) / DeepSeek Sparse Attention (DSA) decoder, MoE routing, VLM processing, Hugging Face checkpoint conversion, and packed-sequence context parallelism.

The included MedPix recipe uses the validated EP72/CP2 topology: TP1, PP1, 9 nodes / 72 GPUs, packed sequence length 2048, and 100 optimizer steps. It now selects the shared FlashMLA-forward/cuDNN-backward kernel for sparse MLA layers; KDA layers remain on FLA.

Important

Full-model single-GPU checkpoint loading and training are not supported. The supported base-checkpoint initialization path is distributed DCP loading.

User training does not require Transformers 5.16 or a nightly build. This implementation supports AutoModel's pinned transformers==5.12.1, provides its own GLM-5.3 configuration and processor registration, and does not require trust_remote_code.

Changelog

• Register the GLM-5.3-Flash Transformers architecture and native AutoModel configuration.
• Implement the multimodal processor, vision tower, mHC decoder blocks, dense/MoE layers, and HF state-dict adapter with FP8 checkpoint dequantization.
• Run KDA with FLA production kernels and preserve its left-to-right recurrence with contiguous packed CP shards.
• Implement packed DSA CP with document isolation and differentiable full-sequence gather, including empty-local-shard backward participation.
• Reuse a model-independent FlashMLA-forward/cuDNN-backward latent sparse-attention adapter for GLM-5.2 and GLM-5.3; preserve GLM-5.2's model-specific D=576, top-k <= 2048 contract.
• Absorb GLM-5.3's K projection into the query, attend over its shared 512-wide latent K/V, then apply the model-owned V projection.
• Align MoE routing with the HF reference: correction bias selects experts while routing weights use unbiased sigmoid scores.
• Add a MedPix packed-2K EP72/CP2 recipe plus unit, numerical, distributed, and functional parity coverage.

Validation

Hugging Face numerical parity

Reference setup:

• Validation-only truncated model with the first 4 decoder layers: KDA, KDA, KDA, DSA
• Sequence length: 4096
• Full vocabulary: 154,880 logits for all 4096 token positions
• BF16 compute; checkpoint FP8 weights dequantized before compute
• Transformers reference: 5.16.0.dev0
• State-dict audit: 135 expected, 135 loaded, 0 missing, 0 unexpected, 0 shape mismatches

Metric	Result
Mean KL, HF to AutoModel	0.0002477953
P95 KL	0.0005324535
Max KL	0.0271428488
Mean JSD	0.0000619027
Logits cosine similarity	0.9999251899
Top-1 token agreement	98.33984375%
Mean absolute logit difference	0.0115956810

Packed CP1 / CP8 full-model training parity

Both full-model runs use commit e3c3cacad2e8eaadda6c38c19d859a973587f604, mmoukouba/MedPix-VQA, packed sequence length 2048, global batch size 144, local batch size 1, TP1, PP1, and the SDPA sparse-attention reference. Both completed 100/100 optimizer steps with clean Slurm exits.

Run	Parallelism	Resources	Steps	Peak logged memory	W&B
CP1 reference	EP144 / CP1	18 nodes / 144 H100s	100/100	38.89 GiB	xcu2nc2k
CP8 validation	EP144 / CP8	18 nodes / 144 H100s	100/100	41.04 GiB	21yl4mj9

Across all 100 matched steps, CP8 versus CP1:

• Mean absolute loss difference: 0.001879
• RMSE: 0.002800
• Maximum absolute loss difference: 0.0130 at step 93
• Step 99 loss: CP1 1.2344, CP8 1.2328 (absolute difference 0.0016)

cuDNN sparse MLA parity and performance

The released checkpoint uses 64 query heads, a 512-wide latent K/V, index_topk=2048, and index_kpool=4, producing 2051 raw sparse slots before FlashMLA alignment.

H100 shared-kernel stress test (D=512, raw K=2063, larger than the released model's 2051 slots):

Tensor	Cosine	Relative RMSE
Output	0.9999978542	0.0020721
dQ	0.9999986291	0.0016682
dKV	0.9999930263	0.0037219

Full native GLM-5.3 sparse layer versus the SDPA reference (seq=64, D=512, raw K=2051, BF16):

Tensor	Cosine	Relative RMSE
Output	0.9999925494	0.0038936
Input gradient	0.9999864101	0.0052155
Projection weight gradients	0.9999790192 minimum	0.0064959 maximum

Single-node distributed training smoke test:

• First 4 decoder layers, EP8 / CP2 on 8x H100
• MedPix, packed sequence length 2048, global batch size 4
• backend.attn=cudnn, activation checkpointing enabled
• Completed 2/2 optimizer steps; losses 10.7592, 11.2339
• Logged memory 13.29 GiB, then 17.26 GiB
• Step 0 including compilation: 188.79 s; warm step 1: about 4 s / 1157.58 tokens/s

Full-model, same-commit end-to-end comparison:

• Commit: e3c3cacad2e8eaadda6c38c19d859a973587f604
• MedPix packed sequence length 2048, GBS 144, local batch size 1
• EP72 / CP2, TP1 / PP1, 9 nodes / 72 H100s
• 100/100 optimizer steps; both jobs completed with exit code 0
• TPS statistics use the same steady-state window, steps 10-99

Backend	Mean TPS	Median TPS	Mean / peak memory	Training loop	Slurm elapsed	W&B runtime	W&B
SDPA	8,378.15	8,417.33	57.160 / 57.68 GiB	37:15	39:29	2,278.36 s	yn0z3ym3
cuDNN	9,489.04	9,538.12	57.185 / 57.78 GiB	33:18	35:35	2,040.68 s	kbp425qx

cuDNN versus SDPA:

• Mean TPS: +13.26%; median TPS: +13.32%
• Training-loop elapsed: -3:57 / -10.60%
• Slurm elapsed: -3:54 / -9.88%
• W&B runtime: -237.69 s / -10.43%
• Mean steady-state memory difference: +0.025 GiB/GPU for cuDNN

The two 100-step loss curves remain aligned: mean absolute difference 0.002648, RMSE 0.003076, maximum absolute difference 0.0066 at step 73, and final loss SDPA 1.2359 versus cuDNN 1.2337.

Runtime used nvidia-cudnn-frontend[cutedsl]==1.27.0 and FlashMLA b7643bd on H100 with CUDA 13.3.

GLM-5.2 cuDNN regression: PR base versus head

This checks the GLM-5.2 path touched by the shared sparse-attention utility refactor. Both runs use the same nightly image without rebuilding, nvidia-cudnn-frontend[cutedsl]==1.27.0, FlashMLA b7643bd, and the same training configuration:

• zai-org/GLM-5.2 on the Tulu-3 SFT mixture (train[:100000])
• Packed THD sequence length 4096; 19,830 generated packs
• TP1 / CP1 / PP4 / EP64 on 32 nodes / 256 H100s
• Global batch size 256, local batch size 4, and 100 optimizer steps
• backend.attn=cudnn; W&B disabled for this scoped NeMo-CI regression
• Steady-state TPS statistics use the same steps 10-99 window

Revision	Mean TPS	Median TPS	Mean / peak memory	Step 0 incl. compilation	Training loop	Final loss	NeMo-CI
PR base 8c954f67	103,095.64	103,116.44	46.968 / 47.11 GiB	548.49 s	23:06	0.5357	parent / leaf
PR head e3c3caca	103,074.03	103,068.09	46.972 / 47.11 GiB	501.61 s	22:19	0.5356	parent / leaf
Head versus base	-0.021%	-0.047%	+0.004 / +0.000 GiB	-46.88 s	-0:47	-0.0001	Both 100/100, success

Across all 100 matched steps, head versus base loss parity is:

• Mean absolute difference: 0.000091
• RMSE: 0.000136
• Maximum absolute difference: 0.0006 at compilation step 0
• Step 99: base 0.5357, head 0.5356 (absolute difference 0.0001)

The steady-state throughput and memory deltas are effectively neutral. The shorter head training-loop time is explained by run-to-run first-compilation variation; it is not treated as a steady-state speedup.

Unit tests

pytest -q \
  tests/unit_tests/models/glm5_next/test_model.py \
  tests/unit_tests/models/glm_moe_dsa/test_glm_moe_dsa_cudnn.py \
  tests/unit_tests/recipes/test_glm5_next_medpix_recipes.py

Result: 35 passed, 3 skipped in the interactive development container. The skipped cases require real optional GPU kernels; the H100 results above exercise those kernels directly.

Before your PR is "Ready for review"

Pre checks:

• Make sure you read and followed Contributor guidelines
• Did you write any new necessary tests?
• Did you add or update any necessary documentation? (validated runnable recipe included)

Additional Information

• HF model: zai-org/GLM-5.3-Flash
• Transformers runtime: AutoModel's pinned transformers==5.12.1; no upgrade to 5.16/nightly and no trust_remote_code are required.
• Supported checkpoint loading: distributed DCP only; full-model single-GPU loading is unsupported.
• The validated recipe uses EP and CP only; TP and PP remain disabled.
• backend.attn=cudnn separately requires both nvidia-cudnn-frontend[cutedsl] and a compatible FlashMLA build; these kernel dependencies are independent of the Transformers version. Environments without them can select backend.attn=sdpa.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[github-actions]

🌿 Preview your docs: https://nvidia-preview-preview-636bf0bc2ef1.docs.buildwithfern.com/nemo/automodel

[yuhezhang-ai]

Non-blocking — sharing two findings from poking at this PR. I'll follow up with the checkpoint-robustness CI enablement separately.

1. The checkpoint loading path probably needs further optimization. Loading OOMs at 1 node (ep8) and 2 nodes (ep16) with the same ~77 GB/rank footprint — it doesn't shrink with EP, so something dense seems to materialize per-rank before sharding. Everything works at your 9-node EP72 layout, but right now ~9 nodes is the floor even for forward-only use. Happy to help dig if useful.

2. GLM-5.3-Flash does not have the GLM-4.7-Flash numerical-sensitivity issue. Context: 4.7-Flash's released checkpoint amplifies tiny numerical noise through its MoE routing, which is why its parity CI needed a special profile — so I was curious whether the new Flash inherits that. I ran the same probes on 5.3-Flash at EP72 on the SDPA path. These are self-sensitivity measurements (the model compared against itself under sequence-shape changes and forced routing flips), not HF parity. Result: divergence stays at the healthy floor (mean KL 0.005–0.008 across 128–4096 tokens, worst token 0.35), and an injected routing flip dies out with depth (0.0065 at layer 2 → 0.0006 at layer 38) instead of amplifying. Same healthy class as GLM-5.2 — good news.

Minor note: the cuDNN-DSA/FlashMLA runtimes aren't in the image build, so CI will take the SDPA path (same path as your CP1/CP8 validation runs). If the cudnn path should ever be CI-covered, the runtimes would need adding to the Dockerfile the way DeepEP/TileLang are.

[HuiyingLi]

/ok to test e3c3cac

[HuiyingLi]

/claude review

[claude]

Review summary

Thorough, well-structured onboarding of the GLM-5.3-Flash VLM. Registry wiring (MODEL_ARCH_MAPPING + _CUSTOM_CONFIG_REGISTRATIONS), ModelCapabilities (cp/ep/thd, tp/pp off — consistent with the EP72/CP2 recipe), the TieSupport.UNTIED_ONLY constructor guard, checkpoint-free initialize_weights, and the RoPE low-precision hazard are all handled correctly: the vision inv_freq/rotary_pos_emb is registered in _keep_in_fp32_modules_strict and the whole-model cast goes through cast_model_to_dtype, not a raw self.to(dtype). The shared cudnn_sparse_attention refactor keeps glm_moe_dsa's narrower public contract intact and re-points its tests. Good CP parity harness and state-dict round-trip coverage.

Two findings:

1. Missing tensor-layout docstrings on new autograd functions (components/models/glm5_next/cp.py, inline). The new _AllGatherSequence / _AllGatherBackwardAnchor forward/backward methods take tensors but do not document their shapes/axis order, which the tensor-contract rule requires for custom autograd forward/backward.

2. image_processing.py has no direct unit test. smart_resize (including the max_pixels binary-search branch), Glm5NextImageProcessor.resize/patchify, and get_number_of_image_patches contain non-trivial resize/patchification math on tensor inputs, but the only processing tests (test_processing.py) cover the chat-template rewrite and the torchvision-absent import guard — not the image math. A small CPU unit test that feeds a fixed-size dummy image tensor through patchify / get_number_of_image_patches and asserts the resulting image_grid_thw and patch count would guard the token-budget contract the recipe relies on. Consider adding one.

Neither is blocking.

[claude]

These new torch.autograd.Function methods accept tensors but have no docstring documenting their layout. Per the tensor-contract rule, custom autograd forward/backward methods must document each tensor input's shape/axis order. Please add Google-style Args/Returns sections:

• _AllGatherSequence.forward — local_tensor is the per-rank sequence shard (e.g. [batch, local_sequence, ...], gathered along dim); state which axis dim indexes and that the return is [batch, global_sequence, ...].
• _AllGatherSequence.backward — grad_output is the gathered-tensor gradient [batch, global_sequence, ...]; return is the local shard [batch, local_sequence, ...].
• _AllGatherBackwardAnchor.forward/backward (below) — gathered/grad_output layouts.

The all_gather_sequence/all_gather_backward_anchor wrappers also carry only a one-line docstring with no tensor layout.

[HuiyingLi]

/ok to test 636bf0b