Signed-off-by: Youngeun Kwon <youngeunk@nvidia.com>

code refactoring

Signed-off-by: Youngeun Kwon <youngeunk@nvidia.com>

code refactoring

Signed-off-by: Youngeun Kwon <youngeunk@nvidia.com>

Signed-off-by: Youngeun Kwon <youngeunk@nvidia.com>

Signed-off-by: Youngeun Kwon <youngeunk@nvidia.com>

- get_placements: skip TP shard for expert params when EP>1 (current
  scope: expert_tensor_parallel_size=1 only); use len(dim_map) for
  mesh dim count.
- Remove vestigial fields _moe_fused, _moe_num_experts_local,
  _moe_gate_entries, _moe_up_entries, _moe_down_entries; replace with a
  single _moe_kind: "w13" | "w2" discriminator.
- Drop metadata_ep_gathered flag and ep_size parameter; the EP-gathered
  invariant on input metadata is now asserted in the docstring (Megatron
  meets it via export_hf_weights' EP all-gather; DTensor doesn't use EP).

Net diff: -64 / +34 lines.

Signed-off-by: Youngeun Kwon <youngeunk@nvidia.com>

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Signed-off-by: Youngeun Kwon <youngeunk@nvidia.com>

These per-user SLURM submission scripts are local conveniences and shouldn't
live in the repo.  Keep them on disk (still gitignore-able) but stop tracking.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Signed-off-by: Youngeun Kwon <youngeunk@nvidia.com>

`_param_map` (HF-name → Megatron-param view), `_expert_groups` (regex
grouping of expert params by layer/proj), and `_hf_to_vllm` (HF→vLLM
merge-rule mapping) all depend only on model topology, not weight
values.  Build them once in the respective prepare_nccl_reshard_refit_info
calls and reuse for every refit step.  Saves one full Megatron-param
iterator pass + one regex scan + one vLLM named_parameters walk per
refit step.

Also: dedup the per-expert regex (reuse _INDIVIDUAL_EXPERT_RE from
nccl_reshard_utils), drop leftover bringup print()s in
init_pp_comm_groups, and remove a WHAT-narrating comment.

Verified on the full L1 matrix (11 small + dsv3 + qwen3-235b); KL
values unchanged within run-to-run variance.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Signed-off-by: Youngeun Kwon <youngeunk@nvidia.com>

Add disaggregated weight-refit support for NemotronH models (e.g.
Nemotron-3-Nano-30B-A3B): hybrid Mamba2 / attention / MoE-FFN.

Key fix: non-gated (ReLU^2) MoE experts silently dropped their
up-projection. NemotronH MoE has up_proj + down_proj (no gate_proj), and
vLLM builds it as SharedFusedMoE(is_act_and_mul=False) with
w13_weight = [E, intermediate, hidden] (up only). fuse_expert_params_in_metadata
did `if not gate_entries: continue`, so it never created w13 and dropped
up_proj -> vLLM's w13_weight stayed stale -> every MoE layer half-corrupt
-> 100% importance-sampling masking, reward ~0. Now build w13 from up_proj
alone when there is no gate, on both the metadata and value-fusion sides.
The gated SwiGLU path is unchanged (if gate_entries / elif up_entries).

Also:
- is_misc_param: route NemotronH Mamba mixer params (in_proj, conv1d,
  A_log, D, dt_bias, norm) and the MoE router (mixer.gate.weight) to the
  misc/load_weights path; out_proj stays on the bulk xferdtensor path.
- placement: out_proj -> row-parallel; NemotronH token embedding ->
  vocab-parallel.
- vllm_backend: translate backbone.* -> model.* and embeddings ->
  embed_tokens at vLLM-param lookup so bulk mapping resolves.

Validated at baseline parity (reward + KL match the broadcast refit) on
NemotronH (non-gated) and Qwen3-30B-A3B (gated SwiGLU regression).

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Signed-off-by: Youngeun Kwon <youngeunk@nvidia.com>

Three related refactors to the nccl_reshard refit path:

