`nemo_rl.weight_sync.nccl_reshard_utils`#

Refit metadata builders and lightweight wrapper types for nccl_reshard.

This module provides:

MeshInfo: lightweight DeviceMesh-compatible wrapper that doesn’t require a shared torch.distributed process group (needed for cross-world transfers)
Placement rules: mapping param names to TP/EP sharding strategies
build_nccl_reshard_refit_info: compute per-layer param metadata for refit
make_nccl_reshard_refit_info_wire_safe: convert placements and meshes into plain dicts/lists before the refit metadata crosses a process boundary
restore_refit_info_placements: undo msgspec dict-flattening of placements and meshes on the receiving side

The transfer kernel (xferdtensor) and its DTensorRef src/dst wrapper live in nemo_rl/weight_sync/xferdtensor.py — import both from there.

Module Contents#

Classes#

`MeshInfo`	Lightweight mesh metadata compatible with xferdtensor.
`RefitCtx`	Handoff between a param’s `pre` and `post` refit hooks.
`LocalParamSpec`	A backend’s recipe for transferring one HF param via `xferdtensor`.
`HFToLocalParamMap`	`hf_name -> LocalParamSpec` container returned by build_hf_to_local_param_map.
`RefitBuilderInterface`	Structural contract for an nccl_reshard refit backend (train src / gen dst).

Functions#

`get_tp_shard_dim`	Return the TP shard dim for an FFN weight, or None if not TP-sharded.
`is_expert_param`	Return True if the parameter is a MoE expert weight (sharded by EP).
`is_nccl_reshard_param`	Return True iff the param takes the xferdtensor bulk reshard path.
`_get_expert_tp_shard_dim`	Like get_tp_shard_dim but does NOT skip .experts. params.
`_restore_placement`	Reconstruct a single placement object after msgspec round-trip.
`make_nccl_reshard_refit_info_wire_safe`	Copy refit metadata into types safe for vLLM’s subprocess RPC.
`restore_refit_info_placements`	Restore placements and meshes in `refit_info` after msgspec transit.
`build_mesh_info`	Build a `MeshInfo` and dim_map from a parallelism config.
`get_placements`	Determine DTensor placements for a parameter given a dim_map.
`group_expert_params_in_metadata`	Group per-expert MoE params into backend-agnostic grouped HF entries.
`_extract_layer_prefix`	Return the module prefix before `layers.N`.
`_extract_layer_name`	Extract the per-layer group name from a parameter name.
`check_nccl_reshard_refit_support`	Validate `master_config` against every precondition of nccl_reshard_refit.
`build_nccl_reshard_refit_info`	Build per-layer parameter info for nccl_reshard-based refit.

Data#

`COLUMN_PARALLEL_SUFFIXES`
`ROW_PARALLEL_SUFFIXES`
`FFN_PROJ_WEIGHT_SUFFIXES`
`FFN_GROUPED_EXPERT_SUFFIXES`
`_STR_TO_DTYPE`
`_INDIVIDUAL_EXPERT_RE`
`_LAYER_RE`
`_MODEL_PREFIX_RE`

API#

class nemo_rl.weight_sync.nccl_reshard_utils.MeshInfo(rank_tensor: torch.Tensor)#

Lightweight mesh metadata compatible with xferdtensor.

Provides the same .mesh / ._mesh interface as DeviceMesh but without requiring torch.distributed process groups – allowing xferdtensor to read mesh topology across separate torch.distributed worlds.

Initialization

property ndim#

class nemo_rl.weight_sync.nccl_reshard_utils.RefitCtx#

Handoff between a param’s pre and post refit hooks.

The transfer API (xferdtensor) reads only buf. extra is provided for flexible, backend-specific state.

Use case:

vLLM merged params tracks the merged param slice in extra["region"]

buf: torch.Tensor#: None

extra: dict[str, Any]#: ‘field(…)’

