`nemo_rl.weight_sync.vllm_remote_sparse_weight_synchronizer`#

Shared S3/ZeroMQ sparse synchronizer for remote non-colocated vLLM refit.

Module Contents#

Classes#

VllmRemoteSparseWeightSynchronizer

Functions#

validate_vllm_remote_sparse_refit

Validate the optional config and return its internal transport name.

Data#

_REMOTE_SPARSE_TRANSPORTS

API#

nemo_rl.weight_sync.vllm_remote_sparse_weight_synchronizer._REMOTE_SPARSE_TRANSPORTS#: None

nemo_rl.weight_sync.vllm_remote_sparse_weight_synchronizer.validate_vllm_remote_sparse_refit( config: nemo_rl.models.generation.vllm.config.VllmConfig, *, colocated: bool, megatron_enabled: bool, ) → str | None#: Validate the optional config and return its internal transport name.

class nemo_rl.weight_sync.vllm_remote_sparse_weight_synchronizer.VllmRemoteSparseWeightSynchronizer( policy: Any, generation: Any, *, transport: str, api_key_env_var: str | None = None, request_timeout_s: float = 600.0, baseline_init_refs: list[Any] | None = None, )#

Bases: nemo_rl.weight_sync.interfaces.WeightSynchronizer

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

property is_stale: bool#

_run_policy_workers(

method_name: str,

**kwargs: Any,

) → list[Any]#

_run_generation_workers(

method_name: str,

**kwargs: Any,

) → list[Any]#

_request_receivers( path: str, body: dict[str, Any], *, timeout_s: float | None = None, ) → list[dict[str, Any]]#

static start_baseline(policy: Any, transport: str) → list[Any]#

static _merge_refit_info( parts: list[dict[str, Any]], ) → dict[str, Any]#

init_communicator() → None#

shutdown() → None#

nemo_rl.weight_sync.vllm_remote_sparse_weight_synchronizer#