`nemo_rl.models.generation.interfaces`#

Module Contents#

Classes#

`ResourcesConfig`
`OptionalResourcesConfig`
`ColocationConfig`
`CheckpointEngineConfig`	Normalized internal configuration for checkpoint-engine refit.
`GenerationConfig`	Configuration for generation.
`GenerationSamplingParams`	Sampling profile threaded explicitly through rollout entry points.
`GenerationDatumSpec`	Specification for input data required by generation models.
`GenerationOutputSpec`	Specification for output data returned by generation models.
`CollectiveSenderSpec`	Policy-side protocol and packing geometry for NCCL weight transfer.
`GenerationInterface`	Abstract base class defining the interface for RL policies.

Functions#

`_warn_unsupported_in_flight_refit_pause_once`	Warn once per backend type when native refit pause is unavailable.
`get_num_routed_experts`	Best-effort read of the routed-expert count from a HF model config.
`resolve_routed_experts_dtype`	Return the narrowest signed dtype that fits expert ids and the -1 sentinel.
`resolve_routed_experts_dtype_name_for_model`	Resolve the routed-experts carry dtype name (“int8”/”int16”/”int32”) for a model.
`verify_right_padding`	Verify that a tensor is right-padded according to the provided lengths.
`should_use_async_rollouts`	Determine whether a generation backend uses asynchronous rollouts.
`reject_unenforceable_refit_deadline`	Refuse a refit deadline the backend cannot actually apply.

Data#

`ROUTED_EXPERTS_FALLBACK_DTYPE`
`ROUTED_EXPERTS_MISSING_ROUTE_SENTINEL`
`_ROUTED_EXPERTS_DTYPE_NAMES`

API#

nemo_rl.models.generation.interfaces.ROUTED_EXPERTS_FALLBACK_DTYPE#: None

nemo_rl.models.generation.interfaces.ROUTED_EXPERTS_MISSING_ROUTE_SENTINEL#: None

nemo_rl.models.generation.interfaces._ROUTED_EXPERTS_DTYPE_NAMES#: None

nemo_rl.models.generation.interfaces._warn_unsupported_in_flight_refit_pause_once( backend_name: str, ) → None#: Warn once per backend type when native refit pause is unavailable.

nemo_rl.models.generation.interfaces.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.

nemo_rl.models.generation.interfaces.resolve_routed_experts_dtype( num_experts: Optional[int], ) → torch.dtype#: Return the narrowest signed dtype that fits expert ids and the -1 sentinel.

nemo_rl.models.generation.interfaces.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.

nemo_rl.models.generation.interfaces.verify_right_padding( data: Union[nemo_rl.distributed.batched_data_dict.BatchedDataDict[GenerationDatumSpec], nemo_rl.distributed.batched_data_dict.BatchedDataDict[GenerationOutputSpec]], pad_value: int = 0, raise_error: bool = True, ) → tuple[bool, Union[str, None]]#

Verify that a tensor is right-padded according to the provided lengths.

Parameters:

data –
The BatchedDataDict to check, containing either:
- For GenerationDatumSpec: input_ids and input_lengths
- For GenerationOutputSpec: output_ids and unpadded_sequence_lengths
pad_value – The expected padding value (default: 0)
raise_error – Whether to raise an error if wrong padding is detected

Returns:

Tuple of (is_right_padded, error_message)

is_right_padded: True if right padding confirmed, False otherwise
error_message: None if properly padded, otherwise a description of the issue

class nemo_rl.models.generation.interfaces.ResourcesConfig#

Bases: typing.TypedDict

gpus_per_node: int#: None

num_nodes: int#: None

class nemo_rl.models.generation.interfaces.OptionalResourcesConfig#

Bases: typing.TypedDict

gpus_per_node: int | None#: None

num_nodes: int | None#: None

class nemo_rl.models.generation.interfaces.ColocationConfig#

Bases: typing.TypedDict

enabled: bool#: None

resources: nemo_rl.models.generation.interfaces.OptionalResourcesConfig#: None

class nemo_rl.models.generation.interfaces.CheckpointEngineConfig#

Bases: typing.TypedDict

Normalized internal configuration for checkpoint-engine refit.

Initialization

Initialize self. See help(type(self)) for accurate signature.

backend: str#: None

