feat(megatron): add nccl_reshard M-to-N refit for Megatron generation

docs/design-docs/nccl-reshard-refit.md

@@ -28,26 +28,76 @@ single `ValueError` listing every violation. The current requirements are:
   path uses IPC and is unaffected by this feature.
 * **Megatron training backend** — `policy.megatron_cfg.enabled=true` (the DTensor
   training backend is not supported yet.).
-* **vLLM generation backend** — `policy.generation.backend=vllm` (SGLang and TRTLLM
-  backend is not supported yet.).
-* Megatron `expert_tensor_parallel_size` (i.e., ETP) must be 1; custom PP layouts
-  (`pipeline_model_parallel_layout`, virtual PP > 1, embedding/loss pipeline-split
-  accounting) are not supported yet.
-* **Precision** must match end to end for BF16 train ↔ BF16 gen and blockwise-FP8
-  train (`fp8_param=true` + blockwise recipe) ↔ FP8 gen
-  (`vllm_cfg.precision=fp8`). BF16 train → MXFP8 gen is also supported with
-  `vllm_cfg.precision=fp8` and `vllm_cfg.is_mx=true`; the generation ranks
-  quantize each received BF16 shard before installing it. Blockwise-FP8 train →
+* **vLLM or Megatron generation backend** — `policy.generation.backend` must be
+  `vllm` or `megatron` (SGLang and TRTLLM are not supported yet).
+* Training-side Megatron supports expert tensor parallelism when the generation
+  destination is Megatron. A vLLM destination still requires
+  `expert_tensor_parallel_size: 1`. Custom PP layouts
+  (`pipeline_model_parallel_layout`, virtual PP > 1, embedding/loss
+  pipeline-split accounting) are not supported yet.
+* **Generation-side ETP with `inference_optimized` is pinned to 1.** Those MoE
+  layers do not implement expert tensor parallelism and raise whenever the
+  *resolved* ETP exceeds 1 — and an omitted ETP resolves to TP, not 1. So
+  `merged_inference_megatron_cfg` pins generation-side
+  `expert_tensor_parallel_size` to 1 for that transformer implementation, which
+  is what lets generation-side TP > 1 work; the reshard then handles the
+  train-ETP → gen-ETP=1 gather. Other transformer implementations retain their
+  configured ETP. An explicitly requested generation-side ETP > 1 with
+  `inference_optimized` is rejected by config key name rather than surfacing as
+  a raw MCore assert at model build.
+* **Precision** for vLLM supports BF16 train ↔ BF16 gen, blockwise-FP8 train
+  (`fp8_param=true` + blockwise recipe) ↔ FP8 gen, and BF16 train → MXFP8 gen
+  (`vllm_cfg.precision=fp8`, `vllm_cfg.is_mx=true`). Blockwise-FP8 train →
   MXFP8 gen is not supported.
+* Megatron generation accepts BF16 or supported Transformer Engine FP8 training
+  parameter storage, including blockwise FP8 and MXFP8 with `fp8_param=true`.
+  Quantized sources are materialized as logical BF16 for transport; the
+  destination either stores BF16 or quantizes each complete local weight into
+  MXFP8. Before every refit, an explicit parameter sync materializes optimizer
+  updates that would otherwise wait for the next overlapped all-gather. When
+  MXFP8 parameter all-gather reuses the gradient buffer, that aliased allocation
+  stays GPU-resident across refit so
+  persistent DDP/autograd views remain valid; ordinary gradient buffers and
+  optimizer state are offloaded only when
+  `policy.generation.mcore_generation_config.offload_policy_before_refit` is true.
+* **The wire format is always BF16, even for MXFP8 train → MXFP8 gen.** This is
+  forced by the upstream API, not a shortcut, and is worth stating because it
+  means an MXFP8 trainer does *not* get a smaller refit (expect ~2x the
+  theoretical MXFP8 wire size, plus a dequantize on the source and a re-quantize
+  on the destination). Three reasons it cannot currently be otherwise:
+  * TE MXFP8 and MCore MXFP8 are not byte-compatible. MCore itself dequantizes
+    and re-quantizes when converting between them, deliberately, "to avoid any
+    numerical differences between TE and mcore MXFP8 formats"
+    (`megatron/core/inference/quantization/utils.py`).
+  * `MXFP8Tensor`'s only data constructor is `from_bf16`; `copy_` delegates to
+    `quantize_`, which calls `from_bf16`. There is no relayout entry point.
+  * MCore stores *swizzled* scales, padded to multiples of 128 rows and 4
+    columns, so a shard of the swizzled scales is not a shard of the logical
+    scales. An alignment-aware MXFP8 transport would have to unswizzle,
+    re-slice, and re-swizzle — most of the cost of a requantize anyway.
+  The benefit of this path is capability (M-to-N reshard into a Megatron
+  engine), not bandwidth.
 * BF16 FlashInfer TRTLLM MoE is supported through vLLM's native
   layerwise-reload path. Its grouped expert weights must use expert-parallel
   destination sharding with linear expert placement; tensor-sharded expert
   destinations and round-robin placement are rejected. This path does not
   support an FP8 KV cache or a co-trained MTP drafter; setup rejects both
   combinations.
 * vLLM expert parallelism is supported with the NeMo RL convention