class nemo_rl.weight_sync.nccl_reshard_utils.LocalParamSpec#

A backend’s recipe for transferring one HF param via xferdtensor.

base: base form of the param tensor.
pre:  ``base -> RefitCtx``; materializes the transfer subject/object
    in RefitCtx.buf. If ``None``, ``RefitCtx(buf=base)`` is used.
    e.g., stack grouped MoE expert params.
    e.g., Create a temporary buffer for the merged param.
post: ``RefitCtx -> None``; runs after xferdtensor
    e.g., copy back the received buffer into the merged param.

TODO: A layout that block-permutes the assembled param (e.g. FlashInfer TRTLLM w13) would need a group-level finalize run once after all components land — a future loop-level addition, not a per-param field. pre/post covers today’s backends (Triton, FlashInfer CUTLASS, Megatron).

base: Any#: None

pre: Optional[Callable[[Any], nemo_rl.weight_sync.nccl_reshard_utils.RefitCtx]]#: None

post: Optional[Callable[[nemo_rl.weight_sync.nccl_reshard_utils.RefitCtx], None]]#: None

class nemo_rl.weight_sync.nccl_reshard_utils.HFToLocalParamMap#

hf_name -> LocalParamSpec container returned by build_hf_to_local_param_map.

Holds LocalParamSpec for each HF param name.

specs: dict[str, nemo_rl.weight_sync.nccl_reshard_utils.LocalParamSpec]#: ‘field(…)’

get( hf_name: str, default: Optional[nemo_rl.weight_sync.nccl_reshard_utils.LocalParamSpec] = None, ) → Optional[nemo_rl.weight_sync.nccl_reshard_utils.LocalParamSpec]#: Spec for hf_name or default (None); loops assert non-None.

class nemo_rl.weight_sync.nccl_reshard_utils.RefitBuilderInterface#

Bases: typing.Protocol

Structural contract for an nccl_reshard refit backend (train src / gen dst).

A backend builds its hf_name -> LocalParamSpec map once via

Meth:: build_hf_to_local_param_map, then its nccl_reshard_refit loop drives each param’s transfer through the spec’s pre/post hooks.

build_hf_to_local_param_map( refit_info: dict, ) → nemo_rl.weight_sync.nccl_reshard_utils.HFToLocalParamMap#: Build the unified hf_name -> LocalParamSpec map for nccl_reshard refit.

nemo_rl.weight_sync.nccl_reshard_utils.COLUMN_PARALLEL_SUFFIXES#: (‘gate_proj.weight’, ‘up_proj.weight’)

nemo_rl.weight_sync.nccl_reshard_utils.ROW_PARALLEL_SUFFIXES#: (‘down_proj.weight’,)

nemo_rl.weight_sync.nccl_reshard_utils.get_tp_shard_dim(param_name: str) → Optional[int]#

Return the TP shard dim for an FFN weight, or None if not TP-sharded.

gate/up are column-parallel (dim 0), down is row-parallel (dim 1). MoE experts shard on EP not TP, so they return None here — get_placements routes experts through _get_expert_tp_shard_dim instead.

nemo_rl.weight_sync.nccl_reshard_utils.is_expert_param(param_name: str) → bool#: Return True if the parameter is a MoE expert weight (sharded by EP).

nemo_rl.weight_sync.nccl_reshard_utils.FFN_PROJ_WEIGHT_SUFFIXES#: (‘gate_proj.weight’, ‘up_proj.weight’, ‘down_proj.weight’)

nemo_rl.weight_sync.nccl_reshard_utils.FFN_GROUPED_EXPERT_SUFFIXES#: (‘experts.gate_up_proj’, ‘experts.down_proj’)

nemo_rl.weight_sync.nccl_reshard_utils.is_nccl_reshard_param(param_name: str) → bool#

Return True iff the param takes the xferdtensor bulk reshard path.

