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.

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 / CP2 training parity

Both full-model runs use mmoukouba/MedPix-VQA, packed sequence length 2048, global batch size 144, local batch size 1, TP1, PP1, and complete 100/100 optimizer steps. These runs validate the packed CP implementation and EP72/CP2 topology using the SDPA sparse-attention reference; the cuDNN sparse path is validated separately below.

Run	Parallelism	Resources	Steps	Peak logged memory	W&B
CP1 reference	EP144 / CP1	18 nodes / 144 GPUs	100/100	38.91 GiB	pod2gbnf
CP2 validation	EP72 / CP2	9 nodes / 72 GPUs	100/100	57.68 GiB	g7b4wi14

Across all 100 matched steps:

• Mean absolute loss difference: 0.0018
• Maximum absolute loss difference: 0.0075
• Step 99 loss: CP1 1.3615, CP2 1.3624

cuDNN sparse MLA parity

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

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

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
• 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 requires both nvidia-cudnn-frontend[cutedsl] and a compatible FlashMLA build; 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-e3c3cacad2e8.docs.buildwithfern.com/nemo/automodel