`nemo_rl.utils.checkpoint_engines.base`#

Module Contents#

Classes#

`TensorMeta`
`CheckpointEngine`

Functions#

`create_checkpoint_engine`
`split_weight_chunks`
`merge_weight_chunk_batches`

API#

class nemo_rl.utils.checkpoint_engines.base.TensorMeta#

name: str#: None

shape: torch.Size#: None

dtype: torch.dtype#: None

chunk_offset: int#: None

chunk_size: int#: None

offset: int | None#: None

property nbytes: int#

class nemo_rl.utils.checkpoint_engines.base.CheckpointEngine#

Bases: abc.ABC

shard_expert_weights: bool#: False

get_target_weight_layout() → dict[str, Any] | None#: Return the destination-local layout for this policy rank, if any.

abstractmethod prepare() → Any#: Allocate or register backend resources and return serializable metadata.

abstractmethod init_policy_process_group( *, worker_rank: int, train_world_size: int, rollout_world_size: int, metadata: list[Any], ) → None#: Connect a policy worker to its transfer peer.

abstractmethod init_rollout_process_group( *, rollout_rank: int, train_world_size: int, rollout_world_size: int, metadata: list[Any], ) → None#: Connect a rollout worker to its transfer peer.

finalize() → None#: Release per-refit backend state.

abstractmethod async send_weights( weights: collections.abc.Generator[tuple[str, torch.Tensor], None, None], ) → None#: Send (name, tensor) weights from the policy side.

abstractmethod receive_weight_batches() → collections.abc.AsyncGenerator[list[tuple[str, torch.Tensor]], None]#: Yield (name, tensor) batches on the generation side.

nemo_rl.utils.checkpoint_engines.base.create_checkpoint_engine( backend: str, *, bucket_size_bytes: int, engine_kwargs: dict[str, Any], ) → nemo_rl.utils.checkpoint_engines.base.CheckpointEngine#

nemo_rl.utils.checkpoint_engines.base.split_weight_chunks( weights: collections.abc.Generator[tuple[str, torch.Tensor], None, None], bucket_size: int, ) → collections.abc.Generator[tuple[nemo_rl.utils.checkpoint_engines.base.TensorMeta, torch.Tensor], None, None]#

async nemo_rl.utils.checkpoint_engines.base.merge_weight_chunk_batches( chunk_batches: collections.abc.AsyncGenerator[list[tuple[nemo_rl.utils.checkpoint_engines.base.TensorMeta, torch.Tensor]], None], *, merge_device: torch.device | str | None = None, ) → collections.abc.AsyncGenerator[list[tuple[str, torch.Tensor]], None]#