nemo_rl.models.policy.workers.base_policy_worker#

Module Contents#

Classes#

AbstractPolicyWorker

Base class for policy workers with shared functionality.

API#

class nemo_rl.models.policy.workers.base_policy_worker.AbstractPolicyWorker#

Base class for policy workers with shared functionality.

model_update_group: Optional[Any]#

None

pp_comm_group: Optional[Any]#

None

init_collective(
ip: str,
port: int,
world_size: int,
*,
train_world_size: int,
nccl_peer: str = 'nemo',
) → None#

Initialize the collective communication.

Parameters:

  • ip – IP address for the process group

  • port – Port for the process group

  • world_size – Total world size (train_world_size + inference_world_size)

  • train_world_size – Number of training workers (used in inference cluster)

  • nccl_peer – NCCL initialization protocol used by the inference workers

init_nccl_reshard_comm_group(
pp_ips: list[str],
pp_ports: list[int],
pp_size: int,
my_pp_stage: int,
sub_world_size: int,
my_rank_in_group: int,
) → None#

Bootstrap this train worker’s nccl_reshard comm group.

One comm group per PP stage; each train worker joins exactly its own stage’s group (pp_ips[my_pp_stage] / pp_ports[my_pp_stage]). Non-PP is simply pp_size == 1 that contains all the train ranks.

abstractmethod prepare_nccl_reshard_refit_info(
train_parallelism: dict[str, int],
gen_parallelism: dict[str, int],
train_world_size: int,
gen_world_size: int,
) → dict[str, Any]#

Prepare parameter metadata for NCCL reshard refit.

abstractmethod nccl_reshard_refit(
kv_scales: Optional[dict[str, float]] = None,
refit_timeout_s: Optional[float] = None,
) → None#

Transfer policy weights with NCCL reshard refit.

_refit_transport_state(key: str) → dict#

Return the mutable state dict for one refit transport, e.g. "sglang_ipc".

Transports stash their bookkeeping here (gather groups, engine handles, weight versions) instead of each backend adding fields to the worker constructors. Created on first use because this base class has no __init__ — the same lazy shape as maybe_init_zmq on the ZMQ path.

connect_sglang_rollout_engines(
*,
engine_gpu_counts: list[int],
engine_gpu_offsets: Optional[list[int]] = None,
) → None#

Set up the colocate Gloo gather topology for SGLang weight refit.

Must be called collectively by every trainer rank when SGLang engines are added or recovered. Subsequent calls with the same layout are no-ops.

stand_down_refit_watchdog() → int#

Cancel any in-flight refit deadline in this worker; returns how many.

Called by the controller when a generation shard’s process is confirmed GONE. A dead peer closes its sockets and NCCL unblocks the survivors on its own, so the deadline has nothing useful to do and firing it is actively harmful – see stand_down_armed_watchdogs.

Deliberately trivial: this runs on the actor’s event loop while the refit itself runs off it, so it must not block or touch CUDA.

is_alive() → bool#

Check if the worker is alive.

reset_peak_memory_stats() → None#

Reset peak memory statistics.

get_gpu_info() → dict[str, Any]#

Return information about the GPU being used by this worker.

report_device_id() → str#

Report the UUID of the current CUDA device using NVML.

Returns:

UUID of the device in the format “GPU-xxxxx”

Return type:

str

get_zmq_address() → str#

Get the ZMQ address for the current device.

maybe_init_zmq() → None#

Initialize the ZMQ socket if it doesn’t exist.

get_free_memory_bytes() → int#

Get the available free memory.

shutdown() → bool#

Shutdown the policy.

start_gpu_profiling() → None#

Start GPU profiling.

stop_gpu_profiling() → None#

Stop GPU profiling.

report_node_ip_and_gpu_id() → tuple[str, int]#

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

get_reference_policy_logprobs(
*,
data: nemo_rl.distributed.batched_data_dict.BatchedDataDict[Any],
micro_batch_size: Optional[int] = None,
) → nemo_rl.distributed.batched_data_dict.BatchedDataDict[nemo_rl.models.policy.interfaces.ReferenceLogprobOutputSpec]#

Get the logprobs from the reference policy for a batch of data.

If micro_batch_size is provided, it will be used instead of the configured logprob_batch_size.

Returns:

a BatchedDataDict with key “reference_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.

finalize_async_save() → None#

Block until any in-flight async checkpoint write completes. No-op by default.

finish_training(*args: Any, **kwargs: Any) → None#