-  `expert_parallel_size == tensor_parallel_size`. 
-* Generation-side, PP > 1 is not supported. 
+  `expert_parallel_size == tensor_parallel_size`.
+* Megatron generation supports expert parallelism; generation-side expert tensor
+  parallelism is available only when `transformer_impl` is not `inference_optimized`.
+* Megatron generation uses the same top-level selector as other backends:
+  `refit_transport=null` selects NeMo-RL's packed collective,
+  `refit_transport=mcore` selects Megatron Core's native refit, and
+  `refit_transport=nccl_reshard` selects M-to-N. `refit_backend` is consulted
+  only for `refit_transport=mcore`. Colocated Megatron generation requires
+  `refit_transport=mcore`, because its refit is carried by the in-place
+  wake-reshard; the other transports are rejected rather than silently ignored.
+* **Generation-side PP > 1 is not supported by this refit transport yet.**
+  Megatron-Core and vLLM can run generation with PP, and training-side Megatron
+  PP is supported here; the missing piece is generation-stage-aware destination
+  routing in `nccl_reshard`.
 * **No ModelOpt real quantization** — `policy.generation.real_quant=false`. Real-quant
   rollouts refit through vLLM's layerwise-reload weight loaders, which the bulk
   `xferdtensor` writes bypass.
@@ -90,16 +140,28 @@ nccl-reshard-refit implementation:
   shared `model_update_group` and are loaded on the generation side through the
   backend's regular `load_weights` machinery.
 
-The feature is integrated into the `nemo_rl/weight_sync/` framework:
-`create_weight_synchronizer(..., nccl_reshard_refit=True)` returns a
-`NcclReshardWeightSynchronizer` whose `init_communicator()` performs the one-time setup
-and whose `sync_weights()` runs one refit.
+The feature is integrated into the `nemo_rl/weight_sync/` framework. For vLLM,
+`create_weight_synchronizer(...)` returns an `NcclReshardWeightSynchronizer` directly.
+For Megatron generation, the existing `MegatronWeightSynchronizer` retains ownership of
+the inference-engine lifecycle and delegates only the transfer to an
+`NcclReshardWeightSynchronizer`.
 
 ### Execution Flow: Setup Time
 
 `NcclReshardWeightSynchronizer.init_communicator()` runs three steps once, before
 training starts:
 
+The generation backend declares whether it needs Bridge's physical export or logical
+weights. The synchronizer passes that payload requirement to the source worker; the
+source worker does not inspect or branch on the destination backend's name. Requesting
+logical weights is a Megatron-inference-specific exception: vLLM keeps the universal
+Bridge-export representation, while Megatron inference requests logical weights because
+its destination storage is built by MCore rather than Bridge. Megatron workers are assigned
+an explicit source or destination refit role and expose the same
+`prepare_refit_info`, `build_hf_to_local_param_map`,
+`prepare_nccl_reshard_refit_info`, and `nccl_reshard_refit` entry points in either
+role.
+
 1. **`init_collective()`** — creates the `model_update_group`, a NCCL group spanning all
    training and generation ranks. The bulk path does not use it; it carries the misc
    packed-broadcast, including FP8 KV-cache scales for backend combinations that support
@@ -216,10 +278,16 @@ generation side maps those HF names onto whatever its own storage layout is.
   supported generation parallelism; the comm bootstrap methods; the
   `nccl_reshard_refit()` receive loop; the misc consumer feeding `load_weights`; and
   backend-specific finalization after all weights arrive.
