nemo_rl.models.policy.lm_policy#

Module Contents#

Classes#

Policy

Functions#

_aggregate_megatron_flops_metrics

Aggregate FLOPS metrics from Megatron worker results.

Data#

PathLike

API#

nemo_rl.models.policy.lm_policy.PathLike#

None

nemo_rl.models.policy.lm_policy._aggregate_megatron_flops_metrics(
results: list[dict],
world_size: int,
) → dict#

Aggregate FLOPS metrics from Megatron worker results.

Called when the Megatron worker returns total_flops directly (no FLOPTracker).

class nemo_rl.models.policy.lm_policy.Policy(
cluster: nemo_rl.distributed.virtual_cluster.RayVirtualCluster,
config: nemo_rl.models.policy.PolicyConfig,
tokenizer: transformers.PreTrainedTokenizerBase,
name_prefix: str = 'lm_policy',
workers_per_node: Optional[Union[int, list[int]]] = None,
init_optimizer: bool = True,
weights_path: Optional[nemo_rl.models.policy.lm_policy.PathLike] = None,
optimizer_path: Optional[nemo_rl.models.policy.lm_policy.PathLike] = None,
init_reference_model: bool = True,
processor: Optional[transformers.AutoProcessor] = None,
worker_extension_cls_fqn: Optional[str] = None,
)#

Bases: nemo_rl.models.policy.interfaces.ColocatablePolicyInterface, nemo_rl.models.generation.interfaces.GenerationInterface

property data_parallel_size: int#

Data-parallel degree, read from the policy’s sharding annotations.

run_all_workers_single_data(
method_name: str,
*args,
**kwargs,
) → Any#

Run a method on all workers in parallel with the same data.

Mainly used for worker extension classes.

Parameters:

  • method_name – The name of the method to run.

  • *args – The positional arguments to pass to the method.

  • **kwargs – The keyword arguments to pass to the method.

Returns:

The results of the method run on all workers.

run_all_workers_multiple_data(
method_name: str,
*args,
**kwargs,
) → Any#

Run a method on all workers in parallel with different data.

Mainly used for worker extension classes.

Parameters:

  • method_name – The name of the method to run.

  • *args – The positional arguments to pass to the method.

  • **kwargs – The keyword arguments to pass to the method.

Returns:

The results of the method run on all workers.

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

Initialize the collective communication.

init_collective_mcore_generation(
ip: str,
port: int,
world_size: int,
*,
rank_offset: int,
refit_backend: str = 'gloo',
) → list[ray.ObjectRef]#

Initialize the megatron refit collective on this policy’s workers.

preinit_nvshmem() → list[ray.ObjectRef]#

Pre-initialize NVSHMEM on this policy’s workers (no-op when not using nvshmem).

swap_weights_via_reshard(*, is_source: bool) → list[ray.ObjectRef]#

Send (is_source=True) or receive (is_source=False) weights via megatron reshard.

_shard_for_logprob(
data: nemo_rl.distributed.batched_data_dict.BatchedDataDict[Any],
) → tuple[list[nemo_rl.distributed.batched_data_dict.SlicedDataDict], Optional[list[int]]]#

Shard inputs for get_logprobs / get_reference_policy_logprobs.

Mirrors the legacy shard block (lines 426-450 / 503-530). Returns (sharded_data, unsorted_data_indices) where the second element is the inverse permutation needed to undo seqpack/dynbatch reorder (None when neither is enabled).

_shard_for_train(
data: nemo_rl.distributed.batched_data_dict.BatchedDataDict[Any],
batch_size: int,
) → list[nemo_rl.distributed.batched_data_dict.SlicedDataDict]#

Shard inputs for train.

Mirrors the legacy shard block (lines 706-729). Note vs. _shard_for_logprob: uses train_mb_tokens (not logprob_mb_tokens), passes batch_size (not None), and does not return unsorted_data_indices because train returns scalar metrics (no per-row outputs to reorder).

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 the logprobs of the model for a data dict.

Returns:

a BatchedDataDict with key “logprobs” and shape [batch_size, sequence_length]. We use the convention that the logprob of the first token is 0 so that the sequence length is maintained. The logprob of input token i is specified at position i in the output logprobs tensor.

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 the logprobs of the reference policy for a data dict.

Returns: Identical to get_logprobs.

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]#

Dispatch get_topk_logits to workers (no CP/packed support initially).

get_full_logits_ipc(
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,
) → list[dict[str, Any]]#

Ship the teacher’s full-vocab logits to the student via CUDA IPC.

Used by cross-tokenizer distillation; supports heterogeneous teacher TP/CP. Gathers each worker’s {"dp_rank", "per_sample_handles"} and returns the global-batch-ordered list produced by

Func:

aggregate_per_sample_handles: a length-gbs list where element i is {"teacher_shards": [shard, ...]} holding every TP×CP shard of global sample i. Each shard carries the IPC payload plus the buf_idx / sample_index_in_buf slot index and the TP/CP shard metadata; the loss consumer reassembles the full [T_t, V_t] teacher logits (or its CP-local window) from these shards.

