`nemo_rl.models.policy.utils`#

Module Contents#

Classes#

IPCProtocol

IPC protocol constants for ZMQ weight streaming.

Functions#

`resolve_policy_worker_cls`	Return the quantized policy worker FQN if `quant_cfg` is set, else `default_cls`.
`resolve_model_class`	Resolve the appropriate model class for a given model name.
`is_vllm_v1_engine_enabled`	Check if vLLM V1 engine is enabled.
`get_gpu_info`	Return information about the GPU being used by this worker.
`configure_dynamo_cache`	Disable dynamo autotune_local_cache.
`get_runtime_env_for_policy_worker`	Get runtime environment configuration for policy workers.
`get_megatron_checkpoint_dir`	Gets the default megatron checkpoint directory for initial HF -> Mcore conversion.
`get_handle_from_tensor`	Get IPC handle from a tensor.
`ensure_teacher_ipc_buffer`	Lazy-alloc / grow `[N_mb, B, T, V]` teacher-logits IPC storage.
`aggregate_per_sample_handles`	Flatten teacher per-sample IPC handles into a global-batch-ordered list.
`calculate_aligned_size`	Calculate aligned size for memory alignment.
`stream_weights_via_ipc_zmq_impl`	Shared implementation for streaming weights via IPC ZMQ with improved memory management.
`rebuild_cuda_tensor_from_ipc`	Rebuild a CUDA tensor from an IPC handle.
`_derive_engine_gpu_offsets`	Cumulative-sum offsets for a dense engine layout.
`connect_colocate_topology`	Generalized colocate rollout-engine connect for FSDP and Megatron.
`_check_weight_sync_results`
`iter_named_tensor_buckets`	Group `(name, tensor)` pairs into buckets of at most `buffer_size_bytes`.
`send_hf_buckets_via_ipc_actor_impl`	Send finalized HF tensor buckets to colocated SGLang engines via Ray IPC.
`init_process_group`	Create a side-by-side `ProcessGroup` without touching the default world.
`connect_rollout_engines_from_distributed`	Set up the SGLang NCCL weight-update group with trainer rank 0 as rank 0.
`disconnect_rollout_engines_from_distributed`	Tear down trainer-side and engine-side NCCL state for `group_name`.
`broadcast_hf_buckets_via_distributed_impl`	Broadcast finalized HF tensor buckets to SGLang via NCCL (rank 0 only).

Data#

`AUTOMODEL_FACTORY`
`POLICY_WORKER_OVERRIDES`

API#

nemo_rl.models.policy.utils.AUTOMODEL_FACTORY: Dict[str, Any]#: None

class nemo_rl.models.policy.utils.IPCProtocol(*args, **kwds)#

Bases: enum.Enum

IPC protocol constants for ZMQ weight streaming.

Initialization

COMPLETE#: ‘complete’

ACK#: ‘ack’

nemo_rl.models.policy.utils.POLICY_WORKER_OVERRIDES#: None

nemo_rl.models.policy.utils.resolve_policy_worker_cls(default_cls: str, config: dict) → str#

Return the quantized policy worker FQN if quant_cfg is set, else default_cls.

Safe to call even when ModelOpt is not installed — returns default_cls unchanged whenever quant_cfg is None, so the core policy path stays import-free of ModelOpt.

nemo_rl.models.policy.utils.resolve_model_class(model_name: str) → Any#: Resolve the appropriate model class for a given model name.

nemo_rl.models.policy.utils.is_vllm_v1_engine_enabled() → bool#

Check if vLLM V1 engine is enabled.

Returns:: True if V1 engine is enabled, False otherwise (defaults to True if not set)
Return type:: bool

nemo_rl.models.policy.utils.get_gpu_info(model: torch.nn.Module) → dict[str, Any]#: Return information about the GPU being used by this worker.

nemo_rl.models.policy.utils.configure_dynamo_cache() → None#

Disable dynamo autotune_local_cache.

Dynamo may fail at cached_autotune when there’s already a cache with different order of node_bundles. Disable autotune_local_cache as a workaround. See https://github.com/pytorch/pytorch/issues/153791 for more details.

nemo_rl.models.policy.utils.get_runtime_env_for_policy_worker( policy_worker_name: str, ) → dict[str, Any]#

Get runtime environment configuration for policy workers.

Note: expandable_segments configuration is handled directly in the worker init methods to ensure proper GPU detection after CUDA initialization.

