nemo_rl.models.policy.interfaces

Module Contents

Classes

LogprobOutputSpec	logprobs: Tensor of log probabilities.
ReferenceLogprobOutputSpec	logprobs: Tensor of log probabilities.
ScoreOutputSpec	scores: Tensor of scores.
TopkLogitsOutputSpec	Per-position top-k logits and corresponding global token indices.
PolicyInterface	Abstract base class defining the interface for RL policies.
ColocatablePolicyInterface

API

class nemo_rl.models.policy.interfaces.LogprobOutputSpec

Bases: typing.TypedDict

logprobs: Tensor of log probabilities.

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.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

abstractmethod init_collective(

ip: str,

port: int,

world_size: int,

*,

train_world_size: int,

nccl_peer: str = 'nemo',

) → list[ray.ObjectRef]

abstractmethod offload_before_refit() → None

abstractmethod offload_after_refit() → None

offload_to_cpu() → None

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

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,

) → 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.