`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
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.
`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.

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.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#