nemo_rl.models.policy.utils.get_megatron_checkpoint_dir() → str#

Gets the default megatron checkpoint directory for initial HF -> Mcore conversion.

Megatron initial checkpoint should be saved to a path available on all nodes. The directory used will take this order of precendence:

$NRL_MEGATRON_CHECKPOINT_DIR (if set)
$HF_HOME/nemo_rl (if HF_HOME is set)
~/.cache/huggingface/nemo_rl

HF_HOME is preferred since many users will also have that path mounted and it means one less directory to mount into your runtime environment.

nemo_rl.models.policy.utils.get_handle_from_tensor(tensor: torch.Tensor) → tuple[Any]#: Get IPC handle from a tensor.

nemo_rl.models.policy.utils.ensure_teacher_ipc_buffer( storage: Optional[torch.Tensor], handle: Optional[tuple[Any, ...]], num_microbatches: int, batch_size: int, seq_len: int, vocab_size: int, dtype: torch.dtype, device: torch.device, ) → tuple[torch.Tensor, tuple[Any, ...]]#

Lazy-alloc / grow [N_mb, B, T, V] teacher-logits IPC storage.

Returns the (possibly reallocated) (storage, handle). Reallocates and re-exports the IPC handle whenever any dim of the requested shape exceeds the current storage, or dtype/device changed; otherwise the existing storage and cached handle are returned unchanged.

nemo_rl.models.policy.utils.aggregate_per_sample_handles( worker_results: list[dict[str, Any]], ) → list[dict[str, Any]]#

Flatten teacher per-sample IPC handles into a global-batch-ordered list.

Each worker returns {"dp_rank": int, "per_sample_handles": list} where the handle list is in local sample order; the several workers sharing a dp_rank are TP/CP replicas that each contribute one shard per sample. Concatenating samples in sorted(dp_rank) order reproduces the original global sample order (rank 0 holds the first gbs/dp samples, rank 1 the next, …), so the result is a length-gbs list independent of the teacher’s DP degree. Element i is {"teacher_shards": [shard, ...]} holding all TP×CP shards of global sample i.

nemo_rl.models.policy.utils.calculate_aligned_size(size_bytes: int, alignment: int = 512) → int#

Calculate aligned size for memory alignment.

Parameters:

size_bytes (int) – Size in bytes to align
alignment (int) – Alignment boundary in bytes (default 512)

Returns:

Aligned size in bytes(int).

nemo_rl.models.policy.utils.stream_weights_via_ipc_zmq_impl( params_generator, buffer_size_bytes: int, zmq_socket, rank: int, worker_name: str, ) → None#

Shared implementation for streaming weights via IPC ZMQ with improved memory management.

Uses ping-pong double buffering to enable overlapping communication while reusing buffers to reduce memory allocation overhead and improve stability.

Parameters:

params_generator – Generator yielding (name, tensor) pairs
buffer_size_bytes – total size of buffer in bytes for batching parameters
zmq_socket – ZMQ socket for communication
rank – Worker rank for logging
worker_name – Name of the worker for logging

nemo_rl.models.policy.utils.rebuild_cuda_tensor_from_ipc( cuda_ipc_handle: tuple, device_id: int, ) → torch.Tensor#: Rebuild a CUDA tensor from an IPC handle.

nemo_rl.models.policy.utils._derive_engine_gpu_offsets(engine_gpu_counts: list[int]) → list[int]#: Cumulative-sum offsets for a dense engine layout.

nemo_rl.models.policy.utils.connect_colocate_topology( *, engine_gpu_counts: list[int], engine_gpu_offsets: Optional[list[int]] = None, worker_state: dict, ) → None#

Generalized colocate rollout-engine connect for FSDP and Megatron.

Builds a Gloo gather subgroup for each engine’s GPU rank range and stashes rank-only routing state into worker_state:

worker_state["_ipc_gather_group"]: ProcessGroup covering this trainer rank’s engine, or None if the rank is a placeholder / not covered by any engine.
worker_state["_ipc_gather_groups"]: all subgroup handles created for this layout, retained so a rebuild can destroy every live group.
worker_state["_ipc_gather_src"]: the source rank inside the gather group (the first GPU index of the covering engine), or None.
worker_state["_ipc_engine_index"]: index into the caller’s engine list, or None. The caller is responsible for resolving the actor handle / URL at call time so post-recover actor swaps are picked up.
worker_state["_ipc_layout_key"]: cached topology signature so subsequent connects with the same layout are no-ops.