update_weights_bucket_memory_ratio: float#: None

engine_kwargs: dict[str, dict[str, Any]]#: None

class nemo_rl.models.generation.interfaces.GenerationConfig#

Bases: typing.TypedDict

Configuration for generation.

Initialization

Initialize self. See help(type(self)) for accurate signature.

backend: str#: None

max_new_tokens: int#: None

temperature: float#: None

top_p: float#: None

top_k: int | None#: None

val_temperature: float#: None

val_top_p: float#: None

val_top_k: int | None#: None

model_name: NotRequired[str]#: None

stop_token_ids: list[int] | None#: None

stop_strings: list[str] | None#: None

bad_words: NotRequired[list[str] | None]#: None

colocated: NotRequired[nemo_rl.models.generation.interfaces.ColocationConfig]#: None

port_range_low: NotRequired[int]#: None

port_range_high: NotRequired[int]#: None

use_async_rollouts: NotRequired[bool]#: None

_pad_token_id: NotRequired[int]#: None

_mtp_weights_from_refit: NotRequired[bool]#: None

_debug_payload_metrics: NotRequired[bool]#: None

nemo_rl.models.generation.interfaces.should_use_async_rollouts( generation_config: nemo_rl.models.generation.interfaces.GenerationConfig | None, ) → bool#: Determine whether a generation backend uses asynchronous rollouts.

class nemo_rl.models.generation.interfaces.GenerationSamplingParams#

Sampling profile threaded explicitly through rollout entry points.

Rollout callers construct one from the relevant GenerationConfig fields (train or validation) so the sampling used for a rollout is visible at the call site instead of flowing through config side-channels. Named to distinguish it from TrainingSamplingParams (train-time logit filtering) and vLLM’s own SamplingParams.

temperature: float#: None

top_p: float#: None

top_k: int | None#: None

classmethod from_generation_config( generation_config: nemo_rl.models.generation.interfaces.GenerationConfig, ) → nemo_rl.models.generation.interfaces.GenerationSamplingParams#: Build the train-time sampling profile from a generation config.

class nemo_rl.models.generation.interfaces.GenerationDatumSpec#

Bases: typing.TypedDict

Specification for input data required by generation models.

input_ids: Tensor of token IDs representing the input sequences (right padded)
input_lengths: Tensor containing the actual length of each sequence (without padding)
stop_strings: Optional list of strings to stop generation (per sample)
extra: Additional model-specific data fields

Example of a batch with 4 entries with different sequence lengths:

# Batch of 4 sequences with lengths [3, 5, 2, 4]

input_ids (padded):
[
  [101, 2054, 2003,    0,    0],  # Length 3
  [101, 2054, 2003, 2001, 1996],  # Length 5
  [101, 2054,    0,    0,    0],  # Length 2
  [101, 2054, 2003, 2001,    0],  # Length 4
]

input_lengths:
[3, 5, 2, 4]

All functions receiving or returning GenerationDatumSpec should ensure right padding is maintained. Use verify_right_padding() to check.

Initialization

Initialize self. See help(type(self)) for accurate signature.

input_ids: torch.Tensor#: None

input_lengths: torch.Tensor#: None

stop_strings: NotRequired[list[str]]#: None

__extra__: Any#: None

class nemo_rl.models.generation.interfaces.GenerationOutputSpec#

Bases: typing.TypedDict

Specification for output data returned by generation models.

output_ids: Tensor of token IDs representing the generated sequences (right padded)
generation_lengths: Tensor containing the actual length of each generated sequence
unpadded_sequence_lengths: Tensor containing the actual length of each input + generated sequence (without padding)
logprobs: Tensor of log probabilities for each generated token (right padded with zeros)
truncated: Boolean tensor indicating if each sequence was truncated (hit max_tokens limit)
extra: Additional model-specific data fields

Example of a batch with 2 sequences:

# Sample batch with 2 examples
# - Example 1: Input length 3, generated response length 4
# - Example 2: Input length 5, generated response length 2

output_ids (right-padded):
[
  [101, 2054, 2003, 2023, 2003, 1037, 2200,    0],  # 7 valid tokens (3 input + 4 output)
  [101, 2054, 2003, 2001, 1996, 3014, 2005,    0],  # 7 valid tokens (5 input + 2 output)
]

