`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.
`_ensure_ipc_topology`	Lazily create a per-engine Gloo subgroup and cache rank-only routing state.
`_flush_bucket`	Flatten `named_tensors` per dtype, gather to `gather_src`, and POST to the engine.
`stream_weights_via_http_impl`	Stream FSDP weights to colocated SGLang engines via CUDA IPC over HTTP.

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._ensure_ipc_topology( num_engines: int, num_gpus_per_engine: int, worker_state: dict, ) → None#

Lazily create a per-engine Gloo subgroup and cache rank-only routing state.

Every FSDP rank must call dist.new_group for every engine’s rank range (collective). Only the ranks inside a given range stash gather_src and gather_group into worker_state. The engine handle itself is resolved at call time from the caller-provided rollout_engines list so that post-recover actor swaps are picked up without cache invalidation.

Note: callers must have already applied monkey_patch_torch_reductions once during worker setup; this function no longer applies it.

nemo_rl.models.policy.utils._flush_bucket( named_tensors, gather_src: int, gather_group, engine_url: str, weight_version: int, flattened_tensor_bucket_cls, multiprocessing_serializer_cls, ) → None#: Flatten named_tensors per dtype, gather to gather_src, and POST to the engine.

nemo_rl.models.policy.utils.stream_weights_via_http_impl( params_generator: Iterable[tuple[str, torch.Tensor]], rollout_engine_urls: Iterable[str], num_gpus_per_engine: int, rank: int, world_size: int, worker_name: str, buffer_size_bytes: int, worker_state: dict, ) → None#

Stream FSDP weights to colocated SGLang engines via CUDA IPC over HTTP.

Parameters:

params_generator – Iterable yielding (name, tensor) pairs to stream. Caller is responsible for any pre-processing (LoRA merge, HF adaptation, dtype cast).
rollout_engine_urls – http://host:port base URLs of each engine’s node_rank=0 SGLang HTTP server. One entry per engine, in TP rank-range order: engine i owns global ranks [i * num_gpus_per_engine, (i + 1) * num_gpus_per_engine).
num_gpus_per_engine – TP size per SGLang engine.
rank – Global FSDP rank.
world_size – Global FSDP world size.
worker_name – Human label for logs.
buffer_size_bytes – Max bucket size in bytes.
worker_state – Mutable dict on the worker used to cache topology and weight version across refits.

nemo_rl.models.policy.utils#

Module Contents#

Classes#

Functions#

Data#

API#

`nemo_rl.models.policy.utils`#