All trainer ranks must enter this function collectively (each call to dist.new_group is collective). When the layout changes (e.g. a recovered engine resizes the topology) the cached subgroup is destroyed and rebuilt for the new layout.

nemo_rl.models.policy.utils._check_weight_sync_results(results: list) → None#

nemo_rl.models.policy.utils.iter_named_tensor_buckets( params_generator: Iterable[tuple[str, torch.Tensor]], buffer_size_bytes: int, ) → Iterable[list[tuple[str, torch.Tensor]]]#

Group (name, tensor) pairs into buckets of at most buffer_size_bytes.

Waits on async DTensor redistributes (.wait()) before sizing, so the yielded tensors are always materialized and safe to serialize.

nemo_rl.models.policy.utils.send_hf_buckets_via_ipc_actor_impl( *, bucket_iterator: Iterable[list[tuple[str, torch.Tensor]]], rollout_engines: list, worker_state: dict, weight_version: Optional[int] = None, ) → None#

Send finalized HF tensor buckets to colocated SGLang engines via Ray IPC.

Per bucket: group by dtype, serialize a FlattenedTensorBucket per dtype, dist.gather_object to the gather source rank, then on the source rank call ipc_engine.update_weights_from_tensor.remote(...) once per dtype, block on ray.get(refs) per chunk, validate engine return values, synchronize all trainer ranks, then drop the trainer-side flattened_tensor references before moving on.

The trainer-side topology (_ipc_gather_group / _ipc_gather_src / _ipc_engine_index) must already have been set up by

Func:: connect_colocate_topology. Placeholder ranks (no covering engine) return immediately — they must not call gather_object. Non-source trainer ranks participate in the gather and completion broadcast; they don’t issue Ray RPCs and don’t ray.get.

Returns None. Raises RuntimeError if any chunk fails on the engine side.

nemo_rl.models.policy.utils.init_process_group( backend: str | dist.Backend | None = None, init_method: Optional[str] = None, timeout: Optional[datetime.timedelta] = None, world_size: int = -1, rank: int = -1, store: dist.Store | None = None, group_name: Optional[str] = None, pg_options: Any = None, ) → torch.distributed.ProcessGroup#

Create a side-by-side ProcessGroup without touching the default world.

torch.distributed.init_process_group initializes the default world process group. Once the Megatron trainer has stood up its own world during Policy construction, calling it again to talk to SGLang either errors with “trying to initialize the default process group twice” or — depending on torch version — silently hangs in rendezvous against a peer that has already finished its own custom-group setup.

Same approach as SGLang’s sglang.srt.utils.common.init_custom_process_group: replay the public API’s wiring (rendezvous → PrefixStore → _new_process_group_helper) but skip the “set as default PG” step, so multiple independent groups can coexist in the same process.

Only one of init_method and store may be set; otherwise the rendezvous source is ambiguous.

nemo_rl.models.policy.utils.connect_rollout_engines_from_distributed( *, group_name: str, rollout_engines: list, engine_gpu_counts: list[int], ) → torch.distributed.ProcessGroup#

Set up the SGLang NCCL weight-update group with trainer rank 0 as rank 0.

Only trainer rank 0 broadcasts because the AutoBridge path restores full HF weights, not per-PP slices.

The caller (a trainer) must invoke this only on rank 0; other ranks must not call it.

nemo_rl.models.policy.utils.disconnect_rollout_engines_from_distributed( *, group_name: str, model_update_group: torch.distributed.ProcessGroup, rollout_engines: list, ) → None#: Tear down trainer-side and engine-side NCCL state for group_name.

nemo_rl.models.policy.utils.broadcast_hf_buckets_via_distributed_impl( *, bucket_iterator: Iterable[list[tuple[str, torch.Tensor]]], rollout_engines: list, rollout_engine_lock, group_name: str, model_update_group: torch.distributed.ProcessGroup, weight_version: int, ) → None#

Broadcast finalized HF tensor buckets to SGLang via NCCL (rank 0 only).

Per-bucket protocol: trainer rank 0 sends per-tensor metadata to every engine via Ray (update_weights_from_distributed), then issues one dist.broadcast per tensor over the NCCL group, then waits for the Ray refs to confirm engines finished loading the bucket.

The rollout-engine lock wraps each bucket’s broadcast so concurrent SGLang NCCL operations (e.g. health-check pings) cannot collide with the weight-update broadcast.

nemo_rl.models.policy.utils#