nemo_rl.models.generation.vllm.checkpoint_engine

Module Contents

Classes

VllmCheckpointEngineMixin	Checkpoint-engine lifecycle for vLLM workers.
VllmCheckpointEngineRpcMixin	Dispatch checkpoint-engine calls through a synchronous vLLM engine.
VllmAsyncCheckpointEngineRpcMixin	Dispatch checkpoint-engine calls through an asynchronous vLLM engine.

Functions

configure_nixl_worker	Configure vLLM’s worker hook for early NIXL initialization.
preinit_nixl_from_vllm_config	Create the NIXL preinit agent carried by a vLLM internal worker.
resolve_rollout_rank

Data

NIXL_VLLM_WORKER
_NIXL_CONFIG_KEY

API

nemo_rl.models.generation.vllm.checkpoint_engine.NIXL_VLLM_WORKER

‘nemo_rl.models.generation.vllm.vllm_backend.NixlVllmWorker’

nemo_rl.models.generation.vllm.checkpoint_engine._NIXL_CONFIG_KEY

‘nemo_rl_checkpoint_engine’

nemo_rl.models.generation.vllm.checkpoint_engine.configure_nixl_worker(

config: nemo_rl.models.generation.vllm.config.VllmConfig,

vllm_kwargs: dict[str, Any],

) → None

Configure vLLM’s worker hook for early NIXL initialization.

nemo_rl.models.generation.vllm.checkpoint_engine.preinit_nixl_from_vllm_config(vllm_config: Any) → Any

Create the NIXL preinit agent carried by a vLLM internal worker.

nemo_rl.models.generation.vllm.checkpoint_engine.resolve_rollout_rank(rank_prefix: int, rollout_world_size: int) → int

class nemo_rl.models.generation.vllm.checkpoint_engine.VllmCheckpointEngineMixin

Bases: nemo_rl.models.generation.vllm.refit_loader.VllmShardedExpertRefitMixin

Checkpoint-engine lifecycle for vLLM workers.

checkpoint_engine: nemo_rl.utils.checkpoint_engines.base.CheckpointEngine

None

checkpoint_engine_total_memory_bytes() → int

_load_hf_weights(

policy_weights: list[tuple[str, torch.Tensor]],

) → None

init_checkpoint_engine(

backend: str,

bucket_size_bytes: int,

engine_kwargs: dict[str, Any],

) → None

prepare_checkpoint_engine() → Any

init_checkpoint_engine_process_group(

rank_prefix: int,

train_world_size: int,

rollout_world_size: int,

metadata: list[Any],

) → None

finalize_checkpoint_engine() → None

async _update_weights_from_checkpoint_engine_async() → bool

update_weights_from_checkpoint_engine() → bool

class nemo_rl.models.generation.vllm.checkpoint_engine.VllmCheckpointEngineRpcMixin

Dispatch checkpoint-engine calls through a synchronous vLLM engine.

checkpoint_engine_rpc(

checkpoint_method: str,

method_args: tuple[Any, ...] = (),

) → Any

class nemo_rl.models.generation.vllm.checkpoint_engine.VllmAsyncCheckpointEngineRpcMixin

Dispatch checkpoint-engine calls through an asynchronous vLLM engine.

async checkpoint_engine_rpc_async(

checkpoint_method: str,

method_args: tuple[Any, ...] = (),

) → Any