`nemo_rl.models.generation.interfaces`#

Module Contents#

Classes#

`ResourcesConfig`
`OptionalResourcesConfig`
`DisaggConfig`	Prefill/decode disaggregation, in the terms every backend shares.
`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.
`agreed_media_columns`	The one media column spec the capture workers pinned, or `None`.
`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.DisaggConfig#

Bases: pydantic.BaseModel

Prefill/decode disaggregation, in the terms every backend shares.

enabled: bool#: False

num_prefill_engines: int#: 1

num_decode_engines: int#: 1

num_frontend_workers: int#: 1

frontend_tokenize: bool#: False

frontend_base_port: int#: 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

ignore_eos: NotRequired[bool]#: 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

_draft_weights_from_refit: NotRequired[bool]#: None

_mtp_weights_from_refit: NotRequired[bool]#: None

_draft_full_refit: NotRequired[bool]#: None

_debug_payload_metrics: NotRequired[bool]#: None

disaggregation: NotRequired[nemo_rl.models.generation.interfaces.DisaggConfig]#: 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

use_legacy_packing: bool#: False

nemo_rl.models.generation.interfaces.agreed_media_columns( reported: list[MediaColumnSpec | None], ) → MediaColumnSpec | None#

The one media column spec the capture workers pinned, or None.

Each setup_token_capture fan-out target reports the spec it pinned its staging media column to (None from text-only workers and ranks that host no capture). The reassembler mints empty media rows from this spec, so two workers disagreeing would publish two pixel dtypes or widths into one column; that is refused here, at setup, rather than found by TransferQueue mid-run.

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 transports whose workers own an abortable collective can enforce it.

None disables the deadline and passes through untouched.

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.

abstractmethod restart_shard(shard_idx: int) → Optional[str]#

Rebuild one data-parallel shard’s workers and bring its engine back up.

Declared here because EngineSupervisor calls it by name on whatever backend it was handed. Undeclared, a backend that does not implement it – or one that loses the method to a merge, which has happened once already – degrades to an AttributeError swallowed by the supervisor’s except, and the shard is recorded as a failed restart rather than as an unsupported one. Raising here says which it is.

Blocking and slow – it reloads the model – so callers run it off the control loop.

Returns:: The replacement’s OpenAI base URL, or None for an engine that exposes no HTTP server. The URL is expected to differ from the old one: the new engine binds its own port, so callers must publish it rather than assume the fleet’s URL list is still accurate.

abstractmethod shard_liveness_ref(shard_idx: int) → ray.ObjectRef#

Liveness of the worker leading one data-parallel shard.

Which Ray worker leads shard N depends on how shards are laid out across workers, which is the backend’s business. Asking for it by shard index keeps that here rather than in the control loop, where the same arithmetic had a second copy that also assumed every backend has a worker_group – an assumption that has already broken a lane once ('DynamoGeneration' object has no attribute 'worker_group').

log_shard_gpu_state( shard_idx: int, *, label: str, timeout_s: float = 30.0, ) → None#

Print the state of the GPU one shard holds, from that shard’s own node.

A no-op by default rather than NotImplementedError, unlike restart_shard above: this is a diagnostic taken on the restart path, and a backend that cannot provide it should cost the caller nothing. Failing a restart over a missing log line would be worse than the missing log line.

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

abstractmethod prepare_refit_info( state_dict_info: Optional[dict[str, Any]], ) → Optional[list[str]]#

Prepare the info for refit.

Returns:: Optionally, the parameter names the backend wants pre-quantized on the trainer before streaming (e.g. vllm_cfg.refit_prequantize); None when no trainer-side pre-quantization is requested. The vLLM implementation also installs its FP8 configuration on the workers.

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.

abstractmethod setup_token_capture( dp_cfg: nemo_rl.data_plane.interfaces.DataPlaneConfig, staging_partition: str, *, capture_media: bool = False, ) → MediaColumnSpec | None#

Install token capture in the serving workers (token_capture.enabled).

Declared here for the same reason as :meth:attach_fleet_health: the single-controller setup calls this on whichever backend is configured, so an unsupported backend says so itself instead of failing with AttributeError.

Parameters:

dp_cfg – Data-plane config the workers use to build their in-worker client.
staging_partition – Data-plane partition that captured rows are staged in.
capture_media – Whether the staging partition carries media columns and the workers must stage the media the engine consumed beside each call’s tokens.

Returns:

With capture_media, the pixel dtype and patch size the workers pinned the staging media column to (every worker reports the same spec); the reassembler publishes empty media rows in that geometry for groups without media. None for text-only capture.

abstractmethod set_rollout_weight_version(version: int) → None#

Rotate the weight version workers stamp on captured model calls.

Parameters:: version – Trainer weight version now being served, applied to all subsequent captured requests.

invalidate_kv_cache() → bool#

Confirm stale KV state cannot be reused after a weight update.

Backends may implement this by invalidating active-request KV state directly or by performing equivalent engine-side handling atomically with refit. Clearing only unreferenced prefix-cache entries is insufficient while requests are active.

Returns:: True only when the backend guarantees the recompute contract; False when it is unsupported or could not be completed.

pause_generation_for_refit( *, clear_cache: bool, timeout_s: Optional[float] = None, ) → bool#

Arrange a safe generation pause around the upcoming weight refit.

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. Native pause implementations freeze requests until :meth:resume_generation_after_refit. A backend may instead acknowledge a lifecycle owned by its weight synchronizer, which must pause and handle the requested cache clearing before transfer and resume before returning. Such backends document their request handling and timeout semantics; for example, SGLang may retract or abort requests.

Parameters:

clear_cache – Also clear the engine’s reusable caches at pause time so preserved requests recompute their KV after the weight update.
timeout_s – Maximum time to wait for engines to pause; None waits indefinitely.

Returns:

True if every engine paused or the weight synchronizer guarantees the lifecycle above; False when neither is supported. Pause failures raise from this hook or the synchronizer that owns the pause.

resume_generation_after_refit( *, timeout_s: Optional[float] = None, ) → 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.
Parameters:: timeout_s – Maximum time to wait for engines to resume; None waits indefinitely.
Returns:: True if every surviving engine resumed, including when the weight synchronizer already performed the resume; False when neither path is supported. With fleet health, confirmed dead engines may be excluded if enough survivors remain. Other resume failures raise.

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.

snapshot_step_metrics() → None#

Begin a per-training-step generation metric window.

Backends without per-step generation metrics may use this default no-op.

get_step_metrics() → dict[str, float]#

Finish the current metric window and return generation metrics.

Returns:: Metrics accumulated since the matching snapshot_step_metrics call, not running totals. Backends without per-step generation metrics return an empty dictionary.

drain_latest_logger_metrics() → dict[str, Any]#

Consume a bounded latest-value snapshot for frequent telemetry polls.

Implementations may clear or compact their accumulated metric histories. Callers must not assume that a later get_logger_metrics includes values observed before this drain. Backends supporting raw rollout throughput should return cumulative sampled-token counters under generation_tokens as data_parallel_worker_id -> list[counter]. The controller computes per-worker deltas before summing them, so counter resets are detectable.

nemo_rl.models.generation.interfaces#

Module Contents#

Classes#

Functions#

Data#

API#

`nemo_rl.models.generation.interfaces`#