nemo_rl.utils.routed_experts_codec#
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:x
This module must stay importable inside the NeMo Gym actor, so it may only depend on torch.
Module Contents#
Functions#
Encode a [tokens, num_moe_layers, topk] tensor as a base64 envelope. |
|
Decode routed experts into a tensor of the requested dtype. |
Data#
API#
- nemo_rl.utils.routed_experts_codec._MAGIC#
‘nrlre1’
- nemo_rl.utils.routed_experts_codec._WIRE_TORCH_DTYPES#
None
- nemo_rl.utils.routed_experts_codec._TORCH_DTYPE_NAMES#
None
- nemo_rl.utils.routed_experts_codec.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.
- nemo_rl.utils.routed_experts_codec.decode_routed_experts(
- payload: Union[str, Any],
- dtype: torch.dtype,
Decode routed experts into a tensor of the requested dtype.
Accepts the base64 envelope produced by
encode_routed_expertsor the legacy nested-list format.