FFN projection weights take the bulk path: the split gate_proj / up_proj / down_proj (dense MLP + per-expert MoE) and the grouped gate-up-fused MoE experts (experts.gate_up_proj / experts.down_proj). Everything else falls back to the misc packed_broadcast + vLLM load_weights path.

Shared-expert FFN weights (*.shared_expert.*) are routed to misc path. Bare mtp.-prefixed HF names are routed to misc too, because vLLM keeps the MTP drafter separate and updates it through load_weights. That only covers families whose HF names keep the mtp. prefix. DeepSeek exports MTP under model.layers.N HF names (the mtp. appears only on the Megatron side), so those return True here; the caller drops them instead, using the layer set from _collect_mtp_hf_layer_names.

nemo_rl.weight_sync.nccl_reshard_utils._get_expert_tp_shard_dim(param_name: str) → Optional[int]#: Like get_tp_shard_dim but does NOT skip .experts. params.

nemo_rl.weight_sync.nccl_reshard_utils._STR_TO_DTYPE#: None

nemo_rl.weight_sync.nccl_reshard_utils._restore_placement(p)#

Reconstruct a single placement object after msgspec round-trip.

vLLM’s collective_rpc serializes Shard(N) to {"dim": N} and Replicate() to {}; this restores them to real instances. No-op if the input is already a Shard/Replicate.

nemo_rl.weight_sync.nccl_reshard_utils.make_nccl_reshard_refit_info_wire_safe(refit_info: dict) → dict#

Copy refit metadata into types safe for vLLM’s subprocess RPC.

Importing megatron.core replaces torch.storage._load_from_bytes with a Megatron function, so Tensor pickles require Megatron at unpickle time. vLLM subprocesses may not have it importable; convert the metadata to the plain representation accepted by restore_refit_info_placements.

nemo_rl.weight_sync.nccl_reshard_utils.restore_refit_info_placements(refit_info: dict) → dict#

Restore placements and meshes in refit_info after msgspec transit.

vLLM’s collective_rpc encodes Shard/Replicate as plain dicts and MeshInfo as a dict with a nested mesh list. This rebuilds the original Python objects in place so that the canonical xferdtensor (which relies on isinstance(p, Shard)) works correctly. Idempotent — safe to call on already-restored refit_info.

nemo_rl.weight_sync.nccl_reshard_utils.build_mesh_info( num_gpus: int, rank_offset: int, tp_size: int = 1, ep_size: int = 1, pp_size: int = 1, ) → tuple#

Build a MeshInfo and dim_map from a parallelism config.

Dims are emitted in the order (tp, ep, dp, pp), size-1 dims are dropped, and the survivors are reversed into the row-major rank tensor (outer->inner). So the first surviving dim in that order becomes the innermost (rightmost, fastest-varying) axis — consecutive global ranks differ in it.

Callers never activate EP and TP in the same mesh: _build_train_src_meshes builds a separate non-expert mesh (ep_size=1) and expert mesh (tp_size=1). So the innermost active dim — and the coord a modulo recovers — is:

TP in the non-expert mesh (EP dropped): global_rank % tp_size recovers the TP coord, the standard Megatron non-expert layout.
EP in the expert mesh (TP dropped): global_rank % ep_size recovers the EP coord, Megatron-Core’s MoE rank layout.

(Were EP and TP ever active together, the emit order would make TP — not EP — innermost; that case does not arise today.)

Returns:: (MeshInfo, dim_map) where dim_map maps "tp"/"ep"/"dp"/"pp" to the corresponding mesh-tensor axis index.

nemo_rl.weight_sync.nccl_reshard_utils.get_placements(param_name: str, dim_map: dict, ndim: int) → list#

Determine DTensor placements for a parameter given a dim_map.

1-D params (layernorm, bias) are always fully replicated. Expert params shard dim 0 on EP; their TP shard dims are shifted by +1.

