nemo_rl.models.policy.interfaces

Module Contents

Classes

LogprobOutputSpec	Token log probabilities; filtered unsupported tokens retain -inf.
ReferenceLogprobOutputSpec	logprobs: Tensor of log probabilities.
ScoreOutputSpec	scores: Tensor of scores.
TopkLogitsOutputSpec	Per-position top-k logits and corresponding global token indices.
TeacherFullPayloadOutputSpec	Sampled-token logprobs plus the full-vocabulary MOPD teacher payload.
PolicyInterface	Abstract base class defining the interface for RL policies.
ColocatablePolicyInterface

API

class nemo_rl.models.policy.interfaces.LogprobOutputSpec

Bases: typing.TypedDict

Token log probabilities; filtered unsupported tokens retain -inf.

Initialization

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

logprobs: torch.Tensor

None

class nemo_rl.models.policy.interfaces.ReferenceLogprobOutputSpec

Bases: typing.TypedDict

logprobs: Tensor of log probabilities.

Initialization

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

reference_logprobs: torch.Tensor

None

class nemo_rl.models.policy.interfaces.ScoreOutputSpec

Bases: typing.TypedDict

scores: Tensor of scores.

Initialization

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

scores: torch.Tensor

None

class nemo_rl.models.policy.interfaces.TopkLogitsOutputSpec

Bases: typing.TypedDict

Per-position top-k logits and corresponding global token indices.

Initialization

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

topk_logits: torch.Tensor

None

topk_indices: torch.Tensor

None

class nemo_rl.models.policy.interfaces.TeacherFullPayloadOutputSpec

Bases: typing.TypedDict

Sampled-token logprobs plus the full-vocabulary MOPD teacher payload.

teacher_full_payload is [B, S, hidden_size] or [B, S, vocab_size] depending on on_policy_distillation.full.teacher_payload. It is None off the last pipeline stage: the payload is far too large to broadcast, so only the stage that produced it writes it back. It comes back on CPU – each microbatch is moved off the device inside the forward schedule, which is what bounds the resident payload to one microbatch.

Initialization

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

logprobs: torch.Tensor

None

teacher_full_payload: Optional[torch.Tensor]

None

class nemo_rl.models.policy.interfaces.PolicyInterface

Bases: abc.ABC

Abstract base class defining the interface for RL policies.

abstractmethod get_logprobs(

data: nemo_rl.distributed.batched_data_dict.BatchedDataDict[nemo_rl.models.generation.interfaces.GenerationDatumSpec],

timer: Optional[nemo_rl.utils.timer.Timer] = None,

) → nemo_rl.distributed.batched_data_dict.BatchedDataDict[nemo_rl.models.policy.interfaces.LogprobOutputSpec]

Get logprobs of actions from observations.

Parameters:

data – BatchedDataDict containing rollouts (tokens)

Returns:

• logprobs: Tensor of logprobs of actions

Return type:

BatchedDataDict containing

abstractmethod get_reference_policy_logprobs(

data: nemo_rl.distributed.batched_data_dict.BatchedDataDict[nemo_rl.models.generation.interfaces.GenerationDatumSpec],

micro_batch_size: Optional[int] = None,

timer: Optional[nemo_rl.utils.timer.Timer] = None,

) → nemo_rl.distributed.batched_data_dict.BatchedDataDict[nemo_rl.models.policy.interfaces.ReferenceLogprobOutputSpec]

Get logprobs of actions from observations.

Parameters:

data – BatchedDataDict containing rollouts (tokens)

Returns:

• logprobs: Tensor of logprobs of actions

Return type:

BatchedDataDict containing

abstractmethod get_topk_logits(

data: nemo_rl.distributed.batched_data_dict.BatchedDataDict[nemo_rl.models.generation.interfaces.GenerationDatumSpec],

k: int,

micro_batch_size: Optional[int] = None,

timer: Optional[nemo_rl.utils.timer.Timer] = None,

) → nemo_rl.distributed.batched_data_dict.BatchedDataDict[nemo_rl.models.policy.interfaces.TopkLogitsOutputSpec]

Get per-position top-k logits and global indices for a batch of inputs.

.. rubric:: Notes

• Aligns to next-token positions → returns S-1 positions.

abstractmethod train(

data: nemo_rl.distributed.batched_data_dict.BatchedDataDict,

loss_fn: nemo_rl.algorithms.loss.interfaces.LossFunction,

eval_mode: bool = False,

*,

gbs: Optional[int] = None,

mbs: Optional[int] = None,

timer: Optional[nemo_rl.utils.timer.Timer] = None,

) → dict[str, Any]

Train the policy on a global batch of data.