1. Tag MoE expert fusion by role + explicit gated flag.
   Replace the overloaded fused_expert_param_type "w13"/"w2" with the
   semantic role "up_proj"/"down_proj" plus a gated:bool (single source of
   truth for the up-projection layout; the value-side branch asserts it
   against gate_proj presence).  Fused param names stay w13_weight/w2_weight
   (vLLM's FusedMoE convention).

2. Move DTensorRef to xferdtensor.py; rename full_tensor() -> local_tensor().
   DTensorRef is the wrapper xferdtensor reads as its src/dst, so it belongs
   next to the kernel.  full_tensor() only ever returned the local shard and
   had no callers (the kernel reads the ._local_tensor attribute).  MeshInfo
   stays in nccl_reshard_utils.

3. Drop the DTensor train backend (Megatron + vLLM only for this initial
   version).  check_nccl_reshard_refit_support states the Megatron-only scope
   intentionally; the grpo.py refit setup collapses to the Megatron branch;
   the unreachable DTensor-worker nccl-refit methods are removed.

Validated at parity across the L0 matrix (4B + 30B parallelism shapes and
NemotronH) on both gated (Qwen3-30B-A3B) and non-gated (NemotronH) MoE.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Signed-off-by: Youngeun Kwon <youngeunk@nvidia.com>

Two NemotronH-specific gaps blocked the disaggregated refit of
NVIDIA-Nemotron-3-Super-120B-A12B (Latent-MoE, GQA num_kv=2) at
TP8 EP8 PP4; both are now fixed and validated end-to-end (refit at
reward ~0.50-0.52 / gen-KL ~0.004, parity with the broadcast baseline).

1. GQA QKV at tp_size > num_query_groups.  The bulk (xferdtensor) path
   splits the fused QKV on each TP-local shard, which needs every rank
   to own whole K/V heads.  When TP exceeds the KV-head count Megatron
   channel-splits the KV heads, so no rank holds a whole one and the
   local split is impossible.  is_misc_param now takes a qkv_to_misc
   flag (set when num_query_groups < train tp): QKV then takes the misc
   path, which gathers full heads via export_hf_weights like the
   broadcast refit.  _iter_local_hf_params skips QKV in that case so the
   bulk path never attempts the impossible split.  No-op for every other
   model (num_query_groups >= tp_size keeps QKV on the bulk fast path).

2. PP-stage mapping was Llama/Qwen-only.  _LAYER_RE / _MODEL_PREFIX_RE
   and _build_layer_to_pp_stage only knew model.layers.N naming, so
   NemotronH's backbone.* params all defaulted to stage 0 and PP rank 0
   tried to transfer the whole model.  Both are now model|backbone
   naming-agnostic.

Keeps the super-v3 BF16 L1 test (config + baseline/refit scripts).

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Signed-off-by: Youngeun Kwon <youngeunk@nvidia.com>

Decouple build_nccl_reshard_refit_info from vLLM-specific layout rules so
another generation backend could supply its own, and simplify the gen-side
merged-param handling.

- MoE-expert fusion is no longer hardcoded in the builder:
  fuse_expert_params_in_metadata -> vllm_fuse_expert_params_in_metadata,
  passed in as fuse_expert_param_in_metadata_fn (None = no fusion).
  VllmGeneration.expert_fusion_fn() supplies it via the driver.

- Route QKV to the misc/load_weights path whenever KV heads can't be cleanly
  1/tp-sharded on EITHER side:
  qkv_to_misc = num_query_groups < max(train_tp, gen_tp). This subsumes the
  gen-side KV-head-replication case, so the vLLM-specific KV-head slice
  mapping (_compute_tp_local_slice) is removed from the bulk path.

- get_dst_dtensor's merged-param path collapses to a single placement-honoring
  block: receive this component's slice of the merged vLLM param (Shard dst ->
  the 1/gen_tp shard; all-Replicate dst / disable_tp -> the full tensor, which
  is the slice) and copy it in. Golden and the real reshard op now agree on
  dst placements with no path-specific branch.

Validated on the BF16 L0+L1 golden matrix, 12/12: 4b dense (TP/PP/DP), 30b MoE
(EP/TP/PP incl. qkv->misc at gen TP8), 235b, and dsv3 (disable_tp).

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Signed-off-by: Youngeun Kwon <youngeunk@nvidia.com>

