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,
skip_weight_load: bool = False,
)#

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

_report_sharded_payload(
sharded_data: list[nemo_rl.distributed.batched_data_dict.SlicedDataDict],
boundary: str,
) → None#

Measure the exact unique per-DP-shard Ray arguments.

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.

init_nccl_reshard_comm_group(
pp_ips: list[str],
pp_ports: list[int],
pp_size: int,
pp_stages: list[int],
sub_world_size: int,
ranks_in_group: list[int],
) → list[ray.ObjectRef]#

Initialize the nccl_reshard bulk-path comm group on all train workers.

prepare_nccl_reshard_refit_info(
train_parallelism,
gen_parallelism,
train_world_size,
gen_world_size,
)#

Prepare per-layer param metadata for nccl_reshard refit.

nccl_reshard_refit(kv_scales=None) → list[ray.ObjectRef]#

Transfer weights to gen workers via nccl_reshard (xferdtensor).

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 shutdown() and the pointer to the object is lost due to leaving a function scope. It’s always recommended that the user calls 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.