refactor(trainer): add iter_moe_blocks for the MoE layer lookup

src/prime_rl/trainer/model.py

@@ -1,7 +1,7 @@
 import logging
 import os
 import time
-from collections.abc import Callable
+from collections.abc import Callable, Iterator
 from pathlib import Path
 from typing import cast
 
@@ -125,18 +125,22 @@ def freeze_vision_encoder(model: nn.Module, override_attr: str | None = None) ->
     logger.info(f"Froze {num_frozen} parameters in vision encoder")
 
 
+def iter_moe_blocks(model: nn.Module) -> Iterator[MoE]:
+    """Yield the MoE MLP of each decoder layer, skipping dense layers."""
+    for layer in get_language_model(model).layers:
+        mlp = getattr(layer, "mlp", None)
+        if isinstance(mlp, MoE):
+            yield mlp
+
+
 def freeze_moe_router(model: nn.Module) -> None:
     """Freeze MoE router parameters to maintain stable routing during training."""
     logger = get_logger()
-    language_model = get_language_model(model)
     num_frozen = 0
-
-    for layer in language_model.layers:
-        mlp = layer.mlp if hasattr(layer, "mlp") else layer.feed_forward if hasattr(layer, "feed_forward") else None
-        if isinstance(mlp, MoE):
-            for param in mlp.router.parameters():
-                param.requires_grad = False
-                num_frozen += 1
+    for moe in iter_moe_blocks(model):
+        for param in moe.router.parameters():
+            param.requires_grad = False
+            num_frozen += 1
 
     # No-op for non-MoE models: freeze_moe_router=True is the RL default.
     if num_frozen > 0:
@@ -149,16 +153,12 @@ def apply_fp32_moe_router(model: nn.Module) -> None:
     The FSDP bf16 cast exemption is applied separately in `setup_fsdp`.
     """
     logger = get_logger()
-    language_model = get_language_model(model)
     num_routers = 0
-
-    for layer in language_model.layers:
-        mlp = layer.mlp if hasattr(layer, "mlp") else layer.feed_forward if hasattr(layer, "feed_forward") else None
-        if isinstance(mlp, MoE):
-            mlp.router.to(torch.float32)
-            if isinstance(mlp.router, TokenChoiceTopKRouter):
-                mlp.router.fp32_gate = True
-            num_routers += 1
+    for moe in iter_moe_blocks(model):
+        moe.router.to(torch.float32)
+        if isinstance(moe.router, TokenChoiceTopKRouter):
+            moe.router.fp32_gate = True
+        num_routers += 1
 
     # No-op for non-MoE models: moe_router_dtype='float32' is the default,
     # so absence of MoE routers is the common case, not an error.
@@ -176,13 +176,10 @@ def get_full_offload_dtype_policy(
     if config.moe_router_dtype != "float32":
         return policy
 
-    language_model = get_language_model(model)
-    for layer in language_model.layers:
-        mlp = layer.mlp if hasattr(layer, "mlp") else layer.feed_forward if hasattr(layer, "feed_forward") else None
-        if isinstance(mlp, MoE):
-            for param in mlp.router.parameters():
-                if param.is_floating_point():
-                    policy[id(param)] = (torch.float32, torch.float32)
+    for moe in iter_moe_blocks(model):
+        for param in moe.router.parameters():
+            if param.is_floating_point():
+                policy[id(param)] = (torch.float32, torch.float32)
     return policy
 
 
@@ -214,14 +211,10 @@ def freeze_sparse_indexer(model: nn.Module) -> None:
 def apply_force_balanced_routing(model: nn.Module) -> None:
     """Force MoE token-choice routers into round-robin assignment for fake-data smoke tests."""
     logger = get_logger()
-    language_model = get_language_model(model)
     num_routers = 0
-
-    for layer in language_model.layers:
-        mlp = layer.mlp if hasattr(layer, "mlp") else layer.feed_forward if hasattr(layer, "feed_forward") else None
-        if isinstance(mlp, MoE):
-            mlp.router.force_balanced = True
-            num_routers += 1
+    for moe in iter_moe_blocks(model):
+        moe.router.force_balanced = True
+        num_routers += 1
 
     if num_routers == 0:
         raise ValueError("No MoE routers found to force-balance. Is this a MoE model?")
@@ -247,13 +240,7 @@ def get_load_balance_stats(
     """
     per_layer_max_vio = []
     per_layer_routing_confidence = []
-    language_model = get_language_model(model)
-    block_mlps = []
-    for transformer_block in language_model.layers:
-        # This is necessary for models that have mixed dense layers
-        block_mlp = getattr(transformer_block, "mlp", None)
-        if block_mlp is not None and hasattr(block_mlp, "tokens_per_expert"):
-            block_mlps.append(block_mlp)
+    block_mlps = list(iter_moe_blocks(model))
     if not block_mlps:
         return {"max_vio": None, "routing_confidence": None, "tokens_per_expert": torch.empty(0, 0)}
 
@@ -332,7 +319,7 @@ def compute_expert_load_stats(tokens_per_expert: Tensor) -> dict[str, Tensor]:
 
 def get_expert_load_stats(tokens_per_expert: Tensor, group: dist.ProcessGroup) -> dict[str, float]:
     """Sum a step's per-layer expert counts across the group (every rank routes distinct tokens) and compute load stats."""
-    if tokens_per_expert.numel() == 0:  # MoE model without prime-rl MoE layers (HF impl)
+    if tokens_per_expert.numel() == 0:
         return {}
     dist.all_reduce(tokens_per_expert, op=dist.ReduceOp.SUM, group=group)
     return {name: value.item() for name, value in compute_expert_load_stats(tokens_per_expert).items()}