nemo_rl.weight_sync.nccl_reshard_utils._INDIVIDUAL_EXPERT_RE#: ‘compile(…)’

nemo_rl.weight_sync.nccl_reshard_utils.group_expert_params_in_metadata( state_dict_metadata: dict[str, dict[str, Any]], ) → dict[str, dict[str, Any]]#

Group per-expert MoE params into backend-agnostic grouped HF entries.

This function replaces expert entries in the state_dict_metadata with grouped-expert entries.

For each MoE projection, stack every expert’s HF param into ONE entry along a new leading expert dim, keeping the HF projection name:

gate_proj : [E, intermediate, hidden]
up_proj : [E, intermediate, hidden]
down_proj : [E, hidden, intermediate] Each grouped entry is tagged grouped_expert_proj (“gate_proj” / “up_proj” / “down_proj”).

The input state_dict_metadata has a global view of the parameters. Non-expert params are passed through unchanged.

nemo_rl.weight_sync.nccl_reshard_utils._LAYER_RE#: ‘compile(…)’

nemo_rl.weight_sync.nccl_reshard_utils._MODEL_PREFIX_RE#: ‘compile(…)’

nemo_rl.weight_sync.nccl_reshard_utils._extract_layer_prefix(param_name: str) → Optional[str]#

Return the module prefix before layers.N.

model / model.language_model / backbone for the usual layouts, "" for a bare layers.N, or None if the name has no layers.N.

nemo_rl.weight_sync.nccl_reshard_utils._extract_layer_name(param_name: str) → str#

Extract the per-layer group name from a parameter name.

.. rubric:: Examples

model.layers.0.mlp.gate_proj.weight -> model.layers.0 model.language_model.layers.1.mlp.up_proj.weight -> model.language_model.layers.1 backbone.layers.3.mixer.experts.0.down_proj.weight -> backbone.layers.3 layers.2.ffn.shared_experts.w2.weight -> layers.2

nemo_rl.weight_sync.nccl_reshard_utils.check_nccl_reshard_refit_support(master_config: dict) → None#

Validate master_config against every precondition of nccl_reshard_refit.

Collects all violations and raises a single ValueError listing them, so a user fixing their config can address everything in one pass rather than re-running after each individual failure. No-op on success.

Conditions checked here are everything that can be decided from master_config alone (no model loading, no GPU work). The following additional constraint cannot be checked from config and is enforced at runtime by the MoE fusion regex:

MoE experts must use the ...experts.N.{up_proj,down_proj}.weight naming, optionally with a gate_proj sibling. Both gated SwiGLU (gate_proj + up_proj) and non-gated ReLU^2 (up_proj only) are fused; an expert naming the fusion regex doesn’t recognize falls through and vLLM’s w13_weight / w2_weight consumers will then reject it.

Raises:: ValueError – if any precondition is violated.

nemo_rl.weight_sync.nccl_reshard_utils.build_nccl_reshard_refit_info( state_dict_metadata: dict[str, dict[str, Any]], train_parallelism: dict[str, int], gen_parallelism: dict[str, int], train_world_size: int, gen_world_size: int, layer_to_pp_stage: Optional[dict[str, int]] = None, ) → dict[str, Any]#

Build per-layer parameter info for nccl_reshard-based refit.

Parameters:

state_dict_metadata – {hf_param_name: {"shape": list, "dtype": str}} The input state_dict_metadata has a global view of the parameters
gen_parallelism (train_parallelism /) – {"tp_size", "ep_size", "pp_size"}
gen_world_size (train_world_size /) – number of GPUs per side
layer_to_pp_stage – optional mapping from layer name to PP stage index. When provided (PP>1), per-stage meshes are built so each PP stage’s train ranks + all gen ranks form an independent sub-group.

Returns:

[…], “per_layer_params”: {layer: [param_info, …]}, “pp_size”: int}``

Return type:

``{“layer_names”

nemo_rl.weight_sync.nccl_reshard_utils#