From c7f9bcf39849a91b6f5209e68467860e4571ef07 Mon Sep 17 00:00:00 2001 From: Zeyu Zhou Date: Mon, 20 Jul 2026 10:28:13 -0700 Subject: [PATCH 1/4] perf: encode gym routed_experts as base64 envelope instead of JSON int lists At long context lengths the router-replay routes ([tokens, layers, topk]) attached to each chat response via .tolist() become multi-MB nested JSON int lists (~40MB at 30k tokens). Every hop on the NeMo Gym HTTP path (vLLM OpenAI server, gym model server, agent, resources server, gym actor) re-parses, pydantic-validates, and re-serializes them at ~1s of single-threaded CPU per hop, which throttles async-GRPO rollout production while generation GPUs sit idle. Encode the routes as a single self-describing base64 string (nrlre1:::, preserving the resolved routed_experts_dtype) so intermediate hops handle one opaque object. The decoder accepts the legacy nested-list format for compatibility. Validated on a 16-node swe1 gym smoke (Qwen3-30B-A3B, prompts 15-28k tokens, 5 steps, seed 42): mean rollout time r3off 30.0s / r3on unpatched 132.2s (4.4x) / r3on patched 35.4s (1.18x), with per-step gen_kl_error identical to the unpatched R3 run (0.0005-0.0008). Requires the NeMo Gym RoutedExperts type alias to accept str (one-line change in nemo_gym/openai_utils.py, submitted separately). Co-Authored-By: Claude Fable 5 Signed-off-by: Zeyu Zhou --- nemo_rl/environments/nemo_gym.py | 3 +- nemo_rl/models/generation/vllm/utils.py | 11 ++- nemo_rl/utils/routed_experts_codec.py | 99 +++++++++++++++++++ .../unit/models/generation/test_vllm_utils.py | 13 ++- tests/unit/utils/test_routed_experts_codec.py | 62 ++++++++++++ 5 files changed, 181 insertions(+), 7 deletions(-) create mode 100644 nemo_rl/utils/routed_experts_codec.py create mode 100644 tests/unit/utils/test_routed_experts_codec.py 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..5f060e0107f --- /dev/null +++ b/nemo_rl/utils/routed_experts_codec.py @@ -0,0 +1,99 @@ +# 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 numpy and torch. +""" + +import base64 +from typing import Any, Union + +import numpy as np +import torch + +_MAGIC = "nrlre1" +_NP_DTYPES = {"int8": np.int8, "int16": np.int16, "int32": np.int32} +_TORCH_DTYPE_NAMES = { + torch.int8: "int8", + torch.int16: "int16", + torch.int32: "int32", +} + + +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(_NP_DTYPES)}." + ) + arr = routed_experts.detach().cpu().numpy() + tokens, num_layers, topk = arr.shape + data = base64.b64encode(np.ascontiguousarray(arr).tobytes()).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 + if dtype_name not in _NP_DTYPES: + 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." + ) + arr = np.frombuffer(base64.b64decode(data), dtype=_NP_DTYPES[dtype_name]) + expected = shape[0] * shape[1] * shape[2] + if arr.size != expected: + raise ValueError( + f"routed_experts envelope has {arr.size} elements, expected " + f"{expected} for shape {shape}." + ) + # frombuffer views are read-only; copy() yields a writable array. + return torch.from_numpy(arr.reshape(shape).copy()).to(dtype) 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)) From d6be3659377e9186a92fcde5ba6fc35822c141a8 Mon Sep 17 00:00:00 2001 From: Zeyu Zhou Date: Mon, 20 Jul 2026 15:46:58 -0700 Subject: [PATCH 2/4] ci: add routed_experts_codec.py to pyrefly project-includes Co-Authored-By: Claude Fable 5 Signed-off-by: Zeyu Zhou --- pyrefly.toml | 1 + 1 file changed, 1 insertion(+) 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", From 0cbe0ac120f5c19d579e8077b31567a52413d3ce Mon Sep 17 00:00:00 2001 From: Zeyu Zhou Date: Mon, 20 Jul 2026 17:57:27 -0700 Subject: [PATCH 3/4] perf: encode routed_experts buffer zero-copy via memoryview Review suggestion from @ZhiyuLi-Nvidia: tobytes() materializes a full-size intermediate copy of the payload; memoryview on the contiguous array feeds b64encode directly. Byte-identical output. Co-Authored-By: Claude Fable 5 Signed-off-by: Zeyu Zhou --- nemo_rl/utils/routed_experts_codec.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/nemo_rl/utils/routed_experts_codec.py b/nemo_rl/utils/routed_experts_codec.py index 5f060e0107f..8fbf85cbbb9 100644 --- a/nemo_rl/utils/routed_experts_codec.py +++ b/nemo_rl/utils/routed_experts_codec.py @@ -60,9 +60,11 @@ def encode_routed_experts(routed_experts: torch.Tensor) -> str: f"Unsupported routed_experts dtype {routed_experts.dtype}; " f"expected one of {sorted(_NP_DTYPES)}." ) - arr = routed_experts.detach().cpu().numpy() + arr = routed_experts.detach().cpu().contiguous().numpy() tokens, num_layers, topk = arr.shape - data = base64.b64encode(np.ascontiguousarray(arr).tobytes()).decode("ascii") + # 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}" From 79fb96bd7666db4189a5d21e12a52ae5e74c5775 Mon Sep 17 00:00:00 2001 From: Zeyu Zhou Date: Mon, 20 Jul 2026 18:08:24 -0700 Subject: [PATCH 4/4] refactor: decode routed_experts via torch.frombuffer on a writable bytearray Review suggestion from @ZhiyuLi-Nvidia: decode into a bytearray and let torch.frombuffer share its memory instead of round-tripping through a read-only numpy view plus copy(). Copy count is unchanged in the common case (the bytearray construction is the one unavoidable copy), but this drops the numpy dependency, avoids a second copy when the target dtype differs from the wire dtype, and validates base64 strictly. Signed-off-by: Zeyu Zhou --- nemo_rl/utils/routed_experts_codec.py | 35 +++++++++++++++------------ 1 file changed, 20 insertions(+), 15 deletions(-) diff --git a/nemo_rl/utils/routed_experts_codec.py b/nemo_rl/utils/routed_experts_codec.py index 8fbf85cbbb9..e7cc41be2c5 100644 --- a/nemo_rl/utils/routed_experts_codec.py +++ b/nemo_rl/utils/routed_experts_codec.py @@ -25,22 +25,21 @@ "nrlre1::xx:" This module must stay importable inside the NeMo Gym actor, so it may only -depend on numpy and torch. +depend on torch. """ import base64 from typing import Any, Union -import numpy as np import torch _MAGIC = "nrlre1" -_NP_DTYPES = {"int8": np.int8, "int16": np.int16, "int32": np.int32} -_TORCH_DTYPE_NAMES = { - torch.int8: "int8", - torch.int16: "int16", - torch.int32: "int32", +_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: @@ -58,7 +57,7 @@ def encode_routed_experts(routed_experts: torch.Tensor) -> str: if dtype_name is None: raise ValueError( f"Unsupported routed_experts dtype {routed_experts.dtype}; " - f"expected one of {sorted(_NP_DTYPES)}." + f"expected one of {sorted(_WIRE_TORCH_DTYPES)}." ) arr = routed_experts.detach().cpu().contiguous().numpy() tokens, num_layers, topk = arr.shape @@ -83,19 +82,25 @@ def decode_routed_experts(payload: Union[str, Any], dtype: torch.dtype) -> torch f"'{_MAGIC}:::' envelope." ) _, dtype_name, shape_str, data = parts - if dtype_name not in _NP_DTYPES: + 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." ) - arr = np.frombuffer(base64.b64decode(data), dtype=_NP_DTYPES[dtype_name]) + # 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 arr.size != expected: + if len(raw) != expected * wire_dtype.itemsize: raise ValueError( - f"routed_experts envelope has {arr.size} elements, expected " - f"{expected} for shape {shape}." + f"routed_experts envelope has {len(raw) // wire_dtype.itemsize} " + f"elements, expected {expected} for shape {shape}." ) - # frombuffer views are read-only; copy() yields a writable array. - return torch.from_numpy(arr.reshape(shape).copy()).to(dtype) + 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)