diff --git a/nemo_rl/environments/nemo_gym.py b/nemo_rl/environments/nemo_gym.py index 77a21b9613a..bbd11b869d7 100644 --- a/nemo_rl/environments/nemo_gym.py +++ b/nemo_rl/environments/nemo_gym.py @@ -27,6 +27,7 @@ _get_node_ip_local, ) from nemo_rl.environments.interfaces import EnvironmentInterface +from nemo_rl.utils.routed_experts_codec import decode_routed_experts from nemo_rl.utils.timer import Timer # Kept local (not imported from models.generation) so the gym actor stays free of @@ -366,7 +367,7 @@ def _postprocess_nemo_gym_to_nemo_rl_result( routed_experts_dtype = _ROUTED_EXPERTS_DTYPES[ self.cfg.get("routed_experts_dtype", "int16") ] - routed_experts = torch.as_tensor( + routed_experts = decode_routed_experts( routed_experts_raw, dtype=routed_experts_dtype ) if routed_experts.dim() != 3: diff --git a/nemo_rl/models/generation/vllm/utils.py b/nemo_rl/models/generation/vllm/utils.py index 4d389b3aedb..176c3d6c56c 100644 --- a/nemo_rl/models/generation/vllm/utils.py +++ b/nemo_rl/models/generation/vllm/utils.py @@ -22,6 +22,7 @@ ROUTED_EXPERTS_FALLBACK_DTYPE, GenerationDatumSpec, ) +from nemo_rl.utils.routed_experts_codec import encode_routed_experts R3_MISSING_ROUTE_SENTINEL = -1 @@ -292,9 +293,13 @@ def attach_routed_experts_to_chat_response_choices( r3_stats["actual_routes"], r3_stats["expected_routes"], ) - choice.message.routed_experts = routed_experts.to( - dtype=routed_experts_dtype - ).tolist() + # Base64 envelope instead of .tolist(): nested JSON int lists cost + # ~1s of CPU per serialize/parse hop at long context lengths and get + # re-validated at every gym HTTP hop; a single string passes through + # the gym chain opaquely. + choice.message.routed_experts = encode_routed_experts( + routed_experts.to(dtype=routed_experts_dtype) + ) if len(attached_choice_indices) != len(choices): missing_choice_indices = sorted( diff --git a/nemo_rl/utils/routed_experts_codec.py b/nemo_rl/utils/routed_experts_codec.py new file mode 100644 index 00000000000..e7cc41be2c5 --- /dev/null +++ b/nemo_rl/utils/routed_experts_codec.py @@ -0,0 +1,106 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Compact wire codec for router-replay routed-expert indices. + +Routed experts for a long-context sample are millions of ints (shape +[tokens, num_moe_layers, topk]). Serialized as nested JSON lists they cost +~1s of single-threaded CPU per serialize/parse hop, and every HTTP hop on +the NeMo Gym path (model server, agent, resources server) pays that again +for pydantic validation and re-serialization. Encoded as a single base64 +string the payload stays one opaque Python object end to end, so +intermediate hops only pay a string copy. + +Envelope format (version 1): + "nrlre1::xx:" + +This module must stay importable inside the NeMo Gym actor, so it may only +depend on torch. +""" + +import base64 +from typing import Any, Union + +import torch + +_MAGIC = "nrlre1" +_WIRE_TORCH_DTYPES = { + "int8": torch.int8, + "int16": torch.int16, + "int32": torch.int32, +} +_TORCH_DTYPE_NAMES = {v: k for k, v in _WIRE_TORCH_DTYPES.items()} + + +def encode_routed_experts(routed_experts: torch.Tensor) -> str: + """Encode a [tokens, num_moe_layers, topk] tensor as a base64 envelope. + + The tensor's own dtype (int8/int16/int32, as resolved by + ``resolve_routed_experts_dtype``) is preserved on the wire. + """ + if routed_experts.dim() != 3: + raise ValueError( + "routed_experts must have shape [tokens, num_moe_layers, topk], " + f"got {tuple(routed_experts.shape)}." + ) + dtype_name = _TORCH_DTYPE_NAMES.get(routed_experts.dtype) + if dtype_name is None: + raise ValueError( + f"Unsupported routed_experts dtype {routed_experts.dtype}; " + f"expected one of {sorted(_WIRE_TORCH_DTYPES)}." + ) + arr = routed_experts.detach().cpu().contiguous().numpy() + tokens, num_layers, topk = arr.shape + # memoryview feeds the buffer to b64encode zero-copy; tobytes() would + # materialize a full-size intermediate copy of the payload. + data = base64.b64encode(memoryview(arr)).decode("ascii") + return f"{_MAGIC}:{dtype_name}:{tokens}x{num_layers}x{topk}:{data}" + + +def decode_routed_experts(payload: Union[str, Any], dtype: torch.dtype) -> torch.Tensor: + """Decode routed experts into a tensor of the requested dtype. + + Accepts the base64 envelope produced by ``encode_routed_experts`` or the + legacy nested-list format. + """ + if not isinstance(payload, str): + return torch.as_tensor(payload, dtype=dtype) + parts = payload.split(":", 3) + if len(parts) != 4 or parts[0] != _MAGIC: + raise ValueError( + "routed_experts string payload is not a valid " + f"'{_MAGIC}:::' envelope." + ) + _, dtype_name, shape_str, data = parts + wire_dtype = _WIRE_TORCH_DTYPES.get(dtype_name) + if wire_dtype is None: + raise ValueError(f"Unsupported routed_experts dtype '{dtype_name}'.") + shape = tuple(int(dim) for dim in shape_str.split("x")) + if len(shape) != 3: + raise ValueError( + f"routed_experts envelope shape '{shape_str}' is not 3-dimensional." + ) + # Decode into a writable bytearray so torch.frombuffer can share its + # memory directly (frombuffer writability follows the underlying buffer; + # the bytearray construction is the single unavoidable copy). + raw = bytearray(base64.b64decode(data, validate=True)) + expected = shape[0] * shape[1] * shape[2] + if len(raw) != expected * wire_dtype.itemsize: + raise ValueError( + f"routed_experts envelope has {len(raw) // wire_dtype.itemsize} " + f"elements, expected {expected} for shape {shape}." + ) + if expected == 0: + return torch.empty(shape, dtype=dtype) + tensor = torch.frombuffer(raw, dtype=wire_dtype).reshape(shape) + return tensor if tensor.dtype == dtype else tensor.to(dtype) diff --git a/pyrefly.toml b/pyrefly.toml index 75315f5d862..ef0d04e4771 100644 --- a/pyrefly.toml +++ b/pyrefly.toml @@ -195,6 +195,7 @@ project-includes = [ "nemo_rl/utils/packed_tensor.py", "nemo_rl/utils/prefetch_venvs.py", "nemo_rl/utils/r3_trace.py", + "nemo_rl/utils/routed_experts_codec.py", "nemo_rl/utils/timer.py", "nemo_rl/utils/venvs.py", "nemo_rl/utils/weight_transfer_http.py", diff --git a/tests/unit/models/generation/test_vllm_utils.py b/tests/unit/models/generation/test_vllm_utils.py index e4e033e0a99..835cfc12e2e 100644 --- a/tests/unit/models/generation/test_vllm_utils.py +++ b/tests/unit/models/generation/test_vllm_utils.py @@ -35,6 +35,11 @@ model_dump_chat_response_with_routed_experts, pad_and_align_routed_expert_indices, ) +from nemo_rl.utils.routed_experts_codec import decode_routed_experts + + +def _decoded_routes(payload: str) -> list: + return decode_routed_experts(payload, dtype=torch.int32).tolist() def _mk_inputs(batch_size: int = 2, seq_len: int = 5): @@ -449,13 +454,15 @@ def test_attach_routed_experts_to_chat_response_choices_reassociates_by_choice_i device=torch.device("cpu"), ) - assert response.choices[0].message.routed_experts == [ + # Routes travel as a base64 string envelope, one opaque object per choice. + assert isinstance(response.choices[0].message.routed_experts, str) + assert _decoded_routes(response.choices[0].message.routed_experts) == [ [[10]], [[11]], [[30]], [[0]], ] - assert response.choices[1].message.routed_experts == [ + assert _decoded_routes(response.choices[1].message.routed_experts) == [ [[10]], [[11]], [[31]], @@ -517,7 +524,7 @@ def test_attach_routed_experts_to_chat_response_choices_warns_on_missing_routes( 2, 4, ) - assert response.choices[0].message.routed_experts == [ + assert _decoded_routes(response.choices[0].message.routed_experts) == [ [[10]], [[11]], [[R3_MISSING_ROUTE_SENTINEL]], diff --git a/tests/unit/utils/test_routed_experts_codec.py b/tests/unit/utils/test_routed_experts_codec.py new file mode 100644 index 00000000000..725660b717b --- /dev/null +++ b/tests/unit/utils/test_routed_experts_codec.py @@ -0,0 +1,62 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import pytest +import torch + +from nemo_rl.utils.routed_experts_codec import ( + decode_routed_experts, + encode_routed_experts, +) + + +@pytest.mark.parametrize("dtype", [torch.int8, torch.int16, torch.int32]) +def test_round_trip_preserves_values_and_wire_dtype(dtype): + routes = torch.randint(0, 100, (7, 3, 2), dtype=torch.int32) + routes[2, 1, 0] = -1 # missing-route sentinel survives signed dtypes + routes = routes.to(dtype) + + payload = encode_routed_experts(routes) + assert isinstance(payload, str) + assert payload.startswith(f"nrlre1:{str(dtype).removeprefix('torch.')}:7x3x2:") + + decoded = decode_routed_experts(payload, dtype=torch.int32) + assert decoded.dtype == torch.int32 + assert torch.equal(decoded, routes.to(torch.int32)) + + +def test_decode_accepts_legacy_nested_lists(): + decoded = decode_routed_experts([[[1, 2]], [[3, 4]]], dtype=torch.int16) + assert decoded.dtype == torch.int16 + assert decoded.tolist() == [[[1, 2]], [[3, 4]]] + + +def test_decode_rejects_malformed_payloads(): + with pytest.raises(ValueError, match="envelope"): + decode_routed_experts("not-an-envelope", dtype=torch.int32) + + good = encode_routed_experts(torch.zeros(2, 3, 2, dtype=torch.int16)) + magic, dtype_name, _, data = good.split(":", 3) + with pytest.raises(ValueError, match="expected"): + # Element count does not match the declared shape. + decode_routed_experts(f"{magic}:{dtype_name}:5x3x2:{data}", dtype=torch.int32) + with pytest.raises(ValueError, match="dtype"): + decode_routed_experts(f"{magic}:int64:2x3x2:{data}", dtype=torch.int32) + + +def test_encode_rejects_bad_inputs(): + with pytest.raises(ValueError, match="shape"): + encode_routed_experts(torch.zeros(3, 2, dtype=torch.int16)) + with pytest.raises(ValueError, match="dtype"): + encode_routed_experts(torch.zeros(2, 3, 2, dtype=torch.int64))