xferdtensor() now dispatches to the real point-to-point reshard
(nccl.xfer.api.XdtensorRedistribute) when it is available and not overridden,
otherwise to the existing broadcast-based golden reference.

- Guarded import: golden-only containers (no nccl.xfer module) auto-fall-back
  to golden; NRL_XFERDTENSOR_GOLDEN=1 forces golden on any container.
- Each rank holds only one side (train ranks own the src shard, gen ranks the
  dst shard); the op takes a one-sided None and derives the absent side from
  the present side's global shape + mesh/placements. PyTorch Shard/Replicate
  placements pass through; meshes go in as nccl.xfer Mesh via _to_xfer_mesh.

Safe to merge ahead of a real-op-capable container thanks to the guarded
import. The xfer-capable nccl4py + matching custom NCCL are supplied by the
container/deps, not by this change. The real op's ~100x per-call latency is a
known nccl_xfer-library issue (handed to the NCCL team); until it's fixed the
golden path stays the default for the async refit loop.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Signed-off-by: Youngeun Kwon <youngeunk@nvidia.com>

Validated green across the full golden L0/L1 BF16+FP8 matrix plus Nemotron
Nano-v3 (L0) and Super-v3 (L1); no behavior change on tested configs.

- vllm_backend: size merged-param slices from refit_info["gen_tp_size"]
  instead of torch.distributed.get_world_size(), which equals gen TP only
  while gen PP/EP are pinned to 1 -- robust if gen PP is ever added.
- nccl_reshard_utils: reject unsupported gen precision values up front in
  check_nccl_reshard_refit_support (allow only bf16/bfloat16/auto/fp8/unset),
  so a typo or mismatched precision fails the gate instead of deadlocking the
  bulk collective.
- contiguity guards: .view -> .reshape in the MoE w13 gate/up fusion and
  .contiguous().clone() in the golden all-Replicate fast path, so a
  non-contiguous param view can't crash the reshard.
- rename ne_mesh/ex_mesh -> non_expert_mesh/expert_mesh (and dim_maps),
  unifying three naming styles into one.
- docstrings + worked examples for _compute_shard_slices, _fuse_one_moe_param,
  _build_expert_groups, and the get_dst_dtensor unmapped-discard path
  (clarifying it is collective-symmetry participation, not stale-weight
  corruption).

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Signed-off-by: Youngeun Kwon <youngeunk@nvidia.com>

The low-level op is finalized as "NCCL Xfer", so the refit feature is renamed
from nccl_reshard_* to nccl_xfer_* throughout:
- config knob policy.nccl_reshard_refit -> policy.nccl_xfer_refit
- prepare_/build_/check_nccl_reshard_refit_* -> nccl_xfer_* (+ _async variants)
- the nccl_reshard_refit method -> nccl_xfer_refit; nccl_reshard_refit_info /
  _enabled attrs -> nccl_xfer_*
- module nemo_rl/distributed/nccl_reshard_utils.py -> nccl_xfer_utils.py

Also rename init_pp_comm_groups -> init_per_pp_refit_comm_group (+ _async): it
sets up the per-PP-stage refit comm groups, not the PP group itself.

Dead-code cleanup folded in:
- VOCAB suffix matching uses endswith (was substring `in`)
- drop the unreachable scale_inv MERGE_RULES (scale_inv takes the misc path)
- drop the dropped-DTensor-backend to_local() else-branch and the unused
  DTensorRef.local_tensor() accessor

Behavior-preserving; validated on the golden matrix (BF16 L0 10/10, FP8 L0 9/9,
235b L1 full) -- config-knob read, module import, and method dispatch all
confirmed under the new names.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Signed-off-by: Youngeun Kwon <youngeunk@nvidia.com>

…sion gen-side)

Mirror the QKV pattern for MoE experts: the train worker and the shared
nccl_xfer_refit_info now use backend-agnostic grouped HF names
(gate_proj/up_proj/down_proj, each [E, ...]); the vLLM w13/w2 fusion moves
entirely to the gen side.  A new gen backend now needs only a gen-side mapping,
never a train-worker change.