-
-**To extend to a new backend**, the only piece with genuinely new logic is
-`build_hf_to_local_param_map`. Everything else is boilerplate that follows a fixed
-contract and can be copied from the existing backend almost verbatim.
+* **Megatron generation side** (`megatron_worker.py`): mapping the same canonical
+  HF FFN shards to local fused dense/expert views. BF16 destinations receive in
+  place; MXFP8 destinations use short-lived BF16 staging buffers and quantize into
+  their persistent MCore storage. Misc weights continue through Megatron Bridge's
+  packed-broadcast import path.
+
+**To extend to a new backend**, provide a destination map from canonical HF weights
+to that backend's local storage. Both backends implement this as
+`build_hf_to_local_param_map`; Megatron derives its targets from Bridge conversion tasks.
+Everything else follows the fixed transport contract.
 
 **The one backend-specific implementation — `build_hf_to_local_param_map`:** resolve
 each bulk HF name to your local storage as a `LocalParamSpec` — `base` for tensors

docs/guides/refit.md

@@ -1,23 +1,27 @@
 # Weight Refit: Choosing a Transport
 
 Weight refit copies updated policy weights into the rollout model. Choose the
-topology first, then select one non-colocated transport with
+topology first, then select one transport with
 `policy.generation.refit_transport`.
 
 ## Pick a Transport
 
 | Topology | `refit_transport` | Transport | Use when |
 |---|---|---|---|
 | `colocated.enabled: true` | `null` | CUDA IPC | Policy and rollout workers share GPUs; vLLM uses ZMQ plus CUDA IPC handles and SGLang uses Ray CUDA IPC. |
-| `colocated.enabled: false` | `null` | NCCL broadcast | vLLM uses the default full-weight collective; SGLang uses its own weight-update group. |
+| `colocated.enabled: true` | `mcore` | MCore wake-reshard | Megatron policy and generation share GPUs. |
+| `colocated.enabled: false` | `null` | NCCL broadcast | vLLM and Megatron use the default full-weight collective; SGLang uses its own weight-update group. |
+| `colocated.enabled: false` | `mcore` | MCore native refit | Megatron policy sends directly to Megatron generation through an MCore copy service. |
 | `colocated.enabled: false` | `nccl_reshard` | NCCL reshard | Provides the best performance for large models (>100s B models). |
 | `colocated.enabled: false` | `vllm_zmq_sparse` | Sparse delta over ZeroMQ | The link is bandwidth-limited and workers can reach a relay over TCP. |
 | `colocated.enabled: false` | `vllm_s3_sparse` | Sparse delta through S3 | Workers communicate through shared object storage. |
 | `colocated.enabled: false` | `nixl` | NIXL checkpoint engine | The cluster has a fast UCX/RDMA fabric for full-weight refit. |
 
-`null` is the default. The sparse transports read only `refit_cfg.sparse`; NIXL
-reads only `refit_cfg.nixl`. Because one selector chooses the transport, sparse
-delta and NIXL cannot both be active.
+`null` is the default. For non-colocated Megatron generation it selects the
+packed collective; Megatron's native refit must be selected explicitly with
+`mcore`. The sparse transports read only `refit_cfg.sparse`; NIXL reads only
+`refit_cfg.nixl`. Because one selector chooses the transport, sparse delta and
+NIXL cannot both be active.
 
 ## vLLM Reload API
 
@@ -49,8 +53,9 @@ for those cases.
 | Colocated CUDA IPC | vLLM or SGLang | DTensor or Megatron | Uses the generation backend's standard loader. |
 | NCCL | vLLM or Megatron | DTensor or Megatron | Uses the standard full-weight loader. |
 | NCCL + reload API | vLLM | DTensor or Megatron | Default non-colocated NCCL only. Rejects colocated refit, `nccl_reshard`, sparse delta, NIXL/checkpoint-engine transports, ModelOpt quantization (`real_quant` or `quant_cfg`), Eagle/MTP trainer-refit draft weights, and grouped-MoE MXFP8 slabs. |
+| MCore native refit | Megatron | Megatron | Supports MCore's configured copy-service backend. |
 | SGLang NCCL weight-update group | SGLang | Megatron | Trainer rank 0 broadcasts finalized HF tensors to the engine leaders. |
-| NCCL reshard | vLLM | Megatron | Requires matching BF16 or blockwise FP8 precision; Megatron ETP must be 1. Currently supporting Megatron+vLLM backends. |
+| NCCL reshard | vLLM or Megatron | Megatron | Supports BF16 and the documented FP8/MXFP8 combinations; see the [NCCL Reshard Refit design](../design-docs/nccl-reshard-refit.md). |
 | Sparse delta | vLLM | Megatron | BF16/FP16, unquantized rollout only. |
 | NIXL, full weights | vLLM | DTensor or Megatron | Supports the standard full-weight FP8 loader. DTensor FP8 KV-cache scale transfer is not yet supported. |
 | NIXL, sharded experts | vLLM | DTensor or Megatron | Unquantized BF16/FP16 Triton MoE only; FP8/MXFP8 and dynamic expert placement are rejected. |
