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

README.md

@@ -49,6 +49,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

@@ -21,6 +21,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 | ✅ |

src/prime_rl/inference/patches.py

@@ -21,6 +21,7 @@ def apply_shared_vllm_patches():
     monkey_patch_kv_xfer_finished_tolerate_freed()
     monkey_patch_online_fp8_parameter_cast()
     monkey_patch_deepseek_v4_allowed_layer_types()
+    monkey_patch_glm5_next_allowed_layer_types()
     monkey_patch_tokenize_params_validation()
     monkey_patch_strip_routed_experts_from_chat()
     monkey_patch_dp_coordinator_startup_timeout()
@@ -47,6 +48,13 @@ def monkey_patch_deepseek_v4_allowed_layer_types():
     allow_deepseek_v4_layer_types()
 
 
+def monkey_patch_glm5_next_allowed_layer_types():
+    """Let vLLM's `Glm5NextConfig` construct under the pinned transformers (see `allow_glm5_next_layer_types`)."""
+    from prime_rl.utils.transformers_compat import allow_glm5_next_layer_types
+
+    allow_glm5_next_layer_types()
+
+
 def monkey_patch_online_fp8_parameter_cast():
     """Pass plain tensors to vLLM's compiled block-FP8 caster.
 

src/prime_rl/trainer/model.py

@@ -41,6 +41,7 @@
 from prime_rl.trainer.lora import apply_lora_to_model, freeze_all_except_lora_and_specified, strip_lora_from_state_dict
 from prime_rl.trainer.models import (
     AutoModelForCausalLMPrimeRL,
+    PreTrainedModelPrimeRL,
     PrimeLmOutput,
     cast_float_and_contiguous,
     get_custom_causal_lm_cls,
@@ -53,6 +54,7 @@
     get_fsdp_shard_placement_fn,
     write_back_loaded_packed_parameters,
 )
+from prime_rl.trainer.models.glm5_next.sparse_attention import KPoolIndexer
 from prime_rl.trainer.models.glm_moe_dsa.sparse_mla_attention import Indexer
 from prime_rl.trainer.models.layers.fp8_linear import replace_linear_with_fp8_blockwise_linear
 from prime_rl.trainer.models.layers.lm_head import inject_prime_lm_head
@@ -170,18 +172,24 @@ def apply_fp32_moe_router(model: nn.Module) -> None:
         logger.info(f"Running {num_routers} MoE router gates in fp32")
 
 
+def get_fp32_modules(model: nn.Module) -> list[nn.Module]:
+    """Modules the model keeps in fp32 for compute (see `PreTrainedModelPrimeRL.fp32_modules`)."""
+    return model.fp32_modules() if isinstance(model, PreTrainedModelPrimeRL) else []
+
+
 def get_full_offload_dtype_policy(
     model: nn.Module,
     config: ModelConfig,
 ) -> dict[int, tuple[torch.dtype, torch.dtype]]:
     """Return persistent parameter and reduced-gradient dtypes for full offload."""
     # FSDP casts reduced gradients back to the persistent sharded-parameter dtype.
     policy = {id(param): (torch.bfloat16, torch.bfloat16) for param in model.parameters() if param.is_floating_point()}
-    if config.moe_router_dtype != "float32":
-        return policy
+    fp32_modules = get_fp32_modules(model)
+    if config.moe_router_dtype == "float32":
+        fp32_modules += [moe.router for moe in iter_moe_blocks(model)]
 
-    for moe in iter_moe_blocks(model):
-        for param in moe.router.parameters():
+    for module in fp32_modules:
+        for param in module.parameters():
             if param.is_floating_point():
                 policy[id(param)] = (torch.float32, torch.float32)
     return policy
@@ -203,7 +211,7 @@ def freeze_sparse_indexer(model: nn.Module) -> None:
     num_frozen = 0
 
     for module in model.modules():
-        if isinstance(module, (Indexer, DeepseekV4Indexer, SparseAttentionIndexer)):
+        if isinstance(module, (Indexer, DeepseekV4Indexer, SparseAttentionIndexer, KPoolIndexer)):
             for param in module.parameters():
                 param.requires_grad = False
                 num_frozen += 1
@@ -548,6 +556,25 @@ def shard_placement_fn(parameter: nn.Parameter) -> Shard | None:
     transformer_layers = language_model.layers
 
     fullgraph = config.compile is not None and config.compile.fullgraph
+    fp32_modules = get_fp32_modules(model)
+    if fp32_modules and fullgraph:
+        raise ValueError(
+            "model.compile.fullgraph=true is not supported for this model: its fp32 modules are FSDP units "
+            "inside the compiled blocks, and dynamo cannot trace FSDP hooks."
+        )
+    for module in fp32_modules:
+        # The module converts its inputs itself, and the code around it may use the parameters its
+        # forward returns, so they stay unsharded until the backward.
+        fully_shard(
+            module,
+            mesh=hsdp_mesh,
+            mp_policy=MixedPrecisionPolicy(
+                param_dtype=torch.float32, reduce_dtype=torch.float32, cast_forward_inputs=False
+            ),
+            offload_policy=offload_policy,
+            reshard_after_forward=False,
+            shard_placement_fn=shard_placement_fn,
+        )
     for transformer_block in transformer_layers:
         for module in transformer_block.modules():
             if isinstance(module, NGramEmbedding) and parallel_dims.get_mesh("head").size() > 1:
@@ -852,15 +879,19 @@ def apply_quantization(model: nn.Module, config: ModelConfig) -> None:
     if quant is None:
         return
 
+    ignore_patterns = quant.ignore_patterns
+    if isinstance(model, PreTrainedModelPrimeRL):
+        ignore_patterns = [*ignore_patterns, *model.quantization_ignore_patterns()]
+
     if isinstance(quant, FP8Config):
-        replace_linear_with_fp8_blockwise_linear(model, ignore_modules=quant.ignore_patterns)
+        replace_linear_with_fp8_blockwise_linear(model, ignore_modules=ignore_patterns)
     elif isinstance(quant, MXFP8Config):
         capability = torch.cuda.get_device_capability()
         if capability[0] < 10:
             raise ValueError(
                 f"MXFP8 quantization requires Blackwell (SM100+), but device is SM{capability[0]}{capability[1]}."
             )
-        replace_linear_with_mxfp8_linear(model, recipe=quant.recipe, ignore_modules=quant.ignore_patterns)
+        replace_linear_with_mxfp8_linear(model, recipe=quant.recipe, ignore_modules=ignore_patterns)
 
 
 def configure_trainable_parameters(model: nn.Module, config: ModelConfig) -> nn.Module | None:
@@ -968,6 +999,9 @@ def setup_model(
     if config.moe_router_dtype == "float32":
         apply_fp32_moe_router(model)
 
+    for module in get_fp32_modules(model):
+        module.to(torch.float32)
+
     # A sparse-attention indexer runs its forward under torch.no_grad(), so it is never
     # trainable. Freeze it so optimizer state stays symmetric across checkpoint save/resume.
     # No-op for models without a sparse indexer.

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
@@ -40,6 +41,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)
@@ -64,6 +67,8 @@
     (AfmoeConfig, AfmoeForCausalLM),
     (DeepseekV4Config, DeepseekV4ForCausalLM),
     (Glm4MoeConfig, Glm4MoeForCausalLM),
+    (Glm5NextTextConfig, Glm5NextForCausalLM),
+    (Glm5NextConfig, Glm5NextForCausalLM),
     (GlmMoeDsaConfig, GlmMoeDsaForCausalLM),
     (LagunaConfig, LagunaForCausalLM),
     (MiniMaxM2Config, MiniMaxM2ForCausalLM),
@@ -116,6 +121,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,
 }
 
 

src/prime_rl/trainer/models/base.py

@@ -1,7 +1,7 @@
 from dataclasses import dataclass
 from typing import Literal
 
-from torch import Tensor
+from torch import Tensor, nn
 from transformers.modeling_utils import PreTrainedModel
 
 from prime_rl.utils.cp import CPContext
@@ -47,6 +47,22 @@ def keep_in_fp32_for_weight_transfer(cls, name: str) -> bool:
         """
         return False
 
+    def fp32_modules(self) -> list[nn.Module]:
+        """Small modules whose parameters train and compute in fp32 instead of the bf16 compute dtype.
+
+        Each becomes its own FSDP unit with an fp32 policy (fp32 compute dtype under full offload), so
+        it must own only fp32 parameters. FSDP does not cast its inputs, and its parameters stay
+        unsharded from its forward until the backward, so its forward may return them.
+        """
+        return []
+
+    def quantization_ignore_patterns(self) -> list[str]:
+        """Dense linear modules that stay BF16 under ``model.quantization``, beyond the configured patterns.
+
+        For linears that inference keeps in BF16 when it quantizes the rest of the model.
+        """
+        return []
+
     @classmethod
     def from_config(cls, config, trust_remote_code: bool = False, **kwargs):
         """Public from_config that mirrors the Auto class API."""

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",
+]