The producer-side IPC storage is persistent and reused across calls (via copy_); the caller releases it once via

Meth:

release_ipc_buffer at the end of training / validation (and on error), not per call.

v0 limitation: no dynamic batching, no sequence packing.

release_ipc_buffer() → None#

Tell all workers to drop their stashed IPC tensors.

train(
data: nemo_rl.distributed.batched_data_dict.BatchedDataDict[Any],
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,
check_dim_skip_keys: Optional[Iterable[str]] = None,
) → dict[str, Any]#

Train the policy on a batch of data with a given loss function.

Parameters:

check_dim_skip_keys – Keys whose tensors are not student-sequence-aligned at dim 1 and must be excluded from the worker’s sequence-dim pre-flight check. Used by cross-tokenizer distillation to pass through teacher / alignment auxiliaries that ride on the same data dict.

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

Generate a batch of data using the policy.

score(
data: nemo_rl.distributed.batched_data_dict.BatchedDataDict[nemo_rl.models.generation.interfaces.GenerationDatumSpec],
) → nemo_rl.distributed.batched_data_dict.BatchedDataDict[nemo_rl.models.policy.interfaces.ScoreOutputSpec]#

Score a batch of data using the policy.

prepare_for_generation(
*args: Any,
**kwargs: Any,
) → bool#
finish_generation(*args: Any, **kwargs: Any) → bool#
prepare_for_training(*args: Any, **kwargs: Any) → None#
prepare_for_lp_inference(
*args: Any,
**kwargs: Any,
) → None#
invalidate_kv_cache(*args: Any, **kwargs: Any) → bool#
prepare_refit_info() → Optional[dict[str, Any]]#

Prepare the info for refit.

Returns:

A dictionary containing the info for refit.

Return type:

dict

finish_inference() → None#

Offload policy model to CPU after inference.

finish_training(*args: Any, **kwargs: Any) → None#
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]#

Trigger KV-cache FP8 scale calibration across Megatron workers and return results.

Note: The backend MegatronPolicyWorker.calibrate_qkv_fp8_scales already implements distributed reduction, returning results merged across ranks. Therefore, we shard the input by DP and call in parallel, then take the result from the first worker.

get_free_memory_bytes() → int#

Get the available free memory.

stream_weights_via_ipc_zmq(
buffer_size_bytes: int,
kv_scales: Optional[dict[str, float]] = None,
) → list[ray.ObjectRef]#

Send the weights for IPC handles via ZMQ socket.

stream_weights_via_http(
rollout_engine_urls: list[str],
buffer_size_bytes: int,
) → list[ray.ObjectRef]#

Send the weights to colocated SGLang engines via CUDA IPC over HTTP.

Parameters:

  • rollout_engine_urls – http://host:port base URLs of each engine’s node_rank=0 SGLang HTTP server. The caller resolves these once (via engine.get_base_url) and passes them in, so every FSDP rank doesn’t redo the Ray RPC.

  • buffer_size_bytes – Max bucket size in bytes before flushing.

The rollout TP size is captured once via set_rollout_num_gpus_per_engine and reused by each worker.

set_rollout_num_gpus_per_engine(num_gpus_per_engine: int) → None#

Broadcast the rollout engine TP size to every policy worker.

broadcast_weights_for_collective(
kv_scales: Optional[dict[str, float]] = None,
) → list[ray.ObjectRef]#

Broadcast the weights for collective communication.

offload_before_refit() → None#

Offload the optimizer and buffers to the CPU.

offload_after_refit() → None#

Offload the optimizer and buffers to the CPU.

offload_to_cpu() → None#

Offload to CPU to free GPU memory; currently only used by PPO.

save_checkpoint(
weights_path: str,
optimizer_path: Optional[str] = None,
tokenizer_path: Optional[str] = None,
checkpointing_cfg: Optional[nemo_rl.utils.checkpoint.CheckpointingConfig] = None,
) → None#

Save a checkpoint of the model.

With Megatron async_save=True, this returns after D2H staging. The caller must call finalize_async_save() before renaming the checkpoint directory.

finalize_async_save() → None#

Block until all workers’ in-flight async checkpoint writes complete.

No-op when async_save is disabled. Must be called before the checkpoint directory is renamed from tmp_step_N/ to step_N/.

shutdown() → bool#

Shut down all HF workers and clean up resources.

__del__() → None#

Shuts down the worker groups when the object is deleted or is garbage collected.

This is an extra safety net in case the user forgets to call worker_group.shutdown() and the pointer to the object is lost due to leaving a function scope. It’s always recommended that the user calls worker_group.shutdown().

start_gpu_profiling() → None#

Start GPU profiling.

stop_gpu_profiling() → None#

Stop GPU profiling.

print_node_ip_and_gpu_id() → list[tuple[str, int]]#

Print the node IP and GPU ID of the current worker.