`nemo_rl.models.generation.vllm.refit_loader`#

Module Contents#

Classes#

VllmShardedExpertRefitMixin

Load destination-local expert shards into vLLM storage.

API#

class nemo_rl.models.generation.vllm.refit_loader.VllmShardedExpertRefitMixin#

Load destination-local expert shards into vLLM storage.

_is_sharded_refit_weight(name: str, tensor: torch.Tensor) → bool#

_sharded_refit_param_names(name: str) → list[str]#

_checkpoint_engine_weight_layout() → nemo_rl.models.generation.vllm.refit_layout.VllmWeightLayout#

static _validate_expert_storage( param_name: str, param: torch.nn.Parameter, ) → None#: Reject incompatible vLLM storage before advertising shards.

_local_expert_id(param: torch.nn.Parameter, expert_id: int) → int#

_load_destination_local_expert_group( param_name: str, param: torch.nn.Parameter, shard_id: str, items: list[tuple[int, torch.Tensor]], ) → None#: Copy destination-local experts into canonical vLLM storage.

_load_sharded_expert_weight_groups( weights: list[tuple[str, torch.Tensor]], ) → list[tuple[str, torch.Tensor]]#

_load_sharded_expert_weights( policy_weights: list[tuple[str, torch.Tensor]], ) → None#