feat: size routed_experts carry dtype from the model's expert count

nemo_rl/algorithms/grpo.py

@@ -93,7 +93,10 @@
     run_async_nemo_gym_rollout,
     run_multi_turn_rollout,
 )
-from nemo_rl.models.generation.interfaces import GenerationInterface
+from nemo_rl.models.generation.interfaces import (
+    GenerationInterface,
+    resolve_routed_experts_dtype_name_for_model,
+)
 from nemo_rl.models.generation.megatron import MegatronGeneration
 from nemo_rl.models.generation.sglang.config import SGLangConfig
 from nemo_rl.models.generation.sglang.sglang_generation import SGLangGeneration
@@ -588,6 +591,11 @@ def _spinup_nemo_gym(base_urls, model_name):
             invalid_tool_call_patterns=invalid_tool_call_patterns,
             thinking_tags=thinking_tags,
             require_routed_experts=router_replay_enabled(policy_config),
+            routed_experts_dtype=(
+                resolve_routed_experts_dtype_name_for_model(model_name)
+                if router_replay_enabled(policy_config)
+                else "int16"
+            ),
             initial_global_config_dict=nemo_gym_dict,
         )
         nemo_gym_opts = {}

nemo_rl/environments/nemo_gym.py

@@ -29,6 +29,15 @@
 from nemo_rl.environments.interfaces import EnvironmentInterface
 from nemo_rl.utils.timer import Timer
 
