feat(models): GLM-5.3 Flash text model (glm5_next)

README.md

@@ -46,6 +46,7 @@ The trainer runs the PrimeRL model implementations under `src/prime_rl/trainer/m
 | Family | Example IDs | MoE | EP | CP |
 |--------|-------------|-----|----|-----|
 | GLM-5 (`glm_moe_dsa`) | `zai-org/GLM-5`, `zai-org/GLM-5-FP8` | yes | ✅ | ✅ |
+| GLM-5.3 Flash, text only (`glm5_next`) | `zai-org/GLM-5.3-Flash-BF16` | yes | ✅ | ✅ |
 | Qwen3 MoE (`qwen3_moe`) | `Qwen/Qwen3-30B-A3B`, … | yes | ✅ | ✅ |
 | Qwen3.5 MoE (`qwen3_5_moe`) | `Qwen/Qwen3.5-35B-A3B`, … | yes | ✅ | ✅ |
 | Qwen3.5 VLMs | see [advanced.md](docs/advanced.md#multimodal-training) (`qwen3_5`, `qwen3_5_moe`) | MoE only on MoE VLMs | MoE only | ✅ |

docs/advanced.md

@@ -22,6 +22,7 @@ The trainer only runs `prime-rl`'s own model implementations, selected from the
 | Family | HF config types | EP | CP |
 |---|---|---|---|
 | GLM-5 / GLM-5.2 (`glm_moe_dsa`) | `zai-org/GLM-5`, `zai-org/GLM-5-FP8`, `zai-org/GLM-5.2`, `zai-org/GLM-5.2-FP8` | ✅ | ✅ |
+| GLM-5.3 Flash, text only (`glm5_next`) | `zai-org/GLM-5.3-Flash-BF16` | ✅ | ✅ |
 | Qwen3 MoE | `Qwen/Qwen3-30B-A3B`, … | ✅ | ✅ |
 | Qwen3.5 MoE | `Qwen/Qwen3.5-35B-A3B`, … | ✅ | ✅ |
 | Qwen3.5 VLMs | see [Multimodal training](#multimodal-training) | MoE only | ✅ |

examples/advanced/glm-5.3-flash/swe.toml

@@ -0,0 +1,106 @@
+# GLM-5.3 Flash SWE RL on SWE-rebench-V2 (20-turn cap) at 128K context:
+# 2 trainer nodes + 4 vLLM replicas (TP 8, one node each), H200.
+#
+#   uv run rl @ examples/advanced/glm-5.3-flash/swe.toml
+
+max_steps = 200
+seq_len = 131072
+output_dir = "/shared/outputs" # FILL IN: point at your shared filesystem
+clean = true
+
+[run]
+name = "glm53-flash-swe"
+
+[ckpt]
+interval = 10
+keep_last = 1
+
+[log]
+level = "debug"
+
+[model]
+name = "zai-org/GLM-5.3-Flash-BF16"
+
+[monitors.wandb]
+project = "glm53-flash"
+name = "swe"
+
+[weight_broadcast]
+type = "nccl"
+timeout = 12000
+
+[deployment]
+type = "multi_node"
+num_train_nodes = 2
+num_infer_nodes = 1
+num_infer_replicas = 4
+# Trainer nodes peak near 3 TB host RAM with AdamW offloaded; the orchestrator's 128K rollout buffer tips them over.
+orchestrator_on_inference = true
+
+[slurm]
+job_name = "glm53-flash-swe"
+partition = "all" # FILL IN: set to your cluster's partition
+
+[inference.vllm]
+tensor_parallel_size = 8
+max_model_len = 131072
+tool_call_parser = "glm47"
+reasoning_parser = "glm45"
+# vLLM sizes this hybrid model's pages at 640 tokens; DeepGEMM's k-pool indexer on SM90 needs a multiple of 256.
+block_size = 768
+# Higher values OOM on the first weight update while receiving an MoE layer.
+gpu_memory_utilization = 0.75
+# vLLM's KDA decode kernel puts num_seqs * 64 on a grid axis capped at 65535.
+max_num_seqs = 512
+
+[trainer]
+dist_timeout_seconds = 1800
+
+[trainer.ckpt]
+# Weights only: AdamW state would make each checkpoint ~4.5 TB instead of ~0.65 TB.
+skip_optimizer = true
+
+[trainer.model]
+ep = 8
+# Keeps 16K tokens per GPU at 128K context.
+cp = 8
+offload = "full"
+# Host RAM sits at ~99% with AdamW state offloaded; activation offloading tipped a trainer node into a stall.
+ac_offloading = "None"
+
+[trainer.model.ac]
+mode = "full"
+
+[trainer.model.compile]
+
+[trainer.optim]
+type = "adamw"
+lr = 1e-6
+
+[orchestrator]
+batch_size = 256
+
+[orchestrator.train]
+group_size = 16
+
+[orchestrator.concurrency]
+initial_inflight = 256
+max_inflight = 512
+
+[orchestrator.renderer]
+# GLM-5.3 Flash's text template is GLM-5.3's.
+name = "glm-5.3"
+clear_thinking = false
+
+[orchestrator.train.sampling.extra_body.chat_template_kwargs]
+clear_thinking = false
+
+[[orchestrator.train.source]]
+name = "swerebench"
+env.taskset.id = "swerebench-v2"
+env.taskset.dataset_name = "PrimeIntellect/SWE-rebench-V2-Filtered-Verified"
+env.agent.harness.id = "rlm"
+env.agent.max_turns = 20
+# Fail rollouts stuck on a dead sandbox stream instead of letting them hold in-flight slots forever.
+env.timeout.episode = 7200
+env.agent.runtime.labels = ["glm53-flash", "swe"]

src/prime_rl/inference/patches.py

@@ -24,6 +24,7 @@ def apply_shared_vllm_patches():
     monkey_patch_routed_experts_cut_loaded_prefix_hits()
     monkey_patch_kv_xfer_finished_tolerate_freed()
     monkey_patch_online_fp8_parameter_cast()
+    monkey_patch_glm5_next_kda_conv_weights()
     monkey_patch_online_mxfp8_moe_keep_kernel_on_reload()
     monkey_patch_strip_routed_experts_from_chat()
     monkey_patch_skip_prompt_routing_for_chat()
@@ -297,6 +298,34 @@ def _setup_kernel(self, layer, w13, w2, w13_scale, w2_scale, w13_input_scale, w2
     Mxfp8OnlineMoEMethod._setup_kernel = _setup_kernel
 
 
+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.
+    After loading, store the three weights as slices of one buffer and point the cache at it;
+    reloads copy into parameter storage in place, so updates land in the buffer.
+
+    Remove when vLLM keeps the q/k/v conv weights in one merged parameter
+    (https://github.com/vllm-project/vllm/pull/60991).
+    """
+    from vllm.models.glm5next.common.kda import Glm5NextLinearAttention
+    from vllm.models.glm5next.common.model import Glm5NextForCausalLM, Glm5NextForConditionalGeneration
+
+    @torch.no_grad()
+    def process_weights_after_loading(self) -> None:
+        for module in self.modules():
+            if isinstance(module, Glm5NextLinearAttention):
+                convs = (module.q_conv1d, module.k_conv1d, module.v_conv1d)
+                merged = torch.cat([conv.weight.data for conv in convs])
+                for conv, part in zip(convs, merged.chunk(3)):
+                    conv.weight.data = part
+                # merged shape needs to be [3C, K] instead of [3C, 1, K]
+                module._merged_conv_weight = merged.flatten(1)
+
+    Glm5NextForCausalLM.process_weights_after_loading = process_weights_after_loading
+    Glm5NextForConditionalGeneration.process_weights_after_loading = process_weights_after_loading
+
+
 def monkey_patch_online_fp8_parameter_cast():
     """Pass plain tensors to vLLM's compiled block-FP8 caster.
 

src/prime_rl/trainer/models/__init__.py

@@ -13,6 +13,7 @@
 from prime_rl.trainer.models.base import PreTrainedModelPrimeRL
 from prime_rl.trainer.models.deepseek_v4 import DeepseekV4Config, DeepseekV4ForCausalLM
 from prime_rl.trainer.models.glm4_moe import Glm4MoeConfig, Glm4MoeForCausalLM
+from prime_rl.trainer.models.glm5_next import Glm5NextConfig, Glm5NextForCausalLM, Glm5NextTextConfig
 from prime_rl.trainer.models.glm_moe_dsa import GlmMoeDsaConfig, GlmMoeDsaForCausalLM
 from prime_rl.trainer.models.gpt_oss import GptOssConfig, GptOssForCausalLM
 from prime_rl.trainer.models.laguna import LagunaConfig, LagunaForCausalLM
@@ -41,6 +42,8 @@
 AutoConfig.register("afmoe", AfmoeConfig, exist_ok=True)
 AutoConfig.register("deepseek_v4", DeepseekV4Config, exist_ok=True)
 AutoConfig.register("glm4_moe", Glm4MoeConfig, exist_ok=True)
+AutoConfig.register("glm5_next", Glm5NextConfig, exist_ok=True)
+AutoConfig.register("glm5_next_text", Glm5NextTextConfig, exist_ok=True)
 AutoConfig.register("glm_moe_dsa", GlmMoeDsaConfig, exist_ok=True)
 AutoConfig.register("gpt_oss", GptOssConfig, exist_ok=True)
 AutoConfig.register("laguna", LagunaConfig, exist_ok=True)
@@ -67,6 +70,8 @@
     (AfmoeConfig, AfmoeForCausalLM),
     (DeepseekV4Config, DeepseekV4ForCausalLM),
     (Glm4MoeConfig, Glm4MoeForCausalLM),
+    (Glm5NextTextConfig, Glm5NextForCausalLM),
+    (Glm5NextConfig, Glm5NextForCausalLM),
     (GlmMoeDsaConfig, GlmMoeDsaForCausalLM),
     (LagunaConfig, LagunaForCausalLM),
     (MiniMaxM2Config, MiniMaxM2ForCausalLM),
@@ -121,6 +126,7 @@ def supports_custom_impl(model_config: PretrainedConfig) -> bool:
     "qwen3_5": Qwen3_5ForCausalLM,
     "qwen3_5_moe": Qwen3_5ForCausalLM,
     "qwen4_exp": Qwen3_8FlashNextForCausalLM,
+    "glm5_next": Glm5NextForCausalLM,
     "muse_glimmer": MuseGlimmerForCausalLM,
 }
 

src/prime_rl/trainer/models/glm4_moe/converting_glm4_moe.py

@@ -33,17 +33,17 @@
 )
 
 
-def glm_moe_layer_ops(layer_idx: int) -> list[ConvOp]:
+def glm_moe_layer_ops(layer_idx: int, layers_prefix: str = "model.layers") -> list[ConvOp]:
     """The HF <-> PrimeRL MoE ops for one GLM-style layer.
 
-    Shared by GLM-4 MoE and GLM-MoE-DSA (their HF <-> PrimeRL MoE conversion is
-    identical; DSA's attention/MLA weights pass through untouched)."""
-    p = f"model.layers.{layer_idx}.mlp"
+    Shared by GLM-4 MoE, GLM-MoE-DSA and GLM-5.3 Flash (their HF <-> PrimeRL MoE
+    conversion is identical; attention weights pass through untouched)."""
+    p = f"{layers_prefix}.{layer_idx}.mlp"
     ops: list[ConvOp] = [
         Rename(f"{p}.gate.weight", f"{p}.router.gate.weight"),
         Rename(f"{p}.gate.e_score_correction_bias", f"{p}.router.selection_bias"),
         routed_experts_op(
-            f"model.layers.{layer_idx}",
+            f"{layers_prefix}.{layer_idx}",
             hf_experts="mlp.experts",
             prime_experts="mlp.experts",
             fused=True,

src/prime_rl/trainer/models/glm5_next/__init__.py

@@ -0,0 +1,16 @@
+from prime_rl.trainer.models.glm5_next.configuration_glm5_next import Glm5NextConfig, Glm5NextTextConfig
+from prime_rl.trainer.models.glm5_next.modeling_glm5_next import (
+    Glm5NextForCausalLM,
+    Glm5NextModel,
+    Glm5NextPreTrainedModel,
+    Glm5NextTextModel,
+)
+
+__all__ = [
+    "Glm5NextConfig",
+    "Glm5NextForCausalLM",
+    "Glm5NextModel",
+    "Glm5NextPreTrainedModel",
+    "Glm5NextTextConfig",
+    "Glm5NextTextModel",
+]

src/prime_rl/trainer/models/glm5_next/configuration_glm5_next.py

@@ -0,0 +1,155 @@
+from transformers.configuration_utils import PretrainedConfig
+
+
+class Glm5NextTextConfig(PretrainedConfig):
+    model_type = "glm5_next_text"
+    base_config_key = "text_config"
+    keys_to_ignore_at_inference = ["past_key_values"]
+    attribute_map = {"num_local_experts": "n_routed_experts"}
+
+    def __init__(
+        self,
+        vocab_size: int = 154880,
+        hidden_size: int = 4096,
+        intermediate_size: int = 12288,
+        moe_intermediate_size: int = 2048,
+        num_hidden_layers: int = 45,
+        num_attention_heads: int = 64,
+        num_key_value_heads: int = 64,
+        hidden_act: str = "silu",
+        max_position_embeddings: int = 1048576,
+        initializer_range: float = 0.02,
+        rms_norm_eps: float = 1e-5,
+        attention_bias: bool = False,
+        attention_dropout: float = 0.0,
+        layer_types: list[str] | None = None,
+        mlp_layer_types: list[str] | None = None,
+        first_k_dense_replace: int = 3,
+        linear_attn_config: dict | None = None,
+        q_lora_rank: int = 1536,
+        kv_lora_rank: int = 512,
+        qk_nope_head_dim: int = 256,
+        qk_rope_head_dim: int = 0,
+        v_head_dim: int = 256,
+        index_n_heads: int = 32,
+        index_head_dim: int = 128,
+        index_topk: int = 2048,
+        index_kpool: int = 4,
+        index_kpool_compress: bool = True,
+        index_kpool_always_select_tail: bool = True,
+        indexer_types: list[str] | None = None,
+        n_routed_experts: int = 288,
+        num_experts_per_tok: int = 8,
+        n_shared_experts: int = 1,
+        routed_scaling_factor: float = 2.5,
+        norm_topk_prob: bool = True,
+        scoring_func: str = "sigmoid",
+        n_group: int = 1,
+        topk_group: int = 1,
+        swiglu_limit: float = 10.0,
+        hc_mult: int = 4,
+        hc_sinkhorn_iters: int = 20,
+        hc_eps: float = 1e-6,
+        num_nextn_predict_layers: int = 1,
+        pad_token_id: int | None = 154820,
+        bos_token_id: int | None = None,
+        eos_token_id: int | list[int] | None = None,
+        tie_word_embeddings: bool = False,
+        use_cache: bool = False,
+        **kwargs,
+    ) -> None:
+        self.vocab_size = vocab_size
+        self.hidden_size = hidden_size
+        self.intermediate_size = intermediate_size
+        self.moe_intermediate_size = moe_intermediate_size
+        self.num_hidden_layers = num_hidden_layers
+        self.num_attention_heads = num_attention_heads
+        self.num_key_value_heads = num_key_value_heads
+        self.hidden_act = hidden_act
+        self.max_position_embeddings = max_position_embeddings
+        self.initializer_range = initializer_range
+        self.rms_norm_eps = rms_norm_eps
+        self.attention_bias = attention_bias
+        self.attention_dropout = attention_dropout
+
+        self.layer_types = layer_types or ["linear_attention"] * num_hidden_layers
+        self.first_k_dense_replace = first_k_dense_replace
+        self.mlp_layer_types = mlp_layer_types or [
+            "dense" if layer_index < first_k_dense_replace else "sparse" for layer_index in range(num_hidden_layers)
+        ]
+
+        linear_attn_config = linear_attn_config or {}
+        self.linear_attn_config = linear_attn_config
+        self.linear_num_heads = linear_attn_config.get("num_heads", 64)
+        self.linear_head_dim = linear_attn_config.get("head_dim", 128)
+        self.linear_conv_kernel_dim = linear_attn_config.get("short_conv_kernel_size", 4)
+        self.linear_lower_bound = linear_attn_config.get("gate_lower_bound", -5.0)
+
+        self.q_lora_rank = q_lora_rank
+        self.kv_lora_rank = kv_lora_rank
+        self.qk_nope_head_dim = qk_nope_head_dim
+        self.qk_rope_head_dim = qk_rope_head_dim
+        self.qk_head_dim = qk_nope_head_dim + qk_rope_head_dim
+        self.v_head_dim = v_head_dim
+
+        self.index_n_heads = index_n_heads
+        self.index_head_dim = index_head_dim
+        self.index_topk = index_topk
+        self.index_kpool = index_kpool
+        self.index_kpool_compress = index_kpool_compress
+        self.index_kpool_always_select_tail = index_kpool_always_select_tail
+        self.indexer_types = indexer_types or ["full"] * num_hidden_layers
+
+        self.n_routed_experts = n_routed_experts
+        self.num_experts_per_tok = num_experts_per_tok
+        self.n_shared_experts = n_shared_experts
+        self.routed_scaling_factor = routed_scaling_factor
+        self.norm_topk_prob = norm_topk_prob
+        self.scoring_func = scoring_func
+        self.n_group = n_group
+        self.topk_group = topk_group
+        self.swiglu_limit = swiglu_limit
+
+        self.hc_mult = hc_mult
+        self.hc_sinkhorn_iters = hc_sinkhorn_iters
+        self.hc_eps = hc_eps
+        self.num_nextn_predict_layers = num_nextn_predict_layers
+        self.use_cache = use_cache
+
+        # The remaining checkpoint keys pass through untouched so an export's config.json still
+        # carries them: vLLM reads topk_method (selection bias), moe_router_dtype (fp32 routing)
+        # and indexer_rope_interleave off it. head_dim is 0 in the checkpoint, which generic
+        # head_dim readers (perf accounting) would take literally.
+        for unused in ("dtype", "head_dim"):
+            kwargs.pop(unused, None)
+        super().__init__(
+            pad_token_id=pad_token_id,
+            bos_token_id=bos_token_id,
+            eos_token_id=eos_token_id,
+            tie_word_embeddings=tie_word_embeddings,
+            **kwargs,
+        )
+
+
+class Glm5NextConfig(PretrainedConfig):
+    model_type = "glm5_next"
+    sub_configs = {"text_config": Glm5NextTextConfig}
+
+    def __init__(
+        self,
+        text_config: Glm5NextTextConfig | dict | None = None,
+        tie_word_embeddings: bool = False,
+        **kwargs,
+    ) -> None:
+        if isinstance(text_config, dict):
+            text_config = Glm5NextTextConfig(**text_config)
+        self.text_config = text_config or Glm5NextTextConfig()
+
+        # Vision keys are not modelled but pass through, so an export's config.json keeps them.
+        kwargs.setdefault("pad_token_id", self.text_config.pad_token_id)
+        kwargs.setdefault("bos_token_id", self.text_config.bos_token_id)
+        kwargs.setdefault("eos_token_id", self.text_config.eos_token_id)
+        super().__init__(tie_word_embeddings=tie_word_embeddings, **kwargs)
+
+
+__all__ = ["Glm5NextConfig", "Glm5NextTextConfig"]