nccl_xfer_utils.py:
- vllm_fuse_expert_params_in_metadata -> agnostic group_expert_params_in_metadata:
  stack per-expert HF params into one [E, *per_expert_shape] entry per projection,
  tagged grouped_expert_proj.  No w13/w2, no gated, no gen_tp.  The builder calls
  it unconditionally (grouping is universal); dropped the
  fuse_expert_param_in_metadata_fn injection param.
- Removed w13_weight/w2_weight from COLUMN/ROW_PARALLEL_SUFFIXES; grouped
  gate/up/down resolve via _get_expert_tp_shard_dim -> Shard(1)/Shard(2).

megatron_policy_worker.py:
- _fuse_one_moe_param -> _group_experts: plain torch.stack of the per-expert
  tensors; deleted the gate/up interleave + gen_tp_size + gated logic.
- _get_src_local_tensor dispatches on grouped_expert_proj; dropped gen_tp_size.
- prepare_nccl_xfer_refit_info: dropped the injection param.

vllm_backend.py:
- Generalized merged_param_slice to a multi-dim index tuple.
- New flag-dispatched grouped-expert branch in _build_hf_to_vllm_mapping: gate/up
  -> w13_weight dim-1 halves ([:, :P]/[:, P:2P]) when gated (has_gate pre-pass),
  up/down -> direct (whole w13 / w2).  get_dst_dtensor's existing buffer +
  post_refit_hook places the received Shard(1)/Shard(2) shard.

vllm_generation.py / grpo.py / lm_policy.py: removed expert_fusion_fn + all
fuse_expert_param_in_metadata_fn threading.

w13_weight/w2_weight no longer appear in the train worker or the shared metadata
— only in vllm_backend (where vLLM-specific layout belongs).

Validated on the golden path (NRL_XFERDTENSOR_GOLDEN=true): byte-identical by
construction (the gen-side slice reproduces the old interleave layout) + reward/KL
parity across Qwen3-30B-A3B gated (gen_tp=8), NemotronH-nano non-gated, FP8 30B
(scale_inv stays misc), and 30B PP2+TP (per-stage) — all with zero grouped-mapping
warnings and KL in the expected BF16/FP8 ranges.

Signed-off-by: Youngeun Kwon <youngeunk@nvidia.com>
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

…Store workaround)

FP8 KV cache now works on the nccl_xfer refit path: the per-layer k/v(/q)
scales ride the existing misc packed-broadcast as plain scale tensors, instead
of an out-of-band TCPStore scalar channel.  This works because is_misc_param
already routes .k_scale/.v_scale/.q_scale to the misc path and their names are
already in misc_meta (built from _iter_params_with_optional_kv_scales at prepare
time), so producer/gen agree and the scales broadcast like any other misc param.

Train (megatron_policy_worker.py):
- nccl_xfer_refit: drop the NotImplementedError on kv_scales; forward kv_scales
  to _broadcast_misc_params_packed.
- _broadcast_misc_params_packed(kv_scales): forward to
  _iter_params_with_optional_kv_scales(kv_scales=...) so the real calibrated
  scales (grpo's calibrate_qkv_fp8_scales) broadcast instead of the 1.0
  placeholder.

Gen (vllm_backend.py):
- nccl_xfer_refit: after the misc receive, call _maybe_process_fp8_kv_cache()
  (guarded on kv_cache_dtype=fp8) to finalize the per-layer scales — exactly the
  process_weights_after_loading step the prior comment said this work would need.

Test: script/new_refit/fp8/4b_tp2_tp4_fp8kv.sh (FP8 4b TP2DP4->TP4DP2 +
kv_cache_dtype=fp8).  Note FP8 KV cache requires SYNC rollouts
(async_grpo=false, async_engine=false) per grpo.py's guard; disaggregated +
sync is valid (only the inverse — async requires non-colocated — is enforced).

Validated golden (job 12718430): refit reached, recompute_kv_scales runs each
step, Gen KL 0.0086-0.0122 (sane FP8-KV range, ~3x FP8-weights-only as expected
from KV quantization), no grouped-mapping / kv-scale / process_weights errors.
Also confirms sync-mode non-colocated nccl_xfer refit works.

Signed-off-by: Youngeun Kwon <youngeunk@nvidia.com>
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>