nemo_rl.models.policy.workers.megatron_remote_sparse_refit

Canonical Hugging Face sparse-refit state for a Megatron policy worker.

Module Contents

Classes

MegatronRemoteSparseRefit

API

class nemo_rl.models.policy.workers.megatron_remote_sparse_refit.MegatronRemoteSparseRefit(

worker: Any,

refit_config: nemo_rl.models.generation.vllm.config.VllmSparseRefitConfig,

)

Initialization

initialize_baseline(

*,

shard_rank: int,

shard_count: int,

transport: nemo_rl.models.generation.vllm.config.VllmRefitTransportName,

) → dict[str, tuple[tuple[int, ...], torch.dtype]]

stream(

transport: nemo_rl.models.generation.vllm.config.VllmRefitTransportName,

targets: list[str],

*,

transfer_id: str,

api_key_env_var: str | None,

timeout_s: float,

shard_rank: int,

shard_count: int,

overwrite_names: list[str],

) → dict[str, int]

finish(succeeded: bool) → None