+# Kept local (not imported from models.generation) so the gym actor stays free of
+# generation-module imports. Must cover every name resolve_routed_experts_dtype
+# can produce.
+_ROUTED_EXPERTS_DTYPES = {
+    "int8": torch.int8,
+    "int16": torch.int16,
+    "int32": torch.int32,
+}
+
 DEFAULT_INVALID_TOOL_CALL_PATTERNS = [
     "<tool_call>",
     "</tool_call>",
@@ -80,6 +89,9 @@ class NemoGymConfig(TypedDict):
     require_routed_experts: NotRequired[
         bool
     ]  # Require Gym output items to carry R3 routed_experts
+    routed_experts_dtype: NotRequired[
+        str
+    ]  # Carry dtype name for routed_experts tensors ("int8"/"int16"/"int32"), resolved from the model's expert count
 
 
 def _detect_invalid_tool_call_and_malformed_thinking(
@@ -351,7 +363,12 @@ def _postprocess_nemo_gym_to_nemo_rl_result(
 
             routed_experts = None
             if routed_experts_raw is not None:
-                routed_experts = torch.as_tensor(routed_experts_raw, dtype=torch.int32)
+                routed_experts_dtype = _ROUTED_EXPERTS_DTYPES[
+                    self.cfg.get("routed_experts_dtype", "int16")
+                ]
+                routed_experts = torch.as_tensor(
+                    routed_experts_raw, dtype=routed_experts_dtype
+                )
                 if routed_experts.dim() != 3:
                     raise ValueError(
                         "NeMo Gym returned routed_experts with invalid shape. "

nemo_rl/models/generation/interfaces.py

@@ -12,13 +12,75 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 from abc import ABC, abstractmethod
-from typing import Any, NotRequired, TypedDict, Union
+from typing import Any, NotRequired, Optional, TypedDict, Union
 
 import ray
 import torch
 
 from nemo_rl.distributed.batched_data_dict import BatchedDataDict
 
+# Routed-expert index tensors ([seq, layers, topk]) are carried in the narrowest
+# signed dtype that fits ids 0..num_experts-1 plus the -1 missing-route sentinel:
+# int8 for <=128 experts (e.g. Qwen3-MoE), int16 for <=32768 (e.g. DeepSeek-V3),
+# int32 beyond. This shrinks message logs, transports, replay buffers, and
+# checkpoints 2-4x vs int32. The Megatron replay install converts to int64 at the
+# gather site, so training math is unaffected. When the expert count cannot be
+# determined, fall back to int16.
+ROUTED_EXPERTS_FALLBACK_DTYPE = torch.int16
+
+_ROUTED_EXPERTS_DTYPE_NAMES = {
+    torch.int8: "int8",
+    torch.int16: "int16",
+    torch.int32: "int32",
+}
+
+
+def get_num_routed_experts(hf_config: Any) -> Optional[int]:
+    """Best-effort read of the routed-expert count from a HF model config.
+
+    Checks the attribute names used by the common MoE architectures (Qwen-MoE,
+    DeepSeek, Mixtral), including nested ``text_config`` for VLMs. Returns None
+    for dense models or unrecognized configs.
+    """
+    for owner in (hf_config, getattr(hf_config, "text_config", None)):
+        if owner is None:
+            continue
+        for attr in ("num_experts", "n_routed_experts", "num_local_experts"):
+            value = getattr(owner, attr, None)
+            if isinstance(value, int) and value > 0:
+                return value
+    return None
+
+
+def resolve_routed_experts_dtype(num_experts: Optional[int]) -> torch.dtype:
+    """Return the narrowest signed dtype that fits expert ids and the -1 sentinel."""
+    if num_experts is None:
+        return ROUTED_EXPERTS_FALLBACK_DTYPE
+    if num_experts - 1 <= torch.iinfo(torch.int8).max:
+        return torch.int8
+    if num_experts - 1 <= torch.iinfo(torch.int16).max:
+        return torch.int16
+    return torch.int32
+
+
+def resolve_routed_experts_dtype_name_for_model(model_name: str) -> str:
+    """Resolve the routed-experts carry dtype name ("int8"/"int16"/"int32") for a model.
+
+    Used where only the model name is available (e.g. building the NeMo-Gym env
+    config on the driver). Falls back to the default dtype name if the config
+    cannot be loaded.
+    """
+    # Deferred import: transformers config loading is only needed for this sizing.
+    from transformers import AutoConfig
+
+    try:
+        hf_config = AutoConfig.from_pretrained(model_name, trust_remote_code=True)
+    except (OSError, ValueError):
+        return _ROUTED_EXPERTS_DTYPE_NAMES[ROUTED_EXPERTS_FALLBACK_DTYPE]
+    return _ROUTED_EXPERTS_DTYPE_NAMES[
+        resolve_routed_experts_dtype(get_num_routed_experts(hf_config))
+    ]
+
 
 def verify_right_padding(
     data: Union[

nemo_rl/models/generation/vllm/utils.py

@@ -18,10 +18,42 @@
 import torch
 
 from nemo_rl.distributed.batched_data_dict import BatchedDataDict
-from nemo_rl.models.generation.interfaces import GenerationDatumSpec
+from nemo_rl.models.generation.interfaces import (
+    ROUTED_EXPERTS_FALLBACK_DTYPE,
+    GenerationDatumSpec,
+)
 
 R3_MISSING_ROUTE_SENTINEL = -1
 
+# The expert-id range vs carry dtype is model-constant, so it is verified on the
+# first non-empty routed-experts tensor per process and skipped afterwards.
+G_ROUTED_EXPERTS_RANGE_CHECKED = False
+
+
+def _as_routed_experts_tensor(
+    value: Any, *, device: torch.device, dtype: torch.dtype
+) -> torch.Tensor:
+    """Convert backend routed-expert ids to the resolved carry dtype.
+
+    Guards against expert ids overflowing ``dtype`` before the narrowing cast,
+    which would otherwise wrap silently (e.g. if the expert count was
+    mis-detected when resolving the dtype).
+    """
+    global G_ROUTED_EXPERTS_RANGE_CHECKED
+    tensor = torch.as_tensor(value, device=device)
+    if not G_ROUTED_EXPERTS_RANGE_CHECKED and tensor.numel() > 0:
+        max_id = int(tensor.max())
+        limit = torch.iinfo(dtype).max
+        if max_id > limit:
+            raise ValueError(
+                f"routed expert id {max_id} exceeds the resolved carry dtype "
+                f"{dtype} (max {limit}); the model's expert count was likely "
+                "mis-detected (see resolve_routed_experts_dtype in "
+                "nemo_rl.models.generation.interfaces)."
+            )
+        G_ROUTED_EXPERTS_RANGE_CHECKED = True
+    return tensor.to(dtype=dtype)
+
 
 def format_prompt_for_vllm_generation(
     data: BatchedDataDict[GenerationDatumSpec], sample_idx: Optional[int] = None
@@ -114,15 +146,20 @@ def pad_and_align_routed_expert_indices(
     require_complete_routed_experts: bool = False,
     allow_missing_routed_experts_fallback: bool = True,
     return_stats: bool = False,
+    routed_experts_dtype: torch.dtype = ROUTED_EXPERTS_FALLBACK_DTYPE,
 ) -> Optional[torch.Tensor] | tuple[Optional[torch.Tensor], dict[str, int]]:
-    """Return full-sequence-aligned routed experts as ``[S, L, topk]`` int32."""
+    """Return full-sequence-aligned routed experts as ``[S, L, topk]`` in ``routed_experts_dtype``."""
     routed = getattr(completion_output, "routed_experts", None)
     prompt_routed = getattr(request_output, "prompt_routed_experts", None)
 
     if prompt_routed is not None:
-        prompt_routed = torch.as_tensor(prompt_routed, dtype=torch.int32, device=device)
+        prompt_routed = _as_routed_experts_tensor(
+            prompt_routed, device=device, dtype=routed_experts_dtype
+        )
     if routed is not None:
-        routed = torch.as_tensor(routed, dtype=torch.int32, device=device)
+        routed = _as_routed_experts_tensor(
+            routed, device=device, dtype=routed_experts_dtype
+        )
 
     if prompt_routed is not None and routed is not None:
         routed = torch.cat((prompt_routed, routed), dim=0)
@@ -180,15 +217,14 @@ def pad_and_align_routed_expert_indices(
 
     default_route = torch.arange(
         routed.shape[2],
-        dtype=torch.int32,
+        dtype=routed_experts_dtype,
         device=device,
     )
     full = (
         default_route.view(1, 1, -1)
         .expand(padded_length, routed.shape[1], routed.shape[2])
         .clone()
     )
-    full = full.to(dtype=torch.int32)
     routes_to_copy = min(expected_routes, routed.shape[0])
     if routes_to_copy > 0:
         full[:routes_to_copy] = routed[:routes_to_copy].to(device=device)
@@ -203,6 +239,7 @@ def attach_routed_experts_to_chat_response_choices(
     *,
     device: torch.device,
     logger: Any = None,
+    routed_experts_dtype: torch.dtype = ROUTED_EXPERTS_FALLBACK_DTYPE,
 ) -> Any:
     """Attach aligned routed experts to OpenAI chat response choices."""
     outputs_by_index = {
@@ -229,6 +266,7 @@ def attach_routed_experts_to_chat_response_choices(
             device=device,
             require_complete_routed_experts=True,
             return_stats=True,
+            routed_experts_dtype=routed_experts_dtype,
         )
         if not isinstance(routed_result, tuple):
             raise RuntimeError(
@@ -254,7 +292,9 @@ def attach_routed_experts_to_chat_response_choices(
                 r3_stats["actual_routes"],
                 r3_stats["expected_routes"],
             )
-        choice.message.routed_experts = routed_experts.to(dtype=torch.int32).tolist()
+        choice.message.routed_experts = routed_experts.to(
+            dtype=routed_experts_dtype
+        ).tolist()
 
     if len(attached_choice_indices) != len(choices):
         missing_choice_indices = sorted(

nemo_rl/models/generation/vllm/vllm_worker.py

@@ -30,8 +30,11 @@
 )
 from nemo_rl.distributed.worker_group_utils import get_nsight_config_if_pattern_matches
 from nemo_rl.models.generation.interfaces import (
+    ROUTED_EXPERTS_FALLBACK_DTYPE,
     GenerationDatumSpec,
     GenerationOutputSpec,
+    get_num_routed_experts,
+    resolve_routed_experts_dtype,
     verify_right_padding,
 )
 from nemo_rl.models.generation.vllm.config import VllmConfig
@@ -203,6 +206,8 @@ def _init_config(
         """Lightweight config setup. No model loading, no heavy imports."""
         self.cfg = config
         self.model_name = self.cfg["model_name"]
+        # Refined from the model's expert count in _load_model.
+        self.routed_experts_dtype = ROUTED_EXPERTS_FALLBACK_DTYPE
         self.tensor_parallel_size = self.cfg["vllm_cfg"]["tensor_parallel_size"]
         self.pipeline_parallel_size = self.cfg["vllm_cfg"]["pipeline_parallel_size"]
         self.expert_parallel_size = self.cfg["vllm_cfg"]["expert_parallel_size"]
@@ -349,6 +354,9 @@ def _load_model(self, bundle_indices, seed):
         # Override HF config for gpt-oss models to ensure compatibility with megatron
         # The megatron --> hf export is done in bf16, so we disable quantization
         hf_config = AutoConfig.from_pretrained(self.model_name, trust_remote_code=True)
+        self.routed_experts_dtype = resolve_routed_experts_dtype(
+            get_num_routed_experts(hf_config)
+        )
         if "GptOssForCausalLM" in getattr(hf_config, "architectures", []):
             if "quantization_config" in hf_config:
                 assert load_format == "dummy", (
@@ -734,6 +742,7 @@ def generate(
                 device=input_ids.device,
                 require_complete_routed_experts=return_routed_experts,
                 return_stats=True,
+                routed_experts_dtype=self.routed_experts_dtype,
             )
             if return_routed_experts and full_routed_experts is None:
                 raise RuntimeError(

nemo_rl/models/generation/vllm/vllm_worker_async.py

@@ -717,6 +717,7 @@ async def capture_result_generator():
                     final_res,
                     device=torch.device("cpu"),
                     logger=LOGGER,
+                    routed_experts_dtype=worker_self.routed_experts_dtype,
                 )
 
         class NeMoRLOpenAIServingChat(NeMoRLOpenAIServingChatMixin, OpenAIServingChat):
@@ -1212,6 +1213,7 @@ async def process_single_sample(sample_idx):
                 device=original_input_ids_single_row.device,
                 require_complete_routed_experts=return_routed_experts,
                 return_stats=True,
+                routed_experts_dtype=self.routed_experts_dtype,
             )
             if return_routed_experts and routed_experts is None:
                 raise RuntimeError(

tests/unit/environments/test_nemo_gym_router_replay.py

@@ -88,3 +88,34 @@ class _MockSelf:
         NemoGym.__ray_metadata__.modified_class._postprocess_nemo_gym_to_nemo_rl_result(
             _MockSelf(), nemo_gym_result, _Tokenizer()
         )
+
+
+def test_nemo_gym_postprocess_casts_routed_experts_to_configured_dtype():
+    import torch
+
+    nemo_gym_result = {
+        "response": {
+            "output": [
+                {
+                    "prompt_token_ids": [1, 2],
+                    "generation_token_ids": [3],
+                    "generation_log_probs": [-0.1],
+                    "routed_experts": _routes(3),
+                },
+            ]
+        },
+        "responses_create_params": {"input": []},
+    }
+
+    class _MockSelf:
+        cfg = {"require_routed_experts": True, "routed_experts_dtype": "int8"}
+
+    result = (
+        NemoGym.__ray_metadata__.modified_class._postprocess_nemo_gym_to_nemo_rl_result(
+            _MockSelf(), nemo_gym_result, _Tokenizer()
+        )
+    )
+
+    for message in result["message_log"]:
+        if "routed_experts" in message:
+            assert message["routed_experts"].dtype == torch.int8