Parameters:

• data – BatchedDataDict containing rollouts (tokens)

• loss_fn – Loss function to use for training

• eval_mode – Whether to run in evaluation mode (no gradient updates)

• gbs – Global batch size override (if None, uses config default)

• mbs – Micro batch size override (if None, uses config default)

abstractmethod calibrate_qkv_fp8_scales(

data: nemo_rl.distributed.batched_data_dict.BatchedDataDict[nemo_rl.models.generation.interfaces.GenerationDatumSpec],

micro_batch_size: Optional[int] = None,

percentile: float = 99.9,

margin: float = 1.05,

include_q: bool = False,

) → dict[str, Any]

Calibrate FP8 scales for Q/K/V activations used by KV cache.

Parameters:

• data – BatchedDataDict containing input_ids and input_lengths.

• micro_batch_size – Optional override for micro batch size during calibration.

• percentile – Percentile for per-tensor amax estimation.

• margin – Safety margin multiplier applied to amax.

• include_q – Whether to also compute scale for Q in addition to K/V.

Returns:

Dict with overall configuration and per-layer scales.

abstractmethod prepare_for_training(*args: Any, **kwargs: Any) → None

abstractmethod finish_training(*args: Any, **kwargs: Any) → None

abstractmethod save_checkpoint(*args: Any, **kwargs: Any) → None

abstractmethod shutdown() → bool

class nemo_rl.models.policy.interfaces.ColocatablePolicyInterface

Bases: nemo_rl.models.policy.interfaces.PolicyInterface

cfg: nemo_rl.models.policy.PolicyConfig

None

abstractmethod init_collective(

ip: str,

port: int,

world_size: int,

*,

train_world_size: int,

rank_offset: int = 0,

nccl_peer: str = 'nemo',

) → list[ray.ObjectRef]

abstractmethod sync_params_before_refit() → None

Materialize the latest policy parameters before weight transfer.

abstractmethod offload_before_refit() → None

abstractmethod offload_after_refit() → None

offload_to_cpu() → None

abstractmethod prepare_refit_info() → Optional[dict[str, Any]]

abstractmethod enable_refit_prequantize(

param_names: list[str],

) → Optional[dict[str, Any]]

Quantize the listed params on the trainer during refit streaming.

Returns:

Refit info updated with the quantized dtypes and scale entries.

abstractmethod stream_weights_via_ipc_zmq(

buffer_size_bytes: int,

kv_scales: Optional[dict[str, float]] = None,

) → list[ray.ObjectRef]

abstractmethod connect_sglang_rollout_engines(

*,

engine_gpu_counts: list[int],

engine_gpu_offsets: Optional[list[int]] = None,

) → None

abstractmethod update_weights_to_sglang_colocated(

*,

rollout_engines: list[ray.actor.ActorHandle],

buffer_size_bytes: int,

target_precision: str = 'bf16',

sglang_quantization_cfg: Optional[dict[str, Any]] = None,

) → list[ray.ObjectRef]

abstractmethod connect_sglang_rollout_engines_distributed(

*,

rollout_engines: list[ray.actor.ActorHandle],

engine_gpu_counts: list[int],

group_name: Optional[str] = None,

) → None

abstractmethod update_weights_to_sglang_distributed(

*,

rollout_engines: list[ray.actor.ActorHandle],

rollout_engine_lock: ray.actor.ActorHandle,

buffer_size_bytes: int,

target_precision: str = 'bf16',

sglang_quantization_cfg: Optional[dict[str, Any]] = None,

) → list[ray.ObjectRef]

abstractmethod broadcast_weights_for_collective(

kv_scales: Optional[dict[str, float]] = None,

refit_timeout_s: Optional[float] = None,

*,

buffer_size_bytes: Optional[int] = None,

num_buffers: Optional[int] = None,

use_legacy_packing: bool = False,

) → list[ray.ObjectRef]

abstractmethod prepare_nccl_reshard_refit_info(

train_parallelism: dict[str, int],

gen_parallelism: dict[str, int],

train_world_size: int,

gen_world_size: int,

) → Any

Prepare per-layer param metadata for nccl_reshard-based refit.

abstractmethod nccl_reshard_refit(

kv_scales: Optional[dict[str, float]] = None,

refit_timeout_s: Optional[float] = None,

) → list[ray.ObjectRef]

Sync weights to generation workers via the NCCL-reshard path.

abstractmethod prepare_for_lp_inference(keep_train_buffers: bool = False) → None

Put the policy in eval mode for logprob inference.

Parameters:

keep_train_buffers – Leave grad buffers and optimizer state on CUDA because a train step is already open and its accumulated gradients must survive this call.