nemo_rl.weight_sync.factory#

Factory for creating WeightSynchronizer instances.

Selects the appropriate weight synchronizer based on the deployment topology (colocated vs. non-colocated) and the generation backend (vLLM uses IPC/ZMQ, SGLang uses HTTP, non-colocated uses NCCL).

Module Contents#

Functions#

create_weight_synchronizer

Create the appropriate WeightSynchronizer for the given deployment.

API#

nemo_rl.weight_sync.factory.create_weight_synchronizer(
policy: Any,
generation: Any,
generation_backend: str,
colocated: bool,
train_cluster: Optional[Any] = None,
inference_cluster: Optional[Any] = None,
refit_buffer_size_gb: Optional[int] = None,
) → nemo_rl.weight_sync.interfaces.WeightSynchronizer#

Create the appropriate WeightSynchronizer for the given deployment.

Parameters:

  • policy – Policy object (ColocatablePolicyInterface).

  • generation – Generation object (GenerationInterface).

  • generation_backend – Name of the generation backend (“vllm”, “sglang”, “megatron”).

  • colocated – Whether policy and generation share the same GPUs.

  • train_cluster – RayVirtualCluster for training workers (required for non-colocated).

  • inference_cluster – RayVirtualCluster for inference workers (required for non-colocated).

  • refit_buffer_size_gb – Optional fixed buffer size for IPC weight staging.

Returns:

A WeightSynchronizer instance appropriate for the deployment topology.

Raises:

  • NotImplementedError – If the requested configuration is not supported.

  • ValueError – If required arguments are missing.