nemo_rl.weight_sync.checkpoint_engine_weight_synchronizer#

Module Contents#

Classes#

CheckpointEngineWeightSynchronizer

Coordinate checkpoint-engine setup and policy-to-rollout transfers.

Functions#

_flatten_metadata

_sort_ranked_metadata

_ordered_generation_metadata

Order vLLM generation metadata by global rollout rank.

Data#

_MEBIBYTE

API#

nemo_rl.weight_sync.checkpoint_engine_weight_synchronizer._MEBIBYTE#

None

nemo_rl.weight_sync.checkpoint_engine_weight_synchronizer._flatten_metadata(results: list[Any]) → list[Any]#
nemo_rl.weight_sync.checkpoint_engine_weight_synchronizer._sort_ranked_metadata(metadata: list[Any]) → list[Any]#
nemo_rl.weight_sync.checkpoint_engine_weight_synchronizer._ordered_generation_metadata(
generation_results: list[Any],
) → list[Any]#

Order vLLM generation metadata by global rollout rank.

Each result belongs to one vLLM data-parallel group. Engine-local ranks are unique only within a group, so sort each group before concatenating them in worker-group order.

class nemo_rl.weight_sync.checkpoint_engine_weight_synchronizer.CheckpointEngineWeightSynchronizer#

Bases: nemo_rl.weight_sync.interfaces.WeightSynchronizer

Coordinate checkpoint-engine setup and policy-to-rollout transfers.

_policy: Any#

None

_generation: Any#

None

_checkpoint_engine_config: nemo_rl.models.generation.interfaces.CheckpointEngineConfig#

None

_stale: bool#

True

_checkpoint_engine_ready: bool#

False

_bucket_size_bytes: int | None#

None

init_communicator() → None#
property is_stale: bool#
_release_after_refit() → bool#
_run_policy(
checkpoint_method: str,
**method_kwargs: Any,
) → list[ray.ObjectRef]#
_generation_rpc() → str#
_run_generation(
checkpoint_method: str,
method_args: tuple[Any, ...] = (),
) → list[ray.ObjectRef]#
_resolve_bucket_size_bytes() → int#
_ensure_checkpoint_engine_ready() → None#
sync_weights(
*,
timer: Optional[nemo_rl.utils.timer.Timer] = None,
kv_scales: Optional[dict[str, float]] = None,
) → None#
shutdown() → None#