generation_lengths:
[4, 2]  # Length of just the generated response part

unpadded_sequence_lengths:
[7, 7]  # Length of full valid sequence (input + generated response)

logprobs (right-padded with zeros):
[
  [0.0, 0.0, 0.0, -1.2, -0.8, -2.1, -1.5, 0.0],  # First 3 are 0 (input tokens), next 4 are actual logprobs
  [0.0, 0.0, 0.0, 0.0, 0.0, -0.9, -1.7, 0.0],     # First 5 are 0 (input tokens), next 2 are actual logprobs
]

truncated:
[False, True]  # Example 2 was truncated (hit max_tokens limit without EOS)

All functions receiving or returning GenerationOutputSpec should ensure right padding is maintained. Use verify_right_padding() to check.

Initialization

Initialize self. See help(type(self)) for accurate signature.

output_ids: torch.Tensor#: None

generation_lengths: torch.Tensor#: None

unpadded_sequence_lengths: torch.Tensor#: None

logprobs: torch.Tensor#: None

routed_experts: NotRequired[torch.Tensor]#: None

r3_routed_experts_missing_routes: NotRequired[torch.Tensor]#: None

r3_routed_experts_expected_routes: NotRequired[torch.Tensor]#: None

r3_routed_experts_actual_routes: NotRequired[torch.Tensor]#: None

truncated: NotRequired[torch.Tensor]#: None

__extra__: Any#: None

class nemo_rl.models.generation.interfaces.CollectiveSenderSpec#

Policy-side protocol and packing geometry for NCCL weight transfer.

nccl_peer: str#: ‘nemo’

buffer_size_bytes: int | None#: None

num_buffers: int | None#: None

nemo_rl.models.generation.interfaces.reject_unenforceable_refit_deadline( backend: str, refit_timeout_s: Optional[float], ) → None#

Refuse a refit deadline the backend cannot actually apply.

Accepting it and doing nothing would be worse than refusing. The deadline exists so that a generation rank dying mid-refit cannot hang the weight-sync collective forever; a user who sets it on a backend that ignores it gets exactly that hang, while believing they are protected. Only vLLM threads the deadline down to the collective today.

None – every path that does not configure a deadline, which is all of them by default – passes through untouched, so this is inert unless someone opts in.

class nemo_rl.models.generation.interfaces.GenerationInterface#

Bases: abc.ABC

Abstract base class defining the interface for RL policies.

classmethod validate_settings( master_config: nemo_rl.algorithms.single_controller_utils.config.MasterConfig, ) → None#

Backend-specific pure-config validation, run before any build.

Parameters:: master_config – The single-controller MasterConfig.

abstractmethod init_collective( ip: str, port: int, world_size: int, *, train_world_size: int, ) → list[ray.ObjectRef]#: Initialize the collective communication.

abstractmethod generate( data: nemo_rl.distributed.batched_data_dict.BatchedDataDict[nemo_rl.models.generation.interfaces.GenerationDatumSpec], greedy: bool, ) → nemo_rl.distributed.batched_data_dict.BatchedDataDict[nemo_rl.models.generation.interfaces.GenerationOutputSpec]#

abstractmethod prepare_for_generation(

*args: Any,

**kwargs: Any,

) → bool#

Ready the engine for a generation phase (start or wake it).

Idempotent wake: calling this on an already-running engine must be safe and cheap.

abstractmethod finish_generation(*args: Any, **kwargs: Any) → bool#

Wind down after a generation phase.

Callers may pass release_gpu (keyword-only, default True): True means the caller needs the GPUs for itself (a training step or a checkpoint save), so even a colocated engine must fully stand down; False means the phase is merely over, and a colocated engine must keep serving usable with no intervening prepare_for_generation. Only the colocated Megatron backend honors the flag today; other backends ignore it, as do engines on dedicated GPUs.

abstractmethod shutdown() → bool#: Shut down generation resources; repeated calls must be safe.

abstractmethod pause_generation(mode: str) → None#: Pause in-flight generation on the backend.

abstractmethod continue_generation() → None#: Resume previously paused generation on the backend.

property requires_kv_scale_sync: bool#: Whether the generation backend requires KV cache scales synchronization.

abstractmethod prepare_refit_info(state_dict_info: dict[str, Any]) → None#: Prepare the info for refit.

