feat(models): GLM-5.3 Flash text model (glm5_next) by KyleWong288 · Pull Request #4072 · PrimeIntellect-ai/prime-rl

KyleWong288 · 2026-10-13T17:54:49Z

Summary

Adds the text model of zai-org/GLM-5.3-Flash (glm5_next, 313B params / ~18B active). Based on #4036, which this supersedes.

KDA linear attention (3 of every 4 layers) on FLA's chunk_kda with GLM's lower-bounded forget gate.
NoPE sparse MLA in the remaining layers (absorbed form) on the existing sparse_mla kernels, which now accept q/k without a RoPE tail.
DSA indexer with k-pool compression: keys pooled per document in groups of 4, FP8 scoring after a Hadamard rotation (as vLLM), top 512 pools plus the query's own incomplete pool. Frozen, as in GLM-5.
mHC via DeepseekV4HyperConnection; MoE with ClampedSilu(10), built inline in the decoder layer like the other models.
Compile- and AC-safe cu_seqlens handling (one-document rows, FLA's CP context, per-layer cu_seqlens copy for checkpoint replay).
Composite glm5_next / text glm5_next_text configs; checkpoint-only keys (topk_method, moe_router_dtype, vision keys, ...) pass through so an HF export still loads in vLLM. Conversion stacks the experts, moves the flat hc_* tensors onto the hyper-connection modules, and drops the vision tower and the MTP layer.
Merged with main (vLLM 0.32, transformers 5.19): transformers now accepts GLM-5.3 Flash's layer types, so the compat shim is gone. It renames the checkpoint's deepseek_sparse_attention to indexed_attention on load, which the decoder layer now matches. Context parallelism uses the Ulysses-only cp_blocker API (no head constraint: KDA, the indexer and sparse MLA all-gather over the CP group).
vLLM patch for KDA conv weight updates: vLLM caches a torch.cat copy of the q/k/v short-conv weights on the first forward, so weight updates never reached inference. monkey_patch_glm5_next_kda_conv_weights stores the three weights as slices of one buffer and points the cache at a view of it; reload copies in place, so every update lands. Remove when vLLM keeps them in one merged parameter.
Example SWE RL config (examples/advanced/glm-5.3-flash/swe.toml): SWE-rebench-V2 at 128K context on 2 trainer nodes (EP 8, CP 8) + 6 rollout nodes (TP 8).

Text only: the vision tower is not modelled. Multimodal support follows in a separate PR.

Validation

100+ RL steps on SWE-rebench-V2. mismatch_kl is stable, reward goes up. Wandb run: link

Also, initial tests of 20 RL steps on math-env (Hendrycks-Math train). Two runs at cp = 1 and at cp = 2 (Ulysses). All steps have mismatch_kl below 0.015:

Wandb run for cp = 1: link
Wandb run for cp = 2: link

Unit test: config and checkpoint round trip for the text and composite configs, including the pass-through keys.

Follow-ups

Now that we have a baseline RL run, the additional changes from #4036 can be validated against the KL numbers above:

fp32 hyper-connections and KDA decay (FSDP casts them to bf16 while vLLM keeps them in fp32).
FP8 trainer layout matching vLLM (attention stays bf16).
Export tool without a generation config; int64 offsets in the grouped FP8 weight cast (both model-independent fixes).

🤖 Generated with Claude Code

Port of the text model from #4036: KDA linear attention, NoPE sparse MLA with the k-pool DSA indexer, mHC via DeepseekV4HyperConnection, and a clamped-SwiGLU MoE. Registers the glm5_next / glm5_next_text configs, lets the sparse MLA kernels run without a RoPE tail, and allows the deepseek_sparse_attention layer type in the trainer and the vLLM plugin. The MoE is built inline in the decoder layer, matching the other models. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

Builds a tiny glm5_next model from both the text and the composite config and checks that HF <-> PrimeRL weight conversion is lossless. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

…n vLLM The configs dropped keys the trainer does not model, so an exported config.json lost topk_method (vLLM then creates no selection-bias parameter and fails to load the export), moe_router_dtype, indexer_rope_interleave and the vision keys. Pass them through and drop only dtype and head_dim (0 in the checkpoint). From #4036 (bd355ea). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

…t updates vLLM's Glm5NextLinearAttention builds _merged_conv_weight once, on the first forward, as a torch.cat copy of the q/k/v short-conv weights. Weight updates write into the three parameters but never reach that copy, so inference kept the base checkpoint's conv weights for the whole run. After loading, store the three weights as slices of one buffer and point the cache at a view of it. Layerwise reload copies into parameter storage in place, so every update lands in the buffer the kernel reads. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

r2e-gym at 128K context on 2 trainer nodes (EP 8, CP 8, full offload) and 6 single-node TP 8 vLLM replicas, batch 256 x group 16, with a difficulty-pool curriculum that down-weights all-solved tasks. The orchestrator runs on an inference node: trainer host RAM sits near 3 TB with AdamW offloaded, and its 128K rollout buffer tipped rank 0 over. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

…points and episode timeout Switch the SWE RL example from r2e-gym to SWE-rebench-V2 with a 20-turn cap, save weights-only checkpoints every 10 steps, and fail episodes after 2 h. Without an episode timeout, a sandbox outage left rollouts waiting forever on dead sandbox streams; they filled every in-flight slot and froze the run. Weights-only checkpoints keep each save at ~0.65 TB instead of ~4.5 TB with AdamW state. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

Resolve the conflicts from the vLLM 0.32 / transformers 5.19 bump (#4140): - transformers 5.19 knows GLM-5.3 Flash's layer types, so drop the `allow_glm5_next_layer_types` shim with the rest of `transformers_compat.py` (deleted upstream), and keep the KDA conv-weight patch in `patches.py`. - transformers now renames the checkpoint's `deepseek_sparse_attention` layer type to `indexed_attention` on load; match the new name in the decoder layer. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

#4134 removed ring context parallelism and replaced `cp_support()` with `cp_blocker(config, cp)`. KDA, the indexer and sparse MLA all-gather what they need over the CP group themselves, so no head constraint applies; return None as GLM-5 (`glm_moe_dsa`) does. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

samsja · 2026-10-24T10:52:17Z

+num_infer_nodes = 1
+num_infer_replicas = 6
+# Trainer nodes peak near 3 TB host RAM with AdamW offloaded; the orchestrator's 128K rollout buffer tips them over.
+orchestrator_on_inference = true


this is on by default afaik lets remove

it seems false by default

samsja · 2026-10-24T10:52:54Z

+[trainer.model.compile]
+
+[trainer.optim]
+type = "adamw"


Suggested change

type = "adamw"

type = "muon"

muon would require > 2 trainer nodes since we can only use offload=optimizer instead of offload=full, shall we just keep as adamw?

samsja · 2026-10-24T10:54:45Z

+def monkey_patch_glm5_next_kda_conv_weights():
+    """Make GLM-5.3 Flash's merged KDA conv weight a view of the q/k/v conv weights.
+
+    vLLM builds `_merged_conv_weight` once as a `torch.cat` copy, so weight updates never reach it.


do we have a vllm pr that fix this upstream ? if not can we create it and link it here

Seems like upstream vllm PR 60991 exists already

Drop settings that match the defaults, use 4 inference replicas, remove the difficulty-pool sampler, and link the upstream vLLM fix for the KDA conv patch. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

# Conflicts: # src/prime_rl/trainer/models/__init__.py # src/prime_rl/utils/vlm.py

KyleWong288 and others added 9 commits October 13, 2026 17:51

test(models): GLM-5.3 Flash config and checkpoint round trip

cfb5bbc

Builds a tiny glm5_next model from both the text and the composite config and checks that HF <-> PrimeRL weight conversion is lossless. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

docs: add GLM-5.3 Flash (text only) to the supported models

e531758

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

KyleWong288 marked this pull request as ready for review October 23, 2026 02:05

KyleWong288 requested review from garrett361 and samsja October 23, 2026 02:16

faresobeid mentioned this pull request Oct 23, 2026

fix(trainer): activation offloading no longer copies expert weights to CPU #4237

Merged