`nemo_rl.models.generation.vllm.refit_layout`#

Module Contents#

Classes#

`VllmExpertParamLayout`
`VllmWeightLayout`
`HfExpertWeight`

Functions#

`parse_hf_expert_weight`
`select_hf_weight_for_vllm_target`	Return the destination-local weight, or `None` if not owned.

Data#

`_HF_EXPERT_WEIGHT_RE`
`_HF_PROJECTION_SHARDS`

API#

class nemo_rl.models.generation.vllm.refit_layout.VllmExpertParamLayout#

Bases: typing.TypedDict

tp_rank: int#: None

tp_size: int#: None

local_expert_ids: list[int] | None#: None

class nemo_rl.models.generation.vllm.refit_layout.VllmWeightLayout#

Bases: typing.TypedDict

expert_params: dict[str, nemo_rl.models.generation.vllm.refit_layout.VllmExpertParamLayout]#: None

missing_weight_prefixes: list[str]#: None

class nemo_rl.models.generation.vllm.refit_layout.HfExpertWeight#

parameter_name: str#: None

expert_id: int#: None

shard_id: Literal[w1, w2, w3]#: None

tp_shard_dim: int#: None

nemo_rl.models.generation.vllm.refit_layout._HF_EXPERT_WEIGHT_RE#: ‘compile(…)’

nemo_rl.models.generation.vllm.refit_layout._HF_PROJECTION_SHARDS: dict[str, Literal[w1, w2, w3]]#: None

nemo_rl.models.generation.vllm.refit_layout.parse_hf_expert_weight( name: str, ) → nemo_rl.models.generation.vllm.refit_layout.HfExpertWeight | None#

nemo_rl.models.generation.vllm.refit_layout.select_hf_weight_for_vllm_target( name: str, tensor: torch.Tensor, *, target_layout: nemo_rl.models.generation.vllm.refit_layout.VllmWeightLayout, ) → torch.Tensor | None#

Return the destination-local weight, or None if not owned.

Pipeline stages omit complete parameter prefixes. Within an owned stage, tensor-parallel MoE layers shard every expert tensor while expert-parallel layers place complete experts on selected ranks.

nemo_rl.models.generation.vllm.refit_layout#