abstractmethod update_weights_via_ipc_zmq() → list[ray.ObjectRef]#: Update the model weights from the given IPC handles.

abstractmethod update_weights_from_collective( refit_timeout_s: Optional[float] = None, ) → list[ray.ObjectRef]#

Update the model weights from collective communication.

refit_timeout_s bounds the receive side of the refit. It is part of the signature for every backend, not just the ones that can act on it, because the synchronizer calls this polymorphically – a backend that omits the parameter does not fail at import or type-check time, it fails at the Ray boundary during the first refit. Backends that cannot enforce it should say so via reject_unenforceable_refit_deadline rather than accept it silently.

get_collective_sender_spec() → nemo_rl.models.generation.interfaces.CollectiveSenderSpec#: Return policy-side NCCL protocol and packed-buffer requirements.

get_inference_world_size() → int | None#: Return a backend-specific collective world size when required.

abstractmethod prepare_nccl_reshard_refit_info(refit_info: dict) → None#: Prepare per-layer param metadata for nccl_reshard-based refit.

abstractmethod nccl_reshard_refit( refit_timeout_s: Optional[float] = None, ) → list[ray.ObjectRef]#

Receive weights from training workers via nccl_reshard.

Takes the deadline for the same reason its sibling above does, and the reason is worth repeating because this is the hook that was missed: the synchronizer calls both polymorphically, so a backend whose signature omits the parameter does not fail at import or type-check time – it fails at the Ray boundary during the first refit, a long way from the signature that caused it.

Parameters:: refit_timeout_s – Deadline for this collective, after which the worker aborts its own communicator. None leaves the refit path unchanged.

abstractmethod attach_fleet_health(monitor: Any, selector: Any) → None#

Route this backend’s shard selection through fleet health.

Declared here rather than discovered with hasattr at the call site, so an unsupported backend says so itself and the capability is greppable from the interface. Same shape as the refit hooks above.

Parameters:

monitor – GenerationFleetHealth owning shard eligibility, which the backend also reports observed failures and successes to.
selector – HealthyShardSelector picking among the serving shards.

invalidate_kv_cache() → bool#

pause_generation_for_refit(*, clear_cache: bool) → bool#

Pause in-flight generation while preserving request state.

Backends with native in-flight refit support override this hook. The default implementation warns once per backend type and lets the refit continue with the backend’s existing in-flight behavior. On supported backends, in-flight requests are frozen rather than aborted and resume from

Meth:: resume_generation_after_refit; new requests queue until then.
Parameters:: clear_cache – Also clear the engine’s reusable caches at pause time so preserved requests recompute their KV after the weight update.
Returns:: True if every engine paused; False when the backend has no native pause support. Backends with native support raise when pausing fails.

resume_generation_after_refit() → bool#

Resume generation paused by :meth:pause_generation_for_refit.

The default implementation shares the once-per-backend warning emitted by

Meth:: pause_generation_for_refit and lets the refit continue for backends without native pause/resume support.
Returns:: True if every engine resumed; False when the backend has no native resume support. Backends with native support raise when resuming fails.

blocks_training() → bool#

Whether this engine must stand down before a training step.

True when generation shares GPUs with training (colocated): the training loop then pauses collection and winds the engine down before training. Engines on dedicated GPUs never block training.

wake_carries_weight_updates() → bool#

Whether prepare_for_generation alone serves the latest weights.

True when waking the engine suffices for it to serve weights updated while it slept (colocated Megatron: the wake reshards, or the engine shares the training tensors outright). The async loop may then defer a wake past a checkpoint save and advance the collector’s weight version with no explicit transfer. Backends whose wake does not reload weights must return False so the loop refits instead.

clear_logger_metrics() → None#

Clear logger metrics for performance reporting.

This is an optional method that backends can implement to clear telemetry metrics. Default implementation does nothing.

get_logger_metrics() → dict[str, Any]#

Get logger metrics for performance reporting.

This is an optional method that backends can implement to collect telemetry metrics. Default implementation returns empty dict.

Returns:: Dictionary of metrics. Format may vary by backend.

nemo_rl.models.generation.interfaces#

Module Contents#

Classes#

Functions#

Data#

API#

`nemo_rl.models.generation.interfaces`#