@@ -66,7 +71,7 @@ colocated generation.
 
 ## Minimal Configuration
 
-Colocated refit needs no transport configuration:
+Colocated vLLM and SGLang refit need no transport configuration:
 
 ```yaml
 policy:
@@ -86,6 +91,17 @@ policy:
     refit_transport: null
 ```
 
+For native MCore refit, select it explicitly:
+
+```yaml
+policy:
+  generation:
+    backend: megatron
+    refit_transport: mcore
+    mcore_generation_config:
+      refit_backend: nccl  # gloo | nccl | nccl_m2n (nvshmem is broken; see #3646)
+```
+
 For NCCL reshard with Megatron policy training and vLLM generation:
 
 ```yaml

examples/configs/grpo_math_1B.yaml

@@ -392,8 +392,11 @@ policy:
     val_top_k: ${.top_k}
     stop_token_ids: null
     stop_strings: null
-    # null = topology default (IPC colocated, NCCL non-colocated).
+    # null = topology default (IPC colocated, NCCL non-colocated). For
+    # non-colocated Megatron generation it selects the packed collective; use
+    # mcore for native refit.
     # Non-colocated vLLM also supports vllm_s3_sparse, vllm_zmq_sparse, and nixl.
+    # Non-colocated vLLM or Megatron generation supports nccl_reshard (M-to-N).
     refit_transport: null
     # vLLM backend only - sglang/megatron/trtllm/dynamo require this to remain null.
     # Selects the NeMo-RL Ray actor class, NOT vLLM's own `worker_extension_cls`.
@@ -434,8 +437,9 @@ policy:
       async_sched_mode: "async"  # Dynamic-batching scheduler mode. Options: "legacy", "async". "async" hides scheduler CPU time behind GPU compute.
       materialize_only_last_token_logits: true
       num_speculative_tokens: 0
+      refit_backend: null  # Native MCore copy-service backend, used only when refit_transport="mcore". Options: "gloo", "nccl", or "nccl_m2n" ("nvshmem" is currently broken, see https://github.com/NVIDIA-NeMo/RL/issues/3646).
+      refit_execution_batch_bytes: null  # Native MCore refit staging limit. null matches NeMo-RL's collective default (2% of total GPU memory, capped at 5 GiB); set a positive byte count to override.
       logprobs_mode: processed_logprobs  # Return log-probs after sampling processors. Use raw_logprobs for parity with policy recomputation.
-      refit_backend: "nccl"  # Copy-service backend for non-colocated megatron weight refit. Options: "gloo" or "nccl"
       offload_policy_before_refit: false  # Keep policy state on GPU by default. Enable when a non-colocated Megatron refit needs additional GPU headroom for transfer staging.
       parsers: []  # OpenAI tool-call / response parser names exposed by the HTTP server (only used if expose_http_server=true).
       expose_http_server: false  # Start an OpenAI-compatible HTTP server alongside the persistent engine. Required for NeMo-Gym.

examples/configs/grpo_math_1B_megatron_generation_colocated_single_controller.yaml

@@ -2,6 +2,7 @@ defaults: grpo_math_1B_megatron_single_controller.yaml
 policy:
   generation:
     backend: megatron
+    refit_transport: mcore
     mcore_generation_config:
       cuda_graph_sizing_distribution: hybrid
       cuda_graph_max_tokens: 512

examples/configs/ppo_math_1B.yaml

@@ -286,6 +286,8 @@ policy:
       prefix_caching_routing_alpha: 1.0
       unified_memory_level: 0
       max_tokens: 16384
+      refit_backend: null  # Native MCore copy-service backend, used only when refit_transport="mcore". Options: "gloo", "nccl", or "nccl_m2n" ("nvshmem" is currently broken, see https://github.com/NVIDIA-NeMo/RL/issues/3646).
+      refit_execution_batch_bytes: null  # Native MCore refit staging limit. null matches NeMo-RL's collective default (2% of total GPU memory, capped at 5 GiB); set a positive byte count to override.
       logprobs_mode: processed_logprobs
       offload_policy_before_refit: false
       http_server_num_replicas: 8

examples/configs/recipes/llm/grpo-llama3.2-1b-instruct-1n8g-megatron_generation.yaml

@@ -19,6 +19,9 @@ policy:
   make_sequence_length_divisible_by: 1
   generation:
     backend: megatron
+    refit_transport: mcore
+    mcore_generation_config:
+      refit_backend: nccl
     max_new_tokens: 512
     vllm_cfg:
       max_model_len: 512

examples/configs/recipes/llm/grpo-nanov3-30BA3B-2n8g-megatron_generation-colocated-reshard-async-gym.yaml

@@ -17,7 +17,9 @@ policy:
     context_parallel_size: 2
   generation:
     backend: megatron
+    refit_transport: mcore
     mcore_generation_config:
+      refit_backend: nccl
       tensor_model_parallel_size: 4
       expert_model_parallel_size: 4
 logger:

examples/configs/recipes/llm/grpo-nanov3-30BA3B-2n8g-megatron_generation-noncolocated-async-gym.yaml

@@ -17,7 +17,9 @@ policy:
     context_parallel_size: 2
   generation:
     backend: megatron
+    refit_transport: mcore
     mcore_generation_config:
+      refit_backend: nccl
       expert_model_parallel_size: 4
     colocated:
       enabled: false

examples/configs/recipes/llm/grpo-nanov3-30BA3B-4n4g-megatron_generation-noncolocated-mxfp8-rollouts-flashinfer.yaml

@@ -0,0 +1,5 @@
+defaults: ./grpo-nanov3-30BA3B-4n4g-megatron_generation-noncolocated-mxfp8-rollouts-packed-refit.yaml
+policy:
+  generation:
+    mcore_generation_config:
+      inference_grouped_gemm_backend: flashinfer

examples/configs/recipes/llm/grpo-nanov3-30BA3B-4n4g-megatron_generation-noncolocated-mxfp8-rollouts-packed-refit.yaml

@@ -0,0 +1,6 @@
+defaults: ./grpo-nanov3-30BA3B-4n4g-megatron_generation-noncolocated-mxfp8-rollouts.yaml
+policy:
+  generation:
+    refit_transport: null
+    mcore_generation_config:
+      refit_backend: null

examples/configs/recipes/llm/grpo-nanov3-30BA3B-4n4g-megatron_generation-noncolocated-mxfp8-rollouts.yaml

@@ -25,7 +25,9 @@ policy:
     enabled: false
   generation:
     backend: megatron
+    refit_transport: mcore
     mcore_generation_config:
+      refit_backend: nccl
       transformer_impl: inference_optimized
       inference_grouped_gemm_backend: torch
       inference_moe_token_dispatcher_type: nvls

examples/configs/recipes/llm/grpo-qwen2.5-math-1.5b-instruct-1n4g-megatron_generation-noncolocated-single-controller-async-lag4.yaml

@@ -9,8 +9,13 @@ async_rl:
   sampler:
     max_lookahead_versions: 4
   max_buffered_rollouts: 160
+  generation_fleet_health:
+    # Native MCore refit has no worker-side abort hook for a deadline.
+    refit_timeout_s: null
 policy:
   generation:
     backend: megatron
+    refit_transport: mcore
     mcore_generation_config:
+      refit_backend: nccl
       kv_cache_management_mode: recompute

examples/configs/recipes/llm/grpo-qwen2.5-math-1.5b-instruct-1n8g-megatron_generation-noncolocated-single-controller-async-lag4.yaml

@@ -9,8 +9,13 @@ async_rl:
   sampler:
     max_lookahead_versions: 4
   max_buffered_rollouts: 160
+  generation_fleet_health:
+    # Native MCore refit has no worker-side abort hook for a deadline.
+    refit_timeout_s: null
 policy:
   generation:
     backend: megatron
+    refit_transport: mcore
     mcore_generation_config:
+      refit_backend: nccl
       kv_cache_management_mode: recompute

examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-circle-count-1n4g-megatron_generation.v1.yaml

@@ -10,6 +10,10 @@ grpo:
     enabled: true
     max_trajectory_age_steps: 2
     in_flight_weight_updates: true
+async_rl:
+  generation_fleet_health:
+    # Native MCore refit has no worker-side abort hook for a deadline.
+    refit_timeout_s: null
 loss_fn:
   reference_policy_kl_penalty: 0.0
   use_importance_sampling_correction: true
@@ -34,8 +38,10 @@ policy:
     activation_checkpointing: true
   generation:
     backend: megatron
+    refit_transport: mcore
     bad_words: null
     mcore_generation_config:
+      refit_backend: nccl
       expose_http_server: true
       buffer_size_gb: 8
       num_cuda_graphs: -1