nemo_rl.weight_sync.interfaces#

Weight synchronization interface for NeMo-RL.

WeightSynchronizer is a dedicated abstraction that decouples weight transfer logic from both PolicyInterface and GenerationInterface. It owns the transfer of model weights between training and generation components.

Transport-specific implementations (IPC/ZMQ, HTTP, NCCL collectives, checkpoint engines) each encapsulate the transfer lifecycle, so algorithm code never branches on backend type.

Colocated transports (IPC, HTTP) own GPU phase transitions internally (offload, prepare_for_generation, restore) as part of their sync_weights() implementation. The NCCL collective transport is a pure data mover; the orchestrator handles phase transitions externally since policy and generation run on separate GPU clusters.

This interface assumes global weight updates: all generation workers are updated atomically and are always at the same weight version. Per-worker updates (where different replicas could be at different versions) are not supported. In async GRPO, heterogeneous weight ages are handled at the sample level (via replay buffer target_weight_versions tracking), not at the synchronizer level.

Module Contents#

Classes#

WeightSynchronizer

Abstract base class for weight synchronization between policy and generation.

API#

class nemo_rl.weight_sync.interfaces.WeightSynchronizer#

Bases: abc.ABC

Abstract base class for weight synchronization between policy and generation.

Implementations handle the weight transfer for a specific transport mechanism (ZMQ IPC, HTTP, NCCL collectives). The orchestrator calls sync_weights() without knowing which transport is being used or whether components are colocated; per-step staleness bookkeeping is owned by the training loop.

Colocated transports (IPC, HTTP) own phase transitions internally (offload_before_refit, prepare_for_generation, offload_after_refit). Non-colocated collective and checkpoint-engine transports are pure data movers; the orchestrator handles phases externally.

abstractmethod sync_weights(
*,
timer: Optional[nemo_rl.utils.timer.Timer] = None,
kv_scales: Optional[dict[str, float]] = None,
) → Optional[dict[str, float]]#

Transfer the latest policy weights to the generation backend.

This method encapsulates the full sync lifecycle:

  1. Prepare the policy side (e.g., offload optimizer state to free GPU memory)

  2. Prepare the generation side (e.g., allocate weight buffers)

  3. Transfer weights via the transport mechanism

  4. Verify the transfer succeeded

  5. Restore both sides to their ready state

Steps 1-2 and 5 (phase transitions) are only performed by colocated transports (IPC, HTTP). Non-colocated collective and checkpoint-engine transports skip them since policy and generation run on separate GPUs.

Step 4 (verification) is performed explicitly by IPC and NCCL transports, which check update_success and raise on failure. The HTTP transport relies on ray.get() to propagate any server-side errors (matching the existing grpo.py behavior).

Parameters:

  • timer – Optional Timer for profiling individual phases.

  • kv_scales – Optional KV cache scales for FP8 quantization. Honored by the IPC/ZMQ and NCCL collective transports. The HTTP transport ignores this parameter.

Returns:

Optional transport-specific scalar metrics for the current sync.

Raises:

RuntimeError – If the weight transfer fails.

abstract property is_stale: bool#

Whether the generation backend’s weights are out of date.

Returns True until the first successful sync_weights() completes, so a fresh run always performs its initial sync (a synchronizer that seeds current weights at construction may start False to skip it). Per-step staleness is tracked by the training loop, not here.

abstractmethod init_communicator() → None#

Initialize any communication channels needed for weight transfer.

Called once during setup, after policy and generation workers are constructed. For colocated IPC/HTTP transports this may prepare refit metadata. For NCCL collectives this initializes the process group.

abstractmethod shutdown() → None#

Release all communication resources.