From b285451596bd41b57488c5e10557bfed8607530e Mon Sep 17 00:00:00 2001 From: William Dykas Date: Thu, 3 Sep 2026 11:50:33 -0700 Subject: [PATCH 1/8] feat(megatron): add native M-to-N inference refit Signed-off-by: William Dykas --- docs/design-docs/nccl-reshard-refit.md | 115 ++- docs/guides/refit.md | 28 +- examples/configs/grpo_math_1B.yaml | 7 +- examples/configs/ppo_math_1B.yaml | 2 + ...-1b-instruct-1n8g-megatron_generation.yaml | 2 +- ...ron_generation-noncolocated-async-gym.yaml | 1 + ...eneration-noncolocated-mxfp8-rollouts.yaml | 1 + ...olocated-single-controller-async-lag4.yaml | 1 + ...cle-count-1n4g-megatron_generation.v1.yaml | 1 + ...a3b-clevr-8n4g-megatron_generation.v1.yaml | 1 + examples/nemo_gym/grpo_nanov3.yaml | 3 +- ...megatron_generation_single_controller.yaml | 1 + .../nemo_gym/grpo_qwen3_30ba3b_instruct.yaml | 3 +- nemo_rl/algorithms/grpo.py | 20 +- nemo_rl/models/generation/interfaces.py | 15 +- nemo_rl/models/generation/megatron/config.py | 44 +- .../megatron/megatron_generation.py | 230 ++++- .../generation/megatron/megatron_worker.py | 716 +++++++++++++- nemo_rl/models/megatron/setup.py | 10 +- nemo_rl/models/policy/interfaces.py | 15 +- nemo_rl/models/policy/lm_policy.py | 55 +- .../policy/workers/base_policy_worker.py | 14 +- .../policy/workers/checkpoint_engine.py | 2 +- .../policy/workers/megatron_policy_worker.py | 709 ++++++++++---- .../collective_weight_synchronizer.py | 4 +- nemo_rl/weight_sync/factory.py | 8 +- .../megatron_weight_synchronizer.py | 77 +- nemo_rl/weight_sync/nccl_reshard_utils.py | 194 ++-- .../nccl_reshard_weight_synchronizer.py | 55 +- .../configs/grpo_math_1B.yaml | 2 + .../L1_Functional_Tests_GB200_MXFP8.sh | 5 + .../L1_Functional_Tests_Megatron_4.sh | 1 + .../grpo_megatron_generation_colocated.sh | 1 + .../grpo_megatron_generation_colocated_gym.sh | 1 + ...o_megatron_generation_colocated_reshard.sh | 3 + .../grpo_megatron_generation_multiturn.sh | 1 + .../grpo_megatron_generation_non_colocated.sh | 1 + ...tron_generation_non_colocated_async_gym.sh | 1 + .../grpo_megatron_generation_topp_topk.sh | 1 + .../grpo_megatron_nccl_reshard_refit.sh | 126 +++ .../nemotron_omni_clevr_megatron_1n2g.sh | 1 + .../nemotron_omni_gym_video_megatron_1n2g.sh | 1 + tests/unit/algorithms/test_grpo.py | 48 +- .../generation/test_megatron_generation.py | 373 +++++++- .../models/megatron/test_group_experts.py | 43 +- .../models/megatron/test_megatron_setup.py | 58 ++ .../policy/test_lm_policy_collective.py | 81 ++ .../policy/test_megatron_split_state.py | 5 + .../models/policy/test_megatron_worker.py | 873 +++++++++++++++++- .../unit/reference_configs/grpo_math_1B.yaml | 6 +- .../ppo_math_1B_megatron.yaml | 2 + tests/unit/utils/test_checkpoint_engine.py | 13 + ...t_checkpoint_engine_weight_synchronizer.py | 2 - .../weight_sync/test_nccl_reshard_utils.py | 298 +++++- .../unit/weight_sync/test_reshard_rebuild.py | 21 +- .../weight_sync/test_weight_synchronizer.py | 207 ++++- 56 files changed, 4072 insertions(+), 437 deletions(-) create mode 100755 tests/functional/grpo_megatron_nccl_reshard_refit.sh diff --git a/docs/design-docs/nccl-reshard-refit.md b/docs/design-docs/nccl-reshard-refit.md index b2085335729..898d9810ecc 100644 --- a/docs/design-docs/nccl-reshard-refit.md +++ b/docs/design-docs/nccl-reshard-refit.md @@ -28,20 +28,63 @@ 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. Custom PP layouts + (`pipeline_model_parallel_layout`, virtual PP > 1, + embedding/loss pipeline-split accounting) are not supported yet. +* **Generation-side ETP is pinned to 1.** MCore's `inference_optimized` 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, which is what lets generation-side TP > 1 + work; the reshard then handles the train-ETP → gen-ETP=1 gather. An explicitly + requested generation-side ETP > 1 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. 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 still offloaded. +* **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. * 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 and expert tensor + parallelism, including using both together. +* 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. @@ -80,16 +123,29 @@ nccl-reshard-refit implementation: `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. The older generation-specific method names remain compatibility aliases for +external callers. + 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 (and FP8 KV-cache scales), identical to the conventional collective @@ -202,17 +258,22 @@ generation side maps those HF names onto whatever its own storage layout is. deliberately **shape-driven** so the same code handles generation TP and generation EP; the comm bootstrap methods; the `nccl_reshard_refit()` receive loop; the misc consumer feeding `load_weights`. - -**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. - -**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 -sent/received as-is, and `pre`/`post` hooks wherever your layout requires staging -(fused/merged tensors, layout conversions, grouped-expert stacking). This is the *only* -place your backend's parameter layout is encoded; all cross-mesh byte movement is -already handled by the shared metadata and `xferdtensor`. +* **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. vLLM implements this as +`build_hf_to_local_param_map`; Megatron derives it from Bridge conversion tasks. +Everything else follows the fixed transport contract. + +**The backend-specific destination map:** resolve each bulk HF name to local storage +as a `LocalParamSpec` — `base` for tensors sent/received as-is, and `pre`/`post` hooks +wherever the layout requires staging (fused/merged tensors, layout conversions, +grouped-expert stacking). This is the only place the backend's parameter layout is +encoded; shared metadata and `xferdtensor` handle the cross-mesh byte movement. (A new *training* backend additionally has to produce the HF-named metadata — names, global shapes, dtypes, and the parallelism description the agnostic builder consumes — diff --git a/docs/guides/refit.md b/docs/guides/refit.md index bf9d7d39f2e..ec9a3315f71 100644 --- a/docs/guides/refit.md +++ b/docs/guides/refit.md @@ -1,7 +1,7 @@ # 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 @@ -9,15 +9,19 @@ topology first, then select one non-colocated transport with | 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,15 @@ policy: refit_transport: null ``` +For native MCore refit, select it explicitly: + +```yaml +policy: + generation: + backend: megatron + refit_transport: mcore +``` + For NCCL reshard with Megatron policy training and vLLM generation: ```yaml diff --git a/examples/configs/grpo_math_1B.yaml b/examples/configs/grpo_math_1B.yaml index 99d9ee369cc..e1360af44ca 100644 --- a/examples/configs/grpo_math_1B.yaml +++ b/examples/configs/grpo_math_1B.yaml @@ -389,8 +389,10 @@ 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 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 # Keep real-quant export tensors on CPU by default. Set false only for # colocated CUDA-IPC refit to avoid a GPU -> CPU -> GPU round trip. @@ -420,8 +422,9 @@ policy: kv_cache_management_mode: "persist" # KV cache lifecycle across suspend/resume. Options: "persist", "offload", "recompute". materialize_only_last_token_logits: true num_speculative_tokens: 0 + refit_backend: "nccl" # 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" 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. vllm_cfg: diff --git a/examples/configs/ppo_math_1B.yaml b/examples/configs/ppo_math_1B.yaml index c125d1d7930..61766cf0c4c 100644 --- a/examples/configs/ppo_math_1B.yaml +++ b/examples/configs/ppo_math_1B.yaml @@ -278,6 +278,8 @@ policy: enable_chunked_prefill: true unified_memory_level: 0 max_tokens: 16384 + refit_backend: "nccl" # 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 vllm_cfg: async_engine: false diff --git a/examples/configs/recipes/llm/grpo-llama3.2-1b-instruct-1n8g-megatron_generation.yaml b/examples/configs/recipes/llm/grpo-llama3.2-1b-instruct-1n8g-megatron_generation.yaml index bb641388d87..d616b381092 100644 --- a/examples/configs/recipes/llm/grpo-llama3.2-1b-instruct-1n8g-megatron_generation.yaml +++ b/examples/configs/recipes/llm/grpo-llama3.2-1b-instruct-1n8g-megatron_generation.yaml @@ -19,6 +19,7 @@ policy: make_sequence_length_divisible_by: 1 generation: backend: megatron + refit_transport: mcore max_new_tokens: 512 vllm_cfg: max_model_len: 512 @@ -33,4 +34,3 @@ logger: name: grpo-llama3.2-1b-instruct-1n8g-megatron_generation cluster: gpus_per_node: 8 - diff --git a/examples/configs/recipes/llm/grpo-nanov3-30BA3B-2n8g-megatron_generation-noncolocated-async-gym.yaml b/examples/configs/recipes/llm/grpo-nanov3-30BA3B-2n8g-megatron_generation-noncolocated-async-gym.yaml index f171bcd913f..710c141a90a 100644 --- a/examples/configs/recipes/llm/grpo-nanov3-30BA3B-2n8g-megatron_generation-noncolocated-async-gym.yaml +++ b/examples/configs/recipes/llm/grpo-nanov3-30BA3B-2n8g-megatron_generation-noncolocated-async-gym.yaml @@ -17,6 +17,7 @@ policy: context_parallel_size: 2 generation: backend: megatron + refit_transport: mcore mcore_generation_config: expert_model_parallel_size: 4 colocated: diff --git a/examples/configs/recipes/llm/grpo-nanov3-30BA3B-4n4g-megatron_generation-noncolocated-mxfp8-rollouts.yaml b/examples/configs/recipes/llm/grpo-nanov3-30BA3B-4n4g-megatron_generation-noncolocated-mxfp8-rollouts.yaml index 1d78a49b1d1..b9810d1acb4 100644 --- a/examples/configs/recipes/llm/grpo-nanov3-30BA3B-4n4g-megatron_generation-noncolocated-mxfp8-rollouts.yaml +++ b/examples/configs/recipes/llm/grpo-nanov3-30BA3B-4n4g-megatron_generation-noncolocated-mxfp8-rollouts.yaml @@ -28,6 +28,7 @@ policy: enabled: false generation: backend: megatron + refit_transport: mcore mcore_generation_config: transformer_impl: inference_optimized inference_grouped_gemm_backend: "torch" diff --git a/examples/configs/recipes/llm/grpo-qwen2.5-math-1.5b-instruct-1n8g-megatron_generation-noncolocated-single-controller-async-lag4.yaml b/examples/configs/recipes/llm/grpo-qwen2.5-math-1.5b-instruct-1n8g-megatron_generation-noncolocated-single-controller-async-lag4.yaml index 6cf9f869841..5152a0a7ace 100644 --- a/examples/configs/recipes/llm/grpo-qwen2.5-math-1.5b-instruct-1n8g-megatron_generation-noncolocated-single-controller-async-lag4.yaml +++ b/examples/configs/recipes/llm/grpo-qwen2.5-math-1.5b-instruct-1n8g-megatron_generation-noncolocated-single-controller-async-lag4.yaml @@ -12,5 +12,6 @@ async_rl: policy: generation: backend: megatron + refit_transport: mcore mcore_generation_config: kv_cache_management_mode: recompute diff --git a/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-circle-count-1n4g-megatron_generation.v1.yaml b/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-circle-count-1n4g-megatron_generation.v1.yaml index aca1e6015a6..442e9778ba4 100644 --- a/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-circle-count-1n4g-megatron_generation.v1.yaml +++ b/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-circle-count-1n4g-megatron_generation.v1.yaml @@ -34,6 +34,7 @@ policy: activation_checkpointing: true generation: backend: megatron + refit_transport: mcore bad_words: null mcore_generation_config: expose_http_server: true diff --git a/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-8n4g-megatron_generation.v1.yaml b/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-8n4g-megatron_generation.v1.yaml index eb674c74ffa..2fb2744c997 100644 --- a/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-8n4g-megatron_generation.v1.yaml +++ b/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-8n4g-megatron_generation.v1.yaml @@ -37,6 +37,7 @@ policy: store_param_remainders: true generation: backend: megatron + refit_transport: mcore max_new_tokens: 2048 bad_words: null mcore_generation_config: diff --git a/examples/nemo_gym/grpo_nanov3.yaml b/examples/nemo_gym/grpo_nanov3.yaml index c41a8113b9e..02d46d37806 100644 --- a/examples/nemo_gym/grpo_nanov3.yaml +++ b/examples/nemo_gym/grpo_nanov3.yaml @@ -245,8 +245,9 @@ policy: kv_cache_management_mode: "persist" # KV cache lifecycle across suspend/resume. Options: "persist", "offload", "recompute". materialize_only_last_token_logits: true num_speculative_tokens: 0 + refit_backend: "nccl" # 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 - refit_backend: "nccl" # Copy-service backend for non-colocated megatron weight refit. Options: "gloo" or "nccl". expose_http_server: true enable_prefix_caching: true parsers: diff --git a/examples/nemo_gym/grpo_qwen3_0_6b_megatron_generation_single_controller.yaml b/examples/nemo_gym/grpo_qwen3_0_6b_megatron_generation_single_controller.yaml index dd3ac8f54f4..fef0277c994 100644 --- a/examples/nemo_gym/grpo_qwen3_0_6b_megatron_generation_single_controller.yaml +++ b/examples/nemo_gym/grpo_qwen3_0_6b_megatron_generation_single_controller.yaml @@ -41,6 +41,7 @@ policy: generation: backend: "megatron" + refit_transport: "mcore" mcore_generation_config: # NeMo-Gym drives rollouts through the engine's OpenAI server. expose_http_server: true diff --git a/examples/nemo_gym/grpo_qwen3_30ba3b_instruct.yaml b/examples/nemo_gym/grpo_qwen3_30ba3b_instruct.yaml index ec8d4382150..3924344a9e1 100644 --- a/examples/nemo_gym/grpo_qwen3_30ba3b_instruct.yaml +++ b/examples/nemo_gym/grpo_qwen3_30ba3b_instruct.yaml @@ -77,8 +77,9 @@ policy: kv_cache_management_mode: "persist" materialize_only_last_token_logits: true num_speculative_tokens: 0 + refit_backend: "nccl" # 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 - refit_backend: "nccl" parsers: [] expose_http_server: false diff --git a/nemo_rl/algorithms/grpo.py b/nemo_rl/algorithms/grpo.py index 7df8f4db53a..be508153328 100644 --- a/nemo_rl/algorithms/grpo.py +++ b/nemo_rl/algorithms/grpo.py @@ -161,6 +161,7 @@ checkpoint_engine_refit_config, ) from nemo_rl.weight_sync.factory import create_weight_synchronizer +from nemo_rl.weight_sync.nccl_reshard_utils import check_nccl_reshard_refit_support # =============================================================================== # Configuration @@ -1709,19 +1710,18 @@ def init_dynamo(): generation_config.get("refit_transport") == "nccl_reshard" ) if nccl_reshard_refit_enabled: - from nemo_rl.weight_sync.nccl_reshard_utils import ( - check_nccl_reshard_refit_support, - ) - check_nccl_reshard_refit_support(master_config) - if generation_config.get("refit_transport") is not None and backend != "vllm": + refit_transport = generation_config.get("refit_transport") + if refit_transport is not None and not ( + backend == "vllm" + or (backend == "megatron" and refit_transport in ("mcore", "nccl_reshard")) + ): raise NotImplementedError( - "Non-default refit transports are only supported for the vLLM " - f"generation backend, but policy.generation.backend={backend!r}. " - "Set policy.generation.refit_transport=null. Support for other " - "generation backends is tracked in " - "https://github.com/NVIDIA-NeMo/RL/issues/3288." + f"refit_transport={refit_transport!r} is not supported for " + f"policy.generation.backend={backend!r}. " + "Set policy.generation.refit_transport=null. Megatron generation " + "supports refit_transport='mcore' or 'nccl_reshard'." ) if backend == "megatron": diff --git a/nemo_rl/models/generation/interfaces.py b/nemo_rl/models/generation/interfaces.py index 00f95740e7d..3d3a1418c9f 100644 --- a/nemo_rl/models/generation/interfaces.py +++ b/nemo_rl/models/generation/interfaces.py @@ -14,13 +14,22 @@ from abc import ABC, abstractmethod from dataclasses import dataclass from functools import cache -from typing import TYPE_CHECKING, Any, NotRequired, Optional, TypedDict, Union +from typing import TYPE_CHECKING, Any, Literal, NotRequired, Optional, TypedDict, Union import ray import torch from nemo_rl.distributed.batched_data_dict import BatchedDataDict +# ``bridge_export`` preserves the universal refit contract: the source exports +# an HF-named, backend-independent representation and each destination converts +# it into local storage. ``logical_weights`` is an explicit Megatron-to-Megatron +# exception: Megatron generation may assume a Megatron training source and ask +# it to materialize logical weights. Megatron training must continue to use the +# universal mode for other destinations (such as vLLM), and new backends should +# not inherit this coupling implicitly. +RefitPayloadMode = Literal["bridge_export", "logical_weights"] + if TYPE_CHECKING: from nemo_rl.algorithms.single_controller_utils.config import MasterConfig @@ -521,6 +530,10 @@ def get_inference_world_size(self) -> int | None: """Return a backend-specific collective world size when required.""" return None + def get_refit_payload_mode(self) -> RefitPayloadMode: + """Return the backend's required representation for transferred weights.""" + return "bridge_export" + def prepare_nccl_reshard_refit_info(self, refit_info: dict) -> None: """Prepare per-layer param metadata for nccl_reshard-based refit.""" raise NotImplementedError diff --git a/nemo_rl/models/generation/megatron/config.py b/nemo_rl/models/generation/megatron/config.py index eb3a41f0258..81188642e74 100644 --- a/nemo_rl/models/generation/megatron/config.py +++ b/nemo_rl/models/generation/megatron/config.py @@ -16,6 +16,19 @@ from nemo_rl.models.generation.interfaces import GenerationConfig from nemo_rl.models.policy import Fp8Config, PolicyConfig +from nemo_rl.utils.packed_tensor import get_target_packed_tensor_size + + +def resolve_refit_execution_batch_bytes(configured_bytes: int | None) -> int: + """Resolve native MCore staging bytes using NeMo-RL's collective default.""" + if configured_bytes is None: + return get_target_packed_tensor_size() + if configured_bytes <= 0: + raise ValueError( + "policy.generation.mcore_generation_config." + "refit_execution_batch_bytes must be positive or null." + ) + return configured_bytes class MCoreGenerationSpecificArgs(TypedDict): @@ -55,7 +68,11 @@ class MCoreGenerationSpecificArgs(TypedDict): vision_embedding_cache_max_bytes: NotRequired[int] allow_stale_multimodal_embeddings: NotRequired[bool] - refit_backend: Literal["gloo", "nccl", "nvshmem"] + # Copy-service backend used only when refit_transport="mcore". + refit_backend: Literal["gloo", "nccl", "nccl_m2n", "nvshmem"] + # Soft per-rank staging limit for native MCore refit. None uses the same + # dynamic packed-buffer target as NeMo-RL's existing collective refit. + refit_execution_batch_bytes: int | None num_speculative_tokens: int mamba_inference_ssm_states_dtype: NotRequired[str] @@ -99,6 +116,9 @@ class MCoreGenerationSpecificArgs(TypedDict): class MCoreGenerationConfig(GenerationConfig): """Generation config for Megatron Inference.""" + # None uses NeMo-RL's packed collective, mcore uses Megatron Core's native + # refit, and nccl_reshard selects the non-colocated NCCL M-to-N transport. + refit_transport: NotRequired[Literal["mcore", "nccl_reshard"] | None] mcore_generation_config: MCoreGenerationSpecificArgs @@ -132,6 +152,28 @@ def merged_inference_megatron_cfg(policy_config: PolicyConfig) -> dict[str, Any] "with TP>1 on the generation model: set " "policy.generation.mcore_generation_config.sequence_parallel=true." ) + + # inference_optimized MoE layers do not implement expert tensor parallelism: + # MCore raises whenever the *resolved* ETP exceeds 1, and an omitted ETP + # resolves to TP. Pin the generation-side ETP to 1 so that TP>1 generation + # keeps working without the caller having to know that rule (ETP is inert + # for dense models). An explicitly requested gen-side ETP>1 is unsatisfiable, + # so name the config key rather than letting MCore assert at model build. + if merged.get("transformer_impl") == "inference_optimized": + gen_overrides = cast( + dict[str, Any], generation_config.get("mcore_generation_config") or {} + ) + explicit_etp = cast( + Optional[int], gen_overrides.get("expert_tensor_parallel_size") + ) + if explicit_etp is not None and explicit_etp > 1: + raise ValueError( + "transformer_impl=inference_optimized does not support expert " + "tensor parallelism on the generation model: set " + "policy.generation.mcore_generation_config." + f"expert_tensor_parallel_size=1 (got {explicit_etp})." + ) + merged["expert_tensor_parallel_size"] = 1 return merged diff --git a/nemo_rl/models/generation/megatron/megatron_generation.py b/nemo_rl/models/generation/megatron/megatron_generation.py index e1fb93b5e48..22c719025f0 100644 --- a/nemo_rl/models/generation/megatron/megatron_generation.py +++ b/nemo_rl/models/generation/megatron/megatron_generation.py @@ -12,6 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +import warnings from copy import deepcopy from typing import TYPE_CHECKING, Any, AsyncGenerator, Optional, cast @@ -27,6 +28,7 @@ GenerationDatumSpec, GenerationInterface, GenerationOutputSpec, + RefitPayloadMode, reject_unenforceable_refit_deadline, ) from nemo_rl.models.generation.megatron.config import ( @@ -41,6 +43,7 @@ from nemo_rl.algorithms.single_controller_utils.config import MasterConfig from nemo_rl.distributed.worker_groups import RayWorkerGroup from nemo_rl.models.policy.lm_policy import Policy + from nemo_rl.weight_sync.membership import RefitMembership class MegatronGeneration(GenerationInterface): @@ -215,7 +218,10 @@ def __init__( policy: Existing training Policy reused for colocated generation. name_prefix: Prefix for naming the worker group (non-colocated only). processor: Optional processor for VLMs (non-colocated only). - skip_weight_load: Do not load the weights from the checkpoint; refit will do it. + skip_weight_load: Do not load weights from the checkpoint; refit will do it. + Inference-engine initialization is deferred until that first refit so CUDA + graphs capture the final persistent weight buffers rather than placeholder + checkpoint tensors. reserved_http_server_port: Driver-reserved OpenAI server port for non-colocated. """ # Import here to avoid circular imports @@ -237,10 +243,54 @@ def __init__( # inference receives a copy because worker setup may modify it. self._policy_config = config self.cfg: MCoreGenerationConfig = config["generation"] + refit_transport = self.cfg.get("refit_transport") + if refit_transport not in (None, "mcore", "nccl_reshard"): + raise ValueError( + "policy.generation.refit_transport must be null, 'mcore', or " + f"'nccl_reshard' for Megatron generation, got {refit_transport!r}." + ) + if refit_transport == "mcore": + refit_backend = self.cfg["mcore_generation_config"].get( + "refit_backend", "gloo" + ) + if refit_backend not in ("gloo", "nccl", "nccl_m2n", "nvshmem"): + raise ValueError( + "policy.generation.mcore_generation_config.refit_backend " + "must be 'gloo', 'nccl', 'nccl_m2n', or 'nvshmem' when " + f"refit_transport='mcore', got {refit_backend!r}." + ) + if policy is not None and refit_backend == "nccl_m2n": + raise ValueError( + "policy.generation.mcore_generation_config.refit_backend=" + "'nccl_m2n' is only supported with non-colocated generation." + ) + elif self.cfg["mcore_generation_config"].get("refit_backend") in ( + "nccl_m2n", + "nvshmem", + ): + # Only the native MCore copy service reads refit_backend. Warn rather + # than raise: every shipped exemplar sets refit_backend unconditionally, + # so another transport legitimately inherits a non-null value. A user + # who explicitly asked for a non-default transport, though, would + # otherwise silently get neither. + # .get(): unlike the 'mcore' branch above, the other paths never + # requires this key, so reading it must not turn an omitted key into a + # KeyError. + warnings.warn( + "policy.generation.mcore_generation_config.refit_backend=" + f"{self.cfg['mcore_generation_config']['refit_backend']!r} is " + f"ignored when refit_transport={refit_transport!r}; it is only " + "read by the native MCore refit (refit_transport='mcore').", + stacklevel=2, + ) # Populated after the first prepare_for_generation (which starts the HTTP server). self.dp_openai_server_base_urls: list[Optional[str]] = [] # Installed by setup via create_weight_synchronizer. self.weight_synchronizer: Optional["WeightSynchronizer"] = None + # The nccl_reshard synchronizer records its current rank layout before + # dispatching any communicator or refit calls. None is the legacy/full-group + # path used by refit implementations that do not manage membership. + self._refit_membership: Optional["RefitMembership"] = None if policy is not None: # Reuse the existing training policy. @@ -268,6 +318,7 @@ def __init__( init_optimizer=False, init_reference_model=False, skip_weight_load=skip_weight_load, + refit_role="destination", reserved_http_server_port=reserved_http_server_port, ) @@ -277,6 +328,15 @@ def __init__( if not skip_weight_load: self.prepare_for_generation() + @property + def uses_native_refit(self) -> bool: + """Whether non-colocated refit uses Megatron Core's native mechanism.""" + return self.cfg.get("refit_transport") == "mcore" + + def get_refit_payload_mode(self) -> RefitPayloadMode: + """Use the Megatron-to-Megatron logical-weight exception for M-to-N.""" + return "logical_weights" + @property def worker_group(self) -> "RayWorkerGroup": """The underlying policy's worker group (fleet-health probes read dp_size).""" @@ -289,35 +349,170 @@ def init_collective( world_size: int, *, train_world_size: int, - refit_backend: str = "gloo", + refit_backend: Optional[str] = None, ) -> list[ray.ObjectRef]: - """Initialize the refit collective for weight synchronization. + """Join the configured refit collective after the training ranks. Args: ip: IP address for the process group rendezvous. port: Port for the process group rendezvous. world_size: Total world size (train + inference workers). train_world_size: Number of training workers (used to offset ranks). - refit_backend: Copy service backend ("gloo" or "nccl"; - "nvshmem" is currently broken and warns at setup). + refit_backend: Optional override for the native MCore copy-service + backend ("gloo", "nccl", or "nccl_m2n"; "nvshmem" is + currently broken and warns at setup, see + https://github.com/NVIDIA-NeMo/RL/issues/3646). Ignored by + the packed collective and nccl_reshard transports. Returns: List of Ray ObjectRefs for the collective init futures. """ - return self._policy.init_collective_mcore_generation( + if self.uses_native_refit: + backend = ( + refit_backend or self.cfg["mcore_generation_config"]["refit_backend"] + ) + return self._policy.init_collective_mcore_generation( + ip, + port, + world_size, + rank_offset=train_world_size, + refit_execution_batch_bytes=self.cfg["mcore_generation_config"][ + "refit_execution_batch_bytes" + ], + refit_backend=backend, + ) + return self._policy.init_collective( ip, port, world_size, + train_world_size=train_world_size, rank_offset=train_world_size, - refit_backend=refit_backend, ) def update_weights_from_collective( self, refit_timeout_s: Optional[float] = None ) -> list[ray.ObjectRef]: - """Receive updated weights from the training cluster via collective communication.""" - reject_unenforceable_refit_deadline("Megatron", refit_timeout_s) - return self._policy.swap_weights_via_reshard(is_source=False) + """Receive weights through the configured Megatron refit mechanism.""" + if self.uses_native_refit: + reject_unenforceable_refit_deadline("Megatron", refit_timeout_s) + return self._policy.swap_weights_via_reshard(is_source=False) + return self._policy.worker_group.run_all_workers_single_data( + "update_weights_from_collective", refit_timeout_s=refit_timeout_s + ) + + def init_nccl_reshard_comm_group( + self, + *, + pp_ips: list[str], + pp_ports: list[int], + pp_size: int, + train_ranks_per_stage: int, + sub_world_size: int, + ) -> list[ray.ObjectRef]: + """Join every training PP stage's NCCL M-to-N communicator.""" + return self._policy.worker_group.run_all_workers_single_data( + "init_nccl_reshard_comm_groups_generation", + pp_ips=pp_ips, + pp_ports=pp_ports, + pp_size=pp_size, + train_ranks_per_stage=train_ranks_per_stage, + sub_world_size=sub_world_size, + ) + + def set_refit_membership(self, membership: "RefitMembership") -> None: + """Record the inference ranks participating in nccl_reshard refits.""" + self._refit_membership = membership + + def _refit_ranked_workers( + self, membership: Optional["RefitMembership"] = None + ) -> list[tuple[Any, int, int]]: + """Return ``(actor, rebuilt rank, original rank)`` for live workers.""" + active = membership or self._refit_membership + if active is None: + raise RuntimeError("Refit membership has not been initialized.") + + workers = self.worker_group.workers + ranked_workers: list[tuple[Any, int, int]] = [] + for shard_idx, rank_prefix in active.shard_prefixes.items(): + worker_start = shard_idx * active.workers_per_shard + for local_rank in range(active.workers_per_shard): + worker_idx = worker_start + local_rank + if worker_idx >= len(workers): + raise RuntimeError( + f"shard {shard_idx} maps to worker {worker_idx}, but the " + f"group has {len(workers)} workers" + ) + ranked_workers.append( + (workers[worker_idx], rank_prefix + local_rank, worker_idx) + ) + return ranked_workers + + def rebuild_collective( + self, membership: "RefitMembership", ip: str, port: int + ) -> list[ray.ObjectRef]: + """Build the misc-weight communicator over the selected Megatron ranks.""" + futures = [] + for worker, rank, original_rank in self._refit_ranked_workers(membership): + futures.append( + worker.init_collective.remote( + ip=ip, + port=port, + world_size=membership.world_size, + train_world_size=membership.train_world_size, + rank_offset=membership.train_world_size + rank - original_rank, + nccl_peer=self.get_collective_sender_spec().nccl_peer, + ) + ) + return futures + + def rebuild_nccl_reshard_comm_group( + self, + membership: "RefitMembership", + *, + pp_ips: list[str], + pp_ports: list[int], + pp_size: int, + train_ranks_per_stage: int, + sub_world_size: int, + ) -> list[ray.ObjectRef]: + """Build bulk communicators over the selected Megatron ranks.""" + return [ + worker.init_nccl_reshard_comm_groups_generation.remote( + pp_ips=pp_ips, + pp_ports=pp_ports, + pp_size=pp_size, + train_ranks_per_stage=train_ranks_per_stage, + sub_world_size=sub_world_size, + rank_prefix=rank, + ) + for worker, rank, _original_rank in self._refit_ranked_workers(membership) + ] + + def prepare_nccl_reshard_refit_info(self, refit_info: dict[str, Any]) -> None: + """Build each inference worker's HF-to-Megatron M-to-N receive map.""" + if self._refit_membership is None: + futures = self._policy.worker_group.run_all_workers_single_data( + "prepare_nccl_reshard_refit_info", refit_info=refit_info + ) + else: + futures = [ + worker.prepare_nccl_reshard_refit_info.remote(refit_info=refit_info) + for worker, _rank, _original_rank in self._refit_ranked_workers() + ] + ray.get(futures) + + def nccl_reshard_refit( + self, refit_timeout_s: Optional[float] = None + ) -> list[ray.ObjectRef]: + """Receive one NCCL M-to-N refit on every Megatron inference worker.""" + if self._refit_membership is None: + return self._policy.worker_group.run_all_workers_single_data( + "nccl_reshard_refit", refit_timeout_s=refit_timeout_s + ) + return [ + worker.nccl_reshard_refit.remote(refit_timeout_s=refit_timeout_s) + for worker, _rank, _original_rank in self._refit_ranked_workers() + ] def generate( self, data: BatchedDataDict[GenerationDatumSpec], greedy: bool = False @@ -420,6 +615,10 @@ def preinit_nvshmem_collective(self) -> list[ray.ObjectRef]: Must be called simultaneously on both training and inference workers. """ + if not self.uses_native_refit: + raise RuntimeError( + "NVSHMEM pre-initialization is only valid with refit_transport='mcore'." + ) return self._policy.preinit_nvshmem() def suspend_for_refit(self) -> None: @@ -435,8 +634,15 @@ def resume_after_refit(self) -> None: ) def prepare_refit_info(self, state_dict_info: Optional[dict[str, Any]]) -> None: - """Accept the cross-backend refit-prep contract; Megatron needs none of it.""" - pass + """Prepare Bridge conversion tasks on every dedicated inference worker.""" + if not self._owns_policy or self.uses_native_refit: + return + if state_dict_info is None: + raise ValueError("Megatron collective refit requires state_dict_info.") + futures = self._policy.worker_group.run_all_workers_single_data( + "prepare_refit_info", state_dict_info=state_dict_info + ) + ray.get(futures) def start_gpu_profiling(self) -> None: """Start GPU profiling on the dedicated inference workers. diff --git a/nemo_rl/models/generation/megatron/megatron_worker.py b/nemo_rl/models/generation/megatron/megatron_worker.py index 882544e4347..4beea0b295a 100644 --- a/nemo_rl/models/generation/megatron/megatron_worker.py +++ b/nemo_rl/models/generation/megatron/megatron_worker.py @@ -19,6 +19,8 @@ import threading import time import warnings +from collections import Counter, OrderedDict +from dataclasses import dataclass, replace from typing import Any, AsyncGenerator, Optional import requests @@ -31,10 +33,21 @@ PrefixCachingCoordinatorPolicy, ) from megatron.core.inference.engines.dynamic_engine import EngineState +from megatron.core.inference.quantization.mxfp8_tensor import MXFP8Tensor +from megatron.core.inference.quantization.utils import ( + quantize_params_to_mxfp8, + resolve_mxfp8_backend, +) from megatron.core.inference.sampling_params import SamplingParams from megatron.core.inference.utils import set_decode_expert_padding +from megatron.core.models.gpt.gpt_layer_specs import ( + get_gpt_layer_with_inference_spec, +) from megatron.core.resharding.copy_services.gloo_copy_service import GlooCopyService from megatron.core.resharding.copy_services.nccl_copy_service import NCCLCopyService +from megatron.core.resharding.copy_services.nccl_m2n_copy_service import ( + NCCLM2NCopyService, +) from megatron.core.resharding.refit import ( prepare_swap_model_weights, swap_model_weights, @@ -54,6 +67,12 @@ toggle_cuda_graphs, ) from megatron.core.utils import unwrap_model +from torch.distributed.distributed_c10d import ( + PrefixStore, + ProcessGroup, + ProcessGroupGloo, + _world, +) from nemo_rl.data.multimodal_utils import CACHED_VIDEO_FRAME_MANIFEST_MAGIC from nemo_rl.distributed.batched_data_dict import BatchedDataDict @@ -63,6 +82,9 @@ GenerationOutputSpec, verify_right_padding, ) +from nemo_rl.models.generation.megatron.config import ( + resolve_refit_execution_batch_bytes, +) from nemo_rl.models.generation.megatron.utils import ( build_image_preprocessing_config, build_prompt_and_multimodal_data, @@ -77,6 +99,104 @@ resume_inference_weights, ) from nemo_rl.utils.nsys import wrap_with_nvtx_name +from nemo_rl.utils.packed_tensor import packed_broadcast_consumer +from nemo_rl.weight_sync.nccl_reshard_utils import ( + _INDIVIDUAL_EXPERT_RE, + _STR_TO_DTYPE, + HFToLocalParamMap, + LocalParamSpec, + RefitCtx, + is_nccl_reshard_param, + restore_refit_info_placements, +) + + +def _inference_optimized_transformer_layer_spec(config: Any) -> Any: + """Build the generic GPT layer spec backed by MCore inference linears.""" + return get_gpt_layer_with_inference_spec( + qk_layernorm=config.qk_layernorm, + multi_latent_attention=config.multi_latent_attention, + qk_l2_norm=config.qk_l2_norm, + num_experts=config.num_moe_experts, + moe_grouped_gemm=config.moe_grouped_gemm, + moe_use_legacy_grouped_gemm=getattr( + config, "moe_use_legacy_grouped_gemm", False + ), + ) + + +def _configure_inference_optimized_layer_spec(model_provider: Any) -> bool: + """Select MCore inference linears for a Bridge generic-GPT provider. + + Only the generic path needs this. Model-specific providers (DeepSeek, + MiniMax, GLM, ...) install ``get_gpt_decoder_block_spec``, which already + dispatches ``transformer_impl='inference_optimized'`` itself and — unlike + the uniform spec built here — gives dense and MoE layers *different* specs. + Overwriting those would rebuild a ``moe_layer_freq`` model's dense layers + as MoE layers, so leave any non-default spec alone. + + Non-GPT providers are left alone too. Hybrid/Mamba stacks (Nemotron-H via + ``HybridModelProvider``) are not ``GPTModelProvider`` subclasses and resolve + ``inference_optimized`` through their own stack spec, so this is a no-op for + them rather than an error. + + Returns: + True if the provider's layer spec was replaced, False if the provider + carries a model-specific spec (or is not a GPT provider) and handles + inference selection itself. + """ + # Bridge imports ModelOpt plugins that can re-enter MCore while this module + # is still initializing, so keep this cycle-sensitive provider import local. + from megatron.bridge.models.gpt_provider import GPTModelProvider, default_layer_spec + + if not isinstance(model_provider, GPTModelProvider): + return False + if model_provider.transformer_layer_spec is not default_layer_spec: + return False + model_provider.transformer_layer_spec = _inference_optimized_transformer_layer_spec + return True + + +def _resolve_mxfp8_refit_backend(model_config: Any) -> str: + """Resolve the MXFP8 storage required by the grouped-GEMM backend.""" + return resolve_mxfp8_backend(model_config.inference_grouped_gemm_backend) + + +@dataclass +class _MegatronRefitTask: + """A Bridge import task and its mutable inference destination.""" + + conversion_task: Any + destination: Any + target_id: int + + @property + def param_name(self) -> str: + return self.conversion_task.param_name + + @property + def dependencies(self) -> tuple[str, ...]: + return self.conversion_task.hf_param_names + + @property + def expected_shape(self) -> torch.Size: + return torch.Size(self.destination.shape) + + @property + def is_mxfp8(self) -> bool: + return isinstance(self.destination, MXFP8Tensor) + + +@dataclass +class _MegatronBulkRefitPiece: + """One HF-local shard and the Megatron destination region it populates.""" + + task: _MegatronRefitTask + spec: Any + shape: torch.Size + dtype: torch.dtype + device: torch.device + destination: torch.Tensor | None class MegatronGenerationMixin: @@ -1051,13 +1171,41 @@ async def _generate_with_persistent_engine( class MegatronGenerationRefitMixin: """Refit collective, weight transfer, and engine suspend/resume around refits.""" + def _init_generation_refit_state(self) -> None: + """Reset every refit attribute to its unprepared state. + + Declaring these up front keeps ``prepare_refit_info`` readiness a value + check rather than an ``hasattr`` probe, and makes the full set visible + to anyone adding a ``__getstate__``. + """ + # Bridge import plan, populated by prepare_refit_info. + self._generation_refit_state_dict_info = None + self._generation_refit_tasks = None + self._generation_refit_model_chunks = None + self._generation_refit_dependency_counts = None + # Per-refit progress, reset at the start of each transfer. + self._generation_refit_task_index = 0 + self._generation_refit_remaining_dependencies = {} + self._generation_refit_pending_weights = {} + self._generation_refit_pending_streams = {} + # nccl_reshard (M-to-N) transport state. The source and destination + # roles use the same public state names and RefitBuilderInterface. + self._generation_nccl_reshard_groups = None + self.nccl_reshard_refit_info = None + self.hf_to_local_param_map = None + self._generation_m2n_pending = None + # Native MCore copy service, built by init_collective_mcore_generation. + self.refit_copy_service = None + self.refit_execution_batch_bytes: int | None = None + def init_collective_mcore_generation( self, ip: str, port: int, world_size: int, rank_offset: int, - refit_backend: str = "gloo", + refit_execution_batch_bytes: int | None, + refit_backend: str, ) -> None: """Initialize the refit collective for non-colocated weight transfer. @@ -1066,8 +1214,11 @@ def init_collective_mcore_generation( port: Port for the process group rendezvous. world_size: Total world size (train + inference workers). rank_offset: Offset for this side's ranks (`train_world_size` for inference). - refit_backend: Copy-service backend ("gloo" or "nccl"; - "nvshmem" is currently broken, see the issue below). + refit_execution_batch_bytes: Configured native MCore staging limit. + None uses NeMo-RL's dynamic collective-refit default. + refit_backend: Copy-service backend ("gloo", "nccl", or + "nccl_m2n"; "nvshmem" is currently broken, see the issue + below). """ if refit_backend == "nvshmem": warnings.warn( @@ -1075,12 +1226,8 @@ def init_collective_mcore_generation( '"gloo". See https://github.com/NVIDIA-NeMo/RL/issues/3646', stacklevel=2, ) - - from torch.distributed.distributed_c10d import ( - PrefixStore, - ProcessGroup, - ProcessGroupGloo, - _world, + execution_batch_bytes = resolve_refit_execution_batch_bytes( + refit_execution_batch_bytes ) local_rank = torch.distributed.get_rank() @@ -1118,7 +1265,9 @@ def init_collective_mcore_generation( # registered for the cuda device on this cross-world PG. GLOO stays the # default backend so the object collectives in `prepare_swap_model_weights` # (all_gather_object / broadcast_object_list) keep using CPU tensors. - if refit_backend == "nccl": + # NCCLM2NCopyService also uses this backend to coordinate topology and + # communicator IDs before native M-to-N calls. + if refit_backend in ("nccl", "nccl_m2n"): from torch.distributed.distributed_c10d import ProcessGroupNCCL # Ensure the NCCL communicator binds to this rank's own GPU. @@ -1142,6 +1291,7 @@ def init_collective_mcore_generation( pg._set_group_name(group_name) self.refit_pg = pg + self.refit_execution_batch_bytes = execution_batch_bytes # Register in torch.distributed's global state so that high-level ops # (all_gather_object, broadcast_object_list) work with this PG. @@ -1156,6 +1306,8 @@ def init_collective_mcore_generation( ) self.refit_copy_service = NVSHMEMCopyService(group=self.refit_pg) + elif refit_backend == "nccl_m2n": + self.refit_copy_service = NCCLM2NCopyService(group=self.refit_pg) elif refit_backend == "nccl": self.refit_copy_service = NCCLCopyService(group=self.refit_pg) else: @@ -1175,6 +1327,7 @@ def init_collective_mcore_generation( group=self.refit_pg, src_rank_offset=0, dst_rank_offset=self.refit_dst_rank_offset, + execution_batch_bytes=self.refit_execution_batch_bytes, ) def preinit_nvshmem_collective(self) -> None: @@ -1184,7 +1337,7 @@ def preinit_nvshmem_collective(self) -> None: outside CUDA graph capture. Lazy initialization during graph recording or replay can corrupt CUDA graph state. """ - if not hasattr(self, "refit_copy_service"): + if self.refit_copy_service is None: return if not hasattr(self.refit_copy_service, "_ensure_initialized"): return @@ -1201,6 +1354,8 @@ def swap_weights_via_reshard(self, is_source: bool) -> bool: """ src_model = self.model if is_source else None dst_model = None if is_source else self.model + if self.refit_execution_batch_bytes is None: + raise RuntimeError("Native MCore refit was not initialized.") swap_model_weights( src_model, @@ -1209,10 +1364,542 @@ def swap_weights_via_reshard(self, is_source: bool) -> bool: group=self.refit_pg, src_rank_offset=0, dst_rank_offset=self.refit_dst_rank_offset, + execution_batch_bytes=self.refit_execution_batch_bytes, ) return True + def init_nccl_reshard_comm_groups_generation( + self, + *, + pp_ips: list[str], + pp_ports: list[int], + pp_size: int, + train_ranks_per_stage: int, + sub_world_size: int, + rank_prefix: Optional[int] = None, + ) -> None: + """Join every training PP stage's M-to-N communicator as a gen rank.""" + # Keep this local because the module imports optional NCCL4Py bindings. + from nemo_rl.distributed.refit_watchdog import RELEASE_GRACE_S, release_within + from nemo_rl.distributed.stateless_process_group import StatelessProcessGroup + + generation_rank = self.rank if rank_prefix is None else rank_prefix + rank_in_group = train_ranks_per_stage + generation_rank + torch.cuda.empty_cache() + stale_groups = getattr(self, "_generation_nccl_reshard_groups", None) or {} + self._generation_nccl_reshard_groups = {} + for stage in range(pp_size): + group = StatelessProcessGroup( + master_address=pp_ips[stage], + port=pp_ports[stage], + rank=rank_in_group, + world_size=sub_world_size, + ) + group.init_nccl_communicator(device=torch.cuda.current_device()) + self._generation_nccl_reshard_groups[stage] = group + + for previous in stale_groups.values(): + release_within( + previous.abort, RELEASE_GRACE_S, "a previous reshard bulk communicator" + ) + + def _commit_megatron_bulk_refit_piece( + self, piece: _MegatronBulkRefitPiece, tensor: torch.Tensor + ) -> None: + """Commit one received HF-local weight to its inference destination.""" + if tensor.shape != piece.shape: + raise ValueError( + f"Shape mismatch for Megatron bulk refit weight {piece.spec.name!r} " + f"of {piece.task.param_name!r}: expected {tuple(piece.shape)}, " + f"got {tuple(tensor.shape)}." + ) + if not piece.task.is_mxfp8: + assert piece.destination is not None + piece.destination.copy_(tensor) + return + + pending = self._generation_m2n_pending.setdefault(piece.task.target_id, {}) + pending[piece.spec.name] = tensor + required = { + spec.name for spec in piece.task.conversion_task.local_hf_param_specs() + } + if required.issubset(pending): + logical_weight = piece.task.conversion_task.combine_local_hf_weights( + pending + ) + self._write_generation_refit_weight(piece.task, logical_weight) + del self._generation_m2n_pending[piece.task.target_id] + + def _build_destination_hf_to_local_param_map( + self, + refit_info: dict[str, Any], + tasks: list[_MegatronRefitTask], + ) -> HFToLocalParamMap: + """Map canonical HF FFN shards onto local Megatron destinations.""" + pieces: dict[str, _MegatronBulkRefitPiece] = {} + + def add_piece(task: _MegatronRefitTask, spec: Any) -> None: + if spec.name in pieces: + raise RuntimeError( + f"Duplicate Megatron M-to-N target for {spec.name!r}." + ) + destination = None if task.is_mxfp8 else spec.select(task.destination) + # MXFP8 destinations stage through BF16: the wire payload is logical + # BF16 by design, and MXFP8Tensor.dtype is optional metadata that is + # None until the destination's first successful update. Pin the + # staging dtype rather than inferring it from the quantized store. + dtype = torch.bfloat16 if task.is_mxfp8 else task.destination.dtype + pieces[spec.name] = _MegatronBulkRefitPiece( + task=task, + spec=spec, + shape=spec.selected_shape(task.expected_shape), + dtype=dtype, + device=task.destination.device, + destination=destination, + ) + + for task in tasks: + for spec in task.conversion_task.local_hf_param_specs(): + if is_nccl_reshard_param(spec.name): + add_piece(task, spec) + + expert_pieces: dict[ + tuple[str, str], list[tuple[int, _MegatronBulkRefitPiece]] + ] = {} + for name, piece in pieces.items(): + match = _INDIVIDUAL_EXPERT_RE.match(name) + if match: + expert_pieces.setdefault((match.group(1), match.group(3)), []).append( + (int(match.group(2)), piece) + ) + + def grouped_spec( + grouped: tuple[_MegatronBulkRefitPiece, ...], + ) -> LocalParamSpec: + first = grouped[0] + + def pre(_base: Any) -> RefitCtx: + return RefitCtx( + buf=torch.empty( + (len(grouped), *first.shape), + dtype=first.dtype, + device=first.device, + ) + ) + + def post(ctx: RefitCtx) -> None: + for received, piece in zip(ctx.buf.unbind(0), grouped, strict=True): + self._commit_megatron_bulk_refit_piece(piece, received) + + return LocalParamSpec(base=None, pre=pre, post=post) + + def staged_spec(piece: _MegatronBulkRefitPiece) -> LocalParamSpec: + def pre(_base: Any) -> RefitCtx: + return RefitCtx( + buf=torch.empty(piece.shape, dtype=piece.dtype, device=piece.device) + ) + + def post(ctx: RefitCtx) -> None: + self._commit_megatron_bulk_refit_piece(piece, ctx.buf) + + return LocalParamSpec(base=None, pre=pre, post=post) + + specs = {} + for layer_name in refit_info["layer_names"]: + for param_info in refit_info["per_layer_params"][layer_name]: + name = param_info["name"] + grouped_proj = param_info.get("grouped_expert_proj") + if grouped_proj is not None: + prefix = name.rsplit(f".{grouped_proj}.weight", 1)[0] + grouped = tuple( + piece + for _, piece in sorted( + expert_pieces.get((prefix, grouped_proj), []) + ) + ) + if not grouped: + raise ValueError( + f"No local Megatron experts map to M-to-N weight {name!r}." + ) + specs[name] = grouped_spec(grouped) + continue + + piece = pieces.get(name) + if piece is None: + raise ValueError( + f"No local Megatron destination maps to M-to-N weight {name!r}." + ) + specs[name] = ( + staged_spec(piece) + if piece.task.is_mxfp8 + else LocalParamSpec(base=piece.destination) + ) + + return HFToLocalParamMap(specs=specs) + + def _prepare_mxfp8_refit(self, tasks: list[_MegatronRefitTask]) -> None: + """Install persistent MCore MXFP8 destinations before engine initialization.""" + model_chunks = ( + self.model if isinstance(self.model, (list, tuple)) else [self.model] + ) + cores = [unwrap_model(model) for model in model_chunks] + mxfp8_cores = [ + core + for core in cores + if getattr(core.config, "transformer_impl", None) == "inference_optimized" + and getattr(core.config, "fp8_recipe", None) == "mxfp8" + ] + if not mxfp8_cores: + return + if self._inference_engine_initialized: + raise RuntimeError( + "MXFP8 refit buffers must be prepared before inference-engine " + "initialization; construct MegatronGeneration with skip_weight_load=True." + ) + + destination_by_id: dict[int, MXFP8Tensor] = {} + for core in mxfp8_cores: + decoder = core.decoder if hasattr(core, "decoder") else core + param_name_by_id = { + id(param): name for name, param in decoder.named_parameters() + } + persistent_buffers = quantize_params_to_mxfp8( + decoder, backend=_resolve_mxfp8_refit_backend(core.config) + ) + destination_by_id.update( + { + param_id: persistent_buffers[name] + for param_id, name in param_name_by_id.items() + if name in persistent_buffers + } + ) + + for task in tasks: + if task.target_id in destination_by_id: + task.destination = destination_by_id[task.target_id] + + def _build_generation_refit_tasks( + self, + ) -> tuple[list[torch.nn.Module], list[_MegatronRefitTask]]: + """Resolve Bridge import tasks against the persistent inference model.""" + model_chunks = ( + list(self.model) if isinstance(self.model, (list, tuple)) else [self.model] + ) + conversion_tasks = self.megatron_bridge.get_conversion_tasks(model_chunks) + tasks: list[_MegatronRefitTask] = [] + for conversion_task in conversion_tasks: + if conversion_task is None or conversion_task.megatron_module is None: + continue + destination = conversion_task.param_weight + if destination is None: + raise RuntimeError( + f"Bridge task {conversion_task.param_name!r} has no destination." + ) + tasks.append( + _MegatronRefitTask( + conversion_task=replace(conversion_task, param_weight=None), + destination=destination, + target_id=id(destination), + ) + ) + + if not tasks: + raise RuntimeError("Megatron Bridge produced no local refit tasks.") + return model_chunks, tasks + + def _install_generation_refit_plan( + self, + *, + model_chunks: list[torch.nn.Module], + tasks: list[_MegatronRefitTask], + state_dict_info: dict[str, Any], + ) -> None: + """Install the packed-broadcast subset of a Bridge receive plan.""" + if not state_dict_info: + raise ValueError( + "Megatron packed refit requires non-empty state_dict_info." + ) + + required_names = {name for task in tasks for name in task.dependencies} + missing_names = sorted(required_names.difference(state_dict_info)) + if missing_names: + preview = ", ".join(missing_names[:20]) + raise ValueError( + "Megatron refit metadata is missing Bridge inputs: " + f"{preview}{' ...' if len(missing_names) > 20 else ''}" + ) + + self._generation_refit_model_chunks = model_chunks + self._generation_refit_tasks = tasks + self._generation_refit_dependency_counts = Counter( + name for task in tasks for name in task.dependencies + ) + self._generation_refit_state_dict_info = state_dict_info + + @torch.no_grad() + def _prepare_destination_refit_info(self, state_dict_info: dict[str, Any]) -> None: + """Build the local HF-to-Megatron receive plan before the first refit.""" + model_chunks, tasks = self._build_generation_refit_tasks() + self._prepare_mxfp8_refit(tasks) + self._install_generation_refit_plan( + model_chunks=model_chunks, + tasks=tasks, + state_dict_info=state_dict_info, + ) + + @torch.no_grad() + def _prepare_destination_nccl_reshard_refit_info( + self, refit_info: dict[str, Any] + ) -> None: + """Prepare Megatron as the destination of NCCL M-to-N reshard refit.""" + refit_info = restore_refit_info_placements(refit_info) + model_chunks, tasks = self._build_generation_refit_tasks() + self._prepare_mxfp8_refit(tasks) + bulk_map = self.build_hf_to_local_param_map(refit_info, destination_tasks=tasks) + + misc_meta = refit_info.get("misc_meta", {}) + misc_state_dict_info = { + name: (tuple(meta["shape"]), _STR_TO_DTYPE[meta["dtype"]]) + for name, meta in misc_meta.items() + } + misc_tasks = [ + task + for task in tasks + if set(task.dependencies).issubset(misc_state_dict_info) + ] + self._install_generation_refit_plan( + model_chunks=model_chunks, + tasks=misc_tasks, + state_dict_info=misc_state_dict_info, + ) + self.nccl_reshard_refit_info = refit_info + self.hf_to_local_param_map = bulk_map + + @torch.no_grad() + def _destination_nccl_reshard_refit( + self, refit_timeout_s: Optional[float] = None + ) -> bool: + """Receive bulk FFN shards via M-to-N, then import packed misc weights.""" + from nemo_rl.distributed.refit_watchdog import sync_stream_within + + # Keep this transport import local: xferdtensor probes the optional + # nccl.m2n extension at import time. + from nemo_rl.weight_sync.xferdtensor import DTensorRef, xferdtensor + + self._generation_m2n_pending: dict[int, dict[str, torch.Tensor]] = {} + + def receive_one(param_info: dict[str, Any], group: Any, stream: Any) -> None: + spec = self.hf_to_local_param_map.get(param_info["name"]) + if spec is None: + raise RuntimeError( + f"Megatron M-to-N refit has no destination for {param_info['name']!r}." + ) + ctx = ( + spec.pre(spec.base) if spec.pre is not None else RefitCtx(buf=spec.base) + ) + destination = DTensorRef(ctx.buf, param_info["global_shape"]) + xferdtensor( + None, + param_info["src_mesh_info"], + param_info["src_placements"], + destination, + param_info["dst_mesh_info"], + param_info["dst_placements"], + group, + stream, + ) + if spec.post is not None: + spec.post(ctx) + + stage_params: OrderedDict[int, list[dict[str, Any]]] = OrderedDict() + refit_info = self.nccl_reshard_refit_info + for layer_name in refit_info["layer_names"]: + for param_info in refit_info["per_layer_params"][layer_name]: + stage_params.setdefault(param_info.get("pp_stage", 0), []).append( + param_info + ) + + num_streams = max( + 1, + min(int(os.environ.get("NRL_REFIT_NUM_STREAMS", "2")), len(stage_params)), + ) + streams = [torch.cuda.Stream() for _ in range(num_streams)] + events: dict[int, torch.cuda.Event] = {} + try: + for index, (stage, params) in enumerate(stage_params.items()): + previous_index = index - num_streams + if previous_index in events: + events[previous_index].synchronize() + stage_stream = streams[index % num_streams] + with torch.cuda.stream(stage_stream): + group = self._generation_nccl_reshard_groups[stage] + for param_info in params: + receive_one(param_info, group, stage_stream) + event = torch.cuda.Event() + event.record() + events[index] = event + + current_stream = torch.cuda.current_stream() + for event in events.values(): + current_stream.wait_event(event) + sync_stream_within( + current_stream, + refit_timeout_s, + "the Megatron destination bulk parameter transfer", + ) + if self._generation_m2n_pending: + raise RuntimeError( + "Megatron M-to-N refit ended with incomplete fused MXFP8 weights: " + f"{sorted(self._generation_m2n_pending)}" + ) + torch.cuda.empty_cache() + return self._update_destination_weights_from_collective( + refit_timeout_s=refit_timeout_s + ) + finally: + self._generation_m2n_pending.clear() + + def _write_generation_refit_weight( + self, task: _MegatronRefitTask, converted_weight: torch.Tensor + ) -> None: + """Copy one converted weight into its persistent inference destination. + + ``copy_`` rather than rebinding: an MXFP8 destination quantizes in place + and keeps its storage address, which is what already-captured CUDA + graphs point at. + + Args: + task: The Bridge import task owning the destination tensor. + converted_weight: The assembled weight in the destination's layout. + + Raises: + ValueError: If the converted weight does not match the destination + shape, which means the refit plan and the model disagree. + """ + if converted_weight.shape != task.expected_shape: + raise ValueError( + f"Shape mismatch for Megatron parameter {task.param_name!r}: " + f"expected {tuple(task.expected_shape)}, got {tuple(converted_weight.shape)}." + ) + + task.destination.copy_(converted_weight) + + def _load_generation_refit_batch( + self, weights: list[tuple[str, torch.Tensor]] + ) -> None: + """Import every Bridge task whose HF inputs have now all arrived. + + Called once per ``packed_broadcast_consumer`` batch. Weights arrive in + HF order, but a task may fuse several of them (gate/up, stacked + experts), so each tensor is buffered until its task's dependency count + reaches zero, then the task is converted and written in task order. + Buffered tensors record their producing stream so the conversion can + wait on it before reading across streams. + + Args: + weights: ``(hf_name, tensor)`` pairs unpacked from this batch. + """ + batch_stream = ( + torch.cuda.current_stream() if weights and weights[0][1].is_cuda else None + ) + for name, tensor in weights: + if self._generation_refit_remaining_dependencies.get(name, 0) > 0: + self._generation_refit_pending_weights[name] = tensor + if batch_stream is not None: + self._generation_refit_pending_streams[name] = batch_stream + + tasks = self._generation_refit_tasks + while self._generation_refit_task_index < len(tasks): + task = tasks[self._generation_refit_task_index] + if any( + name not in self._generation_refit_pending_weights + for name in task.dependencies + ): + break + + # packed_broadcast_consumer alternates CUDA streams. A compound Bridge + # mapping can retain one input across a batch boundary, so make the + # current batch wait for the stream that received each retained input. + current_stream = ( + torch.cuda.current_stream() if batch_stream is not None else None + ) + for name in task.dependencies: + source_stream = self._generation_refit_pending_streams.get(name) + if current_stream is not None and source_stream is not None: + current_stream.wait_stream(source_stream) + + converted = next( + iter( + self.megatron_bridge.stream_weights_hf_to_megatron( + self._generation_refit_model_chunks, + conversion_tasks=[task.conversion_task], + hf_state_dict=self._generation_refit_pending_weights, + ) + ), + None, + ) + if converted is None: + raise RuntimeError( + f"Bridge produced no local value for {task.param_name!r}." + ) + self._write_generation_refit_weight(task, converted.weight) + + for name in task.dependencies: + self._generation_refit_remaining_dependencies[name] -= 1 + if self._generation_refit_remaining_dependencies[name] == 0: + del self._generation_refit_pending_weights[name] + self._generation_refit_pending_streams.pop(name, None) + self._generation_refit_task_index += 1 + + @torch.no_grad() + def _update_destination_weights_from_collective( + self, refit_timeout_s: Optional[float] = None + ) -> bool: + """Receive packed HF tensors and import them into the Megatron model.""" + from nemo_rl.distributed.refit_watchdog import sync_stream_within + + if self._generation_refit_state_dict_info is None: + raise RuntimeError( + "Megatron refit metadata is not prepared. Call prepare_refit_info first." + ) + + self._generation_refit_task_index = 0 + self._generation_refit_pending_weights: dict[str, torch.Tensor] = {} + self._generation_refit_pending_streams: dict[str, torch.cuda.Stream] = {} + self._generation_refit_remaining_dependencies = Counter( + self._generation_refit_dependency_counts + ) + try: + packed_broadcast_consumer( + iterator=iter(self._generation_refit_state_dict_info.items()), + group=self.model_update_group, + src=0, + post_unpack_func=self._load_generation_refit_batch, + ) + if self._generation_refit_task_index != len(self._generation_refit_tasks): + task = self._generation_refit_tasks[self._generation_refit_task_index] + missing = [ + name + for name in task.dependencies + if name not in self._generation_refit_pending_weights + ] + raise RuntimeError( + f"Megatron refit ended before {task.param_name!r}; " + f"missing inputs: {missing}." + ) + + self.megatron_bridge.finalize_hf_import(self._generation_refit_model_chunks) + sync_stream_within( + torch.cuda.current_stream(), + refit_timeout_s, + "the Megatron destination packed parameter transfer", + ) + return True + finally: + self._generation_refit_pending_weights.clear() + self._generation_refit_pending_streams.clear() + def _onload_inference_model(self) -> None: """Restore the colocated inference weights to GPU before resharding / generation.""" if not self._inference_model_offloaded: @@ -1236,6 +1923,11 @@ def _reshard_into_inference_model(self) -> None: inference_model = self.inference_model if inference_model is None: return + execution_batch_bytes = resolve_refit_execution_batch_bytes( + self.cfg["generation"]["mcore_generation_config"][ + "refit_execution_batch_bytes" + ] + ) # Bring the inference weights back to GPU. self._onload_inference_model() @@ -1258,6 +1950,7 @@ def _reshard_into_inference_model(self) -> None: group=None, src_rank_offset=0, dst_rank_offset=0, + execution_batch_bytes=execution_batch_bytes, ) self._swap_weights_plan_prepared = True @@ -1270,6 +1963,7 @@ def _reshard_into_inference_model(self) -> None: group=None, src_rank_offset=0, dst_rank_offset=0, + execution_batch_bytes=execution_batch_bytes, ) # Offload training model. self.model = self.move_model( diff --git a/nemo_rl/models/megatron/setup.py b/nemo_rl/models/megatron/setup.py index f2928945a4f..8fb276f68be 100644 --- a/nemo_rl/models/megatron/setup.py +++ b/nemo_rl/models/megatron/setup.py @@ -1627,16 +1627,20 @@ def _create_megatron_config( # fp8_param_gather and reuse_grad_buf_for_mxfp8_param_ag are derived: both are # only valid when fp8 is enabled, fp8_param=True, and recipe is mxfp8. Mcore's # DDP __post_init__ asserts they remain in sync, so we centralize the derivation - # rather than exposing two redundant YAML knobs that can disagree. + # rather than exposing two redundant YAML knobs that can disagree. The optimizer + # fp8_recipe comes from the same canonical model config so it selects the + # matching main-parameter representation. fp8_cfg = config["megatron_cfg"].get("fp8_cfg", None) - reuse_grad_buf_for_mxfp8_param_ag = ( - fp8_param_enabled and fp8_cfg.get("fp8_recipe") == "mxfp8" + fp8_recipe = ( + fp8_cfg.get("fp8_recipe") if fp8_cfg and fp8_cfg.get("enabled", False) else None ) + reuse_grad_buf_for_mxfp8_param_ag = fp8_param_enabled and fp8_recipe == "mxfp8" overlap_param_gather = config["megatron_cfg"]["distributed_data_parallel_config"][ "overlap_param_gather" ] optimizer_kwargs = { **_resolve_optimizer_dtype_kwargs(config["megatron_cfg"]["optimizer"]), + "fp8_recipe": fp8_recipe, "overlap_param_gather": overlap_param_gather, "reuse_grad_buf_for_mxfp8_param_ag": reuse_grad_buf_for_mxfp8_param_ag, } diff --git a/nemo_rl/models/policy/interfaces.py b/nemo_rl/models/policy/interfaces.py index ef4139c992b..352e2156215 100644 --- a/nemo_rl/models/policy/interfaces.py +++ b/nemo_rl/models/policy/interfaces.py @@ -12,16 +12,18 @@ # See the License for the specific language governing permissions and # limitations under the License. from abc import ABC, abstractmethod -from typing import Any, Optional, TypedDict +from typing import Any, Literal, Optional, TypedDict import ray import torch from nemo_rl.algorithms.loss.interfaces import LossFunction from nemo_rl.distributed.batched_data_dict import BatchedDataDict -from nemo_rl.models.generation.interfaces import GenerationDatumSpec +from nemo_rl.models.generation.interfaces import GenerationDatumSpec, RefitPayloadMode from nemo_rl.utils.timer import Timer +RefitRole = Literal["source", "destination"] + class LogprobOutputSpec(TypedDict): """logprobs: Tensor of log probabilities.""" @@ -172,6 +174,7 @@ def init_collective( world_size: int, *, train_world_size: int, + rank_offset: int = 0, nccl_peer: str = "nemo", ) -> list[ray.ObjectRef]: pass @@ -188,7 +191,11 @@ def offload_to_cpu(self) -> None: pass @abstractmethod - def prepare_refit_info(self) -> Optional[dict[str, Any]]: + def prepare_refit_info( + self, + *, + refit_payload_mode: Optional[RefitPayloadMode] = None, + ) -> Optional[dict[str, Any]]: pass @abstractmethod @@ -263,6 +270,8 @@ def prepare_nccl_reshard_refit_info( gen_parallelism: dict[str, int], train_world_size: int, gen_world_size: int, + *, + refit_payload_mode: Optional[RefitPayloadMode] = None, ) -> Any: """Prepare per-layer param metadata for nccl_reshard-based refit.""" raise NotImplementedError diff --git a/nemo_rl/models/policy/lm_policy.py b/nemo_rl/models/policy/lm_policy.py index 2441b623bbe..b70ada0aef6 100644 --- a/nemo_rl/models/policy/lm_policy.py +++ b/nemo_rl/models/policy/lm_policy.py @@ -37,12 +37,14 @@ GenerationDatumSpec, GenerationInterface, GenerationOutputSpec, + RefitPayloadMode, ) from nemo_rl.models.policy import PolicyConfig from nemo_rl.models.policy.interfaces import ( ColocatablePolicyInterface, LogprobOutputSpec, ReferenceLogprobOutputSpec, + RefitRole, ScoreOutputSpec, TopkLogitsOutputSpec, ) @@ -102,6 +104,7 @@ def __init__( processor: Optional[AutoProcessor] = None, worker_extension_cls_fqn: Optional[str] = None, skip_weight_load: bool = False, + refit_role: RefitRole = "source", reserved_http_server_port: Optional[int] = None, ): self.debug_payload_metrics = False @@ -308,6 +311,10 @@ def __init__( worker_sharding_annotations=self.sharding_annotations, pre_init_communication_queue=pre_init_queue, ) + if megatron_enable: + worker_kwargs["refit_role"] = refit_role + elif refit_role != "source": + raise ValueError("refit_role='destination' requires the Megatron backend.") if skip_weight_load: worker_kwargs["skip_weight_load"] = True if reserved_http_server_port is not None: @@ -467,6 +474,7 @@ def init_collective( world_size: int, *, train_world_size: int, + rank_offset: int = 0, nccl_peer: str = "nemo", ) -> list[ray.ObjectRef]: """Initialize the collective communication.""" @@ -476,6 +484,7 @@ def init_collective( port=port, world_size=world_size, train_world_size=train_world_size, + rank_offset=rank_offset, nccl_peer=nccl_peer, ) # this function should co-work with vllm, so we should wait for all futures to complete outside @@ -488,7 +497,8 @@ def init_collective_mcore_generation( world_size: int, *, rank_offset: int, - refit_backend: str = "gloo", + refit_execution_batch_bytes: int | None, + refit_backend: str, ) -> list[ray.ObjectRef]: """Initialize the megatron refit collective on this policy's workers.""" return self.worker_group.run_all_workers_single_data( @@ -497,6 +507,7 @@ def init_collective_mcore_generation( port=port, world_size=world_size, rank_offset=rank_offset, + refit_execution_batch_bytes=refit_execution_batch_bytes, refit_backend=refit_backend, ) @@ -1036,13 +1047,30 @@ def invalidate_kv_cache(self, *args: Any, **kwargs: Any) -> bool: # We don't need to do anything here return True - def prepare_refit_info(self) -> Optional[dict[str, Any]]: + def prepare_refit_info( + self, + *, + refit_payload_mode: Optional[RefitPayloadMode] = None, + ) -> Optional[dict[str, Any]]: """Prepare the info for refit. Returns: dict: A dictionary containing the info for refit. """ - futures = self.worker_group.run_all_workers_single_data("prepare_refit_info") + worker_kwargs = {} + if self.cfg.get("megatron_cfg", {}).get("enabled", False): + if refit_payload_mode is None: + generation_cfg = self.cfg.get("generation") + refit_payload_mode = ( + "logical_weights" + if generation_cfg is not None + and generation_cfg.get("backend") == "megatron" + else "bridge_export" + ) + worker_kwargs["refit_payload_mode"] = refit_payload_mode + futures = self.worker_group.run_all_workers_single_data( + "prepare_refit_info", **worker_kwargs + ) results = ray.get(futures) # Only get the first worker's info since all workers will have the same result return results[0] @@ -1255,18 +1283,29 @@ def init_nccl_reshard_comm_group( def prepare_nccl_reshard_refit_info( self, - train_parallelism, - gen_parallelism, - train_world_size, - gen_world_size, - ): + train_parallelism: dict[str, Any], + gen_parallelism: dict[str, Any], + train_world_size: int, + gen_world_size: int, + *, + refit_payload_mode: Optional[RefitPayloadMode] = None, + ) -> dict[str, Any]: """Prepare per-layer param metadata for nccl_reshard refit.""" + if refit_payload_mode is None: + generation_cfg = self.cfg.get("generation") + refit_payload_mode = ( + "logical_weights" + if generation_cfg is not None + and generation_cfg.get("backend") == "megatron" + else "bridge_export" + ) futures = self.worker_group.run_all_workers_single_data( "prepare_nccl_reshard_refit_info", train_parallelism=train_parallelism, gen_parallelism=gen_parallelism, train_world_size=train_world_size, gen_world_size=gen_world_size, + refit_payload_mode=refit_payload_mode, ) results = ray.get(futures) return results[0] diff --git a/nemo_rl/models/policy/workers/base_policy_worker.py b/nemo_rl/models/policy/workers/base_policy_worker.py index 693cfe84e0d..96d710622f8 100644 --- a/nemo_rl/models/policy/workers/base_policy_worker.py +++ b/nemo_rl/models/policy/workers/base_policy_worker.py @@ -18,6 +18,7 @@ import zmq from nemo_rl.distributed.batched_data_dict import BatchedDataDict +from nemo_rl.models.generation.interfaces import RefitPayloadMode from nemo_rl.models.policy.interfaces import ReferenceLogprobOutputSpec from nemo_rl.telemetry.setup import shutdown_telemetry from nemo_rl.utils.nsys import wrap_with_nvtx_name @@ -39,6 +40,7 @@ def init_collective( world_size: int, *, train_world_size: int, + rank_offset: int = 0, nccl_peer: str = "nemo", ) -> None: """Initialize the collective communication. @@ -48,6 +50,7 @@ def init_collective( port: Port for the process group world_size: Total world size (train_world_size + inference_world_size) train_world_size: Number of training workers (used in inference cluster) + rank_offset: Offset added to this worker group's local rank. nccl_peer: NCCL initialization protocol used by the inference workers """ from nemo_rl.distributed.refit_watchdog import RELEASE_GRACE_S, release_within @@ -59,9 +62,9 @@ def init_collective( # about which port to use. rank 0 is a trainer and is this store's master. print( f" refit: collective rendezvous [train] addr={ip}:{port} " - f"rank={self.rank} world_size={world_size} " + f"rank={self.rank + rank_offset} world_size={world_size} " f"train_world_size={train_world_size} peer={nccl_peer} " - f"master={self.rank == 0}", + f"master={self.rank + rank_offset == 0}", flush=True, ) @@ -72,7 +75,10 @@ def init_collective( old_group, self.model_update_group = ( self.model_update_group, StatelessProcessGroup( - master_address=ip, port=port, rank=self.rank, world_size=world_size + master_address=ip, + port=port, + rank=self.rank + rank_offset, + world_size=world_size, ), ) # Rebuilding is the recovery path for a dead generation rank, so this runs more @@ -152,6 +158,8 @@ def prepare_nccl_reshard_refit_info( gen_parallelism: dict[str, int], train_world_size: int, gen_world_size: int, + *, + refit_payload_mode: Optional[RefitPayloadMode] = None, ) -> dict[str, Any]: """Prepare parameter metadata for NCCL reshard refit.""" # This is a placeholder implementation. diff --git a/nemo_rl/models/policy/workers/checkpoint_engine.py b/nemo_rl/models/policy/workers/checkpoint_engine.py index 85e7fdd4c72..5b6abf19c85 100644 --- a/nemo_rl/models/policy/workers/checkpoint_engine.py +++ b/nemo_rl/models/policy/workers/checkpoint_engine.py @@ -33,7 +33,7 @@ def maybe_preinit_nixl_checkpoint_engine(config: dict[str, Any]) -> Any: """Preinitialize NIXL when checkpoint-engine refit is configured.""" generation_config = config.get("generation") - if generation_config is None: + if generation_config is None or generation_config.get("backend") != "vllm": return None checkpoint_config = checkpoint_engine_refit_config(generation_config) if checkpoint_config is None or checkpoint_config["backend"] != "nixl": diff --git a/nemo_rl/models/policy/workers/megatron_policy_worker.py b/nemo_rl/models/policy/workers/megatron_policy_worker.py index b0ba6eb2689..d6253e79456 100644 --- a/nemo_rl/models/policy/workers/megatron_policy_worker.py +++ b/nemo_rl/models/policy/workers/megatron_policy_worker.py @@ -20,6 +20,7 @@ import warnings from collections import OrderedDict, defaultdict from contextlib import AbstractContextManager, contextmanager, nullcontext +from dataclasses import dataclass, replace from typing import Any, Iterable, Iterator, Optional, TypeVar, cast log = logging.getLogger(__name__) @@ -58,10 +59,11 @@ from nemo_rl.data_plane.worker_mixin import TQWorkerMixin from nemo_rl.distributed.batched_data_dict import BatchedDataDict from nemo_rl.distributed.named_sharding import NamedSharding -from nemo_rl.models.generation.interfaces import GenerationDatumSpec +from nemo_rl.models.generation.interfaces import GenerationDatumSpec, RefitPayloadMode from nemo_rl.models.generation.megatron.megatron_worker import ( MegatronGenerationMixin, MegatronGenerationRefitMixin, + _configure_inference_optimized_layer_spec, ) from nemo_rl.models.generation.vllm.config import VllmConfig from nemo_rl.models.megatron.common import ( @@ -104,6 +106,7 @@ ColocatablePolicyInterface, LogprobOutputSpec, ReferenceLogprobOutputSpec, + RefitRole, ) from nemo_rl.models.policy.utils import ( broadcast_hf_buckets_via_distributed_impl, @@ -127,9 +130,14 @@ from nemo_rl.utils.r3_trace import maybe_r3_trace_stage from nemo_rl.utils.timer import Timer from nemo_rl.weight_sync.nccl_reshard_utils import ( + _INDIVIDUAL_EXPERT_RE, HFToLocalParamMap, LocalParamSpec, RefitCtx, + _extract_layer_name, + _extract_layer_prefix, + build_nccl_reshard_refit_info, + is_nccl_reshard_param, ) TokenizerType = TypeVar("TokenizerType", bound=PreTrainedTokenizerBase) @@ -239,6 +247,66 @@ def _estimate_refit_tensor_size_in_bytes( return param.numel() * tp_size * ep_size * element_size +def _is_quantized_refit_source(tensor: torch.Tensor) -> bool: + """Whether a training parameter uses TE quantized physical storage.""" + # Defer optional TE FP8 helpers until a refit source is inspected. + from megatron.core.fp8_utils import ( + is_float8tensor, + is_grouped_tensor_with_quantized_storage, + ) + + return is_float8tensor(tensor) or is_grouped_tensor_with_quantized_storage(tensor) + + +def _dequantize_refit_source(tensor: torch.Tensor) -> torch.Tensor: + """Materialize a logical BF16 weight from TE quantized parameter storage.""" + # Defer optional TE FP8 helpers until a quantized source is materialized. + from megatron.core.fp8_utils import ( + dequantize_fp8_tensor, + get_grouped_quantized_members, + is_float8tensor, + is_grouped_tensor_with_quantized_storage, + ) + + if is_grouped_tensor_with_quantized_storage(tensor): + members = get_grouped_quantized_members(tensor, create_if_missing=True) + return torch.stack( + [dequantize_fp8_tensor(member).to(torch.bfloat16) for member in members] + ) + if is_float8tensor(tensor): + return dequantize_fp8_tensor(tensor).to(torch.bfloat16) + return tensor + + +def _get_refit_task_source(task: Any) -> Optional[torch.Tensor]: + """Return the original TE parameter when Bridge exposes an export payload.""" + source = task.param_weight + module = getattr(task, "megatron_module", None) + param_name = getattr(task, "param_name", "") + if source is None or module is None or not param_name: + return source + + original = getattr(module, param_name.rsplit(".", 1)[-1], None) + if original is not None and _is_quantized_refit_source(original): + return original + return source + + +@dataclass(frozen=True) +class _QuantizedRefitSource: + """A live quantized parameter plus the Bridge view to transfer.""" + + tensor: torch.Tensor + spec: Any + + +@dataclass(frozen=True) +class _GroupedRefitSource: + """Local expert source specs that must be stacked before transfer.""" + + specs: tuple[LocalParamSpec, ...] + + def _collect_mtp_hf_layer_names(conversion_tasks: Optional[list]) -> set[str]: """Return HF layer names whose weights originate from Megatron's MTP module. @@ -251,8 +319,6 @@ def _collect_mtp_hf_layer_names(conversion_tasks: Optional[list]) -> set[str]: Returns: Set of HF layer names, e.g. ``{"model.layers.61", "mtp.layers.0"}``. """ - from nemo_rl.weight_sync.nccl_reshard_utils import _extract_layer_name - mtp_layers: set[str] = set() for task in conversion_tasks or []: if task is None: @@ -266,6 +332,45 @@ def _collect_mtp_hf_layer_names(conversion_tasks: Optional[list]) -> set[str]: return mtp_layers +def _collect_local_refit_hf_names(conversion_tasks: Iterable[Any]) -> set[str]: + """Return HF names whose Bridge tasks expose safe direct local views. + + Every task contributing to a grouped HF tensor must provide local specs. + Mappings that need transpose, interleave, or other grouped-export transforms + return no specs and therefore stay on the normal Bridge conversion path. + + The result is expert-parallel invariant. Bridge conversion tasks are + EP-local -- each EP rank only enumerates its own ``experts.N.*`` -- while the + name stream this set is matched against is EP-global, so a purely local set + would classify the same expert differently on different EP ranks. The + support map is therefore AND-reduced across the EP group before filtering. + """ + support_by_name: dict[str, bool] = {} + for task in conversion_tasks: + if task is None: + continue + has_local_views = bool(task.local_hf_param_specs()) + for name in task.hf_param_names: + support_by_name[name] = support_by_name.get(name, True) and has_local_views + + # check_initialized=False: this helper also runs outside an initialized + # parallel state (unit tests, single-rank preparation), where there is no EP + # group to reduce over and the local map is already the whole picture. + ep_group = parallel_state.get_expert_model_parallel_group(check_initialized=False) + if ep_group is not None and torch.distributed.get_world_size(ep_group) > 1: + gathered: list[Optional[dict[str, bool]]] = [None] * ( + torch.distributed.get_world_size(ep_group) + ) + torch.distributed.all_gather_object(gathered, support_by_name, group=ep_group) + merged: dict[str, bool] = {} + for rank_support in gathered: + for name, supported in (rank_support or {}).items(): + merged[name] = merged.get(name, True) and supported + support_by_name = merged + + return {name for name, supported in support_by_name.items() if supported} + + @contextmanager def _meta_tensor_alloc_context(): """Skip real GPU work during metadata enumeration. @@ -317,6 +422,10 @@ class MegatronPolicyWorkerImpl( AbstractPolicyWorker, ColocatablePolicyInterface, ): + # Tests and extension classes that bypass __init__ retain the historical + # training/source behavior unless they explicitly select destination. + refit_role: RefitRole = "source" + refit_payload_mode: RefitPayloadMode = "bridge_export" # Holds the split-API train-step state between begin/finish or # begin/abort; None when no step is open. Declared at class level so # ``self._train_step_state = None`` after finish/abort type-checks. @@ -440,10 +549,17 @@ def __init__( *, worker_sharding_annotations: NamedSharding, skip_weight_load: bool = False, + refit_role: RefitRole = "source", reserved_http_server_port: Optional[int] = None, **kwargs: Any, ): """Initialize the MegatronPolicyWorker.""" + if refit_role not in ("source", "destination"): + raise ValueError( + f"refit_role must be 'source' or 'destination', got {refit_role!r}." + ) + self.refit_role = refit_role + self.refit_payload_mode: RefitPayloadMode = "bridge_export" # NVML-based and guarded on torch.cuda.is_initialized(), so this does # not initialize a CUDA context ahead of the set_device below. log_gpu_memory_diagnostics( @@ -542,6 +658,14 @@ def __init__( ) self.megatron_cfg = runtime_config.megatron_cfg + if ( + not init_optimizer + and self.megatron_cfg.model.transformer_impl == "inference_optimized" + ): + # Bridge's generic GPT provider defaults to its training layer spec. + # Dedicated inference workers need MCore inference linears instead. + # A model-specific provider already handles this and returns False. + _configure_inference_optimized_layer_spec(self.megatron_cfg.model) self.dtype = runtime_config.dtype self.optimizer_cpu_offload = runtime_config.optimizer_cpu_offload self.offload_optimizer_for_logprob = ( @@ -698,6 +822,7 @@ def __init__( self._held_gather_buffer = None self._init_inference_engine_state() + self._init_generation_refit_state() self._setup_colocated_cuda_graph_managers() log_gpu_memory_diagnostics( @@ -748,6 +873,17 @@ def _restore_first_train_step_param_sync(self, update_successful: bool) -> None: self._first_train_step_forward_pre_hook_disabled = False def _copy_main_params_to_param_buffer(self, zero_grad_buffer: bool = False) -> None: + """Stage FP32 optimizer masters into the shared MXFP8 param buffer. + + Only valid while the forward pre-hooks are still installed: this stages + the buffer but does not all-gather it, and callers rely on the + hook-driven (or explicit) param sync to publish the result. Staging with + hooks already off would zero the aliased param/grad buffer and leave + every other DP rank's slice at zero. + + Args: + zero_grad_buffer: Also zero the aliased gradient buffer first. + """ if not isinstance(self.model, DistributedDataParallel): return @@ -2272,8 +2408,12 @@ def get_topk_logits( @torch.no_grad() @wrap_with_nvtx_name("megatron_policy_worker/prepare_refit_info") - def prepare_refit_info(self) -> None: + def _prepare_source_refit_info( + self, + refit_payload_mode: RefitPayloadMode, + ) -> dict[str, tuple[torch.Size, torch.dtype]]: """Prepare state dict metadata for weight refitting and IPC streaming.""" + self.refit_payload_mode = refit_payload_mode self.refit_param_info_mcore = self._calculate_refit_param_info() # Collect tensor metadata for refit / hf side info. @@ -2283,6 +2423,64 @@ def prepare_refit_info(self) -> None: return refit_param_info_hf + def prepare_refit_info( + self, + state_dict_info: Optional[dict[str, Any]] = None, + *, + refit_payload_mode: RefitPayloadMode = "bridge_export", + ) -> Optional[dict[str, tuple[torch.Size, torch.dtype]]]: + """Prepare refit state for this worker's explicit source/destination role.""" + if self.refit_role == "destination": + if state_dict_info is None: + raise ValueError("Destination refit requires state_dict_info.") + self._prepare_destination_refit_info(state_dict_info) + return None + if state_dict_info is not None: + raise ValueError("Source refit does not accept state_dict_info.") + return self._prepare_source_refit_info(refit_payload_mode) + + async def update_weights_from_collective( + self, refit_timeout_s: Optional[float] = None + ) -> bool: + """Receive a collective refit without blocking this actor's event loop.""" + if self.refit_role != "destination": + raise RuntimeError( + "update_weights_from_collective is only valid for destination-role workers." + ) + + from nemo_rl.distributed.refit_watchdog import await_off_loop + + device = torch.cuda.current_device() + + def _on_this_workers_device() -> bool: + torch.cuda.set_device(device) + return self._update_destination_weights_from_collective_guarded( + refit_timeout_s + ) + + return await await_off_loop(_on_this_workers_device) + + @torch.no_grad() + def _update_destination_weights_from_collective_guarded( + self, refit_timeout_s: Optional[float] + ) -> bool: + """Receive packed Megatron weights under the shared refit deadline.""" + from nemo_rl.distributed.refit_watchdog import ( + RefitAborted, + RefitAbortWatchdog, + ) + + with RefitAbortWatchdog([self.model_update_group], refit_timeout_s) as guard: + result = self._update_destination_weights_from_collective( + refit_timeout_s=refit_timeout_s + ) + if guard.fired: + raise RefitAborted( + f"Megatron collective refit exceeded {refit_timeout_s}s and was " + "aborted; a training rank most likely stopped participating" + ) + return result + def _collect_mtp_metrics( self, metrics: dict[str, Any], @@ -2421,22 +2619,22 @@ def _is_fp8_export(self) -> bool: and self.fp8_cfg.get("fp8_recipe") == "blockwise" ) + def _uses_logical_refit_payload(self) -> bool: + """Return whether the destination requested logical floating-point weights.""" + return self.refit_payload_mode == "logical_weights" + def _build_refit_conversion_tasks(self) -> list: """Build the conversion-task list driving refit (BF16 or FP8 export). - For BF16 / FP8-but-fp8_param=False training: standard ``get_conversion_tasks``. - For FP8-with-fp8_param=True: Bridge's ``build_export_fp8_tasks``, which - emits a *pair* of tasks per FP8 weight (the FP8 data and a ``*_scale_inv`` - scale tensor). + A destination that requests logical weights consumes standard Bridge + tasks. Otherwise, keep Bridge's physical FP8 data and scale tasks. """ # Deferred import to avoid circular import issues. from nemo_rl.models.megatron.draft import draft_model_detached with draft_model_detached([self.model]): - if self._is_fp8_export(): - return self.megatron_bridge._model_bridge.build_export_fp8_tasks( - self.megatron_bridge.hf_pretrained, [self.model] - ) + if self._is_fp8_export() and not self._uses_logical_refit_payload(): + return self.megatron_bridge.get_export_fp8_tasks(self.model) return [ task for task in self.megatron_bridge.get_conversion_tasks([self.model]) @@ -2488,6 +2686,24 @@ def calculate_size_in_bytes(param, tp_size, ep_size): ) return param_info + @staticmethod + def _iter_logical_refit_conversion_tasks( + conversion_tasks: Iterable[Any], + ) -> Iterator[Any]: + """Yield Bridge tasks whose quantized sources are materialized as BF16.""" + for task in conversion_tasks: + if task is None or task.param_weight is None: + yield task + continue + source = _get_refit_task_source(task) + assert source is not None + logical_weight = _dequantize_refit_source(source) + yield ( + replace(task, param_weight=logical_weight) + if logical_weight is not task.param_weight + else task + ) + def _iter_params_with_optional_kv_scales( self, kv_scales: Optional[dict[str, float]] = None, @@ -2515,10 +2731,19 @@ def _iter_params_with_optional_kv_scales( # Default to the full conversion tasks conversion_tasks = self.refit_conversion_tasks + # Megatron generation consumes transient logical BF16, matching MCore's + # native refit wire format; Bridge's training FP8 is not inference MXFP8. + # Other backends keep Bridge's physical FP8 payload and scale_inv sibling; + # mixing that scale with BF16 would corrupt the imported weight. + if self._uses_logical_refit_payload(): + conversion_tasks = self._iter_logical_refit_conversion_tasks( + conversion_tasks + ) + base_iter = self.megatron_bridge.export_hf_weights( [self.model], show_progress=False, - conversion_tasks=conversion_tasks, # used for metadata caching + conversion_tasks=conversion_tasks, ) # Yield the original parameters first. @@ -2568,7 +2793,45 @@ def _iter_params_with_optional_kv_scales( ).reshape(1) yield param_name, scale_tensor - def _iter_local_hf_param_shards(self) -> Iterator[tuple[str, torch.Tensor]]: + def _local_refit_source_spec( + self, tensor: torch.Tensor, spec: Any + ) -> LocalParamSpec: + """Build a live source spec for a BF16 or TE-quantized parameter.""" + if not _is_quantized_refit_source(tensor): + return LocalParamSpec(base=spec.select(tensor)) + + return LocalParamSpec(base=_QuantizedRefitSource(tensor, spec)) + + def _materialize_local_refit_spec( + self, + spec: LocalParamSpec, + logical_source_cache: dict[int, torch.Tensor], + ) -> RefitCtx: + """Materialize one local source, reusing quantized-source dequantization within a layer.""" + base = spec.base + if isinstance(base, _QuantizedRefitSource): + source_id = id(base.tensor) + logical = logical_source_cache.get(source_id) + if logical is None: + logical = _dequantize_refit_source(base.tensor) + logical_source_cache[source_id] = logical + return RefitCtx(buf=base.spec.select(logical).contiguous()) + if isinstance(base, _GroupedRefitSource): + return RefitCtx( + buf=torch.stack( + [ + self._materialize_local_refit_spec( + expert_spec, logical_source_cache + ).buf + for expert_spec in base.specs + ] + ) + ) + if spec.pre is not None: + return spec.pre(base) + return RefitCtx(buf=base) + + def _iter_local_hf_param_shards(self) -> Iterator[tuple[str, LocalParamSpec]]: """Yield (hf_name, local_tp_shard) for this rank's locally owned FFN params. Used by the nccl_reshard_refit bulk path (``build_hf_to_local_param_map``). @@ -2577,67 +2840,40 @@ def _iter_local_hf_param_shards(self) -> Iterator[tuple[str, torch.Tensor]]: and are skipped here (see ``is_nccl_reshard_param``). Unlike ``_iter_params_with_optional_kv_scales`` (PP broadcast + TP gather - via ``export_hf_weights``), this yields TP-local shards directly from the - Megatron params — no collectives. Returned tensors are views and must - not be modified in place. EP: ``refit_conversion_tasks`` already holds - only this rank's local experts; PP non-local params have - ``param_weight is None``. + via ``export_hf_weights``), this yields TP-local source specs directly + from the Megatron params — no collectives. BF16 specs retain live tensor + views; quantized specs materialize logical BF16 during each refit. EP: + ``refit_conversion_tasks`` already holds only this rank's local experts; + PP non-local params have ``param_weight is None``. + + Only a Megatron destination gets the logical-BF16 materialization. Every + other backend keeps Bridge's payload verbatim, which for an FP8 export + task is the physical fp8 view its ``_scale_inv`` sibling describes; + dequantizing it here would ship BF16 bytes under an fp8 scale. """ - from megatron.bridge.models.conversion.param_mapping import ( - FusedExpertMapping, - FusedGatedExpertMapping, - GatedMLPMapping, - ) - - from nemo_rl.weight_sync.nccl_reshard_utils import is_nccl_reshard_param - - def _expert_idx(megatron_name: str) -> str: - # Grouped-GEMM experts are numbered by the megatron param name - m = re.search(r"\d+$", megatron_name) - assert m, f"expected trailing expert index in {megatron_name!r}" - return m.group() - + uses_logical_payload = self._uses_logical_refit_payload() for task in self.refit_conversion_tasks: - local_tensor = task.param_weight # local megatron tensor + if uses_logical_payload: + local_tensor = _get_refit_task_source(task) + else: + local_tensor = task.param_weight if local_tensor is None: - continue # non-local PP rank - # FP8 scale siblings take the misc path. + continue # Non-local PP rank. + # An FP8 export task's scale sibling must not enter the bulk map. + # Bridge wraps it in _HFNameSuffixMapping, which overrides only + # resolve/hf_to_megatron/megatron_to_hf; local_hf_param_specs falls + # through __getattr__ to the base mapping and returns the UNSUFFIXED + # weight name. Without this skip the scale is emitted under the + # weight's key and overwrites it in the param map. if task.global_param_name.endswith("_scale_inv"): continue - if isinstance(task.mapping, GatedMLPMapping): - # FFN gate/up fused in linear_fc1 as [gate_shard; up_shard] (dim 0). - gate, up = torch.chunk(local_tensor, 2, dim=0) - yield task.mapping.hf_param["gate"], gate - yield task.mapping.hf_param["up"], up - continue - - if isinstance(task.mapping, FusedGatedExpertMapping): - # Grouped-GEMM MoE (e.g. Qwen3.5-VL): linear_fc1 fuses gate+up per - # expert [gate; up] (dim 0) — same layout as the dense branch - # above, but the hf_param is a single, index-less string. Un-fuse - # into gate/up AND re-attach the per-expert index. - idx = _expert_idx(task.global_param_name) - prefix = str(task.mapping.hf_param)[: -len(".gate_up_proj")] - gate, up = torch.chunk(local_tensor, 2, dim=0) - yield f"{prefix}.{idx}.gate_proj.weight", gate - yield f"{prefix}.{idx}.up_proj.weight", up - continue - - if isinstance(task.mapping, FusedExpertMapping): - # Grouped-GEMM down (linear_fc2): re-attach the per-expert index + - # ``.weight`` so it matches standard per-expert down_proj. - idx = _expert_idx(task.global_param_name) - prefix = str(task.mapping.hf_param)[: -len(".down_proj")] - yield f"{prefix}.{idx}.down_proj.weight", local_tensor - continue - - # Simple 1:1 mappings: only the FFN down_proj (and any non-gated - # simple gate/up) hits this branch. QKV (a compound mapping) and - # every non-FFN param fall through to misc, so they are skipped. - hf_param = task.mapping.hf_param - if not isinstance(hf_param, dict) and is_nccl_reshard_param(str(hf_param)): - yield str(hf_param), local_tensor + for spec in task.local_hf_param_specs(): + if is_nccl_reshard_param(spec.name): + yield ( + spec.name, + self._local_refit_source_spec(local_tensor, spec), + ) # ------------------------------------------------------------------ # SGLang weight update (colocate IPC + disaggregate broadcast) @@ -2967,12 +3203,13 @@ def _build_layer_to_pp_stage( return layer_to_pp_stage @torch.no_grad() - def prepare_nccl_reshard_refit_info( + def _prepare_source_nccl_reshard_refit_info( self, train_parallelism, gen_parallelism, train_world_size, gen_world_size, + refit_payload_mode: RefitPayloadMode, ): """Prepare per-layer parameter metadata for nccl_reshard-based refit. @@ -2981,11 +3218,7 @@ def prepare_nccl_reshard_refit_info( its own fused layout (e.g., vLLM w13/w2) gen-side, so this train worker stays agnostic to any gen backend's MoE-fusion layout. """ - from nemo_rl.weight_sync.nccl_reshard_utils import ( - build_nccl_reshard_refit_info, - is_nccl_reshard_param, - ) - + self.refit_payload_mode = refit_payload_mode self.refit_param_info_mcore = self._calculate_refit_param_info() # Single pass over Bridge's stream: classify each param as major @@ -3004,11 +3237,6 @@ def prepare_nccl_reshard_refit_info( # state_dict_metadata[hf_name] -> [shape, dtype] # At the same time, filter the params to the misc subset (packed_broadcast path). # misc_meta[hf_name] -> [shape, dtype] - from nemo_rl.weight_sync.nccl_reshard_utils import ( - _extract_layer_name, - _extract_layer_prefix, - ) - # HF layers whose weights come from Megatron's MTP module. The prefix # gate inside is_nccl_reshard_param only catches families whose HF # names keep the bare ``mtp.`` prefix (NemotronH, Qwen3.5); DeepSeek @@ -3016,6 +3244,9 @@ def prepare_nccl_reshard_refit_info( # the only reliable signal. vLLM keeps the MTP drafter separate from # the main model and updates it through load_weights -> misc path. mtp_hf_layers_names = _collect_mtp_hf_layer_names(self.refit_conversion_tasks) + local_refit_hf_names = _collect_local_refit_hf_names( + self.refit_conversion_tasks + ) layer_prefix = None with _meta_tensor_alloc_context(): @@ -3029,6 +3260,7 @@ def prepare_nccl_reshard_refit_info( # nccl-reshard path; everything else -> misc (packed_broadcast). if ( is_nccl_reshard_param(name) + and name in local_refit_hf_names and _extract_layer_name(name) not in mtp_hf_layers_names ): state_dict_metadata[name] = meta @@ -3057,9 +3289,12 @@ def prepare_nccl_reshard_refit_info( pp_size = train_parallelism.get("pp_size", 1) # Construct a dict[layer_name:str] -> pp_stage:int. layer_to_pp_stage = None - assert layer_prefix is not None, "layer_prefix is not set" if pp_size > 1: - layer_to_pp_stage = self._build_layer_to_pp_stage(pp_size, layer_prefix) + layer_to_pp_stage = ( + self._build_layer_to_pp_stage(pp_size, layer_prefix) + if layer_prefix is not None + else {} + ) # The key metadata, which should shared with generation workers self.nccl_reshard_refit_info = build_nccl_reshard_refit_info( @@ -3100,11 +3335,59 @@ def _task_is_misc(task) -> bool: return self.nccl_reshard_refit_info + def prepare_nccl_reshard_refit_info( + self, + train_parallelism: Optional[dict[str, Any]] = None, + gen_parallelism: Optional[dict[str, Any]] = None, + train_world_size: Optional[int] = None, + gen_world_size: Optional[int] = None, + *, + refit_info: Optional[dict[str, Any]] = None, + refit_payload_mode: RefitPayloadMode = "bridge_export", + ) -> Optional[dict[str, Any]]: + """Prepare NCCL-reshard state for the worker's explicit refit role.""" + if self.refit_role == "destination": + if refit_info is None: + raise ValueError("Destination NCCL refit requires refit_info.") + if any( + value is not None + for value in ( + train_parallelism, + gen_parallelism, + train_world_size, + gen_world_size, + ) + ): + raise ValueError( + "Destination NCCL refit does not accept source parallelism arguments." + ) + self._prepare_destination_nccl_reshard_refit_info(refit_info) + return None + + if refit_info is not None: + raise ValueError("Source NCCL refit does not accept refit_info.") + if ( + train_parallelism is None + or gen_parallelism is None + or train_world_size is None + or gen_world_size is None + ): + raise ValueError( + "Source NCCL refit requires train/gen parallelism and world sizes." + ) + return self._prepare_source_nccl_reshard_refit_info( + train_parallelism, + gen_parallelism, + train_world_size, + gen_world_size, + refit_payload_mode, + ) + def _build_expert_groups(self, param_map): - """Group this rank's local expert params into stack-ready views. + """Group this rank's local expert params into stack-ready source specs. Keyed by (prefix, proj_type) and resolved to ordered ``param_map`` - views ready for ``torch.stack``. + specs ready for per-refit materialization and ``torch.stack``. Megatron exposes each expert's projection as a separate param; this bins them so ``_group_experts`` can stack a layer's experts into one grouped @@ -3118,23 +3401,19 @@ def _build_expert_groups(self, param_map): * group 3 = proj type -> ``"gate_proj"`` so the name keys into ``("model.layers.3.mlp.experts", "gate_proj")``. - Returns ``{(prefix, proj): [tensor_0, tensor_1, ...]}`` — the per-expert - ``param_map`` views sorted by expert index. Example — a layer with 2 + Returns ``{(prefix, proj): [spec_0, spec_1, ...]}`` — the per-expert + ``param_map`` specs sorted by expert index. Example — a layer with 2 local experts (gated MoE) yields three keys: - ``(".../experts", "gate_proj"): [view(expert 0), view(expert 1)]`` - ``(".../experts", "up_proj") : [view(expert 0), view(expert 1)]`` - ``(".../experts", "down_proj"): [view(expert 0), view(expert 1)]`` - - Resolving names → views here (rather than per refit in ``_group_experts``) - costs nothing extra — ``param_map`` already owns these views and they - stay valid across refits (weights are updated in place; the name→view - mapping is stable), so ``_group_experts`` only has to ``torch.stack``. - The index sort matters: the views are stacked in this order, so expert 0 + ``(".../experts", "gate_proj"): [spec(expert 0), spec(expert 1)]`` + ``(".../experts", "up_proj") : [spec(expert 0), spec(expert 1)]`` + ``(".../experts", "down_proj"): [spec(expert 0), spec(expert 1)]`` + + Resolving names → specs here costs nothing extra and stays valid across + refits because weights are updated in place. The index sort matters: + source tensors are stacked in this order, so expert 0 must precede expert 1 to match the EP ``Shard(0)`` layout the gen side expects. """ - from nemo_rl.weight_sync.nccl_reshard_utils import _INDIVIDUAL_EXPERT_RE - index_groups: dict[tuple[str, str], list[tuple[int, str]]] = {} for name in param_map: # find all the expert params @@ -3152,26 +3431,25 @@ def _build_expert_groups(self, param_map): } def _group_experts(self, proj, grouped_name, expert_groups): - """Stack this rank's local experts for one projection into ``[E_local, ...]``. - - Using the pre-calculated ``expert_groups`` (from ``_build_expert_groups``) - it is just calling torch.stack of all the local expert params. - """ + """Describe local experts that must be stacked into ``[E_local, ...]``.""" prefix = grouped_name.rsplit(f".{proj}.weight", 1)[0] - expert_tensors = expert_groups.get((prefix, proj)) - assert expert_tensors, ( + expert_specs = expert_groups.get((prefix, proj)) + assert expert_specs, ( f"no local experts for {grouped_name!r} (proj={proj!r}); " "PP-filter / expert-group-metadata inconsistency" ) - return torch.stack(expert_tensors) + return _GroupedRefitSource(tuple(expert_specs)) - def build_hf_to_local_param_map(self, refit_info: dict) -> HFToLocalParamMap: + def _build_source_hf_to_local_param_map( + self, refit_info: dict[str, Any] + ) -> HFToLocalParamMap: """Build the Megatron-backend ``hf_to_local_param_map`` (HFToLocalParamMap). Wraps this rank's local Megatron shards into ``LocalParamSpec``s: - direct: ``base`` is sharded local tensor view, sent as-is. - - grouped MoE expert: ``pre`` stacks the per-expert views into - ``[E_local, ...]`` fresh each refit via ``_group_experts``. + - quantized: ``base`` defers logical view materialization until refit. + - grouped MoE expert: ``base`` holds the ordered per-expert specs, which + are materialized and stacked into ``[E_local, ...]`` each refit. """ # This rank's local TP/EP HF param shards (live views), and the # per-expert views grouped for torch.stack. Build-time only. @@ -3179,23 +3457,44 @@ def build_hf_to_local_param_map(self, refit_info: dict) -> HFToLocalParamMap: expert_groups = self._build_expert_groups(param_map) def _expert_spec(proj, grouped_name): - def pre(_base): - return RefitCtx( - buf=self._group_experts(proj, grouped_name, expert_groups) - ) - - return LocalParamSpec(base=None, pre=pre) + return LocalParamSpec( + base=self._group_experts(proj, grouped_name, expert_groups) + ) mapping = {} for layer_name in refit_info["layer_names"]: for p in refit_info["per_layer_params"][layer_name]: + if p.get("pp_stage", 0) != self.my_pp_stage: + continue name = p["name"] if p.get("grouped_expert_proj"): mapping[name] = _expert_spec(p["grouped_expert_proj"], name) else: - mapping[name] = LocalParamSpec(base=param_map.get(name)) + spec = param_map.get(name) + if spec is None: + raise RuntimeError( + f"No local Megatron refit source maps to {name!r}." + ) + mapping[name] = spec return HFToLocalParamMap(specs=mapping) + def build_hf_to_local_param_map( + self, + refit_info: dict[str, Any], + *, + destination_tasks: Optional[list[Any]] = None, + ) -> HFToLocalParamMap: + """Build the local map for this worker's source or destination role.""" + if self.refit_role == "destination": + if destination_tasks is None: + _, destination_tasks = self._build_generation_refit_tasks() + return self._build_destination_hf_to_local_param_map( + refit_info, destination_tasks + ) + if destination_tasks is not None: + raise ValueError("Source refit does not accept destination_tasks.") + return self._build_source_hf_to_local_param_map(refit_info) + async def nccl_reshard_refit(self, kv_scales=None, refit_timeout_s=None): """Run the refit off this actor's event loop; see _nccl_reshard_refit_guarded. @@ -3250,78 +3549,96 @@ def _nccl_reshard_refit_guarded(self, kv_scales=None, refit_timeout_s=None): RefitAbortWatchdog, ) - groups = [self.pp_comm_group, self.model_update_group] - with RefitAbortWatchdog(groups, refit_timeout_s) as guard: - self._nccl_reshard_refit( - kv_scales=kv_scales, refit_timeout_s=refit_timeout_s + if self.refit_role == "destination": + if kv_scales is not None: + raise ValueError("Destination NCCL refit does not accept kv_scales.") + generation_groups = ( + getattr(self, "_generation_nccl_reshard_groups", None) or {} ) + groups = [*generation_groups.values(), self.model_update_group] + else: + groups = [self.pp_comm_group, self.model_update_group] + + with RefitAbortWatchdog(groups, refit_timeout_s) as guard: + if self.refit_role == "destination": + result = self._destination_nccl_reshard_refit( + refit_timeout_s=refit_timeout_s + ) + else: + result = self._nccl_reshard_refit( + kv_scales=kv_scales, refit_timeout_s=refit_timeout_s + ) if guard.fired: # The aborted transfer returned cleanly, so this is the only signal there is. raise RefitAborted( f"refit nccl_reshard exceeded {refit_timeout_s}s and was aborted; " "a generation rank most likely stopped participating" ) + return result def _nccl_reshard_refit(self, kv_scales=None, refit_timeout_s=None): # hf_to_local_param_map is built once in prepare_nccl_reshard_refit_info; # weight values change but the name → spec mapping is stable across # refits. from nemo_rl.distributed.refit_watchdog import sync_stream_within + + # Keep this local because xferdtensor probes optional NCCL M-to-N bindings. from nemo_rl.weight_sync.xferdtensor import DTensorRef, xferdtensor - # spec.pre (grouped-MoE expert stacking) and spec.post enqueue on this - # worker's current stream; xferdtensor should use the same stream. + # MXFP8 source dequantization, grouped-MoE stacking, and spec.post enqueue + # on this worker's current stream; xferdtensor uses the same stream. nccl_reshard_stream = torch.cuda.current_stream() for layer_name in self.nccl_reshard_refit_info["layer_names"]: - for param_info in self.nccl_reshard_refit_info["per_layer_params"][ - layer_name - ]: - # Each train worker handles only its own PP stage's params - # (non-PP = every param is in pp_stage 0). - if param_info.get("pp_stage", 0) != self.my_pp_stage: - continue - group = self.pp_comm_group + # Gate/up and grouped expert specs in one logical layer can share a + # training parameter. Keep those materializations only until every + # parameter in the layer has been enqueued, rather than retaining a + # model-sized BF16 cache for the full refit. + logical_source_cache: dict[int, torch.Tensor] = {} + try: + for param_info in self.nccl_reshard_refit_info["per_layer_params"][ + layer_name + ]: + # Each train worker handles only its own PP stage's params + # (non-PP = every param is in pp_stage 0). + if param_info.get("pp_stage", 0) != self.my_pp_stage: + continue + group = self.pp_comm_group - spec = self.hf_to_local_param_map.get(param_info["name"]) - assert spec is not None, ( - f"no spec for {param_info['name']!r} in hf_to_local_param_map" - ) - # pre stacks grouped MoE experts fresh each refit; a direct - # param sends its live TP/EP-local view as-is. - ctx = ( - spec.pre(spec.base) # stack grouped MoE experts - if spec.pre is not None - else RefitCtx(buf=spec.base) # send local shard as-is - ) - assert ctx.buf is not None, ( - f"no local tensor for {param_info['name']!r}" - ) - src_tensor = DTensorRef( - local_tensor=ctx.buf, global_shape=param_info["global_shape"] - ) - xferdtensor( - src_tensor, - param_info["src_mesh_info"], - param_info["src_placements"], - None, - param_info["dst_mesh_info"], - param_info["dst_placements"], - group, - nccl_reshard_stream, - ) - if spec.post is not None: - spec.post(ctx) - # Drop refs to the per-iteration grouped MoE tensor so its CUDA - # memory returns to the caching allocator - del ctx, src_tensor + spec = self.hf_to_local_param_map.get(param_info["name"]) + assert spec is not None, ( + f"no spec for {param_info['name']!r} in hf_to_local_param_map" + ) + ctx = self._materialize_local_refit_spec(spec, logical_source_cache) + assert ctx.buf is not None, ( + f"no local tensor for {param_info['name']!r}" + ) + src_tensor = DTensorRef( + local_tensor=ctx.buf, global_shape=param_info["global_shape"] + ) + xferdtensor( + src_tensor, + param_info["src_mesh_info"], + param_info["src_placements"], + None, + param_info["dst_mesh_info"], + param_info["dst_placements"], + group, + nccl_reshard_stream, + ) + if spec.post is not None: + spec.post(ctx) + # Drop refs to per-param views and grouped tensors promptly. + del ctx, src_tensor + finally: + # Never retain stale BF16 materializations across layers or + # optimizer steps. + logical_source_cache.clear() sync_stream_within( nccl_reshard_stream, refit_timeout_s, "the bulk parameter transfer" ) torch.cuda.empty_cache() - import time - misc_t0 = time.perf_counter() self._broadcast_misc_params_packed(kv_scales=kv_scales) sync_stream_within( @@ -3379,6 +3696,11 @@ def prepare_for_lp_inference(self, keep_train_buffers: bool = False) -> None: the accumulated gradients survive. An mcore change that dropped that guard, or that zeroed unconditionally, would silently reinstate this bug. + + MXFP8 overlap aliases the parameter all-gather buffer to grad + storage. That allocation remains resident even when + ``keep_train_buffers`` is false because logprob forward pre-hooks + may need it for parameter all-gather. """ # First worker call after the controller's select() wait returns, so this # is also the "idle wait end" marker. Whether the offload was suppressed @@ -3390,10 +3712,21 @@ def prepare_for_lp_inference(self, keep_train_buffers: bool = False) -> None: keep_train_buffers, ) self._log_gpu_mem("lp_prep_enter") - self.model = self.move_model(self.model, "cuda", move_grads=False) + uses_mxfp8_shared_buffer = self._uses_mxfp8_overlap_shared_param_buffer() + self.model = self.move_model( + self.model, "cuda", move_grads=uses_mxfp8_shared_buffer + ) + if ( + uses_mxfp8_shared_buffer + and not keep_train_buffers + and self.optimizer is not None + ): + # reload_from_cpu() zeros the shared storage. Restage optimizer + # masters before a logprob forward gathers parameters from it. + self._copy_main_params_to_param_buffer(zero_grad_buffer=True) self.model.eval() - if not keep_train_buffers: + if not keep_train_buffers and not uses_mxfp8_shared_buffer: # offload grads to cpu self.model = self.move_model( self.model, "cpu", move_params=False, move_grads=True @@ -3499,13 +3832,39 @@ def _clear_fp8_caches(self): @wrap_with_nvtx_name("megatron_policy_worker/offload_before_refit") def offload_before_refit(self): - """Offload the optimizer and buffers to the CPU.""" + """Offload optimizer state and buffers that are safe to release.""" # An in-flight async checkpoint keeps references to the CUDA tensors in # its sharded state dict until the write is finalized. Offloading swaps # those tensors for CPU storage, so the checkpoint references would keep # the old CUDA storage alive and defeat the offload. self.finalize_async_save() + # With MXFP8 overlap, the optimizer updates FP32 master shards and the + # next parameter all-gather requantizes them into the model weights. A + # refit happens between optimizer steps, before that next training + # forward, so force the gather now. This both gives generation the + # latest weights and leaves hooks disabled while the shared param/grad + # buffer is held across refit. The normal train-step transition + # re-enables them. + # Deliberately conditional on the hooks being enabled. Every state in + # which they are already off is one where the weights are current + # anyway: before the first train step the buffer holds the checkpoint; + # eval entry already forced a sync via disable_forward_pre_hook( + # param_sync=True) and runs no optimizer step; a skipped step leaves the + # masters unchanged; and a successful step re-enables the hooks before + # returning. If a stale case is ever found, note that staging alone does + # NOT fix it - with reuse_grad_buf_for_mxfp8_param_ag param_data aliases + # grad_data, so zero_grad_buffer() wipes the parameters and + # _copy_main_params_to_param_buffer restores only this rank's shard, + # leaving every other DP rank at zero. Upstream pairs that staging with a + # following start_param_sync (DistributedOptimizer. + # prepare_model_params_for_param_sync); any fix needs the sync too. + if ( + self._uses_mxfp8_overlap_shared_param_buffer() + and self._forward_pre_hook_enabled() + ): + self._disable_forward_pre_hook_until_next_train_step(param_sync=True) + no_grad = torch.no_grad() no_grad.__enter__() allocated = torch.cuda.memory_allocated() / (1024**3) # Convert to GB @@ -3513,9 +3872,17 @@ def offload_before_refit(self): print( f"GPU Memory before optimizer offload: {allocated:.2f}GB allocated, {reserved:.2f}GB reserved" ) + # MXFP8 overlap aliases parameter all-gather and gradient storage. Keep + # that allocation resident: resizing it out from under the persistent + # autograd/DDP views makes the next backward accumulate into invalid + # storage. Ordinary independent grad buffers remain safe to offload. + keep_shared_buffer = self._uses_mxfp8_overlap_shared_param_buffer() self.model = self.move_model( - self.model, "cpu", move_params=False, move_grads=True - ) # get rid of grad buffers + self.model, + "cpu", + move_params=False, + move_grads=not keep_shared_buffer, + ) # When True, clear Transformer Engine's per-module _fp8_workspaces scratch # buffers in offload_before_refit (before weight transfer to the inference @@ -3605,6 +3972,7 @@ def offload_after_refit(self): no_grad = torch.no_grad() no_grad.__enter__() + keep_shared_buffer = self._uses_mxfp8_overlap_shared_param_buffer() # Non-reshard colocated serves both models from the same param buffers. generation_cfg = self.cfg.get("generation") keep_params_for_generation = ( @@ -3619,7 +3987,10 @@ def offload_after_refit(self): # before offloading the model parameters (including self.A_log). self.model.eval() self.model = self.move_model( - self.model, "cpu", move_params=not keep_params_for_generation + self.model, + "cpu", + move_params=not (keep_shared_buffer or keep_params_for_generation), + move_grads=not keep_shared_buffer, ) torch.randn(1).cuda() # wake up torch allocator self.offload_before_refit() # rerun the old offload function diff --git a/nemo_rl/weight_sync/collective_weight_synchronizer.py b/nemo_rl/weight_sync/collective_weight_synchronizer.py index d3598361210..160c6b9cc16 100644 --- a/nemo_rl/weight_sync/collective_weight_synchronizer.py +++ b/nemo_rl/weight_sync/collective_weight_synchronizer.py @@ -181,7 +181,9 @@ def init_communicator(self) -> None: # prepare_refit_info is called before init_collective. This matches # distillation.py ordering. Neither call depends on the other today, # but we document this as the canonical ordering for future reference. - state_dict_info = self._policy.prepare_refit_info() + state_dict_info = self._policy.prepare_refit_info( + refit_payload_mode=self._generation.get_refit_payload_mode() + ) self._generation.prepare_refit_info(state_dict_info) ip, port = self._train_cluster.get_master_address_and_port() diff --git a/nemo_rl/weight_sync/factory.py b/nemo_rl/weight_sync/factory.py index 006aabc1bc9..1ff689cd16b 100644 --- a/nemo_rl/weight_sync/factory.py +++ b/nemo_rl/weight_sync/factory.py @@ -84,7 +84,13 @@ def create_weight_synchronizer( f"Supported backends: {sorted(_SUPPORTED_BACKENDS)}" ) - checkpoint_engine_config = checkpoint_engine_refit_config(generation.cfg) + # Megatron owns its refit selectors (including "mcore"); the vLLM-oriented + # checkpoint-engine normalization rejects that valid Megatron value. + checkpoint_engine_config = ( + None + if generation_backend == MEGATRON_BACKEND + else checkpoint_engine_refit_config(generation.cfg) + ) if checkpoint_engine_config is not None: if colocated: raise ValueError( diff --git a/nemo_rl/weight_sync/megatron_weight_synchronizer.py b/nemo_rl/weight_sync/megatron_weight_synchronizer.py index dab017fce56..6d15a97297a 100644 --- a/nemo_rl/weight_sync/megatron_weight_synchronizer.py +++ b/nemo_rl/weight_sync/megatron_weight_synchronizer.py @@ -18,7 +18,13 @@ import ray from nemo_rl.utils.timer import Timer +from nemo_rl.weight_sync.collective_weight_synchronizer import ( + CollectiveWeightSynchronizer, +) from nemo_rl.weight_sync.interfaces import WeightSynchronizer +from nemo_rl.weight_sync.nccl_reshard_weight_synchronizer import ( + NcclReshardWeightSynchronizer, +) class MegatronWeightSynchronizer(WeightSynchronizer): @@ -31,11 +37,8 @@ class MegatronWeightSynchronizer(WeightSynchronizer): transfer, but one the worker performs internally on wake. Sync therefore reduces to dropping training-only buffers and re-entering inference mode. - Non-colocated adds the cross-group collective: the training and inference - workers are disjoint actor groups that rendezvous in mcore's - reshard-capable weight swap, with the engine suspended around the - transfer. That wiring (a joint refit process group over the configured - copy-service backend) is established once in ``init_communicator``. + Non-colocated generation keeps the Megatron engine lifecycle here and + delegates the transfer to native MCore refit, packed collective, or M2N. """ def __init__( @@ -58,6 +61,34 @@ def __init__( self._train_cluster = train_cluster self._inference_cluster = inference_cluster self._refit_backend: Optional[str] = None + self._transport: Optional[WeightSynchronizer] = None + if colocated: + # Colocated refit always uses the in-place wake-reshard, so + # any other transport is inert. Reject rather than silently ignoring it: + # a user asking for packed collective refit would otherwise get the + # native one with no indication their setting did nothing. + if generation.cfg.get("refit_transport") != "mcore": + raise ValueError( + "policy.generation.refit_transport must be 'mcore' with " + "colocated Megatron generation, which always uses the in-place " + "wake-reshard. Set colocated.enabled=false to use the packed " + "collective or nccl_reshard transport." + ) + if not colocated and not generation.uses_native_refit: + if generation.cfg.get("refit_transport") == "nccl_reshard": + self._transport = NcclReshardWeightSynchronizer( + policy=policy, + generation=generation, + train_cluster=train_cluster, + inference_cluster=inference_cluster, + ) + else: + self._transport = CollectiveWeightSynchronizer( + policy=policy, + generation=generation, + train_cluster=train_cluster, + inference_cluster=inference_cluster, + ) self._stale = True def init_communicator(self) -> None: @@ -68,6 +99,9 @@ def init_communicator(self) -> None: """ if self._colocated: return + if self._transport is not None: + self._transport.init_communicator() + return ip, port = self._train_cluster.get_master_address_and_port() print(f"Using ip: {ip}, port: {port} for collective communication", flush=True) train_world_size = self._train_cluster.world_size() @@ -75,11 +109,15 @@ def init_communicator(self) -> None: self._refit_backend = self._generation.cfg["mcore_generation_config"][ "refit_backend" ] + refit_execution_batch_bytes = self._generation.cfg["mcore_generation_config"][ + "refit_execution_batch_bytes" + ] futures_train = self._policy.init_collective_mcore_generation( ip, port, world_size, rank_offset=0, + refit_execution_batch_bytes=refit_execution_batch_bytes, refit_backend=self._refit_backend, ) futures_inference = self._generation.init_collective( @@ -125,17 +163,20 @@ def sync_weights( else nullcontext() ) with timer_context: - futures_train = self._policy.swap_weights_via_reshard(is_source=True) - futures_inference = self._generation.update_weights_from_collective() - ray.get(futures_train) - results = ray.get(futures_inference) - if not all(result for result in results if result is not None): - raise RuntimeError( - "❌ Error: Updating weights for the generation policy failed " - "during refit.\nThis often indicates an issue with the " - "refit copy service or a problem within the generation " - "backend.\n" - ) + if self._transport is not None: + self._transport.sync_weights(kv_scales=kv_scales) + else: + futures_train = self._policy.swap_weights_via_reshard(is_source=True) + futures_inference = self._generation.update_weights_from_collective() + ray.get(futures_train) + results = ray.get(futures_inference) + if not all(result for result in results if result is not None): + raise RuntimeError( + "❌ Error: Updating weights for the generation policy failed " + "during refit.\nThis often indicates an issue with the " + "refit copy service or a problem within the generation " + "backend.\n" + ) self._generation.prepare_for_generation(tags=["kv_cache"]) self._generation.resume_after_refit() @@ -147,4 +188,6 @@ def is_stale(self) -> bool: return self._stale def shutdown(self) -> None: - """Nothing to tear down; the collective lives in the worker groups.""" + """Release any resources owned by the delegated transport.""" + if self._transport is not None: + self._transport.shutdown() diff --git a/nemo_rl/weight_sync/nccl_reshard_utils.py b/nemo_rl/weight_sync/nccl_reshard_utils.py index 1c30d75ca95..049fd3dbbe6 100644 --- a/nemo_rl/weight_sync/nccl_reshard_utils.py +++ b/nemo_rl/weight_sync/nccl_reshard_utils.py @@ -345,18 +345,16 @@ def build_mesh_info( **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: + Callers build separate non-expert (TP) and expert (EP/ETP) meshes. An + expert mesh may activate both EP and expert TP. The active axes use + Megatron's rank order, with TP innermost and EP immediately outside it: * **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.) + * **EP** in an EP-only expert mesh: ``global_rank % ep_size`` recovers + the EP coord. + * **ETP + EP** in a generation expert mesh: ETP is the innermost axis, + matching Megatron-Core's expert rank generator. Returns: ``(MeshInfo, dim_map)`` where *dim_map* maps ``"tp"``/``"ep"``/``"dp"``/``"pp"`` @@ -387,7 +385,7 @@ def 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. + Expert params shard dim 0 on EP; their ETP shard dims are shifted by +1. """ num_mesh_dims = len(dim_map) or 1 placements = [Replicate() for _ in range(num_mesh_dims)] @@ -399,12 +397,12 @@ def get_placements(param_name: str, dim_map: dict, ndim: int) -> list: if is_expert_param(param_name): if "ep" in dim_map: placements[dim_map["ep"]] = Shard(0) - else: - # We currently never shards both EP and TP for the expert params - # so far. If MCore etp is enabled, the logic should be changed - tp_dim = _get_expert_tp_shard_dim(param_name) - if tp_dim is not None and "tp" in dim_map: - placements[dim_map["tp"]] = Shard(tp_dim + 1) + # The grouped HF tensor has a leading expert dimension, so the usual + # projection shard dimension is shifted by one for ETP. This remains + # independent of EP's Shard(0) when both axes are active. + tp_dim = _get_expert_tp_shard_dim(param_name) + if tp_dim is not None and "tp" in dim_map: + placements[dim_map["tp"]] = Shard(tp_dim + 1) else: # for non-expert params tp_dim = get_tp_shard_dim(param_name) @@ -592,6 +590,10 @@ def check_nccl_reshard_refit_support(master_config: Any) -> None: dtensor_cfg = policy.get("dtensor_cfg", {}) or {} vllm_cfg = generation.get("vllm_cfg", {}) or {} vllm_kwargs = generation.get("vllm_kwargs", {}) or {} + mcore_generation_cfg = { + **megatron_cfg, + **(generation.get("mcore_generation_config", {}) or {}), + } violations: list[str] = [] @@ -602,14 +604,14 @@ def check_nccl_reshard_refit_support(master_config: Any) -> None: "(nccl_reshard_refit is only for disaggregated train/gen)." ) - # Gen backend = vLLM — only backend with prepare_nccl_reshard_refit_info. + # Generation backends with a destination-side HFToLocalParamMap adapter. backend = generation.get("backend") - if backend != "vllm": + if backend not in ("vllm", "megatron"): violations.append( - f"policy.generation.backend must be 'vllm' (got {backend!r})." + f"policy.generation.backend must be 'vllm' or 'megatron' (got {backend!r})." ) - if vllm_kwargs.get("enable_eplb"): + if backend == "vllm" and vllm_kwargs.get("enable_eplb"): violations.append( "policy.generation.vllm_kwargs.enable_eplb must be False " "(nccl_reshard_refit fixes the expert->rank mapping at setup; " @@ -634,9 +636,8 @@ def check_nccl_reshard_refit_support(master_config: Any) -> None: "layerwise-reload weight loaders that ModelOpt real quant requires)." ) - # This initial version supports only the Megatron train + vLLM gen - # combination; the DTensor train backend refit path is intentionally - # dropped (Megatron is the path validated end-to-end). + # Only Megatron training currently provides the local Bridge source views; + # DTensor training is not supported by this transport yet. megatron_enabled = megatron_cfg.get("enabled", False) dtensor_enabled = dtensor_cfg.get("enabled", False) if not megatron_enabled: @@ -651,15 +652,6 @@ def check_nccl_reshard_refit_support(master_config: Any) -> None: ) if megatron_enabled: - etp = megatron_cfg.get("expert_tensor_parallel_size", 1) - - # ETP is not supported yet. - if etp not in (1, None): - violations.append( - f"Megatron expert_tensor_parallel_size is not supported yet " - f"(got etp={etp})." - ) - # PP-layout knobs that _build_layer_to_pp_stage doesn't yet handle. if megatron_cfg.get("pipeline_model_parallel_layout") is not None: violations.append( @@ -680,7 +672,11 @@ def check_nccl_reshard_refit_support(master_config: Any) -> None: "policy.megatron_cfg.account_for_loss_in_pipeline_split must be False." ) - # Precision compatibility (train ↔ gen). Supported combinations: + # Precision compatibility (train ↔ gen). vLLM supports byte-compatible + # BF16/BF16 and blockwise-FP8/FP8, plus receiver-side BF16-to-MXFP8. + # Megatron generation sends logical BF16 weights from BF16 or TE-quantized + # training storage, then writes BF16 or performs on-receive MXFP8 + # quantization at the destination. # BF16 train ↔ BF16 gen (default, tested) # FP8 train ↔ FP8 gen (fp8_param=True + blockwise + vllm precision=fp8) # BF16 storage → MXFP8 gen (receiver quantizes the resharded BF16 shard) @@ -697,7 +693,13 @@ def check_nccl_reshard_refit_support(master_config: Any) -> None: # bytes and deadlock/corrupt the bulk collective; reject anything outside # the supported set up front (this also catches typos such as # "fp8_e4m3" that would otherwise skip the FP8 checks below). - if gen_precision not in (None, "auto", "bf16", "bfloat16", "fp8"): + if backend == "vllm" and gen_precision not in ( + None, + "auto", + "bf16", + "bfloat16", + "fp8", + ): violations.append( f"policy.generation.vllm_cfg.precision={gen_precision!r} is not " "supported by nccl_reshard_refit (use 'bf16'/'bfloat16', 'fp8', " @@ -705,7 +707,7 @@ def check_nccl_reshard_refit_support(master_config: Any) -> None: "gen dtype must match the train dtype." ) - if gen_precision == "fp8": + if backend == "vllm" and gen_precision == "fp8": if fp8_param: if vllm_cfg.get("is_mx"): violations.append( @@ -736,17 +738,82 @@ def check_nccl_reshard_refit_support(master_config: Any) -> None: "policy.megatron_cfg.fp8_cfg.fp8_param=True, or " "is_mx=True for BF16-to-MXFP8 refit." ) - elif fp8_param: + elif backend == "vllm" and fp8_param: violations.append( "policy.megatron_cfg.fp8_cfg.fp8_param=True requires " "policy.generation.vllm_cfg.precision='fp8' " "(FP8 storage on train side has no BF16 gen consumer)." ) - # Gen-backend restrictions. The reshard supports gen-side TP, DP, and EP; - # the vLLM backend shards experts by index across its TP ranks, so its EP - # is either 1 (TP-sharded experts) or equal to TP (EP-sharded). PP is not - # yet supported gen-side. + if backend == "megatron": + if policy.get("precision") != "bfloat16": + violations.append( + "policy.precision must be 'bfloat16' for Megatron-generation " + "nccl_reshard refit." + ) + if fp8_param and not fp8_cfg.get("enabled", False): + violations.append( + "policy.megatron_cfg.fp8_cfg.fp8_param=True requires " + "policy.megatron_cfg.fp8_cfg.enabled=True." + ) + gen_pp = mcore_generation_cfg.get("pipeline_model_parallel_size", 1) + if gen_pp != 1: + violations.append( + "policy.generation.mcore_generation_config." + "pipeline_model_parallel_size must be 1 for nccl_reshard refit " + f"(got {gen_pp})." + ) + + gen_fp8_cfg = mcore_generation_cfg.get("fp8_cfg", {}) or {} + if gen_fp8_cfg.get("enabled") and gen_fp8_cfg.get("fp8_recipe") != "mxfp8": + violations.append( + "Megatron-generation nccl_reshard refit supports BF16 or MXFP8 " + "inference weights; policy.generation.mcore_generation_config." + "fp8_cfg.fp8_recipe must be 'mxfp8' when FP8 is enabled." + ) + + # MXFP8 inference quantizes through resolve_mxfp8_backend, which + # only accepts the 'torch' and 'flashinfer' grouped-GEMM backends. + # MCore does reject this pairing at model build, but only against its + # own field names; validating here fails at config time and names the + # NeMo-RL key the user actually sets. + # + # An omitted key is NOT safe: MCore's TransformerConfig default is + # "vllm", which resolve_mxfp8_backend rejects. Resolve the default + # here so the omitted case fails at config time like the explicit one. + gemm_backend = mcore_generation_cfg.get( + "inference_grouped_gemm_backend", "vllm" + ) + if gen_fp8_cfg.get("enabled") and gemm_backend not in ( + "torch", + "flashinfer", + ): + violations.append( + "MXFP8 Megatron generation requires policy.generation." + "mcore_generation_config.inference_grouped_gemm_backend to be " + f"'torch' or 'flashinfer' (got {gemm_backend!r}; MCore's " + "default is 'vllm', so this key must be set explicitly)." + ) + + # _prepare_mxfp8_refit only installs persistent MXFP8 destinations + # for cores whose transformer_impl is 'inference_optimized'. Without + # it the refit silently falls back to copying BF16 into TE FP8 + # params instead of quantizing, so reject the pairing up front. + if ( + gen_fp8_cfg.get("enabled") + and mcore_generation_cfg.get("transformer_impl") + != "inference_optimized" + ): + violations.append( + "MXFP8 Megatron generation requires policy.generation." + "mcore_generation_config.transformer_impl='inference_optimized' " + f"(got {mcore_generation_cfg.get('transformer_impl')!r})." + ) + + # Gen-backend restrictions. The reshard supports gen-side TP, DP, EP, and + # Megatron ETP. The vLLM backend shards experts by index across + # its TP ranks, so its EP is either 1 (TP-sharded experts) or equal to TP + # (EP-sharded). PP is not yet supported gen-side. if generation.get("backend") == "vllm": gen_tp = vllm_cfg.get("tensor_parallel_size", 1) gen_ep = vllm_cfg.get("expert_parallel_size", 1) @@ -781,7 +848,8 @@ def build_nccl_reshard_refit_info( Args: state_dict_metadata: ``{hf_param_name: {"shape": list, "dtype": str}}`` The input ``state_dict_metadata`` has a global view of the parameters - train_parallelism / gen_parallelism: ``{"tp_size", "ep_size", "pp_size"}`` + train_parallelism / gen_parallelism: + ``{"tp_size", "ep_size", "etp_size", "pp_size"}`` train_world_size / gen_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 @@ -804,15 +872,15 @@ def build_nccl_reshard_refit_info( pp_size = train_parallelism.get("pp_size", 1) tp_size = train_parallelism.get("tp_size", 1) ep_size = train_parallelism.get("ep_size", 1) + etp_size = train_parallelism.get("etp_size", 1) use_per_stage = pp_size > 1 if use_per_stage: assert layer_to_pp_stage is not None, ( "layer_to_pp_stage must be provided when pp_size > 1" ) - # Currently we don't support ETP>1 for the nccl_reshard_refit. - # Non-expert params, ranks are partitioned (tp, dp); - # Expert params, ranks are partitioned (ep, edp). + # Non-expert params partition ranks by (tp, dp); expert params partition + # them by (etp, ep, edp), matching MCore's separate expert rank grid. def _build_train_src_meshes(num_gpus: int, rank_offset: int, stage_pp: int): non_expert_mesh, non_expert_dim_map = build_mesh_info( num_gpus, @@ -824,23 +892,22 @@ def _build_train_src_meshes(num_gpus: int, rank_offset: int, stage_pp: int): expert_mesh, expert_dim_map = build_mesh_info( num_gpus, rank_offset=rank_offset, - tp_size=1, + tp_size=etp_size, ep_size=ep_size, pp_size=stage_pp, ) return (non_expert_mesh, non_expert_dim_map), (expert_mesh, expert_dim_map) # Gen (dst) side. Non-expert params are TP-sharded across the gen ranks. - # Expert params follow the gen-side expert-parallel size: - # * gen ep_size == 1 -> experts are TP-sharded like dense (shared mesh). - # * gen ep_size > 1 -> experts are EP-sharded by expert index (Shard(0)) - # over the EP ranks (tp_size=1, ep_size=gen_ep). - # Mirrors the train (src) expert/non-expert mesh split. - # TODO: current logic might be tied to the vllm-backend. The logic might need - # to be revised when we support other gen-backends. - gen_tp = gen_parallelism.get("tp_size", 1) - gen_ep = gen_parallelism.get("ep_size", 1) - gen_pp = gen_parallelism.get("pp_size", 1) + # Expert params follow the generation backend's EP/ETP layout. vLLM uses + # TP for experts when EP=1, whereas Megatron can keep ETP=1 and replicate + # experts across its ordinary TP ranks. + gen_tp = gen_parallelism.get("tp_size", 1) or 1 + gen_ep = gen_parallelism.get("ep_size", 1) or 1 + # Bind to a definite int: _build_dst_meshes closes over this, and a + # narrowed-in-place Optional would still read as int | None inside it. + gen_etp = gen_parallelism.get("etp_size", gen_tp if gen_ep == 1 else 1) + gen_pp = gen_parallelism.get("pp_size", 1) or 1 def _build_dst_meshes(num_gpus: int, rank_offset: int): non_expert = build_mesh_info( @@ -850,16 +917,13 @@ def _build_dst_meshes(num_gpus: int, rank_offset: int): ep_size=1, pp_size=gen_pp, ) - if gen_ep > 1: - expert = build_mesh_info( - num_gpus, - rank_offset=rank_offset, - tp_size=1, - ep_size=gen_ep, - pp_size=gen_pp, - ) - else: - expert = non_expert + expert = build_mesh_info( + num_gpus, + rank_offset=rank_offset, + tp_size=gen_etp, + ep_size=gen_ep, + pp_size=gen_pp, + ) return non_expert, expert if use_per_stage: diff --git a/nemo_rl/weight_sync/nccl_reshard_weight_synchronizer.py b/nemo_rl/weight_sync/nccl_reshard_weight_synchronizer.py index 3ddd36bbb7d..6a759f54f99 100644 --- a/nemo_rl/weight_sync/nccl_reshard_weight_synchronizer.py +++ b/nemo_rl/weight_sync/nccl_reshard_weight_synchronizer.py @@ -14,7 +14,7 @@ """NCCL-xfer (shard-to-shard) weight synchronizer for non-colocated deployments. -Handles disaggregated Megatron-train -> vLLM-gen weight refit via the +Handles disaggregated Megatron-train -> vLLM/Megatron-gen weight refit via the ``xferdtensor`` reshard: bulk FFN/expert params are resharded shard-to-shard between the train and gen parallelism layouts over a dedicated per-PP-stage NCCL communicator, while the remaining "misc" params ride a packed broadcast over the @@ -31,9 +31,8 @@ sync_weights(): policy.nccl_reshard_refit(kv_scales) + generation.nccl_reshard_refit(); verify. -Like the collective transport, this is a pure data mover: policy and generation -run on separate GPU clusters, so the phase transitions (offload / restore) are -owned by the orchestrator, not here. +Like the collective transport, this is a pure data mover. Backend-specific +phase transitions are owned by the caller. """ from collections.abc import Sequence @@ -42,6 +41,7 @@ import ray +from nemo_rl.models.generation.megatron.config import merged_inference_megatron_cfg from nemo_rl.utils.timer import Timer from nemo_rl.weight_sync.interfaces import WeightSynchronizer from nemo_rl.weight_sync.membership import RefitMembership, plan_refit_membership @@ -89,7 +89,7 @@ def _settle_before_propagating(futures, budget_s, what: str) -> None: class NcclReshardWeightSynchronizer(WeightSynchronizer): """Weight synchronizer using the ``xferdtensor`` shard-to-shard reshard. - For non-colocated Megatron-train -> vLLM-gen deployments where weights are + For non-colocated Megatron-train -> vLLM/Megatron-gen deployments where weights are redistributed directly between the two parallelism layouts (bulk path) plus a packed broadcast for the misc params. Mirrors :class:`CollectiveWeightSynchronizer` but additionally bootstraps the @@ -101,7 +101,7 @@ class NcclReshardWeightSynchronizer(WeightSynchronizer): Args: policy: Policy object implementing ColocatablePolicyInterface (Megatron). - generation: Generation object implementing GenerationInterface (vLLM). + generation: Generation object implementing GenerationInterface. train_cluster: RayVirtualCluster for the training workers. Only used by ``init_communicator()``; may be ``None`` for sync-only instances. inference_cluster: RayVirtualCluster for the inference workers. Only @@ -142,19 +142,47 @@ def __init__( def _train_parallelism(self) -> dict[str, int]: megatron_cfg = self._policy.cfg["megatron_cfg"] + tp_size = megatron_cfg.get("tensor_model_parallel_size", 1) + etp_size = megatron_cfg.get("expert_tensor_parallel_size") return { - "tp_size": megatron_cfg.get("tensor_model_parallel_size", 1), + "tp_size": tp_size, "ep_size": megatron_cfg.get("expert_model_parallel_size", 1), + # MCore accepts None and resolves it to TP in ModelParallelConfig; + # normalize at this boundary so the plan builder receives only ints. + "etp_size": tp_size if etp_size is None else etp_size, "pp_size": megatron_cfg.get("pipeline_model_parallel_size", 1), } def _gen_parallelism(self) -> dict[str, int]: - vllm_cfg = self._policy.cfg["generation"].get("vllm_cfg", {}) - return { - "tp_size": vllm_cfg.get("tensor_parallel_size", 1), - "ep_size": vllm_cfg.get("expert_parallel_size", 1), - "pp_size": vllm_cfg.get("pipeline_parallel_size", 1), - } + generation_cfg = self._policy.cfg["generation"] + if generation_cfg["backend"] == "vllm": + vllm_cfg = generation_cfg.get("vllm_cfg", {}) + tp_size = vllm_cfg.get("tensor_parallel_size", 1) + ep_size = vllm_cfg.get("expert_parallel_size", 1) + return { + "tp_size": tp_size, + "ep_size": ep_size, + "etp_size": tp_size if ep_size == 1 else 1, + "pp_size": vllm_cfg.get("pipeline_parallel_size", 1), + } + if generation_cfg["backend"] == "megatron": + # Resolve through the same merge the generation model is built from, + # so the reshard mesh cannot drift from the layout MCore actually + # instantiates (e.g. the inference_optimized ETP pin). + megatron_cfg = merged_inference_megatron_cfg(self._policy.cfg) + tp_size = megatron_cfg["tensor_model_parallel_size"] + etp_size = megatron_cfg.get("expert_tensor_parallel_size") + return { + "tp_size": tp_size, + "ep_size": megatron_cfg["expert_model_parallel_size"], + # Match MCore's effective default: an omitted/None ETP uses TP. + "etp_size": tp_size if etp_size is None else etp_size, + "pp_size": megatron_cfg["pipeline_model_parallel_size"], + } + raise ValueError( + "NCCL M-to-N refit only supports vLLM or Megatron generation, got " + f"{generation_cfg['backend']!r}." + ) def sync_weights( self, @@ -329,6 +357,7 @@ def _build(self, membership: RefitMembership) -> None: gen_parallelism, train_world_size, inference_world_size, + refit_payload_mode=self._generation.get_refit_payload_mode(), ) # nccl_reshard_refit_info holds MeshInfo rank tensors created under diff --git a/research/template_project/configs/grpo_math_1B.yaml b/research/template_project/configs/grpo_math_1B.yaml index e77057620ed..f60ff8e5316 100644 --- a/research/template_project/configs/grpo_math_1B.yaml +++ b/research/template_project/configs/grpo_math_1B.yaml @@ -301,6 +301,8 @@ policy: use_cuda_graphs_for_non_decode_steps: true # Enable CUDA graphs for prefill/context processing enable_chunked_prefill: true # Split long prefills into chunks for better memory management max_tokens: 16384 # Maximum number of tokens to use in a single step. Analogous to vllm's max_num_batched_tokens + refit_backend: "nccl" # 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. vllm_cfg: async_engine: false precision: ${policy.precision} diff --git a/tests/functional/L1_Functional_Tests_GB200_MXFP8.sh b/tests/functional/L1_Functional_Tests_GB200_MXFP8.sh index 34910911f33..f6fd2492173 100644 --- a/tests/functional/L1_Functional_Tests_GB200_MXFP8.sh +++ b/tests/functional/L1_Functional_Tests_GB200_MXFP8.sh @@ -36,6 +36,11 @@ run_test() { run_test uv run --no-sync bash ./tests/functional/grpo_vllm_mxfp8_rollout_gb200.sh +# MXFP8 leg of the Megatron M-to-N reshard refit. Sized for 2 GPUs +# (cluster.gpus_per_node=2, 1 train + 1 gen), so it fits this shard's runner. +# MXFP8 inference is Blackwell-only, which is what this runner provides. +run_test env REFIT_PRECISION=mxfp8 uv run --no-sync bash ./tests/functional/grpo_megatron_nccl_reshard_refit.sh + cd ${PROJECT_ROOT}/tests if compgen -G ".coverage*" > /dev/null; then coverage combine .coverage* diff --git a/tests/functional/L1_Functional_Tests_Megatron_4.sh b/tests/functional/L1_Functional_Tests_Megatron_4.sh index 4c0accc98f1..161e48c4066 100644 --- a/tests/functional/L1_Functional_Tests_Megatron_4.sh +++ b/tests/functional/L1_Functional_Tests_Megatron_4.sh @@ -36,6 +36,7 @@ run_test() { run_test uv run --no-sync bash ./tests/functional/grpo_megatron_generation_colocated.sh run_test uv run --no-sync bash ./tests/functional/grpo_megatron_generation_non_colocated.sh +run_test uv run --no-sync bash ./tests/functional/grpo_megatron_nccl_reshard_refit.sh run_test uv run --no-sync bash ./tests/functional/grpo_megatron_generation_colocated_reshard.sh run_test uv run --no-sync bash ./tests/functional/grpo_megatron_generation_colocated_gym.sh run_test uv run --no-sync bash ./tests/functional/grpo_megatron_generation_non_colocated_gym.sh diff --git a/tests/functional/grpo_megatron_generation_colocated.sh b/tests/functional/grpo_megatron_generation_colocated.sh index b6ed44dc8c3..53e0ac4c6bd 100644 --- a/tests/functional/grpo_megatron_generation_colocated.sh +++ b/tests/functional/grpo_megatron_generation_colocated.sh @@ -42,6 +42,7 @@ uv run coverage run -a --data-file=$PROJECT_ROOT/tests/.coverage --source=$PROJE policy.logprob_batch_size=4 \ policy.train_micro_batch_size=1 \ policy.generation.backend=megatron \ + policy.generation.refit_transport=mcore \ cluster.gpus_per_node=2 \ cluster.segment_size=1 \ grpo.max_num_steps=2 \ diff --git a/tests/functional/grpo_megatron_generation_colocated_gym.sh b/tests/functional/grpo_megatron_generation_colocated_gym.sh index 3326fe2bea3..0d4a01ded3c 100755 --- a/tests/functional/grpo_megatron_generation_colocated_gym.sh +++ b/tests/functional/grpo_megatron_generation_colocated_gym.sh @@ -64,6 +64,7 @@ uv run coverage run -a --data-file=$PROJECT_ROOT/tests/.coverage --source=$PROJE policy.megatron_cfg.context_parallel_size=1 \ policy.megatron_cfg.sequence_parallel=false \ policy.generation.backend=megatron \ + +policy.generation.refit_transport=mcore \ policy.generation.mcore_generation_config.expose_http_server=true \ policy.max_total_sequence_length=512 \ policy.generation.max_new_tokens=128 \ diff --git a/tests/functional/grpo_megatron_generation_colocated_reshard.sh b/tests/functional/grpo_megatron_generation_colocated_reshard.sh index 9af8da7f47d..c2d5ad2163d 100755 --- a/tests/functional/grpo_megatron_generation_colocated_reshard.sh +++ b/tests/functional/grpo_megatron_generation_colocated_reshard.sh @@ -33,6 +33,9 @@ uv run coverage run -a --data-file=$PROJECT_ROOT/tests/.coverage --source=$PROJE policy.train_micro_batch_size=1 \ policy.megatron_cfg.tensor_model_parallel_size=2 \ policy.generation.backend=megatron \ + policy.generation.refit_transport=mcore \ + policy.generation.top_k=1 \ + policy.generation.mcore_generation_config.logprobs_mode=raw_logprobs \ ++policy.generation.mcore_generation_config.transformer_impl=inference_optimized \ ++policy.generation.mcore_generation_config.tensor_model_parallel_size=1 \ policy.generation.mcore_generation_config.refit_backend=nccl \ diff --git a/tests/functional/grpo_megatron_generation_multiturn.sh b/tests/functional/grpo_megatron_generation_multiturn.sh index 7a7867f0964..2c9f4a9a447 100644 --- a/tests/functional/grpo_megatron_generation_multiturn.sh +++ b/tests/functional/grpo_megatron_generation_multiturn.sh @@ -28,6 +28,7 @@ uv run coverage run -a --data-file=$PROJECT_ROOT/tests/.coverage --source=$PROJE policy.dtensor_cfg.enabled=false \ policy.megatron_cfg.enabled=true \ policy.generation.backend=megatron \ + policy.generation.refit_transport=mcore \ cluster.gpus_per_node=2 \ grpo.max_rollout_turns=5 \ grpo.max_num_steps=3 \ diff --git a/tests/functional/grpo_megatron_generation_non_colocated.sh b/tests/functional/grpo_megatron_generation_non_colocated.sh index 5ad43b878cf..ba3449c5994 100755 --- a/tests/functional/grpo_megatron_generation_non_colocated.sh +++ b/tests/functional/grpo_megatron_generation_non_colocated.sh @@ -29,6 +29,7 @@ uv run coverage run -a --data-file=$PROJECT_ROOT/tests/.coverage --source=$PROJE policy.logprob_batch_size=4 \ policy.train_micro_batch_size=1 \ policy.generation.backend=megatron \ + policy.generation.refit_transport=null \ policy.generation.colocated.enabled=false \ policy.generation.colocated.resources.gpus_per_node=1 \ policy.generation.mcore_generation_config.refit_backend=nccl \ diff --git a/tests/functional/grpo_megatron_generation_non_colocated_async_gym.sh b/tests/functional/grpo_megatron_generation_non_colocated_async_gym.sh index 64011ac4d13..6dd60c3f3bd 100755 --- a/tests/functional/grpo_megatron_generation_non_colocated_async_gym.sh +++ b/tests/functional/grpo_megatron_generation_non_colocated_async_gym.sh @@ -64,6 +64,7 @@ uv run coverage run -a --data-file=$PROJECT_ROOT/tests/.coverage --source=$PROJE policy.megatron_cfg.context_parallel_size=1 \ policy.megatron_cfg.sequence_parallel=false \ policy.generation.backend=megatron \ + +policy.generation.refit_transport=mcore \ policy.generation.mcore_generation_config.expose_http_server=true \ policy.generation.mcore_generation_config.enable_prefix_caching=true \ policy.max_total_sequence_length=512 \ diff --git a/tests/functional/grpo_megatron_generation_topp_topk.sh b/tests/functional/grpo_megatron_generation_topp_topk.sh index 11911693582..6da0112a5cb 100644 --- a/tests/functional/grpo_megatron_generation_topp_topk.sh +++ b/tests/functional/grpo_megatron_generation_topp_topk.sh @@ -31,6 +31,7 @@ uv run coverage run -a --data-file=$PROJECT_ROOT/tests/.coverage --source=$PROJE policy.logprob_batch_size=4 \ policy.train_micro_batch_size=1 \ policy.generation.backend=megatron \ + policy.generation.refit_transport=mcore \ policy.generation.temperature=0.8 \ policy.generation.top_p=0.9 \ policy.generation.top_k=50 \ diff --git a/tests/functional/grpo_megatron_nccl_reshard_refit.sh b/tests/functional/grpo_megatron_nccl_reshard_refit.sh new file mode 100755 index 00000000000..f643df003e1 --- /dev/null +++ b/tests/functional/grpo_megatron_nccl_reshard_refit.sh @@ -0,0 +1,126 @@ +#!/bin/bash +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# End-to-end non-colocated Megatron train -> Megatron generation refit over +# nccl_reshard with CUDA graphs. The default CI environment may exercise the +# exact-transfer Python implementation. Set REQUIRE_REAL_NCCL_M2N=1 in an +# environment with nccl-extensions installed to require nccl.m2n.reshard +# instead. Set REFIT_PRECISION=mxfp8 on Blackwell to cover MXFP8 training +# parameters, inference-side MXFP8 refit, and Torch-layout weights; BF16 is the +# portable default. + +set -eou pipefail + +SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd) +PROJECT_ROOT=$(realpath "$SCRIPT_DIR/../..") +git config --global --add safe.directory "$PROJECT_ROOT" + +REFIT_PRECISION=${REFIT_PRECISION:-bf16} +case "$REFIT_PRECISION" in + bf16) + model_name=Qwen/Qwen2.5-0.5B + max_token_mult_prob_error=1.05 + precision_args=() + ;; + mxfp8) + # Qwen2.5 uses QKV bias, which the inference-optimized spec rejects. + # Qwen3 is bias-free and is supported by Megatron Bridge. + model_name=Qwen/Qwen3-0.6B + max_token_mult_prob_error=1.10 + precision_args=( + ++policy.megatron_cfg.fp8_cfg.enabled=true + ++policy.megatron_cfg.fp8_cfg.fp8=e4m3 + ++policy.megatron_cfg.fp8_cfg.fp8_recipe=mxfp8 + ++policy.megatron_cfg.fp8_cfg.fp8_param=true + # Match NeMo-RL's FP8 end-to-end recipes: MXFP8 model parameters + # with the standard FP32-master distributed optimizer. + ++policy.megatron_cfg.optimizer.use_precision_aware_optimizer=false + ++policy.generation.mcore_generation_config.transformer_impl=inference_optimized + ++policy.generation.mcore_generation_config.fp8_cfg.enabled=true + ++policy.generation.mcore_generation_config.fp8_cfg.fp8=e4m3 + ++policy.generation.mcore_generation_config.fp8_cfg.fp8_recipe=mxfp8 + ++policy.generation.mcore_generation_config.fp8_cfg.fp8_param=true + ++policy.generation.mcore_generation_config.inference_grouped_gemm_backend=torch + ) + ;; + *) + echo "REFIT_PRECISION must be bf16 or mxfp8, got $REFIT_PRECISION" >&2 + exit 1 + ;; +esac + +EXP_NAME="$(basename "$0" .sh)-$REFIT_PRECISION" +EXP_DIR="$SCRIPT_DIR/$EXP_NAME" +LOG_DIR="$EXP_DIR/logs" +JSON_METRICS="$EXP_DIR/metrics.json" +RUN_LOG="$EXP_DIR/run.log" +export PYTHONPATH="${PROJECT_ROOT}:${PYTHONPATH:-}" + +rm -rf "$EXP_DIR" +mkdir -p "$LOG_DIR" + +# Exercise the default dispatcher rather than forcing either fallback path. +unset NRL_XFERDTENSOR_GOLDEN +unset NRL_XFERDTENSOR_PYTHON + +cd "$PROJECT_ROOT" +uv run coverage run -a \ + --data-file="$PROJECT_ROOT/tests/.coverage" \ + --source="$PROJECT_ROOT/nemo_rl" \ + "$PROJECT_ROOT/examples/run_grpo.py" \ + --config "$PROJECT_ROOT/examples/configs/grpo_math_1B_megatron.yaml" \ + policy.model_name="$model_name" \ + grpo.num_prompts_per_step=2 \ + grpo.num_generations_per_prompt=4 \ + policy.train_global_batch_size=4 \ + policy.logprob_batch_size=4 \ + policy.train_micro_batch_size=1 \ + policy.generation.backend=megatron \ + policy.generation.colocated.enabled=false \ + policy.generation.colocated.resources.gpus_per_node=1 \ + policy.generation.refit_transport=nccl_reshard \ + policy.generation.mcore_generation_config.cuda_graph_impl=local \ + policy.generation.mcore_generation_config.inference_cuda_graph_scope=block \ + policy.generation.mcore_generation_config.num_cuda_graphs=4 \ + policy.generation.mcore_generation_config.use_cuda_graphs_for_non_decode_steps=true \ + "${precision_args[@]}" \ + cluster.gpus_per_node=2 \ + grpo.max_num_steps=2 \ + logger.tensorboard_enabled=true \ + logger.log_dir="$LOG_DIR" \ + logger.wandb_enabled=false \ + logger.monitor_gpus=true \ + checkpointing.enabled=false \ + policy.megatron_cfg.checkpoint.async_save=false \ + "$@" \ + 2>&1 | tee "$RUN_LOG" + +grep -q "nccl_reshard bulk comm group" "$RUN_LOG" +# NeMo-RL-owned marker: the engine is initialized lazily after the first refit +# so CUDA graphs capture the final (possibly MXFP8) buffers. Do not assert on +# Megatron-LM's "cuda graph warmup" tqdm label - it is upstream-owned and only +# appears when tqdm is not disabled. +grep -q "Initialized persistent inference engine" "$RUN_LOG" +if [[ "${REQUIRE_REAL_NCCL_M2N:-0}" == "1" ]]; then + grep -q "reshard path: real nccl.m2n.reshard" "$RUN_LOG" +else + grep -Eq \ + "reshard path: (real nccl.m2n.reshard|xferdtensor_python \(exact-transfer\))" \ + "$RUN_LOG" +fi + +uv run tests/json_dump_tb_logs.py "$LOG_DIR" --output_path "$JSON_METRICS" +uv run python tests/check_metrics.py "$JSON_METRICS" \ + "max(data[\"train/token_mult_prob_error\"]) < $max_token_mult_prob_error" diff --git a/tests/functional/nemotron_omni_clevr_megatron_1n2g.sh b/tests/functional/nemotron_omni_clevr_megatron_1n2g.sh index 0a61b7b71c6..64d7105697e 100755 --- a/tests/functional/nemotron_omni_clevr_megatron_1n2g.sh +++ b/tests/functional/nemotron_omni_clevr_megatron_1n2g.sh @@ -113,6 +113,7 @@ uv run --no-sync python examples/run_vlm_grpo.py \ ++policy.megatron_cfg.optimizer.exp_avg_sq_dtype=bfloat16 \ ++policy.megatron_cfg.optimizer.store_param_remainders=false \ policy.generation.backend=megatron \ + policy.generation.refit_transport=mcore \ policy.generation.colocated.enabled=true \ policy.generation.colocated.resources.num_nodes=1 \ policy.generation.colocated.resources.gpus_per_node=2 \ diff --git a/tests/functional/nemotron_omni_gym_video_megatron_1n2g.sh b/tests/functional/nemotron_omni_gym_video_megatron_1n2g.sh index 54d84dd1b0d..0a757d4ef04 100755 --- a/tests/functional/nemotron_omni_gym_video_megatron_1n2g.sh +++ b/tests/functional/nemotron_omni_gym_video_megatron_1n2g.sh @@ -104,6 +104,7 @@ uv run --no-sync python examples/nemo_gym/run_grpo_nemo_gym.py \ ++policy.megatron_cfg.optimizer.exp_avg_sq_dtype=bfloat16 \ ++policy.megatron_cfg.optimizer.store_param_remainders=false \ policy.generation.backend=megatron \ + +policy.generation.refit_transport=mcore \ ++policy.generation.bad_words=null \ policy.generation.colocated.enabled=true \ policy.generation.colocated.resources.num_nodes=1 \ diff --git a/tests/unit/algorithms/test_grpo.py b/tests/unit/algorithms/test_grpo.py index 901f986b65f..3fbbfc7886d 100644 --- a/tests/unit/algorithms/test_grpo.py +++ b/tests/unit/algorithms/test_grpo.py @@ -26,6 +26,7 @@ from omegaconf import OmegaConf from torchdata.stateful_dataloader import StatefulDataLoader +from nemo_rl.algorithms import grpo as grpo_mod from nemo_rl.algorithms.advantage_estimator import ( GDPOAdvantageEstimator, GRPOAdvantageEstimator, @@ -219,6 +220,51 @@ def test_refit_policy_generation_forwards_kv_scales_on_colocated_ipc( ) +def test_megatron_m2n_refit_delegates_entirely_to_the_synchronizer() -> None: + """MegatronWeightSynchronizer owns the engine lifecycle; the caller must not duplicate it. + + ``refit_policy_generation`` returns as soon as a weight synchronizer is + present, so suspend/offload/prepare/resume must NOT be driven here — they + live inside ``MegatronWeightSynchronizer.sync_weights`` and are asserted in + ``tests/unit/weight_sync/test_weight_synchronizer.py``. + """ + policy = MagicMock() + generation = object.__new__(MegatronGeneration) + generation.suspend_for_refit = MagicMock() + generation.prepare_for_generation = MagicMock() + generation.resume_after_refit = MagicMock() + generation.weight_synchronizer = MagicMock() + generation.weight_synchronizer.sync_weights.return_value = {"bytes": 16.0} + + metrics = refit_policy_generation( + policy, + generation, + colocated_inference=False, + kv_scales={"layer.0": 0.5}, + ) + + assert metrics == {"bytes": 16.0} + generation.weight_synchronizer.sync_weights.assert_called_once_with( + timer=None, kv_scales={"layer.0": 0.5} + ) + generation.suspend_for_refit.assert_not_called() + generation.prepare_for_generation.assert_not_called() + generation.resume_after_refit.assert_not_called() + policy.offload_before_refit.assert_not_called() + + +def test_refit_returns_empty_metrics_when_synchronizer_returns_none() -> None: + """``sync_weights`` returning None must not propagate as the metrics dict.""" + generation = object.__new__(MegatronGeneration) + generation.weight_synchronizer = MagicMock() + generation.weight_synchronizer.sync_weights.return_value = None + + assert ( + refit_policy_generation(MagicMock(), generation, colocated_inference=False) + == {} + ) + + class TestMaskSampleFilter: def test_masks_env_flagged_samples(self): repeated_batch = BatchedDataDict( @@ -3394,8 +3440,6 @@ def test_setup_refits_noncolocated_megatron_while_nemo_gym_waits( monkeypatch, mock_grpo_components ): """The initial refit must start a skip-load endpoint before Gym can finish.""" - from nemo_rl.algorithms import grpo as grpo_mod - events = [] gym_started = Event() engine_ready = Event() diff --git a/tests/unit/models/generation/test_megatron_generation.py b/tests/unit/models/generation/test_megatron_generation.py index f5ccdcb4b6f..7f02742d694 100644 --- a/tests/unit/models/generation/test_megatron_generation.py +++ b/tests/unit/models/generation/test_megatron_generation.py @@ -15,11 +15,13 @@ import gc from copy import deepcopy from types import SimpleNamespace +from unittest.mock import MagicMock import pytest import ray import torch +import nemo_rl.models.policy.lm_policy as lm_policy from nemo_rl.algorithms.grpo import refit_policy_generation from nemo_rl.algorithms.utils import get_tokenizer from nemo_rl.data.multimodal_utils import PackedTensor @@ -28,6 +30,8 @@ from nemo_rl.models.generation.megatron import MegatronGeneration, megatron_generation from nemo_rl.models.generation.megatron.config import ( dedicated_inference_megatron_cfg, + merged_inference_megatron_cfg, + resolve_refit_execution_batch_bytes, ) from nemo_rl.models.generation.megatron.megatron_worker import MegatronGenerationMixin from nemo_rl.models.generation.megatron.utils import ( @@ -38,11 +42,373 @@ from nemo_rl.weight_sync.megatron_weight_synchronizer import ( MegatronWeightSynchronizer, ) +from nemo_rl.weight_sync.membership import RefitMembership from tests.unit.test_utils import SimpleLossFn model_name = "Qwen/Qwen3-0.6B" +@pytest.mark.mcore +def test_mxfp8_skip_weight_load_defers_http_server_until_refit( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """HTTP URLs must not force engine init before MXFP8 refit buffers exist.""" + config = deepcopy(basic_megatron_test_config) + config["generation"]["colocated"]["enabled"] = False + config["generation"]["mcore_generation_config"]["expose_http_server"] = True + config["generation"]["mcore_generation_config"]["fp8_cfg"] = {"enabled": True} + + prepare_for_generation = MagicMock() + monkeypatch.setattr(lm_policy, "Policy", MagicMock()) + monkeypatch.setattr( + MegatronGeneration, "init_cluster_placement_groups", MagicMock() + ) + monkeypatch.setattr( + MegatronGeneration, "prepare_for_generation", prepare_for_generation + ) + + generation = MegatronGeneration( + config=config, + tokenizer=MagicMock(), + cluster=MagicMock(), + skip_weight_load=True, + ) + + prepare_for_generation.assert_not_called() + assert generation.dp_openai_server_base_urls == [] + assert lm_policy.Policy.call_args.kwargs["refit_role"] == "destination" + + +@pytest.mark.mcore +def test_nccl_m2n_refit_backend_requires_non_colocated_generation() -> None: + config = deepcopy(basic_megatron_test_config) + config["generation"]["colocated"]["enabled"] = True + config["generation"]["refit_transport"] = "mcore" + config["generation"]["mcore_generation_config"]["refit_backend"] = "nccl_m2n" + + with pytest.raises(ValueError, match="only supported with non-colocated"): + MegatronGeneration( + config=config, + tokenizer=MagicMock(), + policy=MagicMock(), + ) + + +@pytest.mark.mcore +def test_null_refit_transport_selects_packed_collective() -> None: + config = deepcopy(basic_megatron_test_config) + config["generation"]["refit_transport"] = None + + generation = MegatronGeneration( + config=config, + tokenizer=MagicMock(), + policy=MagicMock(), + ) + + assert not generation.uses_native_refit + + +@pytest.mark.mcore +def test_inference_optimized_pins_generation_etp_to_one() -> None: + """Generation-side TP>1 must stay usable for MoE. + + MCore's inference_optimized MoE layers reject a *resolved* ETP > 1, and an + omitted ETP resolves to TP rather than 1. Pinning it here is what keeps + TP>1 generation working. + """ + config = deepcopy(basic_megatron_test_config) + config["megatron_cfg"]["tensor_model_parallel_size"] = 2 + config["megatron_cfg"]["expert_tensor_parallel_size"] = 2 + config["generation"]["mcore_generation_config"]["transformer_impl"] = ( + "inference_optimized" + ) + config["generation"]["mcore_generation_config"]["tensor_model_parallel_size"] = 2 + config["generation"]["mcore_generation_config"]["sequence_parallel"] = True + config["generation"]["mcore_generation_config"].pop( + "expert_tensor_parallel_size", None + ) + + merged = merged_inference_megatron_cfg(config) + + assert merged["expert_tensor_parallel_size"] == 1 + # TP>1 is preserved -- the pin narrows ETP only. + assert merged["tensor_model_parallel_size"] == 2 + + +@pytest.mark.mcore +def test_inference_optimized_rejects_explicit_generation_etp_above_one() -> None: + config = deepcopy(basic_megatron_test_config) + config["generation"]["mcore_generation_config"]["transformer_impl"] = ( + "inference_optimized" + ) + config["generation"]["mcore_generation_config"]["sequence_parallel"] = True + config["generation"]["mcore_generation_config"]["expert_tensor_parallel_size"] = 2 + + with pytest.raises(ValueError, match="expert_tensor_parallel_size=1"): + merged_inference_megatron_cfg(config) + + +@pytest.mark.mcore +@pytest.mark.parametrize("refit_transport", [None, "mcore"]) +def test_megatron_generation_dispatches_refit_transport(refit_transport): + generation = object.__new__(MegatronGeneration) + generation.cfg = { + "refit_transport": refit_transport, + "mcore_generation_config": { + "refit_backend": "gloo", + "refit_execution_batch_bytes": 123, + }, + } + generation._policy = MagicMock() + generation._owns_policy = False + + generation.init_collective("127.0.0.1", 1234, 4, train_world_size=2) + generation.update_weights_from_collective() + + if refit_transport == "mcore": + generation._policy.init_collective_mcore_generation.assert_called_once_with( + "127.0.0.1", + 1234, + 4, + rank_offset=2, + refit_execution_batch_bytes=123, + refit_backend="gloo", + ) + generation._policy.swap_weights_via_reshard.assert_called_once_with( + is_source=False + ) + generation._policy.init_collective.assert_not_called() + else: + generation._policy.init_collective.assert_called_once_with( + "127.0.0.1", + 1234, + 4, + train_world_size=2, + rank_offset=2, + ) + generation._policy.worker_group.run_all_workers_single_data.assert_called_once_with( + "update_weights_from_collective", refit_timeout_s=None + ) + generation._policy.init_collective_mcore_generation.assert_not_called() + + +@pytest.mark.mcore +def test_megatron_generation_m2n_transport_uses_packed_collective_api() -> None: + generation = object.__new__(MegatronGeneration) + generation.cfg = { + "refit_transport": "nccl_reshard", + "mcore_generation_config": { + "refit_backend": "nccl", + "refit_execution_batch_bytes": 123, + }, + } + generation._policy = MagicMock() + generation._owns_policy = False + + generation.init_collective("127.0.0.1", 1234, 4, train_world_size=2) + + assert not generation.uses_native_refit + generation._policy.init_collective.assert_called_once_with( + "127.0.0.1", + 1234, + 4, + train_world_size=2, + rank_offset=2, + ) + generation._policy.init_collective_mcore_generation.assert_not_called() + + +@pytest.mark.mcore +def test_megatron_generation_uses_common_refit_worker_api() -> None: + generation = object.__new__(MegatronGeneration) + generation._policy = MagicMock() + generation._owns_policy = False + generation._refit_membership = None + generation._policy.worker_group.run_all_workers_single_data.return_value = [] + refit_info = {"layer_names": [], "per_layer_params": {}} + + generation.prepare_nccl_reshard_refit_info(refit_info) + generation.nccl_reshard_refit() + + calls = generation._policy.worker_group.run_all_workers_single_data.call_args_list + assert calls[0].args == ("prepare_nccl_reshard_refit_info",) + assert calls[0].kwargs == {"refit_info": refit_info} + assert calls[1].args == ("nccl_reshard_refit",) + assert calls[1].kwargs == {"refit_timeout_s": None} + + +@pytest.mark.mcore +def test_megatron_generation_rebuild_dispatches_only_selected_ranks() -> None: + workers = [MagicMock() for _ in range(3)] + for idx, worker in enumerate(workers): + worker.init_collective.remote.return_value = f"misc-{idx}" + worker.init_nccl_reshard_comm_groups_generation.remote.return_value = ( + f"bulk-{idx}" + ) + worker.nccl_reshard_refit.remote.return_value = f"refit-{idx}" + + generation = object.__new__(MegatronGeneration) + generation._policy = SimpleNamespace(worker_group=SimpleNamespace(workers=workers)) + generation._owns_policy = False + generation._refit_membership = None + membership = RefitMembership( + world_size=10, + train_world_size=8, + shard_prefixes={0: 0, 2: 1}, + workers_per_shard=1, + ) + generation.set_refit_membership(membership) + + assert generation.rebuild_collective(membership, "10.0.0.1", 1234) == [ + "misc-0", + "misc-2", + ] + workers[0].init_collective.remote.assert_called_once_with( + ip="10.0.0.1", + port=1234, + world_size=10, + train_world_size=8, + rank_offset=8, + nccl_peer="nemo", + ) + workers[1].init_collective.remote.assert_not_called() + workers[2].init_collective.remote.assert_called_once_with( + ip="10.0.0.1", + port=1234, + world_size=10, + train_world_size=8, + rank_offset=7, + nccl_peer="nemo", + ) + + assert generation.rebuild_nccl_reshard_comm_group( + membership, + pp_ips=["10.0.0.1"], + pp_ports=[1235], + pp_size=1, + train_ranks_per_stage=8, + sub_world_size=10, + ) == ["bulk-0", "bulk-2"] + workers[0].init_nccl_reshard_comm_groups_generation.remote.assert_called_once_with( + pp_ips=["10.0.0.1"], + pp_ports=[1235], + pp_size=1, + train_ranks_per_stage=8, + sub_world_size=10, + rank_prefix=0, + ) + workers[1].init_nccl_reshard_comm_groups_generation.remote.assert_not_called() + workers[2].init_nccl_reshard_comm_groups_generation.remote.assert_called_once_with( + pp_ips=["10.0.0.1"], + pp_ports=[1235], + pp_size=1, + train_ranks_per_stage=8, + sub_world_size=10, + rank_prefix=1, + ) + + assert generation.nccl_reshard_refit(refit_timeout_s=30.0) == [ + "refit-0", + "refit-2", + ] + workers[0].nccl_reshard_refit.remote.assert_called_once_with(refit_timeout_s=30.0) + workers[1].nccl_reshard_refit.remote.assert_not_called() + workers[2].nccl_reshard_refit.remote.assert_called_once_with(refit_timeout_s=30.0) + + +@pytest.mark.mcore +def test_native_refit_batch_bytes_use_collective_default(monkeypatch) -> None: + monkeypatch.setattr( + "nemo_rl.models.generation.megatron.config.get_target_packed_tensor_size", + lambda: 456, + ) + + assert resolve_refit_execution_batch_bytes(None) == 456 + assert resolve_refit_execution_batch_bytes(123) == 123 + with pytest.raises(ValueError, match="must be positive or null"): + resolve_refit_execution_batch_bytes(0) + + +@pytest.mark.mcore +def test_bridge_refit_finalizes_import_before_return( + monkeypatch: pytest.MonkeyPatch, +) -> None: + # Keep the optional MCore/Transformer Engine worker import test-local. + import nemo_rl.models.generation.megatron.megatron_worker as worker_module + + events = [] + worker = object.__new__(worker_module.MegatronGenerationRefitMixin) + worker.model_update_group = object() + worker._generation_refit_state_dict_info = {} + worker._generation_refit_tasks = [] + worker._generation_refit_dependency_counts = {} + worker._generation_refit_model_chunks = [torch.nn.Module()] + worker.megatron_bridge = MagicMock() + worker.megatron_bridge.finalize_hf_import.side_effect = ( + lambda _model_chunks: events.append("finalize") + ) + + monkeypatch.setattr( + worker_module, + "packed_broadcast_consumer", + lambda **_kwargs: events.append("receive"), + ) + monkeypatch.setattr( + "nemo_rl.distributed.refit_watchdog.sync_stream_within", + lambda *_args: events.append("sync"), + ) + + assert worker._update_destination_weights_from_collective() + assert events == ["receive", "finalize", "sync"] + worker.megatron_bridge.finalize_hf_import.assert_called_once_with( + worker._generation_refit_model_chunks + ) + + +@pytest.mark.mcore +def test_bridge_refit_converts_external_state_through_streaming_api() -> None: + # Keep the optional MCore/Transformer Engine worker import test-local. + import nemo_rl.models.generation.megatron.megatron_worker as worker_module + + worker = object.__new__(worker_module.MegatronGenerationRefitMixin) + conversion_task = SimpleNamespace( + param_name="megatron.weight", hf_param_names=("hf.weight",) + ) + task = worker_module._MegatronRefitTask( + conversion_task=conversion_task, + destination=torch.empty(2), + target_id=1, + ) + converted_weight = torch.ones(2) + worker._generation_refit_pending_weights = {} + worker._generation_refit_pending_streams = {} + worker._generation_refit_remaining_dependencies = {"hf.weight": 1} + worker._generation_refit_tasks = [task] + worker._generation_refit_task_index = 0 + worker._generation_refit_model_chunks = [torch.nn.Module()] + worker.megatron_bridge = MagicMock() + streamed_states = [] + + def stream_weights(*_args, hf_state_dict, **_kwargs): + streamed_states.append(dict(hf_state_dict)) + return iter([SimpleNamespace(weight=converted_weight)]) + + worker.megatron_bridge.stream_weights_hf_to_megatron.side_effect = stream_weights + worker._write_generation_refit_weight = MagicMock() + source_weight = torch.zeros(2) + + worker._load_generation_refit_batch([("hf.weight", source_weight)]) + + stream_call = worker.megatron_bridge.stream_weights_hf_to_megatron.call_args + assert stream_call.args == (worker._generation_refit_model_chunks,) + assert stream_call.kwargs["conversion_tasks"] == [conversion_task] + assert streamed_states[0]["hf.weight"] is source_weight + worker._write_generation_refit_weight.assert_called_once_with( + task, converted_weight + ) + assert worker._generation_refit_pending_weights == {} + + @pytest.mark.mcore def test_multimodal_preprocessing_requires_policy_processor(): class _ImageWrapper: @@ -328,6 +694,7 @@ def mock_run_coroutine_threadsafe(call, loop): "max_grad_norm": 1.0, "generation": { "backend": "megatron", + "refit_transport": "mcore", "model_name": model_name, "max_new_tokens": 16, # Small number of tokens for testing "temperature": 1.0, @@ -355,6 +722,7 @@ def mock_run_coroutine_threadsafe(call, loop): "num_speculative_tokens": 0, "logprobs_mode": "processed_logprobs", "refit_backend": "gloo", # not nvshmem: its NVLS multicast init is unavailable in CI + "refit_execution_batch_bytes": None, "parsers": [], "expose_http_server": False, }, @@ -787,7 +1155,10 @@ def test_megatron_generation_colocated( @pytest.mark.timeout(900) @pytest.mark.parametrize("skip_weight_load", [False, True]) def test_megatron_generation_non_colocated_refit( - policy_cluster_separate, test_input_data, tokenizer, skip_weight_load + policy_cluster_separate, + test_input_data, + tokenizer, + skip_weight_load, ): """Non-colocated Megatron generation. diff --git a/tests/unit/models/megatron/test_group_experts.py b/tests/unit/models/megatron/test_group_experts.py index 85335909a17..9e1f1760329 100644 --- a/tests/unit/models/megatron/test_group_experts.py +++ b/tests/unit/models/megatron/test_group_experts.py @@ -12,18 +12,16 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Unit test for the train-side expert stacking (``_group_experts``). +"""Unit tests for train-side expert source grouping and materialization. -``_group_experts`` (``MegatronPolicyWorkerImpl``) stacks this rank's local -per-expert tensors for one projection into ``[E_local, ...]``. It doesn't use -``self`` and operates on plain tensors, so a dummy ``self`` + CPU tensors suffice. +``_group_experts`` (``MegatronPolicyWorkerImpl``) records this rank's local +per-expert sources for one projection. The refit loop materializes and stacks +them into ``[E_local, ...]``. Plain CPU tensors suffice for the BF16 path. Importing ``megatron_policy_worker`` pulls in megatron.core, so this is mcore-marked and skipped where mcore is unavailable. """ -from types import SimpleNamespace - import pytest import torch @@ -37,15 +35,15 @@ from nemo_rl.models.policy.workers.megatron_policy_worker import ( # noqa: E402 MegatronPolicyWorkerImpl, ) +from nemo_rl.weight_sync.nccl_reshard_utils import LocalParamSpec # noqa: E402 pytestmark = pytest.mark.mcore def _group(proj, grouped_name, expert_groups): - # _group_experts ignores self; pass a dummy. - return MegatronPolicyWorkerImpl._group_experts( - SimpleNamespace(), proj, grouped_name, expert_groups - ) + worker = object.__new__(MegatronPolicyWorkerImpl) + grouped = worker._group_experts(proj, grouped_name, expert_groups) + return worker._materialize_local_refit_spec(LocalParamSpec(base=grouped), {}).buf def test_group_experts_stacks_in_order(): @@ -53,7 +51,13 @@ def test_group_experts_stacks_in_order(): e0 = torch.randn(1536, 4096) e1 = torch.randn(1536, 4096) e2 = torch.randn(1536, 4096) - groups = {(prefix, "gate_proj"): [e0, e1, e2]} + groups = { + (prefix, "gate_proj"): [ + LocalParamSpec(base=e0), + LocalParamSpec(base=e1), + LocalParamSpec(base=e2), + ] + } out = _group("gate_proj", f"{prefix}.gate_proj.weight", groups) assert out.shape == (3, 1536, 4096) # Order preserved (expert 0 first). @@ -63,7 +67,7 @@ def test_group_experts_stacks_in_order(): def test_group_experts_missing_group_raises(): - groups = {("other.experts", "gate_proj"): [torch.randn(8, 8)]} + groups = {("other.experts", "gate_proj"): [LocalParamSpec(base=torch.randn(8, 8))]} with pytest.raises(AssertionError): _group("gate_proj", "model.layers.0.mlp.experts.gate_proj.weight", groups) @@ -83,14 +87,15 @@ def test_build_hf_to_local_param_map_train_side(): from nemo_rl.weight_sync.nccl_reshard_utils import HFToLocalParamMap w = object.__new__(MegatronPolicyWorkerImpl) # no __init__ / no megatron state + w.my_pp_stage = 0 prefix = "model.layers.0.mlp.experts" direct = torch.randn(8, 16) # a dense FFN down_proj local shard view e0 = torch.randn(128, 16) # this rank's local expert 0 gate_proj e1 = torch.randn(128, 16) # local expert 1 gate_proj w._iter_local_hf_param_shards = lambda: [ - ("model.layers.0.mlp.down_proj.weight", direct), - (f"{prefix}.0.gate_proj.weight", e0), - (f"{prefix}.1.gate_proj.weight", e1), + ("model.layers.0.mlp.down_proj.weight", LocalParamSpec(base=direct)), + (f"{prefix}.0.gate_proj.weight", LocalParamSpec(base=e0)), + (f"{prefix}.1.gate_proj.weight", LocalParamSpec(base=e1)), ] refit_info = { "layer_names": ["model.layers.0"], @@ -116,10 +121,10 @@ def test_build_hf_to_local_param_map_train_side(): d = pmap.get("model.layers.0.mlp.down_proj.weight") assert d.base is direct and d.pre is None and d.post is None - # Grouped expert: pre stacks this rank's per-expert views into [E_local, ...] - # fresh each refit (base unused — the views are captured in the hook). + # Grouped expert: the base recipe retains the per-expert live views and the + # refit loop stacks them into [E_local, ...] on each refit. g = pmap.get(f"{prefix}.gate_proj.weight") - assert g.pre is not None - ctx = g.pre(g.base) + assert g.pre is None + ctx = w._materialize_local_refit_spec(g, {}) assert ctx.buf.shape == (2, 128, 16) assert torch.equal(ctx.buf[0], e0) and torch.equal(ctx.buf[1], e1) diff --git a/tests/unit/models/megatron/test_megatron_setup.py b/tests/unit/models/megatron/test_megatron_setup.py index 5c060f5fa85..cad33132d32 100644 --- a/tests/unit/models/megatron/test_megatron_setup.py +++ b/tests/unit/models/megatron/test_megatron_setup.py @@ -2695,6 +2695,64 @@ def test_disabled_or_absent_fp8_preserves_optimizer_default( assert optimizer.fp8_recipe is None +@pytest.mark.mcore +class TestCreateMegatronConfigFP8Buffers: + """Tests for MXFP8 parameter-buffer plumbing into optimizer and DDP.""" + + @staticmethod + def _subconfig_kwargs(): + from nemo_rl.models.megatron.setup import _create_megatron_config + + config = { + "megatron_cfg": { + "optimizer": {"use_distributed_optimizer": True}, + "scheduler": {}, + "distributed_data_parallel_config": { + "overlap_param_gather": True, + "grad_reduce_in_fp32": False, + "overlap_grad_reduce": True, + "data_parallel_sharding_strategy": "optim_grads_params", + }, + "fp8_cfg": { + "enabled": True, + "fp8_recipe": "mxfp8", + "fp8_param": True, + }, + "train_iters": 10, + }, + "train_global_batch_size": 8, + } + with ( + patch("nemo_rl.models.megatron.setup.ConfigContainer"), + patch("nemo_rl.models.megatron.setup.TrainingConfig"), + patch("nemo_rl.models.megatron.setup.OptimizerConfig") as mock_optimizer, + patch( + "nemo_rl.models.megatron.setup.DistributedDataParallelConfig" + ) as mock_ddp, + patch("nemo_rl.models.megatron.setup.SchedulerConfig"), + patch("nemo_rl.models.megatron.setup.TokenizerConfig"), + patch("nemo_rl.models.megatron.setup.LoggerConfig"), + ): + _create_megatron_config( + model_cfg=MagicMock(), + checkpoint_config=MagicMock(), + config=config, + hf_model_name="test-model", + dtype=torch.bfloat16, + fp8_param_enabled=True, + ) + + return mock_optimizer.call_args.kwargs, mock_ddp.call_args.kwargs + + def test_mxfp8_recipe_and_param_gather_are_forwarded(self): + optimizer_kwargs, ddp_kwargs = self._subconfig_kwargs() + + assert optimizer_kwargs["fp8_recipe"] == "mxfp8" + assert optimizer_kwargs["reuse_grad_buf_for_mxfp8_param_ag"] is True + assert ddp_kwargs["fp8_param_gather"] is True + assert ddp_kwargs["reuse_grad_buf_for_mxfp8_param_ag"] is True + + @pytest.mark.mcore class TestCreateMegatronConfigOptimizerOffload: """Tests for optimizer CPU-offload plumbing into Megatron Core.""" diff --git a/tests/unit/models/policy/test_lm_policy_collective.py b/tests/unit/models/policy/test_lm_policy_collective.py index 8703370a139..736f0e3a748 100644 --- a/tests/unit/models/policy/test_lm_policy_collective.py +++ b/tests/unit/models/policy/test_lm_policy_collective.py @@ -47,12 +47,93 @@ def shutdown(self, **_kwargs): "port": 1234, "world_size": 4, "train_world_size": 2, + "rank_offset": 0, "nccl_peer": "vllm", }, ) ] +def test_policy_forwards_rank_offset_to_workers(): + """Megatron M-to-N generation ranks join the group after the training ranks.""" + calls = [] + + class WorkerGroup: + def run_all_workers_single_data(self, method_name, **kwargs): + calls.append((method_name, kwargs)) + return ["future"] + + def shutdown(self, **_kwargs): + pass + + policy = Policy.__new__(Policy) + policy.worker_group = WorkerGroup() + + policy.init_collective( + "127.0.0.1", + 1234, + 6, + train_world_size=4, + rank_offset=4, + nccl_peer="nemo", + ) + + assert calls == [ + ( + "init_collective", + { + "ip": "127.0.0.1", + "port": 1234, + "world_size": 6, + "train_world_size": 4, + "rank_offset": 4, + "nccl_peer": "nemo", + }, + ) + ] + + +def test_policy_worker_offsets_its_rank_in_the_collective(monkeypatch): + """A non-zero rank_offset shifts the worker's rank within the shared group.""" + calls = [] + + class ProcessGroup: + def __init__(self, **kwargs): + calls.append(("create", kwargs)) + + def init_nccl_communicator(self, **kwargs): + calls.append(("init", kwargs)) + + monkeypatch.setattr( + "nemo_rl.distributed.stateless_process_group.StatelessProcessGroup", + ProcessGroup, + ) + monkeypatch.setattr( + "nemo_rl.models.policy.workers.base_policy_worker.torch.cuda.current_device", + lambda: 0, + ) + + worker = AbstractPolicyWorker.__new__(AbstractPolicyWorker) + worker.rank = 1 + worker.init_collective( + "127.0.0.1", + 1234, + 6, + train_world_size=4, + rank_offset=4, + ) + + assert calls[0] == ( + "create", + { + "master_address": "127.0.0.1", + "port": 1234, + "rank": 5, + "world_size": 6, + }, + ) + + def test_policy_worker_initializes_requested_nccl_peer(monkeypatch): calls = [] diff --git a/tests/unit/models/policy/test_megatron_split_state.py b/tests/unit/models/policy/test_megatron_split_state.py index 990397b4bc4..ec6f90aeb98 100644 --- a/tests/unit/models/policy/test_megatron_split_state.py +++ b/tests/unit/models/policy/test_megatron_split_state.py @@ -48,6 +48,7 @@ from __future__ import annotations import logging +from types import SimpleNamespace from typing import Any from unittest.mock import MagicMock, patch @@ -149,6 +150,10 @@ def _make_worker(loss_type): }, }, } + w.megatron_cfg = SimpleNamespace( + optimizer=SimpleNamespace(reuse_grad_buf_for_mxfp8_param_ag=False), + ddp=SimpleNamespace(overlap_param_gather=False), + ) w.dp_size = 2 w.cp_size = 1 w.sampling_params = None diff --git a/tests/unit/models/policy/test_megatron_worker.py b/tests/unit/models/policy/test_megatron_worker.py index 78db03db202..fe440c49233 100644 --- a/tests/unit/models/policy/test_megatron_worker.py +++ b/tests/unit/models/policy/test_megatron_worker.py @@ -12,13 +12,15 @@ # See the License for the specific language governing permissions and # limitations under the License. import ast +import asyncio import os import tempfile import time +from dataclasses import dataclass from pathlib import Path from types import SimpleNamespace from typing import Optional -from unittest.mock import MagicMock +from unittest.mock import MagicMock, call import numpy as np import pytest @@ -41,10 +43,677 @@ from nemo_rl.models.policy import PolicyConfig from nemo_rl.models.policy.lm_policy import Policy from nemo_rl.utils.checkpoint import CheckpointManager +from nemo_rl.weight_sync.nccl_reshard_utils import HFToLocalParamMap from tests.unit.test_utils import SimpleLossFn pytestmark = pytest.mark.mcore +# Megatron Core/Bridge and worker-module imports stay test-local so pytest can +# collect this module without the optional MCore/Transformer Engine stack. + + +def _make_refit_task( + *, + param_name: str, + destination: object, + dependencies: tuple[str, ...], + local_specs: tuple[tuple[str, str], ...] = (), + mapping: object | None = None, +): + from megatron.bridge.models.conversion.param_mapping import LocalHFParamSpec + + from nemo_rl.models.generation.megatron.megatron_worker import ( + _MegatronRefitTask, + ) + + specs = tuple( + LocalHFParamSpec( + name, + None if component == "full" else -2, + 1 if component == "up" else 0, + 1 if component == "full" else 2, + ) + for name, component in local_specs + ) + + def combine_local_hf_weights(weights): + if len(specs) == 1: + return weights[specs[0].name] + return torch.cat( + [ + weights[spec.name] + for spec in sorted(specs, key=lambda item: item.split_index) + ], + dim=-2, + ) + + conversion_task = SimpleNamespace( + param_name=param_name, + hf_param_names=dependencies, + mapping=mapping if mapping is not None else MagicMock(), + local_hf_param_specs=lambda: specs, + combine_local_hf_weights=combine_local_hf_weights, + ) + return _MegatronRefitTask( + conversion_task=conversion_task, + destination=destination, + target_id=id(destination), + ) + + +def test_mcore_nccl_m2n_builds_hybrid_group_and_copy_service( + monkeypatch: pytest.MonkeyPatch, +) -> None: + import torch.distributed.distributed_c10d as distributed_c10d + + from nemo_rl.models.generation.megatron import megatron_worker as worker_module + + worker = object.__new__(worker_module.MegatronGenerationRefitMixin) + worker.model = object() + + process_group = MagicMock() + process_group_type = MagicMock(return_value=process_group) + process_group_type.BackendType = worker_module.ProcessGroup.BackendType + gloo_backend = MagicMock() + nccl_backend = MagicMock() + nccl_backend_type = MagicMock(return_value=nccl_backend) + nccl_backend_type.Options.return_value = MagicMock() + m2n_service = MagicMock() + m2n_service_type = MagicMock(return_value=m2n_service) + prepare_swap_model_weights = MagicMock() + + monkeypatch.setattr(worker_module, "ProcessGroup", process_group_type) + monkeypatch.setattr( + worker_module, "ProcessGroupGloo", MagicMock(return_value=gloo_backend) + ) + monkeypatch.setattr(worker_module, "PrefixStore", MagicMock()) + tcp_store_type = MagicMock(return_value=MagicMock()) + monkeypatch.setattr(worker_module.torch.distributed, "TCPStore", tcp_store_type) + monkeypatch.setattr( + worker_module.torch.distributed, "get_rank", MagicMock(return_value=0) + ) + monkeypatch.setattr( + worker_module.torch.cuda, "current_device", MagicMock(return_value=0) + ) + monkeypatch.setattr(worker_module.torch.cuda, "set_device", MagicMock()) + monkeypatch.setattr(distributed_c10d, "ProcessGroupNCCL", nccl_backend_type) + monkeypatch.setattr(worker_module, "NCCLM2NCopyService", m2n_service_type) + monkeypatch.setattr( + worker_module, "prepare_swap_model_weights", prepare_swap_model_weights + ) + monkeypatch.setattr( + worker_module, + "_world", + SimpleNamespace(pg_group_ranks={}, pg_map={}, pg_names={}), + ) + + worker.init_collective_mcore_generation( + "127.0.0.1", + 1234, + 4, + rank_offset=2, + refit_execution_batch_bytes=123, + refit_backend="nccl_m2n", + ) + + assert process_group._register_backend.call_count == 2 + # rank_offset is the whole point of this signature: generation ranks join the + # shared group AFTER the training ranks, so every backend must be built at + # the offset rank (0 + 2), not this world's local rank. + assert process_group_type.call_args.args[1:] == (2, 4) + assert nccl_backend_type.call_args.args[1:3] == (2, 4) + assert tcp_store_type.call_args.kwargs["is_master"] is False + nccl_backend_type.assert_called_once() + m2n_service_type.assert_called_once_with(group=process_group) + assert worker.refit_copy_service is m2n_service + prepare_swap_model_weights.assert_called_once_with( + src_model=None, + target_model=worker.model, + group=process_group, + src_rank_offset=0, + dst_rank_offset=2, + execution_batch_bytes=123, + ) + + +def test_megatron_fp8_refit_tasks_match_payload_mode() -> None: + from nemo_rl.models.policy.workers.megatron_policy_worker import ( + MegatronPolicyWorkerImpl, + ) + + worker = object.__new__(MegatronPolicyWorkerImpl) + worker.fp8_cfg = {"fp8_param": True, "fp8_recipe": "blockwise"} + worker.model = object() + physical_tasks = [object()] + logical_tasks = [None, object()] + worker.megatron_bridge = SimpleNamespace( + get_export_fp8_tasks=MagicMock(return_value=physical_tasks), + get_conversion_tasks=MagicMock(return_value=logical_tasks), + ) + + worker.refit_payload_mode = "logical_weights" + assert worker._build_refit_conversion_tasks() == logical_tasks[1:] + worker.megatron_bridge.get_export_fp8_tasks.assert_not_called() + + worker.refit_payload_mode = "bridge_export" + assert worker._build_refit_conversion_tasks() == physical_tasks + worker.megatron_bridge.get_export_fp8_tasks.assert_called_once_with(worker.model) + + +def test_fp8_export_payload_survives_for_non_megatron_backends() -> None: + """Bridge's FP8 export tasks must reach vLLM as fp8, not dequantized BF16. + + ``_iter_logical_refit_conversion_tasks`` exists so a Megatron inference + engine (BF16 or MXFP8 only) can consume quantized training storage. Applying + it to the vLLM path would replace the physical fp8 ``...weight`` with BF16 + while its ``..._scale_inv`` sibling passed through untouched, so vLLM would + rescale BF16 values with a blockwise scale - wrong weights, no error. + """ + from nemo_rl.models.policy.workers.megatron_policy_worker import ( + MegatronPolicyWorkerImpl, + ) + + fp8_payload = torch.zeros(4, 2, dtype=torch.float8_e4m3fn) + + class _QuantizedParam(torch.Tensor): + """Stands in for a live TE Float8BlockwiseQTensor on the module.""" + + module = SimpleNamespace( + weight=_QuantizedParam(torch.zeros(4, 2, dtype=torch.bfloat16)) + ) + task = SimpleNamespace( + param_weight=fp8_payload, + megatron_module=module, + param_name="linear_fc1.weight", + ) + + worker = object.__new__(MegatronPolicyWorkerImpl) + worker.model = object() + worker.draft_model = None + worker.refit_conversion_tasks = [task] + worker.megatron_bridge = SimpleNamespace( + export_hf_weights=MagicMock(return_value=iter(())) + ) + + worker.cfg = {"generation": {"backend": "vllm"}} + worker.refit_payload_mode = "bridge_export" + list(worker._iter_params_with_optional_kv_scales()) + forwarded = worker.megatron_bridge.export_hf_weights.call_args.kwargs[ + "conversion_tasks" + ] + # The vLLM path forwards Bridge's task list verbatim - not a rewriting generator. + assert list(forwarded) == [task] + assert forwarded[0].param_weight.dtype == torch.float8_e4m3fn + + worker.cfg = {"generation": {"backend": "megatron"}} + worker.refit_payload_mode = "logical_weights" + list(worker._iter_params_with_optional_kv_scales()) + megatron_forwarded = worker.megatron_bridge.export_hf_weights.call_args.kwargs[ + "conversion_tasks" + ] + # Megatron generation gets the materializing generator instead. + assert not isinstance(megatron_forwarded, list) + + +def test_local_hf_shards_follow_bridge_specs() -> None: + from megatron.bridge.models.conversion.param_mapping import LocalHFParamSpec + + from nemo_rl.models.policy.workers.megatron_policy_worker import ( + MegatronPolicyWorkerImpl, + ) + + fused = torch.arange(8, dtype=torch.bfloat16).reshape(4, 2) + task = SimpleNamespace( + param_weight=fused, + megatron_module=None, + param_name="linear_fc1.weight", + global_param_name="decoder.layers.0.mlp.linear_fc1.weight", + local_hf_param_specs=lambda: ( + LocalHFParamSpec("model.layers.0.mlp.gate_proj.weight", -2, 0, 2), + LocalHFParamSpec("model.layers.0.mlp.up_proj.weight", -2, 1, 2), + ), + ) + worker = object.__new__(MegatronPolicyWorkerImpl) + worker.refit_conversion_tasks = [task] + + shards = dict(worker._iter_local_hf_param_shards()) + + assert torch.equal(shards["model.layers.0.mlp.gate_proj.weight"].base, fused[:2]) + assert torch.equal(shards["model.layers.0.mlp.up_proj.weight"].base, fused[2:]) + fused[0, 0] = 99 + assert shards["model.layers.0.mlp.gate_proj.weight"].base[0, 0] == 99 + + +def test_local_refit_names_require_safe_views_from_every_grouped_task() -> None: + from nemo_rl.models.policy.workers.megatron_policy_worker import ( + _collect_local_refit_hf_names, + ) + + grouped_name = "model.layers.0.mlp.experts.down_proj" + safe_name = "model.layers.0.mlp.down_proj.weight" + tasks = [ + SimpleNamespace( + hf_param_names=(grouped_name,), + local_hf_param_specs=lambda: (object(),), + ), + SimpleNamespace( + hf_param_names=(grouped_name,), + local_hf_param_specs=lambda: (), + ), + SimpleNamespace( + hf_param_names=(safe_name,), + local_hf_param_specs=lambda: (object(),), + ), + ] + + assert _collect_local_refit_hf_names(tasks) == {safe_name} + + +def test_local_refit_names_are_expert_parallel_invariant(monkeypatch) -> None: + """Every EP rank must agree on the safe-view set. + + Bridge conversion tasks are EP-local (each rank only enumerates its own + ``experts.N.*``), but the name stream they are matched against is EP-global. + Without an EP reduction rank 0 would omit expert 1 purely because it never + saw a task for it, so the two ranks would disagree on the bulk/misc split. + """ + from nemo_rl.models.policy.workers import megatron_policy_worker as worker_module + from nemo_rl.models.policy.workers.megatron_policy_worker import ( + _collect_local_refit_hf_names, + ) + + expert0 = "model.layers.0.mlp.experts.0.down_proj.weight" + expert1 = "model.layers.0.mlp.experts.1.down_proj.weight" + unsafe = "model.layers.0.self_attn.qkv_proj.weight" + + # Rank 0 owns expert 0 and sees the unsafe attention mapping; rank 1 owns + # expert 1. Neither rank alone can name both experts. + per_rank = [ + {expert0: True, unsafe: False}, + {expert1: True}, + ] + + ep_group = object() + monkeypatch.setattr( + worker_module.parallel_state, + "get_expert_model_parallel_group", + lambda check_initialized=True: ep_group, + ) + monkeypatch.setattr( + worker_module.torch.distributed, "get_world_size", lambda group=None: 2 + ) + + def fake_all_gather_object(out, obj, group=None): + assert group is ep_group + # Pin what this rank contributes, not just what it receives: otherwise a + # wrong local support map would still produce the expected union. + assert obj == {expert0: True, unsafe: False} + out[:] = [obj, per_rank[1]] + + monkeypatch.setattr( + worker_module.torch.distributed, "all_gather_object", fake_all_gather_object + ) + + tasks = [ + SimpleNamespace( + hf_param_names=(expert0,), + local_hf_param_specs=lambda: (object(),), + ), + SimpleNamespace( + hf_param_names=(unsafe,), + local_hf_param_specs=lambda: (), + ), + ] + + # Both experts are present and safe; the unsafe attention name stays out. + assert _collect_local_refit_hf_names(tasks) == {expert0, expert1} + + +@pytest.mark.parametrize("pp_size", [1, 2]) +def test_nccl_reshard_all_misc_refit_supports_empty_bulk(pp_size: int) -> None: + from nemo_rl.models.policy.workers.megatron_policy_worker import ( + MegatronPolicyWorkerImpl, + ) + + misc_name = "model.layers.0.mlp.experts.down_proj" + task = SimpleNamespace( + global_param_name="decoder.layers.0.mlp.experts.linear_fc2.weight", + hf_param_names=(misc_name,), + mapping=SimpleNamespace(hf_param=misc_name), + local_hf_param_specs=lambda: (), + ) + worker = object.__new__(MegatronPolicyWorkerImpl) + worker.refit_conversion_tasks = [task] + worker._calculate_refit_param_info = MagicMock(return_value={}) + worker._iter_params_with_optional_kv_scales = MagicMock( + return_value=iter([(misc_name, torch.empty(2, 2))]) + ) + worker._build_layer_to_pp_stage = MagicMock() + worker.build_hf_to_local_param_map = MagicMock(return_value=HFToLocalParamMap()) + + info = worker.prepare_nccl_reshard_refit_info( + train_parallelism={"tp_size": 1, "ep_size": 1, "pp_size": pp_size}, + gen_parallelism={"tp_size": 1, "ep_size": 1, "pp_size": 1}, + train_world_size=pp_size, + gen_world_size=1, + ) + + assert info["layer_names"] == [] + assert info["per_layer_params"] == {} + assert list(info["misc_meta"]) == [misc_name] + assert worker.hf_to_local_param_map.specs == {} + assert worker._misc_conversion_tasks == [task] + worker._build_layer_to_pp_stage.assert_not_called() + + +def test_destination_refit_role_uses_common_worker_interface( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from nemo_rl.models.policy.workers.megatron_policy_worker import ( + MegatronPolicyWorkerImpl, + ) + + worker = object.__new__(MegatronPolicyWorkerImpl) + worker.refit_role = "destination" + worker.model_update_group = object() + worker._generation_nccl_reshard_groups = {} + tasks = [object()] + expected_map = HFToLocalParamMap() + refit_info = {"layer_names": [], "per_layer_params": {}} + state_dict_info = {"weight": ((2, 2), torch.bfloat16)} + worker._build_generation_refit_tasks = MagicMock( + return_value=([torch.nn.Module()], tasks) + ) + worker._build_destination_hf_to_local_param_map = MagicMock( + return_value=expected_map + ) + worker._prepare_destination_refit_info = MagicMock() + worker._prepare_destination_nccl_reshard_refit_info = MagicMock() + worker._destination_nccl_reshard_refit = MagicMock(return_value=True) + worker._update_destination_weights_from_collective = MagicMock(return_value=True) + set_device = MagicMock() + monkeypatch.setattr(torch.cuda, "current_device", lambda: 3) + monkeypatch.setattr(torch.cuda, "set_device", set_device) + + assert worker.build_hf_to_local_param_map(refit_info) is expected_map + worker.prepare_refit_info(state_dict_info) + worker.prepare_nccl_reshard_refit_info(refit_info=refit_info) + assert asyncio.run(worker.nccl_reshard_refit()) + assert asyncio.run(worker.update_weights_from_collective()) + + worker._build_destination_hf_to_local_param_map.assert_called_once_with( + refit_info, tasks + ) + worker._prepare_destination_refit_info.assert_called_once_with(state_dict_info) + worker._prepare_destination_nccl_reshard_refit_info.assert_called_once_with( + refit_info + ) + worker._destination_nccl_reshard_refit.assert_called_once_with(refit_timeout_s=None) + worker._update_destination_weights_from_collective.assert_called_once_with( + refit_timeout_s=None + ) + assert set_device.call_args_list == [call(3), call(3)] + + +def test_megatron_destination_misc_plan_uses_shipped_misc_metadata() -> None: + from nemo_rl.models.generation.megatron.megatron_worker import ( + MegatronGenerationRefitMixin, + ) + + bulk_task = SimpleNamespace( + dependencies=("model.layers.0.mlp.experts.0.gate_proj.weight",) + ) + misc_task = SimpleNamespace(dependencies=("model.layers.0.input_layernorm.weight",)) + model_chunks = [torch.nn.Module()] + worker = object.__new__(MegatronGenerationRefitMixin) + worker._build_generation_refit_tasks = MagicMock( + return_value=(model_chunks, [bulk_task, misc_task]) + ) + worker._prepare_mxfp8_refit = MagicMock() + worker.build_hf_to_local_param_map = MagicMock(return_value=HFToLocalParamMap()) + worker._install_generation_refit_plan = MagicMock() + refit_info = { + "layer_names": [], + "per_layer_params": {}, + "misc_meta": { + "model.layers.0.input_layernorm.weight": { + "shape": [2], + "dtype": "torch.bfloat16", + } + }, + } + + worker._prepare_destination_nccl_reshard_refit_info(refit_info) + + worker._install_generation_refit_plan.assert_called_once_with( + model_chunks=model_chunks, + tasks=[misc_task], + state_dict_info={ + "model.layers.0.input_layernorm.weight": ((2,), torch.bfloat16) + }, + ) + + +def test_megatron_m2n_stages_fused_mxfp8_weight_until_complete() -> None: + from megatron.core.inference.quantization.mxfp8_tensor import MXFP8Tensor + + from nemo_rl.models.generation.megatron.megatron_worker import ( + MegatronGenerationRefitMixin, + ) + + gate_name = "model.layers.0.mlp.gate_proj.weight" + up_name = "model.layers.0.mlp.up_proj.weight" + destination = MXFP8Tensor( + data=torch.empty((8, 2)), + scale=torch.empty(1), + dtype=torch.bfloat16, + backend="triton", + ) + task = _make_refit_task( + param_name="decoder.layers.0.mlp.linear_fc1.weight", + destination=destination, + dependencies=(gate_name, up_name), + local_specs=((gate_name, "gate"), (up_name, "up")), + ) + refit_info = { + "layer_names": ["model.layers.0"], + "per_layer_params": { + "model.layers.0": [ + {"name": gate_name}, + {"name": up_name}, + ] + }, + } + worker = object.__new__(MegatronGenerationRefitMixin) + worker._generation_m2n_pending = {} + worker._write_generation_refit_weight = MagicMock() + param_map = worker._build_destination_hf_to_local_param_map(refit_info, [task]) + + gate_spec = param_map.get(gate_name) + gate_ctx = gate_spec.pre(None) + gate_ctx.buf.fill_(3) + gate_spec.post(gate_ctx) + worker._write_generation_refit_weight.assert_not_called() + + up_spec = param_map.get(up_name) + up_ctx = up_spec.pre(None) + up_ctx.buf.fill_(5) + up_spec.post(up_ctx) + + fused = worker._write_generation_refit_weight.call_args.args[1] + assert torch.equal(fused[:4], torch.full_like(fused[:4], 3)) + assert torch.equal(fused[4:], torch.full_like(fused[4:], 5)) + assert worker._generation_m2n_pending == {} + + +def test_megatron_m2n_unstacks_grouped_experts_into_local_weights() -> None: + from nemo_rl.models.generation.megatron.megatron_worker import ( + MegatronGenerationRefitMixin, + ) + + tasks = [] + destinations = [] + for expert in range(2): + gate_name = f"model.layers.0.mlp.experts.{expert}.gate_proj.weight" + up_name = f"model.layers.0.mlp.experts.{expert}.up_proj.weight" + destination = torch.zeros((4, 2), dtype=torch.bfloat16) + destinations.append(destination) + tasks.append( + _make_refit_task( + param_name=( + f"decoder.layers.0.mlp.experts.local_experts.{expert}.linear_fc1.weight" + ), + destination=destination, + dependencies=(gate_name, up_name), + local_specs=((gate_name, "gate"), (up_name, "up")), + ) + ) + + grouped_gate = "model.layers.0.mlp.experts.gate_proj.weight" + grouped_up = "model.layers.0.mlp.experts.up_proj.weight" + refit_info = { + "layer_names": ["model.layers.0"], + "per_layer_params": { + "model.layers.0": [ + {"name": grouped_gate, "grouped_expert_proj": "gate_proj"}, + {"name": grouped_up, "grouped_expert_proj": "up_proj"}, + ] + }, + } + worker = object.__new__(MegatronGenerationRefitMixin) + param_map = worker._build_destination_hf_to_local_param_map(refit_info, tasks) + + gate_spec = param_map.get(grouped_gate) + gate_ctx = gate_spec.pre(None) + gate_ctx.buf[0].fill_(1) + gate_ctx.buf[1].fill_(2) + gate_spec.post(gate_ctx) + up_spec = param_map.get(grouped_up) + up_ctx = up_spec.pre(None) + up_ctx.buf[0].fill_(3) + up_ctx.buf[1].fill_(4) + up_spec.post(up_ctx) + + assert torch.equal(destinations[0][:2], torch.ones_like(destinations[0][:2])) + assert torch.equal(destinations[0][2:], torch.full_like(destinations[0][2:], 3)) + assert torch.equal(destinations[1][:2], torch.full_like(destinations[1][:2], 2)) + assert torch.equal(destinations[1][2:], torch.full_like(destinations[1][2:], 4)) + + +@pytest.mark.parametrize( + ("param_info", "expected_message"), + [ + ( + {"name": "model.layers.0.mlp.gate_proj.weight"}, + "No local Megatron destination maps", + ), + ( + { + "name": "model.layers.0.mlp.experts.gate_proj.weight", + "grouped_expert_proj": "gate_proj", + }, + "No local Megatron experts map", + ), + ], +) +def test_megatron_m2n_rejects_weights_with_no_local_destination( + param_info: dict[str, str], expected_message: str +) -> None: + """A train/gen geometry mismatch must fail while the plan is built. + + Every shipped bulk weight has to land somewhere local. If it doesn't, the + receive plan is silently short a destination and the M-to-N collective would + post a mismatched set of transfers at the first refit. + """ + from nemo_rl.models.generation.megatron.megatron_worker import ( + MegatronGenerationRefitMixin, + ) + + worker = object.__new__(MegatronGenerationRefitMixin) + refit_info = { + "layer_names": ["model.layers.0"], + "per_layer_params": {"model.layers.0": [param_info]}, + } + + with pytest.raises(ValueError, match=expected_message): + worker._build_destination_hf_to_local_param_map(refit_info, []) + + +def test_prepare_mxfp8_refit_replaces_only_quantized_parameters( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from nemo_rl.models.generation.megatron import megatron_worker as worker_module + from nemo_rl.models.generation.megatron.megatron_worker import ( + MegatronGenerationRefitMixin, + ) + + core = torch.nn.Module() + core.config = SimpleNamespace( + transformer_impl="inference_optimized", + fp8_recipe="mxfp8", + inference_grouped_gemm_backend="torch", + ) + core.decoder = torch.nn.Module() + core.decoder.weight = torch.nn.Parameter(torch.zeros(2, 2)) + core.decoder.norm = torch.nn.Parameter(torch.zeros(2)) + quantized_destination = object() + quantize = MagicMock(return_value={"weight": quantized_destination}) + # megatron_worker binds this name at import time, so the patch must target + # the consuming module rather than megatron.core...quantization.utils. + monkeypatch.setattr(worker_module, "quantize_params_to_mxfp8", quantize) + + weight_task = _make_refit_task( + param_name="weight", + destination=core.decoder.weight, + dependencies=("weight",), + ) + norm_task = _make_refit_task( + param_name="norm", + destination=core.decoder.norm, + dependencies=("norm",), + ) + worker = object.__new__(MegatronGenerationRefitMixin) + worker.model = core + worker._inference_engine_initialized = False + + worker._prepare_mxfp8_refit([weight_task, norm_task]) + + assert weight_task.destination is quantized_destination + assert norm_task.destination is core.decoder.norm + + +def test_megatron_generation_joins_all_m2n_pipeline_groups( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from nemo_rl.distributed import stateless_process_group + from nemo_rl.models.generation.megatron.megatron_worker import ( + MegatronGenerationRefitMixin, + ) + + groups = [MagicMock(), MagicMock()] + process_group = MagicMock(side_effect=groups) + monkeypatch.setattr(stateless_process_group, "StatelessProcessGroup", process_group) + monkeypatch.setattr(torch.cuda, "current_device", lambda: 0) + monkeypatch.setattr(torch.cuda, "empty_cache", lambda: None) + worker = object.__new__(MegatronGenerationRefitMixin) + worker.rank = 2 + + worker.init_nccl_reshard_comm_groups_generation( + pp_ips=["10.0.0.1", "10.0.0.2"], + pp_ports=[1234, 1235], + pp_size=2, + train_ranks_per_stage=4, + sub_world_size=10, + rank_prefix=1, + ) + + assert process_group.call_args_list == [ + call(master_address="10.0.0.1", port=1234, rank=5, world_size=10), + call(master_address="10.0.0.2", port=1235, rank=5, world_size=10), + ] + for group in groups: + group.init_nccl_communicator.assert_called_once_with(device=0) + def test_model_owned_packing_capability_is_detected(): from nemo_rl.models.policy.workers.megatron_policy_worker import ( @@ -230,6 +899,95 @@ def test_refit_size_estimate_casts_floating_weight_to_export_dtype(): ) +def test_megatron_refit_bridge_tasks_export_logical_quantized_weights( + monkeypatch: pytest.MonkeyPatch, +) -> None: + import nemo_rl.models.policy.workers.megatron_policy_worker as worker_module + + @dataclass(frozen=True) + class Task: + param_weight: torch.Tensor | None + param_name: str = "" + megatron_module: object | None = None + + quantized_source = torch.zeros((4, 2), dtype=torch.uint8) + bridge_payload = torch.ones((4, 2), dtype=torch.uint8) + logical_weight = torch.arange(8, dtype=torch.bfloat16).reshape(4, 2) + bf16_source = torch.ones((2, 2), dtype=torch.bfloat16) + dequantize = MagicMock( + side_effect=lambda tensor: logical_weight + if tensor is quantized_source + else tensor + ) + monkeypatch.setattr( + worker_module, + "_is_quantized_refit_source", + lambda tensor: tensor is quantized_source, + ) + monkeypatch.setattr(worker_module, "_dequantize_refit_source", dequantize) + + tasks = [ + Task( + bridge_payload, + param_name="decoder.layers.0.mlp.weight", + megatron_module=SimpleNamespace(weight=quantized_source), + ), + None, + Task(bf16_source), + ] + exported = list( + worker_module.MegatronPolicyWorkerImpl._iter_logical_refit_conversion_tasks( + tasks + ) + ) + + assert exported[0] is not tasks[0] + assert exported[0].param_weight is logical_weight + assert exported[1] is None + assert exported[2] is tasks[2] + assert dequantize.call_args_list == [call(quantized_source), call(bf16_source)] + + +def test_megatron_refit_quantized_source_dequantizes_once_per_layer( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from megatron.bridge.models.conversion.param_mapping import LocalHFParamSpec + + import nemo_rl.models.policy.workers.megatron_policy_worker as worker_module + + worker = object.__new__(worker_module.MegatronPolicyWorkerImpl) + quantized_source = torch.zeros((2, 4, 2), dtype=torch.uint8) + logical_weight = torch.arange(16, dtype=torch.bfloat16).reshape(2, 4, 2) + dequantize = MagicMock(return_value=logical_weight) + monkeypatch.setattr( + worker_module, + "_is_quantized_refit_source", + lambda tensor: tensor is quantized_source, + ) + monkeypatch.setattr(worker_module, "_dequantize_refit_source", dequantize) + + gate_spec = worker._local_refit_source_spec( + quantized_source, LocalHFParamSpec("gate", -2, 0, 2) + ) + up_spec = worker._local_refit_source_spec( + quantized_source, LocalHFParamSpec("up", -2, 1, 2) + ) + grouped_spec = worker_module.LocalParamSpec( + base=worker_module._GroupedRefitSource((gate_spec, up_spec)) + ) + logical_source_cache = {} + + gate = worker._materialize_local_refit_spec(gate_spec, logical_source_cache).buf + grouped = worker._materialize_local_refit_spec( + grouped_spec, logical_source_cache + ).buf + + assert gate.is_contiguous() + assert torch.equal(grouped[0], logical_weight[:, :2]) + assert torch.equal(grouped[1], logical_weight[:, 2:]) + dequantize.assert_called_once_with(quantized_source) + + def test_qwen3vl_type_fallback_still_delegates_packing(): from megatron.bridge.models.qwen_vl.modelling_qwen3_vl.model import Qwen3VLModel @@ -276,8 +1034,9 @@ def test_megatron_offload_before_refit_finalizes_async_save_first(monkeypatch): worker.fp8_cfg = None worker.cfg = {"megatron_cfg": {"clear_memory_caches_before_refit": False}} worker.finalize_async_save = lambda: events.append("finalize_async_save") + worker._uses_mxfp8_overlap_shared_param_buffer = lambda: False worker.move_model = lambda model, device, move_params, move_grads: ( - events.append("move_model") or model + events.append(("move_model", move_grads)) or model ) class _AllocatorWakeup: @@ -300,7 +1059,48 @@ def cuda(self): MegatronPolicyWorkerImpl.offload_before_refit(worker) assert events[0] == "finalize_async_save" - assert events.index("finalize_async_save") < events.index("move_model") + assert events.index("finalize_async_save") < events.index(("move_model", True)) + + +def test_megatron_offload_before_refit_materializes_latest_mxfp8_weights(monkeypatch): + """Refit must see optimizer updates before the next overlapped train forward.""" + from nemo_rl.models.policy.workers.megatron_policy_worker import ( + MegatronPolicyWorkerImpl, + ) + + events = [] + worker = object.__new__(MegatronPolicyWorkerImpl) + worker.model = object() + worker.optimizer = None + worker.optimizer_cpu_offload = False + worker.fp8_cfg = None + worker.cfg = {"megatron_cfg": {"clear_memory_caches_before_refit": False}} + worker.finalize_async_save = lambda: events.append("finalize_async_save") + worker._uses_mxfp8_overlap_shared_param_buffer = lambda: True + worker._forward_pre_hook_enabled = lambda: True + worker._disable_forward_pre_hook_until_next_train_step = ( + lambda *, param_sync=False: events.append(("disable_hook", param_sync)) + ) + worker.move_model = lambda model, device, move_params, move_grads: ( + events.append(("move_model", move_grads)) or model + ) + + class _AllocatorWakeup: + def cuda(self): + return self + + monkeypatch.setattr(torch.cuda, "memory_allocated", lambda *args, **kwargs: 0) + monkeypatch.setattr(torch.cuda, "memory_reserved", lambda *args, **kwargs: 0) + monkeypatch.setattr(torch.cuda, "empty_cache", lambda: None) + monkeypatch.setattr(torch, "randn", lambda *args, **kwargs: _AllocatorWakeup()) + + MegatronPolicyWorkerImpl.offload_before_refit(worker) + + assert events[:3] == [ + "finalize_async_save", + ("disable_hook", True), + ("move_model", False), + ] @pytest.mark.parametrize("offload_optimizer", [False, True]) @@ -321,6 +1121,7 @@ def test_megatron_offload_before_refit_honors_offload_optimizer_for_refit( worker.fp8_cfg = None worker.cfg = {"megatron_cfg": {"clear_memory_caches_before_refit": False}} worker.finalize_async_save = lambda: None + worker._uses_mxfp8_overlap_shared_param_buffer = lambda: False worker.move_model = lambda model, device, move_params, move_grads: model worker.move_optimizer = lambda device: moved.append(device) @@ -339,24 +1140,34 @@ def cuda(self): @pytest.mark.parametrize( - "generation_backend, colocated, has_inference_model, expect_move_params", + "generation_backend, colocated, has_inference_model, shared_buffer, " + "expect_move_params, expect_move_grads", [ # Plain training worker (no generation config): params always move. - (None, False, False, True), + (None, False, False, False, True, True), # Shared-model colocated Megatron generation: params must stay resident — # inference CUDA graphs replay with capture-time param pointers, and a # CPU round-trip re-allocates their storage. - ("megatron", True, False, False), + ("megatron", True, False, False, False, True), # Colocated reshard (dedicated inference model): params still move; the # graphs live on the dedicated, torch_memory_saver-managed model. - ("megatron", True, True, True), + ("megatron", True, True, False, True, True), + # MXFP8 overlap aliases the parameter and gradient buffers, so neither + # storage can move independently. + (None, False, False, True, False, False), ], ) def test_megatron_offload_after_refit_finalizes_before_model_move( - monkeypatch, generation_backend, colocated, has_inference_model, expect_move_params + monkeypatch, + generation_backend, + colocated, + has_inference_model, + shared_buffer, + expect_move_params, + expect_move_grads, ): """Checkpoint CUDA IPC handles must be dropped before model storage is replaced, - and shared-model colocated generation must keep its params resident.""" + and graph/shared-buffer owners must keep their aliased storage resident.""" from nemo_rl.models.policy.workers.megatron_policy_worker import ( MegatronPolicyWorkerImpl, ) @@ -373,6 +1184,7 @@ def test_megatron_offload_after_refit_finalizes_before_model_move( worker.inference_model = object() if has_inference_model else None worker._colocated_reshard_plan = None worker.finalize_async_save = lambda: events.append("finalize_async_save") + worker._uses_mxfp8_overlap_shared_param_buffer = lambda: shared_buffer worker.move_model = lambda model, device, **kwargs: ( events.append("move_model") or move_kwargs.append(kwargs) or model ) @@ -400,6 +1212,7 @@ def cuda(self): assert events.index("finalize_async_save") < events.index("move_model") assert events.index("eval") < events.index("move_model") assert move_kwargs[0]["move_params"] is expect_move_params + assert move_kwargs[0]["move_grads"] is expect_move_grads def test_megatron_finish_inference_evals_before_model_offload(monkeypatch): @@ -587,7 +1400,6 @@ def test_megatron_prepare_for_training_leaves_native_cpu_optimizer_placement(): worker = object.__new__(MegatronPolicyWorkerImpl) model = _FakeTrainableModel() - worker.model = model worker.optimizer = object() worker.optimizer_cpu_offload = True @@ -602,6 +1414,45 @@ def test_megatron_prepare_for_training_leaves_native_cpu_optimizer_placement(): assert model.train_called +def test_megatron_prepare_for_logprobs_restores_mxfp8_shared_buffer(monkeypatch): + """Restore and restage the shared buffer before MXFP8 parameter gather.""" + from nemo_rl.models.policy.workers.megatron_policy_worker import ( + MegatronPolicyWorkerImpl, + ) + + worker = object.__new__(MegatronPolicyWorkerImpl) + model = _FakeTrainableModel() + move_calls = [] + + worker.rank = 0 + worker.model = model + worker.optimizer = object() + worker.optimizer_cpu_offload = False + worker.offload_optimizer_for_logprob = False + worker._uses_mxfp8_overlap_shared_param_buffer = lambda: True + worker._log_gpu_mem = lambda *_args, **_kwargs: None + worker.move_model = lambda model, device, **kwargs: ( + move_calls.append((device, kwargs)) or model + ) + stage_calls = [] + worker._copy_main_params_to_param_buffer = lambda zero_grad_buffer=False: ( + stage_calls.append(zero_grad_buffer) + ) + + class _AllocatorWakeup: + def cuda(self): + return self + + monkeypatch.setattr(torch, "randn", lambda *args, **kwargs: _AllocatorWakeup()) + monkeypatch.setattr(torch.cuda, "empty_cache", lambda: None) + + MegatronPolicyWorkerImpl.prepare_for_lp_inference(worker) + + assert model.eval_called + assert move_calls == [("cuda", {"move_grads": True})] + assert stage_calls == [True] + + def test_set_moe_grad_scale_func_sets_and_clears_on_model_config(): """_set_moe_grad_scale_func should set/clear moe_grad_scale_func on the config.""" from nemo_rl.models.policy.workers.megatron_policy_worker import ( @@ -848,6 +1699,7 @@ def create_megatron_test_config( "offload_optimizer_for_logprob": False, "generation": { "backend": generation_backend, + "refit_transport": "mcore" if generation_backend == "megatron" else None, "temperature": 1.0, "top_p": 1.0, "top_k": None, @@ -869,6 +1721,7 @@ def create_megatron_test_config( "num_speculative_tokens": 0, "logprobs_mode": "processed_logprobs", "refit_backend": "gloo", # not nvshmem: its NVLS multicast init is unavailable in CI + "refit_execution_batch_bytes": None, "parsers": [], "expose_http_server": False, }, diff --git a/tests/unit/reference_configs/grpo_math_1B.yaml b/tests/unit/reference_configs/grpo_math_1B.yaml index 373ad2bffb2..58128f99e00 100644 --- a/tests/unit/reference_configs/grpo_math_1B.yaml +++ b/tests/unit/reference_configs/grpo_math_1B.yaml @@ -371,7 +371,8 @@ 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 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. refit_transport: null # Keep real-quant export tensors on CPU by default. Set false only for @@ -402,8 +403,9 @@ policy: kv_cache_management_mode: "persist" # KV cache lifecycle across suspend/resume. Options: "persist", "offload", "recompute". materialize_only_last_token_logits: true # Materialize logits only for the last token of each sequence during decode. num_speculative_tokens: 0 + refit_backend: "nccl" # 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" 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. vllm_cfg: diff --git a/tests/unit/reference_configs/ppo_math_1B_megatron.yaml b/tests/unit/reference_configs/ppo_math_1B_megatron.yaml index 9bb901085a5..48b9ba47250 100644 --- a/tests/unit/reference_configs/ppo_math_1B_megatron.yaml +++ b/tests/unit/reference_configs/ppo_math_1B_megatron.yaml @@ -255,6 +255,8 @@ policy: enable_chunked_prefill: true unified_memory_level: 0 max_tokens: 16384 + refit_backend: "nccl" # 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 vllm_cfg: async_engine: false diff --git a/tests/unit/utils/test_checkpoint_engine.py b/tests/unit/utils/test_checkpoint_engine.py index aacd0bb055b..67b980f45c2 100644 --- a/tests/unit/utils/test_checkpoint_engine.py +++ b/tests/unit/utils/test_checkpoint_engine.py @@ -576,6 +576,7 @@ def resolve_nixl_backend_kwargs(nixl_kwargs): maybe_preinit_nixl_checkpoint_engine( { "generation": { + "backend": "vllm", "refit_transport": "nixl", "refit_cfg": {"nixl": {}}, } @@ -585,6 +586,18 @@ def resolve_nixl_backend_kwargs(nixl_kwargs): ) assert calls == [{"backend_name": "UCX", "backend_init_params": None}] + assert ( + maybe_preinit_nixl_checkpoint_engine( + { + "generation": { + "backend": "megatron", + "refit_transport": "mcore", + } + } + ) + is None + ) + def test_merge_weight_chunk_batches_uses_aligned_zero_copy_view(): fp32_w = torch.tensor([7.0, 8.0], dtype=torch.float32) diff --git a/tests/unit/weight_sync/test_checkpoint_engine_weight_synchronizer.py b/tests/unit/weight_sync/test_checkpoint_engine_weight_synchronizer.py index f7e8a7c43e1..732cc5044c6 100644 --- a/tests/unit/weight_sync/test_checkpoint_engine_weight_synchronizer.py +++ b/tests/unit/weight_sync/test_checkpoint_engine_weight_synchronizer.py @@ -19,7 +19,6 @@ import pytest from nemo_rl.models.generation.constants import ( - MEGATRON_BACKEND, SGLANG_BACKEND, VLLM_BACKEND, ) @@ -321,7 +320,6 @@ class TestCheckpointEngineFactory: (VLLM_BACKEND, False, CheckpointEngineWeightSynchronizer), (VLLM_BACKEND, True, ValueError), (SGLANG_BACKEND, False, NotImplementedError), - (MEGATRON_BACKEND, False, NotImplementedError), ], ) def test_checkpoint_engine_factory_routing(self, backend, colocated, expected): diff --git a/tests/unit/weight_sync/test_nccl_reshard_utils.py b/tests/unit/weight_sync/test_nccl_reshard_utils.py index ad3e59d91f7..5cb659505e4 100644 --- a/tests/unit/weight_sync/test_nccl_reshard_utils.py +++ b/tests/unit/weight_sync/test_nccl_reshard_utils.py @@ -159,7 +159,7 @@ def test_check_nccl_reshard_refit_support_rejects_blockwise_fp8_to_mxfp8() -> No ), ( {"backend": "sglang"}, - "policy.generation.backend must be 'vllm' (got 'sglang')", + "policy.generation.backend must be 'vllm' or 'megatron' (got 'sglang')", ), ], ) @@ -175,6 +175,217 @@ def test_check_nccl_reshard_refit_support_rejects_invalid_config( assert expected_violation in str(exc_info.value) +@pytest.mark.parametrize( + ("train_fp8_cfg", "generation_fp8_cfg"), + [ + ({"enabled": False}, {"enabled": False}), + ( + {"enabled": True, "fp8_recipe": "mxfp8", "fp8_param": True}, + {"enabled": True, "fp8_recipe": "mxfp8"}, + ), + ], +) +def test_check_nccl_reshard_refit_support_accepts_megatron_bf16_and_mxfp8( + train_fp8_cfg: dict[str, object], generation_fp8_cfg: dict[str, object] +) -> None: + config = _valid_nccl_reshard_config() + config.policy["precision"] = "bfloat16" + config.policy["megatron_cfg"]["expert_tensor_parallel_size"] = 2 + config.policy["megatron_cfg"]["fp8_cfg"] = train_fp8_cfg + config.policy["generation"] = { + "backend": "megatron", + "refit_transport": "nccl_reshard", + "colocated": {"enabled": False}, + "mcore_generation_config": { + "expert_model_parallel_size": 2, + "expert_tensor_parallel_size": 2, + "pipeline_model_parallel_size": 1, + # MXFP8 destinations quantize through resolve_mxfp8_backend, so the + # grouped-GEMM backend must be one it accepts. MCore's default + # ("vllm") is not, and NeMo-RL supplies no default of its own. + "inference_grouped_gemm_backend": "torch", + # _prepare_mxfp8_refit only installs MXFP8 destinations for + # inference_optimized cores; this mirrors the functional test. + "transformer_impl": "inference_optimized", + "fp8_cfg": generation_fp8_cfg, + }, + } + + check_nccl_reshard_refit_support(config) + + +def test_check_nccl_reshard_accepts_blockwise_megatron_source() -> None: + config = _valid_nccl_reshard_config() + config.policy["precision"] = "bfloat16" + config.policy["megatron_cfg"]["fp8_cfg"] = { + "enabled": True, + "fp8_recipe": "blockwise", + "fp8_param": True, + } + config.policy["generation"] = { + "backend": "megatron", + "refit_transport": "nccl_reshard", + "colocated": {"enabled": False}, + "mcore_generation_config": { + "expert_tensor_parallel_size": 1, + "pipeline_model_parallel_size": 1, + "fp8_cfg": {"enabled": False}, + }, + } + + check_nccl_reshard_refit_support(config) + + +def test_check_nccl_reshard_refit_support_accepts_megatron_transport() -> None: + config = _valid_nccl_reshard_config() + config.policy["precision"] = "bfloat16" + config.policy["generation"] = { + "backend": "megatron", + "refit_transport": "nccl_reshard", + "colocated": {"enabled": False}, + "mcore_generation_config": { + "refit_backend": "nvshmem", + "pipeline_model_parallel_size": 1, + }, + } + + check_nccl_reshard_refit_support(config) + + +def test_check_nccl_reshard_rejects_mxfp8_with_unsupported_grouped_gemm_backend() -> ( + None +): + """MXFP8 inference quantizes only through the torch/flashinfer backends. + + MCore's own guard for this compares ``config.fp8`` (an e4m3/hybrid *format*) + against ``"mxfp8"`` (a *recipe*) and so never fires, and the failure would + otherwise surface deep inside the first refit. + """ + config = _valid_nccl_reshard_config() + config.policy["precision"] = "bfloat16" + config.policy["generation"] = { + "backend": "megatron", + "refit_transport": "nccl_reshard", + "colocated": {"enabled": False}, + "mcore_generation_config": { + "pipeline_model_parallel_size": 1, + "inference_grouped_gemm_backend": "vllm", + "fp8_cfg": {"enabled": True, "fp8_recipe": "mxfp8"}, + }, + } + + with pytest.raises(ValueError, match="inference_grouped_gemm_backend"): + check_nccl_reshard_refit_support(config) + + +@pytest.mark.parametrize("gemm_backend", ["torch", "flashinfer"]) +def test_check_nccl_reshard_accepts_supported_grouped_gemm_backends( + gemm_backend: str, +) -> None: + config = _valid_nccl_reshard_config() + config.policy["precision"] = "bfloat16" + config.policy["generation"] = { + "backend": "megatron", + "refit_transport": "nccl_reshard", + "colocated": {"enabled": False}, + "mcore_generation_config": { + "pipeline_model_parallel_size": 1, + "inference_grouped_gemm_backend": gemm_backend, + "transformer_impl": "inference_optimized", + "fp8_cfg": {"enabled": True, "fp8_recipe": "mxfp8"}, + }, + } + + check_nccl_reshard_refit_support(config) + + +def test_check_nccl_reshard_rejects_mxfp8_with_omitted_grouped_gemm_backend() -> None: + """An omitted key resolves to MCore's default 'vllm', which MXFP8 rejects. + + Treating the omitted case as valid would let exactly the configuration the + guard exists for reach the first refit before failing. + """ + config = _valid_nccl_reshard_config() + config.policy["precision"] = "bfloat16" + config.policy["generation"] = { + "backend": "megatron", + "refit_transport": "nccl_reshard", + "colocated": {"enabled": False}, + "mcore_generation_config": { + "pipeline_model_parallel_size": 1, + "fp8_cfg": {"enabled": True, "fp8_recipe": "mxfp8"}, + }, + } + + with pytest.raises(ValueError, match="inference_grouped_gemm_backend"): + check_nccl_reshard_refit_support(config) + + +def _megatron_gen_config(*, policy_updates=None, **mcore_generation_config): + config = _valid_nccl_reshard_config() + config.policy["precision"] = "bfloat16" + config.policy.update(policy_updates or {}) + config.policy["generation"] = { + "backend": "megatron", + "refit_transport": "nccl_reshard", + "colocated": {"enabled": False}, + "mcore_generation_config": { + "pipeline_model_parallel_size": 1, + **mcore_generation_config, + }, + } + return config + + +@pytest.mark.parametrize( + ("policy_updates", "mcore_generation_config", "expected_violation"), + [ + # Gen-side PP has no stage-aware destination routing yet. + ( + None, + {"pipeline_model_parallel_size": 2}, + "pipeline_model_parallel_size must be 1", + ), + # The transport materializes logical BF16; a non-BF16 trainer would + # silently ship the wrong dtype. + ({"precision": "float32"}, {}, "policy.precision must be 'bfloat16'"), + # fp8_param without fp8_cfg.enabled is a half-configured trainer. + ( + {"megatron_cfg": {"enabled": True, "fp8_cfg": {"fp8_param": True}}}, + {}, + "fp8_cfg.enabled=True", + ), + # Megatron inference weights are BF16 or MXFP8 only. + ( + None, + {"fp8_cfg": {"enabled": True, "fp8_recipe": "blockwise"}}, + "fp8_recipe must be 'mxfp8'", + ), + # Without inference_optimized, _prepare_mxfp8_refit installs no MXFP8 + # destination and the refit silently copies BF16 into TE FP8 params. + ( + None, + { + "fp8_cfg": {"enabled": True, "fp8_recipe": "mxfp8"}, + "inference_grouped_gemm_backend": "torch", + }, + "transformer_impl='inference_optimized'", + ), + ], +) +def test_check_nccl_reshard_rejects_unsupported_megatron_generation( + policy_updates, mcore_generation_config, expected_violation +) -> None: + config = _megatron_gen_config( + policy_updates=policy_updates, **mcore_generation_config + ) + + with pytest.raises(ValueError) as exc_info: + check_nccl_reshard_refit_support(config) + + assert expected_violation in str(exc_info.value) + + # -------------------------------------------------------------------------- # MeshInfo # -------------------------------------------------------------------------- @@ -377,6 +588,24 @@ def test_get_placements_expert_tp_shifts_by_one(): ) +def test_get_placements_expert_ep_and_etp_shard_distinct_axes(): + # Megatron generation activates EP and expert-TP in the SAME mesh (the layout + # build_mesh_info emits for etp>1, ep>1). EP must own the expert dim and ETP + # the shifted projection dim; asserting per-axis rather than on the set of + # shard dims is what catches the two being swapped. + dm = {"ep": 0, "tp": 1} + gate = get_placements("a.mlp.experts.gate_proj.weight", dm, 3) + down = get_placements("a.mlp.experts.down_proj.weight", dm, 3) + assert _shard_dim_at(gate, dm, "ep") == 0 + assert _shard_dim_at(gate, dm, "tp") == 1 + assert _shard_dim_at(down, dm, "ep") == 0 + assert _shard_dim_at(down, dm, "tp") == 2 + # The router gate is not an expert param: it stays replicated on both axes. + assert all( + isinstance(p, Replicate) for p in get_placements("a.mlp.gate.weight", dm, 2) + ) + + # -------------------------------------------------------------------------- # group_expert_params_in_metadata # -------------------------------------------------------------------------- @@ -567,6 +796,73 @@ def test_build_refit_info_groups_experts_and_tags_them(): ) +def test_build_refit_info_replicates_megatron_experts_when_etp_is_one(): + info = build_nccl_reshard_refit_info( + _moe_metadata(num_experts=2), + train_parallelism={"tp_size": 1, "ep_size": 2, "pp_size": 1}, + gen_parallelism={ + "tp_size": 2, + "ep_size": 1, + "etp_size": 1, + "pp_size": 1, + }, + train_world_size=2, + gen_world_size=2, + ) + + gate = _find(info, "model.layers.0.mlp.experts.gate_proj.weight") + assert all(isinstance(placement, Replicate) for placement in gate["dst_placements"]) + + +def test_build_refit_info_shards_megatron_experts_across_ep_and_etp(): + info = build_nccl_reshard_refit_info( + _moe_metadata(num_experts=2), + train_parallelism={ + "tp_size": 1, + "ep_size": 2, + "etp_size": 2, + "pp_size": 1, + }, + gen_parallelism={ + "tp_size": 2, + "ep_size": 2, + "etp_size": 2, + "pp_size": 1, + }, + train_world_size=4, + gen_world_size=4, + ) + + gate = _find(info, "model.layers.0.mlp.experts.gate_proj.weight") + down = _find(info, "model.layers.0.mlp.experts.down_proj.weight") + assert { + placement.dim + for placement in gate["src_placements"] + if isinstance(placement, Shard) + } == {0, 1} + assert { + placement.dim + for placement in down["src_placements"] + if isinstance(placement, Shard) + } == {0, 2} + assert any( + isinstance(placement, Shard) and placement.dim == 0 + for placement in gate["dst_placements"] + ) + assert any( + isinstance(placement, Shard) and placement.dim == 1 + for placement in gate["dst_placements"] + ) + assert any( + isinstance(placement, Shard) and placement.dim == 0 + for placement in down["dst_placements"] + ) + assert any( + isinstance(placement, Shard) and placement.dim == 2 + for placement in down["dst_placements"] + ) + + # -------------------------------------------------------------------------- # make_nccl_reshard_refit_info_wire_safe / restore_refit_info_placements # -------------------------------------------------------------------------- diff --git a/tests/unit/weight_sync/test_reshard_rebuild.py b/tests/unit/weight_sync/test_reshard_rebuild.py index 92fa81e2881..b54211070ff 100644 --- a/tests/unit/weight_sync/test_reshard_rebuild.py +++ b/tests/unit/weight_sync/test_reshard_rebuild.py @@ -81,6 +81,7 @@ def _reshard(dp_size=4, workers_per_shard=1, dead_shards=(), train_world_size=8) # so a stand-in has to answer. Omitting it is how the reshard rebuild went on # passing its tests while silently defaulting to "nemo". get_collective_sender_spec=lambda: SimpleNamespace(nccl_peer="nemo"), + get_refit_payload_mode=lambda: "bridge_export", ) for name in ( "rebuild_collective", @@ -108,12 +109,26 @@ def _reshard(dp_size=4, workers_per_shard=1, dead_shards=(), train_world_size=8) "expert_model_parallel_size": 1, "pipeline_model_parallel_size": 1, }, - "generation": {"vllm_cfg": {"tensor_parallel_size": 1}}, + "generation": { + "backend": "vllm", + "vllm_cfg": {"tensor_parallel_size": 1}, + }, }, init_collective=lambda *a, **k: ["train-f"], init_nccl_reshard_comm_group=lambda **k: ["train-bulk"], - prepare_nccl_reshard_refit_info=lambda tp, gp, tws, iws: ( - plan_calls.append({"train_world_size": tws, "gen_world_size": iws}) + prepare_nccl_reshard_refit_info=lambda tp, + gp, + tws, + iws, + *, + refit_payload_mode: ( + plan_calls.append( + { + "train_world_size": tws, + "gen_world_size": iws, + "refit_payload_mode": refit_payload_mode, + } + ) or {"gen_world_size": iws} ), ) diff --git a/tests/unit/weight_sync/test_weight_synchronizer.py b/tests/unit/weight_sync/test_weight_synchronizer.py index d06d28b6884..02386985090 100644 --- a/tests/unit/weight_sync/test_weight_synchronizer.py +++ b/tests/unit/weight_sync/test_weight_synchronizer.py @@ -82,6 +82,7 @@ def _mock_generation(**overrides): gen.worker_group.workers = [MagicMock()] gen.get_collective_sender_spec.return_value = CollectiveSenderSpec() gen.get_inference_world_size.return_value = None + gen.get_refit_payload_mode.return_value = "bridge_export" for k, v in overrides.items(): setattr(gen, k, v) return gen @@ -122,7 +123,7 @@ class TestIPCWeightSynchronizer: def test_sync_weights_calls_full_lifecycle(self, mock_ray): mock_ray.get.return_value = [True] policy = _mock_policy() - gen = _mock_generation() + gen = _mock_generation(cfg={"backend": "vllm"}) sync = IPCWeightSynchronizer(policy, gen) assert sync.is_stale @@ -578,7 +579,9 @@ def test_init_communicator_sets_up_collective(self, mock_ray): ) sync.init_communicator() - policy.prepare_refit_info.assert_called_once() + policy.prepare_refit_info.assert_called_once_with( + refit_payload_mode="bridge_export" + ) gen.prepare_refit_info.assert_called_once() policy.init_collective.assert_called_once_with( "10.0.0.1", 29500, 6, train_world_size=4, nccl_peer="nemo" @@ -653,12 +656,15 @@ def test_init_communicator_ships_wire_safe_refit_info(self, mock_ray): "expert_model_parallel_size": 1, "pipeline_model_parallel_size": 1, }, - "generation": {"vllm_cfg": {"tensor_parallel_size": 4}}, + "generation": { + "backend": "vllm", + "vllm_cfg": {"tensor_parallel_size": 4}, + }, }, ) policy.init_nccl_reshard_comm_group.return_value = [MagicMock()] policy.prepare_nccl_reshard_refit_info.return_value = refit_info - gen = _mock_generation() + gen = _mock_generation(cfg={"backend": "vllm"}) gen.init_nccl_reshard_comm_group.return_value = [MagicMock()] # tp_size=4 over a 4-GPU generation world -> one DP shard. gen.worker_group.dp_size = 1 @@ -676,7 +682,13 @@ def test_init_communicator_ships_wire_safe_refit_info(self, mock_ray): ) sync.init_communicator() - policy.prepare_nccl_reshard_refit_info.assert_called_once() + policy.prepare_nccl_reshard_refit_info.assert_called_once_with( + {"tp_size": 2, "ep_size": 1, "etp_size": 2, "pp_size": 1}, + {"tp_size": 4, "ep_size": 1, "etp_size": 4, "pp_size": 1}, + 2, + 4, + refit_payload_mode="bridge_export", + ) gen.prepare_nccl_reshard_refit_info.assert_called_once() (shipped,), _ = gen.prepare_nccl_reshard_refit_info.call_args for params in shipped["per_layer_params"].values(): @@ -688,7 +700,10 @@ def test_init_communicator_ships_wire_safe_refit_info(self, mock_ray): def test_shutdown_drops_the_generation_handle(self): sync = NcclReshardWeightSynchronizer( - _mock_policy(), _mock_generation(), _mock_cluster(), _mock_cluster() + _mock_policy(), + _mock_generation(cfg={"backend": "vllm"}), + _mock_cluster(), + _mock_cluster(), ) sync.shutdown() @@ -701,9 +716,23 @@ def test_shutdown_drops_the_generation_handle(self): # --------------------------------------------------------------------------- -def _mock_megatron_generation(refit_backend="nccl", **overrides): +def _mock_megatron_generation( + refit_backend="nccl", + *, + refit_execution_batch_bytes=123, + refit_transport="mcore", + **overrides, +): gen = _mock_generation(**overrides) - gen.cfg = {"mcore_generation_config": {"refit_backend": refit_backend}} + gen.cfg = { + "backend": "megatron", + "refit_transport": refit_transport, + "mcore_generation_config": { + "refit_backend": refit_backend, + "refit_execution_batch_bytes": refit_execution_batch_bytes, + }, + } + gen.uses_native_refit = refit_transport == "mcore" gen.suspend_for_refit.return_value = None gen.resume_after_refit.return_value = None gen.preinit_nvshmem_collective.return_value = [MagicMock()] @@ -719,35 +748,42 @@ def _mock_megatron_policy(**overrides): class TestMegatronWeightSynchronizer: - def test_non_colocated_requires_clusters(self): - with pytest.raises(ValueError): - MegatronWeightSynchronizer( - _mock_megatron_policy(), _mock_megatron_generation(), colocated=False - ) - - def test_colocated_sync_is_offload_and_wake(self): + @patch( + "nemo_rl.weight_sync.megatron_weight_synchronizer.NcclReshardWeightSynchronizer" + ) + def test_m2n_refit_delegates_transfer_and_keeps_megatron_lifecycle( + self, mock_m2n_cls + ): policy = _mock_megatron_policy() - gen = _mock_megatron_generation() - sync = MegatronWeightSynchronizer(policy, gen, colocated=True) + gen = _mock_megatron_generation(refit_transport="nccl_reshard") + transport = mock_m2n_cls.return_value + sync = MegatronWeightSynchronizer( + policy, + gen, + colocated=False, + train_cluster=_mock_cluster(), + inference_cluster=_mock_cluster(), + ) - sync.init_communicator() # no collective to wire - policy.init_collective_mcore_generation.assert_not_called() + sync.init_communicator() + assert sync.sync_weights(kv_scales={"scale": 1.0}) == {} - assert sync.is_stale - assert sync.sync_weights() == {} + transport.init_communicator.assert_called_once() + transport.sync_weights.assert_called_once_with(kv_scales={"scale": 1.0}) + gen.suspend_for_refit.assert_called_once() policy.offload_before_refit.assert_called_once() - # The refit-protocol tag makes the wake bypass the worker's - # engine-awake early-return (the reshard copy rides this wake). - gen.prepare_for_generation.assert_called_once_with(tags=["colocated_refit"]) - gen.suspend_for_refit.assert_not_called() - policy.swap_weights_via_reshard.assert_not_called() - assert not sync.is_stale + assert [ + call.kwargs.get("tags") + for call in gen.prepare_for_generation.call_args_list + ] == [["weights"], ["kv_cache"]] + gen.resume_after_refit.assert_called_once() + @pytest.mark.parametrize("refit_backend", ["nccl", "nccl_m2n"]) @patch("nemo_rl.weight_sync.megatron_weight_synchronizer.ray") - def test_non_colocated_sync_sequence(self, mock_ray): + def test_non_colocated_sync_sequence(self, mock_ray, refit_backend): mock_ray.get.side_effect = lambda futures: [True for _ in futures] policy = _mock_megatron_policy() - gen = _mock_megatron_generation() + gen = _mock_megatron_generation(refit_backend=refit_backend) sync = MegatronWeightSynchronizer( policy, gen, @@ -758,7 +794,18 @@ def test_non_colocated_sync_sequence(self, mock_ray): sync.init_communicator() policy.init_collective_mcore_generation.assert_called_once() + assert ( + policy.init_collective_mcore_generation.call_args.kwargs[ + "refit_execution_batch_bytes" + ] + == 123 + ) + assert ( + policy.init_collective_mcore_generation.call_args.kwargs["refit_backend"] + == refit_backend + ) gen.init_collective.assert_called_once() + assert gen.init_collective.call_args.kwargs["refit_backend"] == refit_backend assert sync.sync_weights() == {} gen.suspend_for_refit.assert_called_once() @@ -773,6 +820,30 @@ def test_non_colocated_sync_sequence(self, mock_ray): policy.preinit_nvshmem.assert_not_called() assert not sync.is_stale + def test_non_colocated_requires_clusters(self): + with pytest.raises(ValueError): + MegatronWeightSynchronizer( + _mock_megatron_policy(), _mock_megatron_generation(), colocated=False + ) + + def test_colocated_sync_is_offload_and_wake(self): + policy = _mock_megatron_policy() + gen = _mock_megatron_generation() + sync = MegatronWeightSynchronizer(policy, gen, colocated=True) + + sync.init_communicator() # no collective to wire + policy.init_collective_mcore_generation.assert_not_called() + + assert sync.is_stale + assert sync.sync_weights() == {} + policy.offload_before_refit.assert_called_once() + # The refit-protocol tag makes the wake bypass the worker's + # engine-awake early-return (the reshard copy rides this wake). + gen.prepare_for_generation.assert_called_once_with(tags=["colocated_refit"]) + gen.suspend_for_refit.assert_not_called() + policy.swap_weights_via_reshard.assert_not_called() + assert not sync.is_stale + @patch("nemo_rl.weight_sync.megatron_weight_synchronizer.ray") def test_non_colocated_nvshmem_preinits(self, mock_ray): mock_ray.get.side_effect = lambda futures: [True for _ in futures] @@ -808,6 +879,30 @@ def test_non_colocated_failed_update_raises(self, mock_ray): sync.sync_weights() assert sync.is_stale + def test_colocated_rejects_packed_collective_refit(self): + """Packed collective refit is unavailable on the colocated wake path.""" + gen = _mock_megatron_generation(refit_transport=None) + + with pytest.raises(ValueError, match="must be 'mcore' with colocated"): + MegatronWeightSynchronizer(_mock_megatron_policy(), gen, colocated=True) + + def test_colocated_allows_mcore_refit_transport(self): + gen = _mock_megatron_generation() + + sync = MegatronWeightSynchronizer(_mock_megatron_policy(), gen, colocated=True) + assert sync._transport is None + + def test_non_colocated_null_transport_uses_collective_transport(self): + gen = _mock_megatron_generation(refit_transport=None) + sync = MegatronWeightSynchronizer( + _mock_megatron_policy(), + gen, + colocated=False, + train_cluster=_mock_cluster(), + inference_cluster=_mock_cluster(), + ) + assert isinstance(sync._transport, CollectiveWeightSynchronizer) + class TestFactory: def test_colocated_vllm_returns_ipc(self): @@ -832,23 +927,32 @@ def test_colocated_sglang_returns_sglang_colocated(self): ) assert isinstance(sync, SGLangColocatedWeightSynchronizer) - def test_colocated_megatron_returns_megatron_synchronizer(self): + def test_non_colocated_vllm_returns_collective(self): policy = _mock_policy() gen = _mock_generation() sync = create_weight_synchronizer( policy=policy, generation=gen, + generation_backend=VLLM_BACKEND, + colocated=False, + train_cluster=_mock_cluster(), + inference_cluster=_mock_cluster(), + ) + assert isinstance(sync, CollectiveWeightSynchronizer) + + def test_colocated_megatron_returns_megatron_synchronizer(self): + sync = create_weight_synchronizer( + policy=_mock_policy(), + generation=_mock_megatron_generation(), generation_backend=MEGATRON_BACKEND, colocated=True, ) assert isinstance(sync, MegatronWeightSynchronizer) def test_non_colocated_megatron_returns_megatron_synchronizer(self): - policy = _mock_policy() - gen = _mock_generation() sync = create_weight_synchronizer( - policy=policy, - generation=gen, + policy=_mock_policy(), + generation=_mock_generation(), generation_backend=MEGATRON_BACKEND, colocated=False, train_cluster=_mock_cluster(), @@ -856,18 +960,43 @@ def test_non_colocated_megatron_returns_megatron_synchronizer(self): ) assert isinstance(sync, MegatronWeightSynchronizer) - def test_non_colocated_vllm_returns_collective(self): + def test_non_colocated_megatron_m2n_uses_effective_parallelism(self): policy = _mock_policy() - gen = _mock_generation() + policy.cfg = { + "megatron_cfg": { + "tensor_model_parallel_size": 2, + "expert_model_parallel_size": 4, + "pipeline_model_parallel_size": 1, + }, + "generation": { + "backend": "megatron", + "refit_transport": "nccl_reshard", + "mcore_generation_config": { + "tensor_model_parallel_size": 8, + "expert_model_parallel_size": 2, + }, + }, + } + generation = _mock_megatron_generation(refit_transport="nccl_reshard") + generation.cfg = policy.cfg["generation"] + sync = create_weight_synchronizer( policy=policy, - generation=gen, - generation_backend=VLLM_BACKEND, + generation=generation, + generation_backend=MEGATRON_BACKEND, colocated=False, train_cluster=_mock_cluster(), inference_cluster=_mock_cluster(), ) - assert isinstance(sync, CollectiveWeightSynchronizer) + + assert isinstance(sync, MegatronWeightSynchronizer) + assert isinstance(sync._transport, NcclReshardWeightSynchronizer) + assert sync._transport._gen_parallelism() == { + "tp_size": 8, + "ep_size": 2, + "etp_size": 8, + "pp_size": 1, + } def test_non_colocated_dynamo_returns_collective(self): sync = create_weight_synchronizer( From e763079ae358cddb0edbf1b27200149ba09faf9e Mon Sep 17 00:00:00 2001 From: William Dykas Date: Mon, 7 Sep 2026 14:10:30 -0700 Subject: [PATCH 2/8] fix(megatron): address refit review feedback Signed-off-by: William Dykas --- docs/design-docs/nccl-reshard-refit.md | 25 ++-- docs/guides/refit.md | 2 + examples/configs/grpo_math_1B.yaml | 2 +- examples/configs/ppo_math_1B.yaml | 2 +- ...-1b-instruct-1n8g-megatron_generation.yaml | 2 + ...ron_generation-noncolocated-async-gym.yaml | 1 + ...oncolocated-mxfp8-rollouts-flashinfer.yaml | 5 + ...colocated-mxfp8-rollouts-packed-refit.yaml | 6 + ...eneration-noncolocated-mxfp8-rollouts.yaml | 1 + ...olocated-single-controller-async-lag4.yaml | 5 + ...cle-count-1n4g-megatron_generation.v1.yaml | 1 + ...a3b-clevr-8n4g-megatron_generation.v1.yaml | 1 + examples/nemo_gym/grpo_nanov3.yaml | 2 +- ...megatron_generation_single_controller.yaml | 4 + .../nemo_gym/grpo_qwen3_30ba3b_instruct.yaml | 2 +- nemo_rl/algorithms/distillation.py | 4 +- nemo_rl/algorithms/grpo.py | 4 +- nemo_rl/algorithms/ppo.py | 4 +- nemo_rl/models/generation/interfaces.py | 30 ++--- nemo_rl/models/generation/megatron/config.py | 2 +- .../megatron/megatron_generation.py | 72 ++++------- .../generation/megatron/megatron_worker.py | 76 +++++++++-- nemo_rl/models/megatron/setup.py | 9 -- nemo_rl/models/policy/interfaces.py | 8 +- nemo_rl/models/policy/lm_policy.py | 34 +---- .../policy/workers/dtensor_policy_worker.py | 6 +- .../workers/dtensor_policy_worker_v2.py | 6 +- .../policy/workers/megatron_policy_worker.py | 41 +++--- .../checkpoint_engine_weight_synchronizer.py | 6 +- nemo_rl/weight_sync/factory.py | 1 + .../weight_sync/ipc_weight_synchronizer.py | 4 +- .../megatron_weight_synchronizer.py | 18 ++- nemo_rl/weight_sync/nccl_reshard_utils.py | 93 +++++++------- .../nccl_reshard_weight_synchronizer.py | 3 +- .../weight_sync/sglang_weight_synchronizer.py | 4 +- .../configs/grpo_math_1B.yaml | 2 +- research/template_project/single_update.py | 4 +- .../grpo_megatron_generation_colocated.sh | 1 + .../grpo_megatron_generation_colocated_gym.sh | 1 + .../grpo_megatron_generation_multiturn.sh | 1 + .../grpo_megatron_generation_non_colocated.sh | 1 - .../grpo_megatron_generation_topp_topk.sh | 1 + ...-noncolocated-mxfp8-rollouts-flashinfer.sh | 15 +++ ...oncolocated-mxfp8-rollouts-packed-refit.sh | 15 +++ tests/test_suites/nightly_gb200.txt | 2 + tests/unit/algorithms/test_distillation.py | 10 +- tests/unit/algorithms/test_grpo.py | 10 +- tests/unit/algorithms/test_ppo.py | 3 +- .../generation/test_megatron_generation.py | 59 +++++++-- .../models/generation/test_vllm_generation.py | 34 +++-- .../generation/test_vllm_quant_backend.py | 8 +- .../models/policy/test_megatron_worker.py | 120 +++++++++++++++--- .../unit/reference_configs/grpo_math_1B.yaml | 2 +- .../ppo_math_1B_megatron.yaml | 2 +- tests/unit/single_controller/test_setup.py | 4 + ...t_checkpoint_engine_weight_synchronizer.py | 5 +- .../unit/weight_sync/test_reshard_rebuild.py | 2 +- .../weight_sync/test_weight_synchronizer.py | 38 +++++- tools/refit_verifier.py | 4 +- 59 files changed, 548 insertions(+), 282 deletions(-) create mode 100644 examples/configs/recipes/llm/grpo-nanov3-30BA3B-4n4g-megatron_generation-noncolocated-mxfp8-rollouts-flashinfer.yaml create mode 100644 examples/configs/recipes/llm/grpo-nanov3-30BA3B-4n4g-megatron_generation-noncolocated-mxfp8-rollouts-packed-refit.yaml create mode 100755 tests/test_suites/llm/grpo-nanov3-30BA3B-4n4g-megatron_generation-noncolocated-mxfp8-rollouts-flashinfer.sh create mode 100755 tests/test_suites/llm/grpo-nanov3-30BA3B-4n4g-megatron_generation-noncolocated-mxfp8-rollouts-packed-refit.sh diff --git a/docs/design-docs/nccl-reshard-refit.md b/docs/design-docs/nccl-reshard-refit.md index 898d9810ecc..136dec2677c 100644 --- a/docs/design-docs/nccl-reshard-refit.md +++ b/docs/design-docs/nccl-reshard-refit.md @@ -33,14 +33,16 @@ single `ValueError` listing every violation. The current requirements are: * Training-side Megatron supports expert tensor parallelism. Custom PP layouts (`pipeline_model_parallel_layout`, virtual PP > 1, embedding/loss pipeline-split accounting) are not supported yet. -* **Generation-side ETP is pinned to 1.** MCore's `inference_optimized` MoE +* **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, which is what lets generation-side TP > 1 - work; the reshard then handles the train-ETP → gen-ETP=1 gather. An explicitly - requested generation-side ETP > 1 is rejected by config key name rather than - surfacing as a raw MCore assert at model build. + `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 → @@ -72,8 +74,8 @@ single `ValueError` listing every violation. The current requirements are: engine), not bandwidth. * vLLM expert parallelism is supported with the NeMo RL convention `expert_parallel_size == tensor_parallel_size`. -* Megatron generation supports expert parallelism and expert tensor - parallelism, including using both together. +* 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 @@ -140,11 +142,10 @@ source worker does not inspect or branch on the destination backend's name. Requ 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 +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. The older generation-specific method names remain compatibility aliases for -external callers. +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 @@ -265,8 +266,8 @@ generation side maps those HF names onto whatever its own storage layout is. packed-broadcast import path. **To extend to a new backend**, provide a destination map from canonical HF weights -to that backend's local storage. vLLM implements this as -`build_hf_to_local_param_map`; Megatron derives it from Bridge conversion tasks. +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 backend-specific destination map:** resolve each bulk HF name to local storage diff --git a/docs/guides/refit.md b/docs/guides/refit.md index ec9a3315f71..368c3c123c5 100644 --- a/docs/guides/refit.md +++ b/docs/guides/refit.md @@ -98,6 +98,8 @@ 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: diff --git a/examples/configs/grpo_math_1B.yaml b/examples/configs/grpo_math_1B.yaml index e1360af44ca..d977940d988 100644 --- a/examples/configs/grpo_math_1B.yaml +++ b/examples/configs/grpo_math_1B.yaml @@ -422,7 +422,7 @@ policy: kv_cache_management_mode: "persist" # KV cache lifecycle across suspend/resume. Options: "persist", "offload", "recompute". materialize_only_last_token_logits: true num_speculative_tokens: 0 - refit_backend: "nccl" # 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_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. parsers: [] # OpenAI tool-call / response parser names exposed by the HTTP server (only used if expose_http_server=true). diff --git a/examples/configs/ppo_math_1B.yaml b/examples/configs/ppo_math_1B.yaml index 61766cf0c4c..d8d27198631 100644 --- a/examples/configs/ppo_math_1B.yaml +++ b/examples/configs/ppo_math_1B.yaml @@ -278,7 +278,7 @@ policy: enable_chunked_prefill: true unified_memory_level: 0 max_tokens: 16384 - refit_backend: "nccl" # 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_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 vllm_cfg: diff --git a/examples/configs/recipes/llm/grpo-llama3.2-1b-instruct-1n8g-megatron_generation.yaml b/examples/configs/recipes/llm/grpo-llama3.2-1b-instruct-1n8g-megatron_generation.yaml index d616b381092..9c0e226343f 100644 --- a/examples/configs/recipes/llm/grpo-llama3.2-1b-instruct-1n8g-megatron_generation.yaml +++ b/examples/configs/recipes/llm/grpo-llama3.2-1b-instruct-1n8g-megatron_generation.yaml @@ -20,6 +20,8 @@ policy: generation: backend: megatron refit_transport: mcore + mcore_generation_config: + refit_backend: nccl max_new_tokens: 512 vllm_cfg: max_model_len: 512 diff --git a/examples/configs/recipes/llm/grpo-nanov3-30BA3B-2n8g-megatron_generation-noncolocated-async-gym.yaml b/examples/configs/recipes/llm/grpo-nanov3-30BA3B-2n8g-megatron_generation-noncolocated-async-gym.yaml index 710c141a90a..ba5b113e684 100644 --- a/examples/configs/recipes/llm/grpo-nanov3-30BA3B-2n8g-megatron_generation-noncolocated-async-gym.yaml +++ b/examples/configs/recipes/llm/grpo-nanov3-30BA3B-2n8g-megatron_generation-noncolocated-async-gym.yaml @@ -19,6 +19,7 @@ policy: backend: megatron refit_transport: mcore mcore_generation_config: + refit_backend: nccl expert_model_parallel_size: 4 colocated: enabled: false diff --git a/examples/configs/recipes/llm/grpo-nanov3-30BA3B-4n4g-megatron_generation-noncolocated-mxfp8-rollouts-flashinfer.yaml b/examples/configs/recipes/llm/grpo-nanov3-30BA3B-4n4g-megatron_generation-noncolocated-mxfp8-rollouts-flashinfer.yaml new file mode 100644 index 00000000000..ba446690674 --- /dev/null +++ b/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 diff --git a/examples/configs/recipes/llm/grpo-nanov3-30BA3B-4n4g-megatron_generation-noncolocated-mxfp8-rollouts-packed-refit.yaml b/examples/configs/recipes/llm/grpo-nanov3-30BA3B-4n4g-megatron_generation-noncolocated-mxfp8-rollouts-packed-refit.yaml new file mode 100644 index 00000000000..428c0d64a92 --- /dev/null +++ b/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 diff --git a/examples/configs/recipes/llm/grpo-nanov3-30BA3B-4n4g-megatron_generation-noncolocated-mxfp8-rollouts.yaml b/examples/configs/recipes/llm/grpo-nanov3-30BA3B-4n4g-megatron_generation-noncolocated-mxfp8-rollouts.yaml index b9810d1acb4..309bf943869 100644 --- a/examples/configs/recipes/llm/grpo-nanov3-30BA3B-4n4g-megatron_generation-noncolocated-mxfp8-rollouts.yaml +++ b/examples/configs/recipes/llm/grpo-nanov3-30BA3B-4n4g-megatron_generation-noncolocated-mxfp8-rollouts.yaml @@ -30,6 +30,7 @@ policy: 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 diff --git a/examples/configs/recipes/llm/grpo-qwen2.5-math-1.5b-instruct-1n8g-megatron_generation-noncolocated-single-controller-async-lag4.yaml b/examples/configs/recipes/llm/grpo-qwen2.5-math-1.5b-instruct-1n8g-megatron_generation-noncolocated-single-controller-async-lag4.yaml index 5152a0a7ace..6df01ce1aeb 100644 --- a/examples/configs/recipes/llm/grpo-qwen2.5-math-1.5b-instruct-1n8g-megatron_generation-noncolocated-single-controller-async-lag4.yaml +++ b/examples/configs/recipes/llm/grpo-qwen2.5-math-1.5b-instruct-1n8g-megatron_generation-noncolocated-single-controller-async-lag4.yaml @@ -14,4 +14,9 @@ policy: backend: megatron refit_transport: mcore mcore_generation_config: + refit_backend: nccl kv_cache_management_mode: recompute +async_rl: + generation_fleet_health: + # Native MCore refit has no worker-side abort hook for a deadline. + refit_timeout_s: null diff --git a/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-circle-count-1n4g-megatron_generation.v1.yaml b/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-circle-count-1n4g-megatron_generation.v1.yaml index 442e9778ba4..8ba4341daa8 100644 --- a/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-circle-count-1n4g-megatron_generation.v1.yaml +++ b/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-circle-count-1n4g-megatron_generation.v1.yaml @@ -37,6 +37,7 @@ policy: refit_transport: mcore bad_words: null mcore_generation_config: + refit_backend: nccl expose_http_server: true buffer_size_gb: 8 num_cuda_graphs: -1 diff --git a/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-8n4g-megatron_generation.v1.yaml b/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-8n4g-megatron_generation.v1.yaml index 2fb2744c997..1a67593604e 100644 --- a/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-8n4g-megatron_generation.v1.yaml +++ b/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-8n4g-megatron_generation.v1.yaml @@ -41,6 +41,7 @@ policy: max_new_tokens: 2048 bad_words: null mcore_generation_config: + refit_backend: nccl buffer_size_gb: 8 async_sched_mode: async num_cuda_graphs: -1 diff --git a/examples/nemo_gym/grpo_nanov3.yaml b/examples/nemo_gym/grpo_nanov3.yaml index 02d46d37806..27432d60bb8 100644 --- a/examples/nemo_gym/grpo_nanov3.yaml +++ b/examples/nemo_gym/grpo_nanov3.yaml @@ -245,7 +245,7 @@ policy: kv_cache_management_mode: "persist" # KV cache lifecycle across suspend/resume. Options: "persist", "offload", "recompute". materialize_only_last_token_logits: true num_speculative_tokens: 0 - refit_backend: "nccl" # 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_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 expose_http_server: true diff --git a/examples/nemo_gym/grpo_qwen3_0_6b_megatron_generation_single_controller.yaml b/examples/nemo_gym/grpo_qwen3_0_6b_megatron_generation_single_controller.yaml index fef0277c994..d91c5984ca5 100644 --- a/examples/nemo_gym/grpo_qwen3_0_6b_megatron_generation_single_controller.yaml +++ b/examples/nemo_gym/grpo_qwen3_0_6b_megatron_generation_single_controller.yaml @@ -43,6 +43,7 @@ policy: backend: "megatron" refit_transport: "mcore" mcore_generation_config: + refit_backend: nccl # NeMo-Gym drives rollouts through the engine's OpenAI server. expose_http_server: true colocated: @@ -63,6 +64,9 @@ data_plane: num_storage_units: ${mul:2, ${cluster.num_nodes}} # TQ wants >= 2 per node async_rl: + generation_fleet_health: + # Native MCore refit has no worker-side abort hook for a deadline. + refit_timeout_s: null sampler: name: in_order # 0 = fully synchronous, so the importance sampling correction above is an diff --git a/examples/nemo_gym/grpo_qwen3_30ba3b_instruct.yaml b/examples/nemo_gym/grpo_qwen3_30ba3b_instruct.yaml index 3924344a9e1..32f5a142b5d 100644 --- a/examples/nemo_gym/grpo_qwen3_30ba3b_instruct.yaml +++ b/examples/nemo_gym/grpo_qwen3_30ba3b_instruct.yaml @@ -77,7 +77,7 @@ policy: kv_cache_management_mode: "persist" materialize_only_last_token_logits: true num_speculative_tokens: 0 - refit_backend: "nccl" # 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_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 parsers: [] diff --git a/nemo_rl/algorithms/distillation.py b/nemo_rl/algorithms/distillation.py index aec53201dde..2aad21c2650 100644 --- a/nemo_rl/algorithms/distillation.py +++ b/nemo_rl/algorithms/distillation.py @@ -592,7 +592,9 @@ def init_nemo_gym(): ) student_generation.weight_synchronizer.init_communicator() elif student_generation is not None: - state_dict_info = student_policy.prepare_refit_info() + state_dict_info = student_policy.prepare_refit_info( + refit_payload_mode=student_generation.get_refit_payload_mode() + ) student_generation.prepare_refit_info(state_dict_info) # if it is not colocated inference, initialize collective communication for update weights diff --git a/nemo_rl/algorithms/grpo.py b/nemo_rl/algorithms/grpo.py index be508153328..a880c3e32a1 100644 --- a/nemo_rl/algorithms/grpo.py +++ b/nemo_rl/algorithms/grpo.py @@ -1828,7 +1828,9 @@ def init_dynamo(): ) is None and _needs_hf_refit_handshake( backend, nccl_reshard_refit_enabled, colocated_inference ): - state_dict_info = policy.prepare_refit_info() + state_dict_info = policy.prepare_refit_info( + refit_payload_mode=policy_generation.get_refit_payload_mode() + ) if policy_generation is not None: policy_generation.prepare_refit_info(state_dict_info) diff --git a/nemo_rl/algorithms/ppo.py b/nemo_rl/algorithms/ppo.py index 25c4b60659c..1f8479fbaff 100644 --- a/nemo_rl/algorithms/ppo.py +++ b/nemo_rl/algorithms/ppo.py @@ -977,7 +977,9 @@ def initialize_generation_with_policy( worker_init_timing_metrics["collective_init_time_s"] = time.perf_counter() - t0 if backend != "sglang": - state_dict_info = policy.prepare_refit_info() + state_dict_info = policy.prepare_refit_info( + refit_payload_mode=policy_generation.get_refit_payload_mode() + ) if policy_generation is not None: policy_generation.prepare_refit_info(state_dict_info) diff --git a/nemo_rl/models/generation/interfaces.py b/nemo_rl/models/generation/interfaces.py index 3d3a1418c9f..ae6f427e88a 100644 --- a/nemo_rl/models/generation/interfaces.py +++ b/nemo_rl/models/generation/interfaces.py @@ -21,14 +21,11 @@ from nemo_rl.distributed.batched_data_dict import BatchedDataDict -# ``bridge_export`` preserves the universal refit contract: the source exports -# an HF-named, backend-independent representation and each destination converts -# it into local storage. ``logical_weights`` is an explicit Megatron-to-Megatron -# exception: Megatron generation may assume a Megatron training source and ask -# it to materialize logical weights. Megatron training must continue to use the -# universal mode for other destinations (such as vLLM), and new backends should -# not inherit this coupling implicitly. -RefitPayloadMode = Literal["bridge_export", "logical_weights"] +# The universal contract is hf_export: the source exports an HF-named, +# backend-independent representation that each destination converts locally. +# logical_weights is a Megatron-to-Megatron exception, read only by the Megatron +# policy worker; new backends should not inherit that coupling implicitly. +RefitPayloadMode = Literal["hf_export", "logical_weights"] if TYPE_CHECKING: from nemo_rl.algorithms.single_controller_utils.config import MasterConfig @@ -421,18 +418,17 @@ def reject_unenforceable_refit_deadline( Accepting it and doing nothing would be worse than refusing. The deadline exists so that a generation rank dying mid-refit cannot hang the weight-sync collective forever; a user who sets it on a backend that ignores it gets exactly that hang, - while believing they are protected. Only vLLM threads the deadline down to the - collective today. + while believing they are protected. Only transports whose workers own an + abortable collective can enforce it. - ``None`` -- every path that does not configure a deadline, which is all of them by - default -- passes through untouched, so this is inert unless someone opts in. + ``None`` disables the deadline and passes through untouched. """ if refit_timeout_s is not None: raise NotImplementedError( - f"{backend} generation cannot enforce a refit deadline " - f"(refit_timeout_s={refit_timeout_s}). Only the vLLM backend threads it " - "into the refit collective. Unset " - "async_rl.generation_fleet_health.refit_timeout_s, or use vLLM generation." + f"{backend} refit cannot enforce a refit deadline " + f"(refit_timeout_s={refit_timeout_s}). Unset " + "async_rl.generation_fleet_health.refit_timeout_s, or select a refit " + "transport with worker-side watchdog support." ) @@ -532,7 +528,7 @@ def get_inference_world_size(self) -> int | None: def get_refit_payload_mode(self) -> RefitPayloadMode: """Return the backend's required representation for transferred weights.""" - return "bridge_export" + return "hf_export" def prepare_nccl_reshard_refit_info(self, refit_info: dict) -> None: """Prepare per-layer param metadata for nccl_reshard-based refit.""" diff --git a/nemo_rl/models/generation/megatron/config.py b/nemo_rl/models/generation/megatron/config.py index 81188642e74..d5294a4ae6a 100644 --- a/nemo_rl/models/generation/megatron/config.py +++ b/nemo_rl/models/generation/megatron/config.py @@ -69,7 +69,7 @@ class MCoreGenerationSpecificArgs(TypedDict): allow_stale_multimodal_embeddings: NotRequired[bool] # Copy-service backend used only when refit_transport="mcore". - refit_backend: Literal["gloo", "nccl", "nccl_m2n", "nvshmem"] + refit_backend: NotRequired[Literal["gloo", "nccl", "nccl_m2n", "nvshmem"] | None] # Soft per-rank staging limit for native MCore refit. None uses the same # dynamic packed-buffer target as NeMo-RL's existing collective refit. refit_execution_batch_bytes: int | None diff --git a/nemo_rl/models/generation/megatron/megatron_generation.py b/nemo_rl/models/generation/megatron/megatron_generation.py index 22c719025f0..4de908ecf82 100644 --- a/nemo_rl/models/generation/megatron/megatron_generation.py +++ b/nemo_rl/models/generation/megatron/megatron_generation.py @@ -12,7 +12,6 @@ # See the License for the specific language governing permissions and # limitations under the License. -import warnings from copy import deepcopy from typing import TYPE_CHECKING, Any, AsyncGenerator, Optional, cast @@ -249,10 +248,8 @@ def __init__( "policy.generation.refit_transport must be null, 'mcore', or " f"'nccl_reshard' for Megatron generation, got {refit_transport!r}." ) - if refit_transport == "mcore": - refit_backend = self.cfg["mcore_generation_config"].get( - "refit_backend", "gloo" - ) + if self.uses_native_refit: + refit_backend = self.cfg["mcore_generation_config"].get("refit_backend") if refit_backend not in ("gloo", "nccl", "nccl_m2n", "nvshmem"): raise ValueError( "policy.generation.mcore_generation_config.refit_backend " @@ -264,25 +261,15 @@ def __init__( "policy.generation.mcore_generation_config.refit_backend=" "'nccl_m2n' is only supported with non-colocated generation." ) - elif self.cfg["mcore_generation_config"].get("refit_backend") in ( - "nccl_m2n", - "nvshmem", - ): - # Only the native MCore copy service reads refit_backend. Warn rather - # than raise: every shipped exemplar sets refit_backend unconditionally, - # so another transport legitimately inherits a non-null value. A user - # who explicitly asked for a non-default transport, though, would - # otherwise silently get neither. - # .get(): unlike the 'mcore' branch above, the other paths never - # requires this key, so reading it must not turn an omitted key into a - # KeyError. - warnings.warn( - "policy.generation.mcore_generation_config.refit_backend=" - f"{self.cfg['mcore_generation_config']['refit_backend']!r} is " - f"ignored when refit_transport={refit_transport!r}; it is only " - "read by the native MCore refit (refit_transport='mcore').", - stacklevel=2, - ) + else: + refit_backend = self.cfg["mcore_generation_config"].get("refit_backend") + if refit_backend is not None: + raise ValueError( + "policy.generation.mcore_generation_config.refit_backend=" + f"{refit_backend!r} is " + f"invalid when refit_transport={refit_transport!r}; it is only " + "read by the native MCore refit (refit_transport='mcore')." + ) # Populated after the first prepare_for_generation (which starts the HTTP server). self.dp_openai_server_base_urls: list[Optional[str]] = [] # Installed by setup via create_weight_synchronizer. @@ -318,7 +305,7 @@ def __init__( init_optimizer=False, init_reference_model=False, skip_weight_load=skip_weight_load, - refit_role="destination", + is_refit_destination=True, reserved_http_server_port=reserved_http_server_port, ) @@ -334,7 +321,7 @@ def uses_native_refit(self) -> bool: return self.cfg.get("refit_transport") == "mcore" def get_refit_payload_mode(self) -> RefitPayloadMode: - """Use the Megatron-to-Megatron logical-weight exception for M-to-N.""" + """Megatron inference receives logical weights on every refit path.""" return "logical_weights" @property @@ -349,7 +336,6 @@ def init_collective( world_size: int, *, train_world_size: int, - refit_backend: Optional[str] = None, ) -> list[ray.ObjectRef]: """Join the configured refit collective after the training ranks. @@ -358,19 +344,11 @@ def init_collective( port: Port for the process group rendezvous. world_size: Total world size (train + inference workers). train_world_size: Number of training workers (used to offset ranks). - refit_backend: Optional override for the native MCore copy-service - backend ("gloo", "nccl", or "nccl_m2n"; "nvshmem" is - currently broken and warns at setup, see - https://github.com/NVIDIA-NeMo/RL/issues/3646). Ignored by - the packed collective and nccl_reshard transports. Returns: List of Ray ObjectRefs for the collective init futures. """ if self.uses_native_refit: - backend = ( - refit_backend or self.cfg["mcore_generation_config"]["refit_backend"] - ) return self._policy.init_collective_mcore_generation( ip, port, @@ -379,7 +357,7 @@ def init_collective( refit_execution_batch_bytes=self.cfg["mcore_generation_config"][ "refit_execution_batch_bytes" ], - refit_backend=backend, + refit_backend=self.cfg["mcore_generation_config"]["refit_backend"], ) return self._policy.init_collective( ip, @@ -429,7 +407,10 @@ def _refit_ranked_workers( """Return ``(actor, rebuilt rank, original rank)`` for live workers.""" active = membership or self._refit_membership if active is None: - raise RuntimeError("Refit membership has not been initialized.") + return [ + (worker, rank, rank) + for rank, worker in enumerate(self.worker_group.workers) + ] workers = self.worker_group.workers ranked_workers: list[tuple[Any, int, int]] = [] @@ -490,25 +471,16 @@ def rebuild_nccl_reshard_comm_group( def prepare_nccl_reshard_refit_info(self, refit_info: dict[str, Any]) -> None: """Build each inference worker's HF-to-Megatron M-to-N receive map.""" - if self._refit_membership is None: - futures = self._policy.worker_group.run_all_workers_single_data( - "prepare_nccl_reshard_refit_info", refit_info=refit_info - ) - else: - futures = [ - worker.prepare_nccl_reshard_refit_info.remote(refit_info=refit_info) - for worker, _rank, _original_rank in self._refit_ranked_workers() - ] + futures = [ + worker.prepare_nccl_reshard_refit_info.remote(refit_info=refit_info) + for worker, _rank, _original_rank in self._refit_ranked_workers() + ] ray.get(futures) def nccl_reshard_refit( self, refit_timeout_s: Optional[float] = None ) -> list[ray.ObjectRef]: """Receive one NCCL M-to-N refit on every Megatron inference worker.""" - if self._refit_membership is None: - return self._policy.worker_group.run_all_workers_single_data( - "nccl_reshard_refit", refit_timeout_s=refit_timeout_s - ) return [ worker.nccl_reshard_refit.remote(refit_timeout_s=refit_timeout_s) for worker, _rank, _original_rank in self._refit_ranked_workers() diff --git a/nemo_rl/models/generation/megatron/megatron_worker.py b/nemo_rl/models/generation/megatron/megatron_worker.py index 4beea0b295a..d453a7c61ba 100644 --- a/nemo_rl/models/generation/megatron/megatron_worker.py +++ b/nemo_rl/models/generation/megatron/megatron_worker.py @@ -125,7 +125,7 @@ def _inference_optimized_transformer_layer_spec(config: Any) -> Any: ) -def _configure_inference_optimized_layer_spec(model_provider: Any) -> bool: +def _configure_inference_optimized_layer_spec(model_provider: Any) -> None: """Select MCore inference linears for a Bridge generic-GPT provider. Only the generic path needs this. Model-specific providers (DeepSeek, @@ -139,22 +139,16 @@ def _configure_inference_optimized_layer_spec(model_provider: Any) -> bool: ``HybridModelProvider``) are not ``GPTModelProvider`` subclasses and resolve ``inference_optimized`` through their own stack spec, so this is a no-op for them rather than an error. - - Returns: - True if the provider's layer spec was replaced, False if the provider - carries a model-specific spec (or is not a GPT provider) and handles - inference selection itself. """ # Bridge imports ModelOpt plugins that can re-enter MCore while this module # is still initializing, so keep this cycle-sensitive provider import local. from megatron.bridge.models.gpt_provider import GPTModelProvider, default_layer_spec if not isinstance(model_provider, GPTModelProvider): - return False + return if model_provider.transformer_layer_spec is not default_layer_spec: - return False + return model_provider.transformer_layer_spec = _inference_optimized_transformer_layer_spec - return True def _resolve_mxfp8_refit_backend(model_config: Any) -> str: @@ -174,6 +168,14 @@ class _MegatronRefitTask: def param_name(self) -> str: return self.conversion_task.param_name + @property + def cache_key(self) -> tuple[int | None, str]: + """Stable identity across Bridge task rebuilds.""" + return ( + self.conversion_task.vp_stage, + self.conversion_task.global_param_name, + ) + @property def dependencies(self) -> tuple[str, ...]: return self.conversion_task.hf_param_names @@ -1183,6 +1185,14 @@ def _init_generation_refit_state(self) -> None: self._generation_refit_tasks = None self._generation_refit_model_chunks = None self._generation_refit_dependency_counts = None + self._generation_mxfp8_destinations: ( + dict[tuple[int | None, str], MXFP8Tensor] | None + ) = None + # MXFP8 conversion removes the original nn.Parameters, so Bridge cannot + # rediscover those tasks through named_parameters() during a later shard + # recovery. Preserve the complete ordered task plan once destinations are + # made persistent. + self._generation_mxfp8_refit_tasks: list[_MegatronRefitTask] | None = None # Per-refit progress, reset at the start of each transfer. self._generation_refit_task_index = 0 self._generation_refit_remaining_dependencies = {} @@ -1540,6 +1550,14 @@ def post(ctx: RefitCtx) -> None: def _prepare_mxfp8_refit(self, tasks: list[_MegatronRefitTask]) -> None: """Install persistent MCore MXFP8 destinations before engine initialization.""" + if self._generation_mxfp8_destinations is not None: + for task in tasks: + destination = self._generation_mxfp8_destinations.get(task.cache_key) + if destination is not None: + task.destination = destination + task.target_id = id(destination) + return + model_chunks = ( self.model if isinstance(self.model, (list, tuple)) else [self.model] ) @@ -1551,6 +1569,7 @@ def _prepare_mxfp8_refit(self, tasks: list[_MegatronRefitTask]) -> None: and getattr(core.config, "fp8_recipe", None) == "mxfp8" ] if not mxfp8_cores: + self._generation_mxfp8_destinations = {} return if self._inference_engine_initialized: raise RuntimeError( @@ -1575,9 +1594,31 @@ def _prepare_mxfp8_refit(self, tasks: list[_MegatronRefitTask]) -> None: } ) + destinations: dict[tuple[int | None, str], MXFP8Tensor] = {} for task in tasks: if task.target_id in destination_by_id: task.destination = destination_by_id[task.target_id] + task.target_id = id(task.destination) + destinations[task.cache_key] = task.destination + self._generation_mxfp8_destinations = destinations + self._generation_mxfp8_refit_tasks = list(tasks) + + def _refresh_flashinfer_mxfp8_weights(self) -> None: + """Refresh derived FlashInfer expert storage after canonical weights change.""" + model_chunks = ( + self.model if isinstance(self.model, (list, tuple)) else [self.model] + ) + refreshed = False + for model in model_chunks: + core = unwrap_model(model) + for module in core.modules(): + refresh = getattr(module, "refresh_flashinfer_mxfp8_weights", None) + if refresh is not None: + refreshed = bool(refresh()) or refreshed + if refreshed: + # Repacking is asynchronous. Finish before another stream replays + # CUDA graphs that read the derived expert buffers. + torch.cuda.synchronize() def _build_generation_refit_tasks( self, @@ -1586,6 +1627,9 @@ def _build_generation_refit_tasks( model_chunks = ( list(self.model) if isinstance(self.model, (list, tuple)) else [self.model] ) + if self._generation_mxfp8_refit_tasks is not None: + return model_chunks, list(self._generation_mxfp8_refit_tasks) + conversion_tasks = self.megatron_bridge.get_conversion_tasks(model_chunks) tasks: list[_MegatronRefitTask] = [] for conversion_task in conversion_tasks: @@ -1754,9 +1798,12 @@ def receive_one(param_info: dict[str, Any], group: Any, stream: Any) -> None: f"{sorted(self._generation_m2n_pending)}" ) torch.cuda.empty_cache() - return self._update_destination_weights_from_collective( - refit_timeout_s=refit_timeout_s + result = self._update_destination_weights_from_collective( + refit_timeout_s=refit_timeout_s, + refresh_mxfp8=False, ) + self._refresh_flashinfer_mxfp8_weights() + return result finally: self._generation_m2n_pending.clear() @@ -1854,7 +1901,10 @@ def _load_generation_refit_batch( @torch.no_grad() def _update_destination_weights_from_collective( - self, refit_timeout_s: Optional[float] = None + self, + refit_timeout_s: Optional[float] = None, + *, + refresh_mxfp8: bool = True, ) -> bool: """Receive packed HF tensors and import them into the Megatron model.""" from nemo_rl.distributed.refit_watchdog import sync_stream_within @@ -1890,6 +1940,8 @@ def _update_destination_weights_from_collective( ) self.megatron_bridge.finalize_hf_import(self._generation_refit_model_chunks) + if refresh_mxfp8: + self._refresh_flashinfer_mxfp8_weights() sync_stream_within( torch.cuda.current_stream(), refit_timeout_s, diff --git a/nemo_rl/models/megatron/setup.py b/nemo_rl/models/megatron/setup.py index 8fb276f68be..7d78072c21e 100644 --- a/nemo_rl/models/megatron/setup.py +++ b/nemo_rl/models/megatron/setup.py @@ -1644,15 +1644,6 @@ def _create_megatron_config( "overlap_param_gather": overlap_param_gather, "reuse_grad_buf_for_mxfp8_param_ag": reuse_grad_buf_for_mxfp8_param_ag, } - # OptimizerConfig.__post_init__ treats fp8_recipe=None as "no fp8 params" and - # lets the precision-aware optimizer keep fp32 masters inside FusedAdam, - # leaving None placeholders in shard_fp32_from_float16_groups; with - # reuse_grad_buf_for_mxfp8_param_ag the shared param buffer must be refilled - # from those masters each step, so the recipe has to be plumbed to the - # optimizer just like Megatron pretrain's get_megatron_optimizer_config does. - if fp8_cfg is not None and fp8_cfg.get("enabled", False): - optimizer_kwargs["fp8_recipe"] = fp8_cfg.get("fp8_recipe") - # Fused linear logprobs run the decoder but read output_layer.weight directly # instead of calling output_layer.forward(). Megatron's distributed-optimizer # overlap_param_gather prefetch chain assumes every param-gather bucket diff --git a/nemo_rl/models/policy/interfaces.py b/nemo_rl/models/policy/interfaces.py index 352e2156215..aa2c3c4309a 100644 --- a/nemo_rl/models/policy/interfaces.py +++ b/nemo_rl/models/policy/interfaces.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. from abc import ABC, abstractmethod -from typing import Any, Literal, Optional, TypedDict +from typing import Any, Optional, TypedDict import ray import torch @@ -22,8 +22,6 @@ from nemo_rl.models.generation.interfaces import GenerationDatumSpec, RefitPayloadMode from nemo_rl.utils.timer import Timer -RefitRole = Literal["source", "destination"] - class LogprobOutputSpec(TypedDict): """logprobs: Tensor of log probabilities.""" @@ -194,7 +192,7 @@ def offload_to_cpu(self) -> None: def prepare_refit_info( self, *, - refit_payload_mode: Optional[RefitPayloadMode] = None, + refit_payload_mode: RefitPayloadMode, ) -> Optional[dict[str, Any]]: pass @@ -271,7 +269,7 @@ def prepare_nccl_reshard_refit_info( train_world_size: int, gen_world_size: int, *, - refit_payload_mode: Optional[RefitPayloadMode] = None, + refit_payload_mode: RefitPayloadMode, ) -> Any: """Prepare per-layer param metadata for nccl_reshard-based refit.""" raise NotImplementedError diff --git a/nemo_rl/models/policy/lm_policy.py b/nemo_rl/models/policy/lm_policy.py index b70ada0aef6..d8ba54f5a40 100644 --- a/nemo_rl/models/policy/lm_policy.py +++ b/nemo_rl/models/policy/lm_policy.py @@ -44,7 +44,6 @@ ColocatablePolicyInterface, LogprobOutputSpec, ReferenceLogprobOutputSpec, - RefitRole, ScoreOutputSpec, TopkLogitsOutputSpec, ) @@ -104,7 +103,7 @@ def __init__( processor: Optional[AutoProcessor] = None, worker_extension_cls_fqn: Optional[str] = None, skip_weight_load: bool = False, - refit_role: RefitRole = "source", + is_refit_destination: bool = False, reserved_http_server_port: Optional[int] = None, ): self.debug_payload_metrics = False @@ -312,9 +311,9 @@ def __init__( pre_init_communication_queue=pre_init_queue, ) if megatron_enable: - worker_kwargs["refit_role"] = refit_role - elif refit_role != "source": - raise ValueError("refit_role='destination' requires the Megatron backend.") + worker_kwargs["is_refit_destination"] = is_refit_destination + elif is_refit_destination: + raise ValueError("is_refit_destination=True requires the Megatron backend.") if skip_weight_load: worker_kwargs["skip_weight_load"] = True if reserved_http_server_port is not None: @@ -1050,26 +1049,15 @@ def invalidate_kv_cache(self, *args: Any, **kwargs: Any) -> bool: def prepare_refit_info( self, *, - refit_payload_mode: Optional[RefitPayloadMode] = None, + refit_payload_mode: RefitPayloadMode, ) -> Optional[dict[str, Any]]: """Prepare the info for refit. Returns: dict: A dictionary containing the info for refit. """ - worker_kwargs = {} - if self.cfg.get("megatron_cfg", {}).get("enabled", False): - if refit_payload_mode is None: - generation_cfg = self.cfg.get("generation") - refit_payload_mode = ( - "logical_weights" - if generation_cfg is not None - and generation_cfg.get("backend") == "megatron" - else "bridge_export" - ) - worker_kwargs["refit_payload_mode"] = refit_payload_mode futures = self.worker_group.run_all_workers_single_data( - "prepare_refit_info", **worker_kwargs + "prepare_refit_info", refit_payload_mode=refit_payload_mode ) results = ray.get(futures) # Only get the first worker's info since all workers will have the same result @@ -1288,17 +1276,9 @@ def prepare_nccl_reshard_refit_info( train_world_size: int, gen_world_size: int, *, - refit_payload_mode: Optional[RefitPayloadMode] = None, + refit_payload_mode: RefitPayloadMode, ) -> dict[str, Any]: """Prepare per-layer param metadata for nccl_reshard refit.""" - if refit_payload_mode is None: - generation_cfg = self.cfg.get("generation") - refit_payload_mode = ( - "logical_weights" - if generation_cfg is not None - and generation_cfg.get("backend") == "megatron" - else "bridge_export" - ) futures = self.worker_group.run_all_workers_single_data( "prepare_nccl_reshard_refit_info", train_parallelism=train_parallelism, diff --git a/nemo_rl/models/policy/workers/dtensor_policy_worker.py b/nemo_rl/models/policy/workers/dtensor_policy_worker.py index b8b09242be9..29e12e3b8ad 100644 --- a/nemo_rl/models/policy/workers/dtensor_policy_worker.py +++ b/nemo_rl/models/policy/workers/dtensor_policy_worker.py @@ -75,6 +75,7 @@ get_grad_norm, to_local_if_dtensor, ) +from nemo_rl.models.generation.interfaces import RefitPayloadMode from nemo_rl.models.huggingface.common import ( get_flash_attention_kwargs, pack_sequences, @@ -1842,8 +1843,11 @@ def return_model_config(self) -> dict[str, Any]: return self.model.config @torch.no_grad() - def prepare_refit_info(self) -> Optional[dict[str, Any]]: + def prepare_refit_info( + self, *, refit_payload_mode: RefitPayloadMode = "hf_export" + ) -> Optional[dict[str, Any]]: """Prepare state dict metadata for weight refitting and IPC streaming.""" + del refit_payload_mode state_dict_info = {} for name, tensor in self.model.state_dict().items(): # all tensor will be casted to self.dtype in stream_weights_via_ipc_zmq/broadcast_weights_for_collective diff --git a/nemo_rl/models/policy/workers/dtensor_policy_worker_v2.py b/nemo_rl/models/policy/workers/dtensor_policy_worker_v2.py index 0ba3be31774..d42a63ce20d 100644 --- a/nemo_rl/models/policy/workers/dtensor_policy_worker_v2.py +++ b/nemo_rl/models/policy/workers/dtensor_policy_worker_v2.py @@ -55,6 +55,7 @@ forward_with_post_processing_fn, prepare_model_forward, ) +from nemo_rl.models.generation.interfaces import RefitPayloadMode from nemo_rl.models.policy import PolicyConfig from nemo_rl.models.policy.interfaces import ( ColocatablePolicyInterface, @@ -1041,8 +1042,11 @@ def return_model_config(self) -> dict[str, Any]: return self.model.config @torch.no_grad() - def prepare_refit_info(self) -> Optional[dict[str, Any]]: + def prepare_refit_info( + self, *, refit_payload_mode: RefitPayloadMode = "hf_export" + ) -> Optional[dict[str, Any]]: """Prepare state dict metadata for weight refitting and IPC streaming.""" + del refit_payload_mode state_dict_info = {} for name, tensor in self.model.state_dict().items(): if name.endswith(".lora_A.weight") or name.endswith(".lora_B.weight"): diff --git a/nemo_rl/models/policy/workers/megatron_policy_worker.py b/nemo_rl/models/policy/workers/megatron_policy_worker.py index d6253e79456..3580bd34100 100644 --- a/nemo_rl/models/policy/workers/megatron_policy_worker.py +++ b/nemo_rl/models/policy/workers/megatron_policy_worker.py @@ -106,7 +106,6 @@ ColocatablePolicyInterface, LogprobOutputSpec, ReferenceLogprobOutputSpec, - RefitRole, ) from nemo_rl.models.policy.utils import ( broadcast_hf_buckets_via_distributed_impl, @@ -424,8 +423,8 @@ class MegatronPolicyWorkerImpl( ): # Tests and extension classes that bypass __init__ retain the historical # training/source behavior unless they explicitly select destination. - refit_role: RefitRole = "source" - refit_payload_mode: RefitPayloadMode = "bridge_export" + is_refit_destination: bool = False + refit_payload_mode: RefitPayloadMode = "hf_export" # Holds the split-API train-step state between begin/finish or # begin/abort; None when no step is open. Declared at class level so # ``self._train_step_state = None`` after finish/abort type-checks. @@ -549,17 +548,13 @@ def __init__( *, worker_sharding_annotations: NamedSharding, skip_weight_load: bool = False, - refit_role: RefitRole = "source", + is_refit_destination: bool = False, reserved_http_server_port: Optional[int] = None, **kwargs: Any, ): """Initialize the MegatronPolicyWorker.""" - if refit_role not in ("source", "destination"): - raise ValueError( - f"refit_role must be 'source' or 'destination', got {refit_role!r}." - ) - self.refit_role = refit_role - self.refit_payload_mode: RefitPayloadMode = "bridge_export" + self.is_refit_destination = is_refit_destination + self.refit_payload_mode: RefitPayloadMode = "hf_export" # NVML-based and guarded on torch.cuda.is_initialized(), so this does # not initialize a CUDA context ahead of the set_device below. log_gpu_memory_diagnostics( @@ -2427,10 +2422,10 @@ def prepare_refit_info( self, state_dict_info: Optional[dict[str, Any]] = None, *, - refit_payload_mode: RefitPayloadMode = "bridge_export", + refit_payload_mode: RefitPayloadMode = "hf_export", ) -> Optional[dict[str, tuple[torch.Size, torch.dtype]]]: """Prepare refit state for this worker's explicit source/destination role.""" - if self.refit_role == "destination": + if self.is_refit_destination: if state_dict_info is None: raise ValueError("Destination refit requires state_dict_info.") self._prepare_destination_refit_info(state_dict_info) @@ -2443,7 +2438,7 @@ async def update_weights_from_collective( self, refit_timeout_s: Optional[float] = None ) -> bool: """Receive a collective refit without blocking this actor's event loop.""" - if self.refit_role != "destination": + if not self.is_refit_destination: raise RuntimeError( "update_weights_from_collective is only valid for destination-role workers." ) @@ -2619,10 +2614,6 @@ def _is_fp8_export(self) -> bool: and self.fp8_cfg.get("fp8_recipe") == "blockwise" ) - def _uses_logical_refit_payload(self) -> bool: - """Return whether the destination requested logical floating-point weights.""" - return self.refit_payload_mode == "logical_weights" - def _build_refit_conversion_tasks(self) -> list: """Build the conversion-task list driving refit (BF16 or FP8 export). @@ -2633,7 +2624,7 @@ def _build_refit_conversion_tasks(self) -> list: from nemo_rl.models.megatron.draft import draft_model_detached with draft_model_detached([self.model]): - if self._is_fp8_export() and not self._uses_logical_refit_payload(): + if self._is_fp8_export() and self.refit_payload_mode != "logical_weights": return self.megatron_bridge.get_export_fp8_tasks(self.model) return [ task @@ -2735,7 +2726,7 @@ def _iter_params_with_optional_kv_scales( # native refit wire format; Bridge's training FP8 is not inference MXFP8. # Other backends keep Bridge's physical FP8 payload and scale_inv sibling; # mixing that scale with BF16 would corrupt the imported weight. - if self._uses_logical_refit_payload(): + if self.refit_payload_mode == "logical_weights": conversion_tasks = self._iter_logical_refit_conversion_tasks( conversion_tasks ) @@ -2851,7 +2842,7 @@ def _iter_local_hf_param_shards(self) -> Iterator[tuple[str, LocalParamSpec]]: task is the physical fp8 view its ``_scale_inv`` sibling describes; dequantizing it here would ship BF16 bytes under an fp8 scale. """ - uses_logical_payload = self._uses_logical_refit_payload() + uses_logical_payload = self.refit_payload_mode == "logical_weights" for task in self.refit_conversion_tasks: if uses_logical_payload: local_tensor = _get_refit_task_source(task) @@ -3343,10 +3334,10 @@ def prepare_nccl_reshard_refit_info( gen_world_size: Optional[int] = None, *, refit_info: Optional[dict[str, Any]] = None, - refit_payload_mode: RefitPayloadMode = "bridge_export", + refit_payload_mode: RefitPayloadMode = "hf_export", ) -> Optional[dict[str, Any]]: """Prepare NCCL-reshard state for the worker's explicit refit role.""" - if self.refit_role == "destination": + if self.is_refit_destination: if refit_info is None: raise ValueError("Destination NCCL refit requires refit_info.") if any( @@ -3485,7 +3476,7 @@ def build_hf_to_local_param_map( destination_tasks: Optional[list[Any]] = None, ) -> HFToLocalParamMap: """Build the local map for this worker's source or destination role.""" - if self.refit_role == "destination": + if self.is_refit_destination: if destination_tasks is None: _, destination_tasks = self._build_generation_refit_tasks() return self._build_destination_hf_to_local_param_map( @@ -3549,7 +3540,7 @@ def _nccl_reshard_refit_guarded(self, kv_scales=None, refit_timeout_s=None): RefitAbortWatchdog, ) - if self.refit_role == "destination": + if self.is_refit_destination: if kv_scales is not None: raise ValueError("Destination NCCL refit does not accept kv_scales.") generation_groups = ( @@ -3560,7 +3551,7 @@ def _nccl_reshard_refit_guarded(self, kv_scales=None, refit_timeout_s=None): groups = [self.pp_comm_group, self.model_update_group] with RefitAbortWatchdog(groups, refit_timeout_s) as guard: - if self.refit_role == "destination": + if self.is_refit_destination: result = self._destination_nccl_reshard_refit( refit_timeout_s=refit_timeout_s ) diff --git a/nemo_rl/weight_sync/checkpoint_engine_weight_synchronizer.py b/nemo_rl/weight_sync/checkpoint_engine_weight_synchronizer.py index 946547adbd6..c014f012eaa 100644 --- a/nemo_rl/weight_sync/checkpoint_engine_weight_synchronizer.py +++ b/nemo_rl/weight_sync/checkpoint_engine_weight_synchronizer.py @@ -67,7 +67,11 @@ class CheckpointEngineWeightSynchronizer(WeightSynchronizer): _bucket_size_bytes: int | None = None def init_communicator(self) -> None: - self._generation.prepare_refit_info(self._policy.prepare_refit_info()) + self._generation.prepare_refit_info( + self._policy.prepare_refit_info( + refit_payload_mode=self._generation.get_refit_payload_mode() + ) + ) self._ensure_checkpoint_engine_ready() @property diff --git a/nemo_rl/weight_sync/factory.py b/nemo_rl/weight_sync/factory.py index 1ff689cd16b..e113c2c68ab 100644 --- a/nemo_rl/weight_sync/factory.py +++ b/nemo_rl/weight_sync/factory.py @@ -130,6 +130,7 @@ def create_weight_synchronizer( colocated=colocated, train_cluster=train_cluster, inference_cluster=inference_cluster, + refit_timeout_s=refit_timeout_s, ) if generation_backend == SGLANG_BACKEND: diff --git a/nemo_rl/weight_sync/ipc_weight_synchronizer.py b/nemo_rl/weight_sync/ipc_weight_synchronizer.py index b23f4906268..a1f4dcb9521 100644 --- a/nemo_rl/weight_sync/ipc_weight_synchronizer.py +++ b/nemo_rl/weight_sync/ipc_weight_synchronizer.py @@ -109,7 +109,9 @@ def is_stale(self) -> bool: return self._stale def init_communicator(self) -> None: - state_dict_info = self._policy.prepare_refit_info() + state_dict_info = self._policy.prepare_refit_info( + refit_payload_mode=self._generation.get_refit_payload_mode() + ) self._generation.prepare_refit_info(state_dict_info) def shutdown(self) -> None: diff --git a/nemo_rl/weight_sync/megatron_weight_synchronizer.py b/nemo_rl/weight_sync/megatron_weight_synchronizer.py index 6d15a97297a..36455e7ab41 100644 --- a/nemo_rl/weight_sync/megatron_weight_synchronizer.py +++ b/nemo_rl/weight_sync/megatron_weight_synchronizer.py @@ -17,6 +17,7 @@ import ray +from nemo_rl.models.generation.interfaces import reject_unenforceable_refit_deadline from nemo_rl.utils.timer import Timer from nemo_rl.weight_sync.collective_weight_synchronizer import ( CollectiveWeightSynchronizer, @@ -49,17 +50,23 @@ def __init__( colocated: bool, train_cluster: Optional[Any] = None, inference_cluster: Optional[Any] = None, + refit_timeout_s: Optional[float] = None, ): if not colocated and (train_cluster is None or inference_cluster is None): raise ValueError( "train_cluster and inference_cluster are required for " "non-colocated Megatron weight synchronization." ) + if not colocated and generation.uses_native_refit: + # Native MCore's copy service does not expose an abortable collective. + # Reject during setup, before either side can enter a refit. + reject_unenforceable_refit_deadline("native MCore", refit_timeout_s) self._policy = policy self._generation = generation self._colocated = colocated self._train_cluster = train_cluster self._inference_cluster = inference_cluster + self._refit_timeout_s = refit_timeout_s self._refit_backend: Optional[str] = None self._transport: Optional[WeightSynchronizer] = None if colocated: @@ -67,7 +74,7 @@ def __init__( # any other transport is inert. Reject rather than silently ignoring it: # a user asking for packed collective refit would otherwise get the # native one with no indication their setting did nothing. - if generation.cfg.get("refit_transport") != "mcore": + if not generation.uses_native_refit: raise ValueError( "policy.generation.refit_transport must be 'mcore' with " "colocated Megatron generation, which always uses the in-place " @@ -81,6 +88,7 @@ def __init__( generation=generation, train_cluster=train_cluster, inference_cluster=inference_cluster, + refit_timeout_s=refit_timeout_s, ) else: self._transport = CollectiveWeightSynchronizer( @@ -88,6 +96,7 @@ def __init__( generation=generation, train_cluster=train_cluster, inference_cluster=inference_cluster, + refit_timeout_s=refit_timeout_s, ) self._stale = True @@ -125,7 +134,6 @@ def init_communicator(self) -> None: port, world_size, train_world_size=train_world_size, - refit_backend=self._refit_backend, ) ray.get(futures_train + futures_inference) @@ -166,8 +174,12 @@ def sync_weights( if self._transport is not None: self._transport.sync_weights(kv_scales=kv_scales) else: + # Dispatch the destination first: unsupported deadlines are rejected + # before any source rank can enter the blocking native transfer. + futures_inference = self._generation.update_weights_from_collective( + refit_timeout_s=self._refit_timeout_s + ) futures_train = self._policy.swap_weights_via_reshard(is_source=True) - futures_inference = self._generation.update_weights_from_collective() ray.get(futures_train) results = ray.get(futures_inference) if not all(result for result in results if result is not None): diff --git a/nemo_rl/weight_sync/nccl_reshard_utils.py b/nemo_rl/weight_sync/nccl_reshard_utils.py index 049fd3dbbe6..e7cd9bcd4ea 100644 --- a/nemo_rl/weight_sync/nccl_reshard_utils.py +++ b/nemo_rl/weight_sync/nccl_reshard_utils.py @@ -693,57 +693,58 @@ def check_nccl_reshard_refit_support(master_config: Any) -> None: # bytes and deadlock/corrupt the bulk collective; reject anything outside # the supported set up front (this also catches typos such as # "fp8_e4m3" that would otherwise skip the FP8 checks below). - if backend == "vllm" and gen_precision not in ( - None, - "auto", - "bf16", - "bfloat16", - "fp8", - ): - violations.append( - f"policy.generation.vllm_cfg.precision={gen_precision!r} is not " - "supported by nccl_reshard_refit (use 'bf16'/'bfloat16', 'fp8', " - "'auto', or leave unset); the refit byte-copies weights, so the " - "gen dtype must match the train dtype." - ) + if backend == "vllm": + if gen_precision not in ( + None, + "auto", + "bf16", + "bfloat16", + "fp8", + ): + violations.append( + f"policy.generation.vllm_cfg.precision={gen_precision!r} is not " + "supported by nccl_reshard_refit (use 'bf16'/'bfloat16', 'fp8', " + "'auto', or leave unset); the refit byte-copies weights, so the " + "gen dtype must match the train dtype." + ) - if backend == "vllm" and gen_precision == "fp8": - if fp8_param: - if vllm_cfg.get("is_mx"): - violations.append( - "policy.generation.vllm_cfg.is_mx=True does not support " - "blockwise-FP8 storage from " - "policy.megatron_cfg.fp8_cfg.fp8_param; use BF16 training " - "storage for receiver-side MXFP8 quantization." - ) - elif fp8_recipe != "blockwise": - violations.append( - "policy.megatron_cfg.fp8_cfg.fp8_recipe must be 'blockwise' " - f"when fp8_param=True (got {fp8_recipe!r}); other recipes " - "don't produce export-ready scale_inv tensors." - ) - elif vllm_cfg.get("is_mx"): - # Policy precision uses the canonical NeMo-RL spelling; unlike - # vLLM precision, it does not accept "bf16", "auto", or None. - if trainer_precision != "bfloat16": + if gen_precision == "fp8": + if fp8_param: + if vllm_cfg.get("is_mx"): + violations.append( + "policy.generation.vllm_cfg.is_mx=True does not support " + "blockwise-FP8 storage from " + "policy.megatron_cfg.fp8_cfg.fp8_param; use BF16 training " + "storage for receiver-side MXFP8 quantization." + ) + elif fp8_recipe != "blockwise": + violations.append( + "policy.megatron_cfg.fp8_cfg.fp8_recipe must be 'blockwise' " + f"when fp8_param=True (got {fp8_recipe!r}); other recipes " + "don't produce export-ready scale_inv tensors." + ) + elif vllm_cfg.get("is_mx"): + # Policy precision uses the canonical NeMo-RL spelling; unlike + # vLLM precision, it does not accept "bf16", "auto", or None. + if trainer_precision != "bfloat16": + violations.append( + "policy.generation.vllm_cfg.is_mx=True with " + "policy.megatron_cfg.fp8_cfg.fp8_param=False requires " + "policy.precision='bfloat16' for receiver-side MXFP8 " + f"quantization (got {trainer_precision!r})." + ) + else: violations.append( - "policy.generation.vllm_cfg.is_mx=True with " - "policy.megatron_cfg.fp8_cfg.fp8_param=False requires " - "policy.precision='bfloat16' for receiver-side MXFP8 " - f"quantization (got {trainer_precision!r})." + "policy.generation.vllm_cfg.precision='fp8' requires " + "policy.megatron_cfg.fp8_cfg.fp8_param=True, or " + "is_mx=True for BF16-to-MXFP8 refit." ) - else: + elif fp8_param: violations.append( - "policy.generation.vllm_cfg.precision='fp8' requires " - "policy.megatron_cfg.fp8_cfg.fp8_param=True, or " - "is_mx=True for BF16-to-MXFP8 refit." + "policy.megatron_cfg.fp8_cfg.fp8_param=True requires " + "policy.generation.vllm_cfg.precision='fp8' " + "(FP8 storage on train side has no BF16 gen consumer)." ) - elif backend == "vllm" and fp8_param: - violations.append( - "policy.megatron_cfg.fp8_cfg.fp8_param=True requires " - "policy.generation.vllm_cfg.precision='fp8' " - "(FP8 storage on train side has no BF16 gen consumer)." - ) if backend == "megatron": if policy.get("precision") != "bfloat16": diff --git a/nemo_rl/weight_sync/nccl_reshard_weight_synchronizer.py b/nemo_rl/weight_sync/nccl_reshard_weight_synchronizer.py index 6a759f54f99..6ead67f30d0 100644 --- a/nemo_rl/weight_sync/nccl_reshard_weight_synchronizer.py +++ b/nemo_rl/weight_sync/nccl_reshard_weight_synchronizer.py @@ -25,7 +25,8 @@ Lifecycle: init_communicator(): 1. policy/generation.init_collective() -- model_update_group (misc) - 2. policy/generation.init_nccl_reshard_comm_group() -- per-PP-stage bulk groups + 2. policy.init_nccl_reshard_comm_group() and + generation.rebuild_nccl_reshard_comm_group() -- per-PP-stage bulk groups 3. policy.prepare_nccl_reshard_refit_info() -> generation.prepare_nccl_reshard_refit_info() -- backend-agnostic metadata sync_weights(): diff --git a/nemo_rl/weight_sync/sglang_weight_synchronizer.py b/nemo_rl/weight_sync/sglang_weight_synchronizer.py index 9560d18b799..63342f6c5e1 100644 --- a/nemo_rl/weight_sync/sglang_weight_synchronizer.py +++ b/nemo_rl/weight_sync/sglang_weight_synchronizer.py @@ -88,7 +88,9 @@ def is_stale(self) -> bool: return self._stale def init_communicator(self) -> None: - state_dict_info = self._policy.prepare_refit_info() + state_dict_info = self._policy.prepare_refit_info( + refit_payload_mode=self._generation.get_refit_payload_mode() + ) self._generation.prepare_refit_info(state_dict_info) def shutdown(self) -> None: diff --git a/research/template_project/configs/grpo_math_1B.yaml b/research/template_project/configs/grpo_math_1B.yaml index f60ff8e5316..5c67734b6e6 100644 --- a/research/template_project/configs/grpo_math_1B.yaml +++ b/research/template_project/configs/grpo_math_1B.yaml @@ -301,7 +301,7 @@ policy: use_cuda_graphs_for_non_decode_steps: true # Enable CUDA graphs for prefill/context processing enable_chunked_prefill: true # Split long prefills into chunks for better memory management max_tokens: 16384 # Maximum number of tokens to use in a single step. Analogous to vllm's max_num_batched_tokens - refit_backend: "nccl" # 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_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. vllm_cfg: async_engine: false diff --git a/research/template_project/single_update.py b/research/template_project/single_update.py index dd3554110a0..c18569eb9c2 100644 --- a/research/template_project/single_update.py +++ b/research/template_project/single_update.py @@ -120,7 +120,9 @@ def main(config: MasterConfig) -> None: print(f" ✓ Results for return_input: {results}") # Prepare refit info once before first refit - state_dict_info = policy.prepare_refit_info() + state_dict_info = policy.prepare_refit_info( + refit_payload_mode=policy_generation.get_refit_payload_mode() + ) policy_generation.prepare_refit_info(state_dict_info or {}) # Create tiny numeric batch and train with NLLLossFn diff --git a/tests/functional/grpo_megatron_generation_colocated.sh b/tests/functional/grpo_megatron_generation_colocated.sh index 53e0ac4c6bd..5639c4868b1 100644 --- a/tests/functional/grpo_megatron_generation_colocated.sh +++ b/tests/functional/grpo_megatron_generation_colocated.sh @@ -43,6 +43,7 @@ uv run coverage run -a --data-file=$PROJECT_ROOT/tests/.coverage --source=$PROJE policy.train_micro_batch_size=1 \ policy.generation.backend=megatron \ policy.generation.refit_transport=mcore \ + policy.generation.mcore_generation_config.refit_backend=nccl \ cluster.gpus_per_node=2 \ cluster.segment_size=1 \ grpo.max_num_steps=2 \ diff --git a/tests/functional/grpo_megatron_generation_colocated_gym.sh b/tests/functional/grpo_megatron_generation_colocated_gym.sh index 0d4a01ded3c..62950717d4c 100755 --- a/tests/functional/grpo_megatron_generation_colocated_gym.sh +++ b/tests/functional/grpo_megatron_generation_colocated_gym.sh @@ -65,6 +65,7 @@ uv run coverage run -a --data-file=$PROJECT_ROOT/tests/.coverage --source=$PROJE policy.megatron_cfg.sequence_parallel=false \ policy.generation.backend=megatron \ +policy.generation.refit_transport=mcore \ + policy.generation.mcore_generation_config.refit_backend=nccl \ policy.generation.mcore_generation_config.expose_http_server=true \ policy.max_total_sequence_length=512 \ policy.generation.max_new_tokens=128 \ diff --git a/tests/functional/grpo_megatron_generation_multiturn.sh b/tests/functional/grpo_megatron_generation_multiturn.sh index 2c9f4a9a447..7e997a6010e 100644 --- a/tests/functional/grpo_megatron_generation_multiturn.sh +++ b/tests/functional/grpo_megatron_generation_multiturn.sh @@ -29,6 +29,7 @@ uv run coverage run -a --data-file=$PROJECT_ROOT/tests/.coverage --source=$PROJE policy.megatron_cfg.enabled=true \ policy.generation.backend=megatron \ policy.generation.refit_transport=mcore \ + policy.generation.mcore_generation_config.refit_backend=nccl \ cluster.gpus_per_node=2 \ grpo.max_rollout_turns=5 \ grpo.max_num_steps=3 \ diff --git a/tests/functional/grpo_megatron_generation_non_colocated.sh b/tests/functional/grpo_megatron_generation_non_colocated.sh index ba3449c5994..d8f17b5ce62 100755 --- a/tests/functional/grpo_megatron_generation_non_colocated.sh +++ b/tests/functional/grpo_megatron_generation_non_colocated.sh @@ -32,7 +32,6 @@ uv run coverage run -a --data-file=$PROJECT_ROOT/tests/.coverage --source=$PROJE policy.generation.refit_transport=null \ policy.generation.colocated.enabled=false \ policy.generation.colocated.resources.gpus_per_node=1 \ - policy.generation.mcore_generation_config.refit_backend=nccl \ cluster.gpus_per_node=2 \ grpo.max_num_steps=2 \ logger.tensorboard_enabled=true \ diff --git a/tests/functional/grpo_megatron_generation_topp_topk.sh b/tests/functional/grpo_megatron_generation_topp_topk.sh index 6da0112a5cb..7eaafaf640e 100644 --- a/tests/functional/grpo_megatron_generation_topp_topk.sh +++ b/tests/functional/grpo_megatron_generation_topp_topk.sh @@ -32,6 +32,7 @@ uv run coverage run -a --data-file=$PROJECT_ROOT/tests/.coverage --source=$PROJE policy.train_micro_batch_size=1 \ policy.generation.backend=megatron \ policy.generation.refit_transport=mcore \ + policy.generation.mcore_generation_config.refit_backend=nccl \ policy.generation.temperature=0.8 \ policy.generation.top_p=0.9 \ policy.generation.top_k=50 \ diff --git a/tests/test_suites/llm/grpo-nanov3-30BA3B-4n4g-megatron_generation-noncolocated-mxfp8-rollouts-flashinfer.sh b/tests/test_suites/llm/grpo-nanov3-30BA3B-4n4g-megatron_generation-noncolocated-mxfp8-rollouts-flashinfer.sh new file mode 100755 index 00000000000..4f0346e1d6c --- /dev/null +++ b/tests/test_suites/llm/grpo-nanov3-30BA3B-4n4g-megatron_generation-noncolocated-mxfp8-rollouts-flashinfer.sh @@ -0,0 +1,15 @@ +#!/bin/bash +SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd) + +# ===== BEGIN CONFIG ===== +# Mirrors the MXFP8 rollout test driver (delegated base). +NUM_NODES=4 +GPUS_PER_NODE=4 +STEPS_PER_RUN=10 +MAX_STEPS=10 +NUM_RUNS=$(( (MAX_STEPS + STEPS_PER_RUN - 1) / STEPS_PER_RUN )) # Round up +NUM_MINUTES=120 +# ===== END CONFIG ===== + +export EXP_NAME="$(basename "$0" .sh)" +bash "$SCRIPT_DIR/grpo-nanov3-30BA3B-4n4g-megatron_generation-noncolocated-mxfp8-rollouts.sh" "$@" diff --git a/tests/test_suites/llm/grpo-nanov3-30BA3B-4n4g-megatron_generation-noncolocated-mxfp8-rollouts-packed-refit.sh b/tests/test_suites/llm/grpo-nanov3-30BA3B-4n4g-megatron_generation-noncolocated-mxfp8-rollouts-packed-refit.sh new file mode 100755 index 00000000000..4f0346e1d6c --- /dev/null +++ b/tests/test_suites/llm/grpo-nanov3-30BA3B-4n4g-megatron_generation-noncolocated-mxfp8-rollouts-packed-refit.sh @@ -0,0 +1,15 @@ +#!/bin/bash +SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd) + +# ===== BEGIN CONFIG ===== +# Mirrors the MXFP8 rollout test driver (delegated base). +NUM_NODES=4 +GPUS_PER_NODE=4 +STEPS_PER_RUN=10 +MAX_STEPS=10 +NUM_RUNS=$(( (MAX_STEPS + STEPS_PER_RUN - 1) / STEPS_PER_RUN )) # Round up +NUM_MINUTES=120 +# ===== END CONFIG ===== + +export EXP_NAME="$(basename "$0" .sh)" +bash "$SCRIPT_DIR/grpo-nanov3-30BA3B-4n4g-megatron_generation-noncolocated-mxfp8-rollouts.sh" "$@" diff --git a/tests/test_suites/nightly_gb200.txt b/tests/test_suites/nightly_gb200.txt index 2c1d6bed868..eeed8ab7c68 100644 --- a/tests/test_suites/nightly_gb200.txt +++ b/tests/test_suites/nightly_gb200.txt @@ -17,6 +17,8 @@ tests/test_suites/llm/grpo-llama3.2-1b-instruct-1n4g-megatron_generation.sh tests/test_suites/llm/grpo-nanov3-30BA3B-4n4g-megatron_generation-noncolocated-async-gym.sh tests/test_suites/llm/grpo-nanov3-30BA3B-4n4g-megatron_generation-colocated-reshard-async-gym.sh tests/test_suites/llm/grpo-nanov3-30BA3B-4n4g-megatron_generation-noncolocated-mxfp8-rollouts.sh +tests/test_suites/llm/grpo-nanov3-30BA3B-4n4g-megatron_generation-noncolocated-mxfp8-rollouts-packed-refit.sh +tests/test_suites/llm/grpo-nanov3-30BA3B-4n4g-megatron_generation-noncolocated-mxfp8-rollouts-flashinfer.sh # TRT-LLM generation backend tests/test_suites/llm/grpo-qwen3-1.7b-2n4g-fsdp2-trtllm.sh diff --git a/tests/unit/algorithms/test_distillation.py b/tests/unit/algorithms/test_distillation.py index a642e47a1f4..6f9e4a61636 100644 --- a/tests/unit/algorithms/test_distillation.py +++ b/tests/unit/algorithms/test_distillation.py @@ -1004,7 +1004,7 @@ class DummyPolicy: def __init__(self, *args, **kwargs): pass - def prepare_refit_info(self): + def prepare_refit_info(self, *, refit_payload_mode): return {} def offload_after_refit(self): @@ -1027,6 +1027,9 @@ def finish_generation(self): def prepare_refit_info(self, *args, **kwargs): return None + def get_refit_payload_mode(self): + return "hf_export" + def init_collective(self, *args, **kwargs): self.collective_calls.append((args, kwargs)) return [MagicMock()] @@ -1153,7 +1156,7 @@ def __init__(self, *args, **kwargs): def offload_after_refit(self): return None - def prepare_refit_info(self): + def prepare_refit_info(self, *, refit_payload_mode): return {} class DummyVllmGeneration: @@ -1175,6 +1178,9 @@ def finish_generation(self): def prepare_refit_info(self, *args, **kwargs): self.prepare_refit_info_called = True + def get_refit_payload_mode(self): + return "hf_export" + nemo_gym_actor = MagicMock() with ( diff --git a/tests/unit/algorithms/test_grpo.py b/tests/unit/algorithms/test_grpo.py index 3fbbfc7886d..2b0d0c47294 100644 --- a/tests/unit/algorithms/test_grpo.py +++ b/tests/unit/algorithms/test_grpo.py @@ -3215,7 +3215,7 @@ def print_node_ip_and_gpu_id(self): def init_collective(self, *_args, **_kwargs): return [] - def prepare_refit_info(self): + def prepare_refit_info(self, *, refit_payload_mode): return {} def legacy_policy_factory( @@ -3253,6 +3253,9 @@ def finish_generation(self): def prepare_refit_info(self, _state): pass + def get_refit_payload_mode(self): + return "hf_export" + def init_collective(self, *_args, **_kwargs): return [] @@ -3354,7 +3357,7 @@ class DummyPolicy: def print_node_ip_and_gpu_id(self): pass - def prepare_refit_info(self): + def prepare_refit_info(self, *, refit_payload_mode): return {} class DummyTrtllmGeneration: @@ -3367,6 +3370,9 @@ def finish_generation(self): def prepare_refit_info(self, _state): pass + def get_refit_payload_mode(self): + return "hf_export" + nemo_gym_actor = object() spinup_nemo_gym_actor = MagicMock(return_value=nemo_gym_actor) monkeypatch.setattr(grpo_mod, "Logger", lambda *_args, **_kwargs: DummyLogger()) diff --git a/tests/unit/algorithms/test_ppo.py b/tests/unit/algorithms/test_ppo.py index 926a4d73ac1..9461183b06a 100644 --- a/tests/unit/algorithms/test_ppo.py +++ b/tests/unit/algorithms/test_ppo.py @@ -1569,6 +1569,7 @@ def get_placement_groups(self): policy.init_collective.return_value = ["policy-future"] value_model = MagicMock() generation = MagicMock() + generation.get_refit_payload_mode.return_value = "hf_export" generation.init_collective.return_value = ["generation-future"] policy_factory = MagicMock(return_value=policy) value_factory = MagicMock(return_value=value_model) @@ -2049,7 +2050,7 @@ def test_noncolocated_vllm_builds_separate_clusters_and_collective(monkeypatch): value_model = result[2] value_model.finish_training.assert_called_once_with() policy.prepare_for_training.assert_called_once_with() - policy.prepare_refit_info.assert_called_once_with() + policy.prepare_refit_info.assert_called_once_with(refit_payload_mode="hf_export") generation.prepare_refit_info.assert_called_once_with({"state": "dict"}) diff --git a/tests/unit/models/generation/test_megatron_generation.py b/tests/unit/models/generation/test_megatron_generation.py index 7f02742d694..75e16b817e2 100644 --- a/tests/unit/models/generation/test_megatron_generation.py +++ b/tests/unit/models/generation/test_megatron_generation.py @@ -76,7 +76,7 @@ def test_mxfp8_skip_weight_load_defers_http_server_until_refit( prepare_for_generation.assert_not_called() assert generation.dp_openai_server_base_urls == [] - assert lm_policy.Policy.call_args.kwargs["refit_role"] == "destination" + assert lm_policy.Policy.call_args.kwargs["is_refit_destination"] is True @pytest.mark.mcore @@ -98,6 +98,7 @@ def test_nccl_m2n_refit_backend_requires_non_colocated_generation() -> None: def test_null_refit_transport_selects_packed_collective() -> None: config = deepcopy(basic_megatron_test_config) config["generation"]["refit_transport"] = None + config["generation"]["mcore_generation_config"]["refit_backend"] = None generation = MegatronGeneration( config=config, @@ -108,6 +109,34 @@ def test_null_refit_transport_selects_packed_collective() -> None: assert not generation.uses_native_refit +@pytest.mark.mcore +def test_native_refit_requires_explicit_backend() -> None: + config = deepcopy(basic_megatron_test_config) + config["generation"]["mcore_generation_config"]["refit_backend"] = None + + with pytest.raises(ValueError, match="got None"): + MegatronGeneration( + config=config, + tokenizer=MagicMock(), + policy=MagicMock(), + ) + + +@pytest.mark.mcore +@pytest.mark.parametrize("refit_transport", [None, "nccl_reshard"]) +def test_non_native_refit_rejects_mcore_backend(refit_transport) -> None: + config = deepcopy(basic_megatron_test_config) + config["generation"]["refit_transport"] = refit_transport + config["generation"]["mcore_generation_config"]["refit_backend"] = "nccl" + + with pytest.raises(ValueError, match="only read by the native MCore refit"): + MegatronGeneration( + config=config, + tokenizer=MagicMock(), + policy=MagicMock(), + ) + + @pytest.mark.mcore def test_inference_optimized_pins_generation_etp_to_one() -> None: """Generation-side TP>1 must stay usable for MoE. @@ -219,22 +248,27 @@ def test_megatron_generation_m2n_transport_uses_packed_collective_api() -> None: @pytest.mark.mcore -def test_megatron_generation_uses_common_refit_worker_api() -> None: +def test_megatron_generation_uses_common_refit_worker_api( + monkeypatch: pytest.MonkeyPatch, +) -> None: + workers = [MagicMock(), MagicMock()] generation = object.__new__(MegatronGeneration) - generation._policy = MagicMock() + generation._policy = SimpleNamespace(worker_group=SimpleNamespace(workers=workers)) generation._owns_policy = False generation._refit_membership = None - generation._policy.worker_group.run_all_workers_single_data.return_value = [] + monkeypatch.setattr(megatron_generation.ray, "get", lambda refs: refs) refit_info = {"layer_names": [], "per_layer_params": {}} generation.prepare_nccl_reshard_refit_info(refit_info) - generation.nccl_reshard_refit() + assert generation.nccl_reshard_refit() == [ + worker.nccl_reshard_refit.remote.return_value for worker in workers + ] - calls = generation._policy.worker_group.run_all_workers_single_data.call_args_list - assert calls[0].args == ("prepare_nccl_reshard_refit_info",) - assert calls[0].kwargs == {"refit_info": refit_info} - assert calls[1].args == ("nccl_reshard_refit",) - assert calls[1].kwargs == {"refit_timeout_s": None} + for worker in workers: + worker.prepare_nccl_reshard_refit_info.remote.assert_called_once_with( + refit_info=refit_info + ) + worker.nccl_reshard_refit.remote.assert_called_once_with(refit_timeout_s=None) @pytest.mark.mcore @@ -347,6 +381,9 @@ def test_bridge_refit_finalizes_import_before_return( worker.megatron_bridge.finalize_hf_import.side_effect = ( lambda _model_chunks: events.append("finalize") ) + worker._refresh_flashinfer_mxfp8_weights = MagicMock( + side_effect=lambda: events.append("refresh") + ) monkeypatch.setattr( worker_module, @@ -359,7 +396,7 @@ def test_bridge_refit_finalizes_import_before_return( ) assert worker._update_destination_weights_from_collective() - assert events == ["receive", "finalize", "sync"] + assert events == ["receive", "finalize", "refresh", "sync"] worker.megatron_bridge.finalize_hf_import.assert_called_once_with( worker._generation_refit_model_chunks ) diff --git a/tests/unit/models/generation/test_vllm_generation.py b/tests/unit/models/generation/test_vllm_generation.py index f08d1990cf4..4277e315ed4 100644 --- a/tests/unit/models/generation/test_vllm_generation.py +++ b/tests/unit/models/generation/test_vllm_generation.py @@ -1434,7 +1434,7 @@ async def test_vllm_policy_generation_async( lm_policy = Policy(cluster, dtensor_config, tokenizer) print("preparing refit info...") - state_dict_info = lm_policy.prepare_refit_info() + state_dict_info = lm_policy.prepare_refit_info(refit_payload_mode="hf_export") async_policy.prepare_refit_info(state_dict_info) print("refitting vllm policy...") @@ -1535,7 +1535,7 @@ def test_vllm_worker_seed_behavior(cluster, tokenizer): dtensor_config = basic_dtensor_test_config lm_policy = Policy(cluster, dtensor_config, tokenizer) - state_dict_info = lm_policy.prepare_refit_info() + state_dict_info = lm_policy.prepare_refit_info(refit_payload_mode="hf_export") policy.prepare_refit_info(state_dict_info) print("refitting vllm policy...") @@ -1872,7 +1872,7 @@ async def test_vllm_generation_with_hf_training_colocated( # Prepare refit info print("Preparing refit info...") - state_dict_info = lm_policy.prepare_refit_info() + state_dict_info = lm_policy.prepare_refit_info(refit_payload_mode="hf_export") vllm_policy.prepare_refit_info(state_dict_info) # Test @@ -1972,7 +1972,7 @@ async def test_vllm_generation_with_hf_training_non_colocated( ray.get(futures_train + futures_inference) # prepare refit info - state_dict_info = lm_policy.prepare_refit_info() + state_dict_info = lm_policy.prepare_refit_info(refit_payload_mode="hf_export") vllm_policy.prepare_refit_info(state_dict_info) # Test @@ -2628,7 +2628,7 @@ def test_vllm_weight_update_and_prefix_cache_reset( vllm_policy = VllmGeneration(cluster, vllm_config) print("preparing refit info...") - state_dict_info = lm_policy.prepare_refit_info() + state_dict_info = lm_policy.prepare_refit_info(refit_payload_mode="hf_export") vllm_policy.prepare_refit_info(state_dict_info) # Prepare input data (batch size 2) @@ -2746,7 +2746,7 @@ def test_vllm_weight_update_memory(cluster, tokenizer, train_backend): lm_policy = Policy(cluster, train_config, tokenizer) print("preparing refit info...") - state_dict_info = lm_policy.prepare_refit_info() + state_dict_info = lm_policy.prepare_refit_info(refit_payload_mode="hf_export") vllm_policy.prepare_refit_info(state_dict_info) print("refitting vllm policy...") @@ -2820,7 +2820,7 @@ def test_vllm_generation_with_stop(cluster, test_input_data, tokenizer, is_eval) lm_policy = Policy(cluster, dtensor_config, tokenizer) print("preparing refit info...") - state_dict_info = lm_policy.prepare_refit_info() + state_dict_info = lm_policy.prepare_refit_info(refit_payload_mode="hf_export") vllm_generation.prepare_refit_info(state_dict_info) print("refitting vllm policy...") @@ -2959,7 +2959,7 @@ async def test_vllm_refit_non_colocated_update_weights( ray.get(futures_train + futures_inference) # prepare refit info - state_dict_info = lm_policy.prepare_refit_info() + state_dict_info = lm_policy.prepare_refit_info(refit_payload_mode="hf_export") vllm_generation.prepare_refit_info(state_dict_info) print("refitting vllm policy...") @@ -3082,7 +3082,9 @@ def test_vllm_generation_with_megatron_training( megatron_policy = Policy(cluster, megatron_config, test_tokenizer) print("preparing refit info...") - state_dict_info = megatron_policy.prepare_refit_info() + state_dict_info = megatron_policy.prepare_refit_info( + refit_payload_mode="hf_export" + ) vllm_policy.prepare_refit_info(state_dict_info) print("Refitting vLLM policy with Megatron weights...") @@ -3244,7 +3246,9 @@ def test_vllm_generation_with_megatron_training_moe_model( megatron_policy = Policy(moe_cluster, megatron_config, test_tokenizer) print("preparing refit info...") - state_dict_info = megatron_policy.prepare_refit_info() + state_dict_info = megatron_policy.prepare_refit_info( + refit_payload_mode="hf_export" + ) vllm_policy.prepare_refit_info(state_dict_info) print("Refitting vLLM policy with Megatron weights...") @@ -3370,7 +3374,7 @@ def test_vllm_megatron_weight_update_memory(cluster, tokenizer): megatron_policy = Policy(cluster, megatron_config, test_tokenizer) print("preparing refit info...") - state_dict_info = megatron_policy.prepare_refit_info() + state_dict_info = megatron_policy.prepare_refit_info(refit_payload_mode="hf_export") vllm_policy.prepare_refit_info(state_dict_info) print("Refitting vLLM policy with Megatron...") @@ -3486,7 +3490,9 @@ def test_vllm_megatron_pipeline_parallel(cluster, tokenizer): megatron_policy = Policy(cluster, megatron_config, test_tokenizer) print("preparing refit info...") - state_dict_info = megatron_policy.prepare_refit_info() + state_dict_info = megatron_policy.prepare_refit_info( + refit_payload_mode="hf_export" + ) vllm_policy.prepare_refit_info(state_dict_info) print("Refitting vLLM with Megatron PP=2 weights...") @@ -3552,7 +3558,9 @@ def test_vllm_megatron_weight_update_with_packing(cluster, test_input_data): vllm_generation = VllmGeneration(cluster, vllm_config) # prepare refit info - state_dict_info = megatron_policy.prepare_refit_info() + state_dict_info = megatron_policy.prepare_refit_info( + refit_payload_mode="hf_export" + ) vllm_generation.prepare_refit_info(state_dict_info) print("refitting vllm policy...") diff --git a/tests/unit/models/generation/test_vllm_quant_backend.py b/tests/unit/models/generation/test_vllm_quant_backend.py index 7e88f4c993a..2738bc58e83 100644 --- a/tests/unit/models/generation/test_vllm_quant_backend.py +++ b/tests/unit/models/generation/test_vllm_quant_backend.py @@ -176,7 +176,9 @@ def test_vllm_quant_refit(cluster, async_engine): ) # Refit: transfer pre-folded weights + input_quantizer amax from Megatron to vLLM - state_dict_info = megatron_policy.prepare_refit_info() + state_dict_info = megatron_policy.prepare_refit_info( + refit_payload_mode="hf_export" + ) vllm_policy.prepare_refit_info(state_dict_info) refit_policy_generation( megatron_policy, vllm_policy, vllm_config["colocated"]["enabled"] @@ -261,7 +263,9 @@ def test_vllm_kv_quant_refit(cluster, recipe, uses_constant_amax, monkeypatch): "have non-positive amax" ) - state_dict_info = megatron_policy.prepare_refit_info() + state_dict_info = megatron_policy.prepare_refit_info( + refit_payload_mode="hf_export" + ) vllm_policy.prepare_refit_info(state_dict_info) refit_policy_generation( megatron_policy, vllm_policy, vllm_config["colocated"]["enabled"] diff --git a/tests/unit/models/policy/test_megatron_worker.py b/tests/unit/models/policy/test_megatron_worker.py index fe440c49233..f10d3662703 100644 --- a/tests/unit/models/policy/test_megatron_worker.py +++ b/tests/unit/models/policy/test_megatron_worker.py @@ -60,21 +60,23 @@ def _make_refit_task( local_specs: tuple[tuple[str, str], ...] = (), mapping: object | None = None, ): - from megatron.bridge.models.conversion.param_mapping import LocalHFParamSpec - from nemo_rl.models.generation.megatron.megatron_worker import ( _MegatronRefitTask, ) - specs = tuple( - LocalHFParamSpec( - name, - None if component == "full" else -2, - 1 if component == "up" else 0, - 1 if component == "full" else 2, + specs = () + if local_specs: + from megatron.bridge.models.conversion.param_mapping import LocalHFParamSpec + + specs = tuple( + LocalHFParamSpec( + name, + None if component == "full" else -2, + 1 if component == "up" else 0, + 1 if component == "full" else 2, + ) + for name, component in local_specs ) - for name, component in local_specs - ) def combine_local_hf_weights(weights): if len(specs) == 1: @@ -89,6 +91,8 @@ def combine_local_hf_weights(weights): conversion_task = SimpleNamespace( param_name=param_name, + global_param_name=param_name, + vp_stage=0, hf_param_names=dependencies, mapping=mapping if mapping is not None else MagicMock(), local_hf_param_specs=lambda: specs, @@ -195,7 +199,7 @@ def test_megatron_fp8_refit_tasks_match_payload_mode() -> None: assert worker._build_refit_conversion_tasks() == logical_tasks[1:] worker.megatron_bridge.get_export_fp8_tasks.assert_not_called() - worker.refit_payload_mode = "bridge_export" + worker.refit_payload_mode = "hf_export" assert worker._build_refit_conversion_tasks() == physical_tasks worker.megatron_bridge.get_export_fp8_tasks.assert_called_once_with(worker.model) @@ -236,7 +240,7 @@ class _QuantizedParam(torch.Tensor): ) worker.cfg = {"generation": {"backend": "vllm"}} - worker.refit_payload_mode = "bridge_export" + worker.refit_payload_mode = "hf_export" list(worker._iter_params_with_optional_kv_scales()) forwarded = worker.megatron_bridge.export_hf_weights.call_args.kwargs[ "conversion_tasks" @@ -406,7 +410,7 @@ def test_nccl_reshard_all_misc_refit_supports_empty_bulk(pp_size: int) -> None: worker._build_layer_to_pp_stage.assert_not_called() -def test_destination_refit_role_uses_common_worker_interface( +def test_refit_destination_uses_common_worker_interface( monkeypatch: pytest.MonkeyPatch, ) -> None: from nemo_rl.models.policy.workers.megatron_policy_worker import ( @@ -414,7 +418,7 @@ def test_destination_refit_role_uses_common_worker_interface( ) worker = object.__new__(MegatronPolicyWorkerImpl) - worker.refit_role = "destination" + worker.is_refit_destination = True worker.model_update_group = object() worker._generation_nccl_reshard_groups = {} tasks = [object()] @@ -455,6 +459,57 @@ def test_destination_refit_role_uses_common_worker_interface( assert set_device.call_args_list == [call(3), call(3)] +def test_m2n_destination_refreshes_flashinfer_after_misc_import( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from nemo_rl.models.generation.megatron import megatron_worker as worker_module + + worker = object.__new__(worker_module.MegatronGenerationRefitMixin) + worker.nccl_reshard_refit_info = { + "layer_names": [], + "per_layer_params": {}, + } + worker._generation_nccl_reshard_groups = {} + worker._update_destination_weights_from_collective = MagicMock(return_value=True) + worker._refresh_flashinfer_mxfp8_weights = MagicMock() + + monkeypatch.setattr(torch.cuda, "Stream", MagicMock) + monkeypatch.setattr(torch.cuda, "current_stream", MagicMock) + monkeypatch.setattr(torch.cuda, "empty_cache", MagicMock()) + monkeypatch.setattr( + "nemo_rl.distributed.refit_watchdog.sync_stream_within", MagicMock() + ) + + assert worker._destination_nccl_reshard_refit(refit_timeout_s=30.0) + worker._update_destination_weights_from_collective.assert_called_once_with( + refit_timeout_s=30.0, + refresh_mxfp8=False, + ) + worker._refresh_flashinfer_mxfp8_weights.assert_called_once_with() + + +@pytest.mark.parametrize("did_refresh", [False, True]) +def test_flashinfer_mxfp8_refresh_synchronizes_only_after_repacking( + monkeypatch: pytest.MonkeyPatch, + did_refresh: bool, +) -> None: + from nemo_rl.models.generation.megatron import megatron_worker as worker_module + + core = torch.nn.Module() + core.experts = torch.nn.Module() + core.experts.refresh_flashinfer_mxfp8_weights = MagicMock(return_value=did_refresh) + worker = object.__new__(worker_module.MegatronGenerationRefitMixin) + worker.model = core + synchronize = MagicMock() + monkeypatch.setattr(worker_module, "unwrap_model", lambda model: model) + monkeypatch.setattr(torch.cuda, "synchronize", synchronize) + + worker._refresh_flashinfer_mxfp8_weights() + + core.experts.refresh_flashinfer_mxfp8_weights.assert_called_once_with() + assert synchronize.call_count == int(did_refresh) + + def test_megatron_destination_misc_plan_uses_shipped_misc_metadata() -> None: from nemo_rl.models.generation.megatron.megatron_worker import ( MegatronGenerationRefitMixin, @@ -639,8 +694,10 @@ def test_megatron_m2n_rejects_weights_with_no_local_destination( worker._build_destination_hf_to_local_param_map(refit_info, []) -def test_prepare_mxfp8_refit_replaces_only_quantized_parameters( +@pytest.mark.parametrize("grouped_gemm_backend", ["torch", "flashinfer"]) +def test_prepare_mxfp8_refit_replaces_only_quantized_parameters_idempotently( monkeypatch: pytest.MonkeyPatch, + grouped_gemm_backend: str, ) -> None: from nemo_rl.models.generation.megatron import megatron_worker as worker_module from nemo_rl.models.generation.megatron.megatron_worker import ( @@ -651,7 +708,7 @@ def test_prepare_mxfp8_refit_replaces_only_quantized_parameters( core.config = SimpleNamespace( transformer_impl="inference_optimized", fp8_recipe="mxfp8", - inference_grouped_gemm_backend="torch", + inference_grouped_gemm_backend=grouped_gemm_backend, ) core.decoder = torch.nn.Module() core.decoder.weight = torch.nn.Parameter(torch.zeros(2, 2)) @@ -675,11 +732,42 @@ def test_prepare_mxfp8_refit_replaces_only_quantized_parameters( worker = object.__new__(MegatronGenerationRefitMixin) worker.model = core worker._inference_engine_initialized = False + worker._generation_mxfp8_destinations = None + worker._generation_mxfp8_refit_tasks = None worker._prepare_mxfp8_refit([weight_task, norm_task]) assert weight_task.destination is quantized_destination assert norm_task.destination is core.decoder.norm + quantize.assert_called_once_with(core.decoder, backend="triton") + + # Rebuilds happen after lazy engine initialization. They must reuse the + # persistent destination map instead of quantizing/rebinding storage again. + rebuilt_weight_task = _make_refit_task( + param_name="weight", + destination=core.decoder.weight, + dependencies=("weight",), + ) + rebuilt_norm_task = _make_refit_task( + param_name="norm", + destination=core.decoder.norm, + dependencies=("norm",), + ) + worker._inference_engine_initialized = True + worker._prepare_mxfp8_refit([rebuilt_weight_task, rebuilt_norm_task]) + + assert rebuilt_weight_task.destination is quantized_destination + assert rebuilt_norm_task.destination is core.decoder.norm + quantize.assert_called_once() + + # Bridge enumerates nn.Parameters, but quantization replaces the weight with + # an MXFP8Tensor attribute. Recovery therefore reuses the complete ordered + # plan rather than asking Bridge to rediscover a now-invisible task. + worker.megatron_bridge = MagicMock() + model_chunks, rebuilt_tasks = worker._build_generation_refit_tasks() + assert model_chunks == [core] + assert rebuilt_tasks == [weight_task, norm_task] + worker.megatron_bridge.get_conversion_tasks.assert_not_called() def test_megatron_generation_joins_all_m2n_pipeline_groups( diff --git a/tests/unit/reference_configs/grpo_math_1B.yaml b/tests/unit/reference_configs/grpo_math_1B.yaml index 58128f99e00..29ec9face5d 100644 --- a/tests/unit/reference_configs/grpo_math_1B.yaml +++ b/tests/unit/reference_configs/grpo_math_1B.yaml @@ -403,7 +403,7 @@ policy: kv_cache_management_mode: "persist" # KV cache lifecycle across suspend/resume. Options: "persist", "offload", "recompute". materialize_only_last_token_logits: true # Materialize logits only for the last token of each sequence during decode. num_speculative_tokens: 0 - refit_backend: "nccl" # 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_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. parsers: [] # OpenAI tool-call / response parser names exposed by the HTTP server (only used if expose_http_server=true). diff --git a/tests/unit/reference_configs/ppo_math_1B_megatron.yaml b/tests/unit/reference_configs/ppo_math_1B_megatron.yaml index 48b9ba47250..1900812de7b 100644 --- a/tests/unit/reference_configs/ppo_math_1B_megatron.yaml +++ b/tests/unit/reference_configs/ppo_math_1B_megatron.yaml @@ -255,7 +255,7 @@ policy: enable_chunked_prefill: true unified_memory_level: 0 max_tokens: 16384 - refit_backend: "nccl" # 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_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 vllm_cfg: diff --git a/tests/unit/single_controller/test_setup.py b/tests/unit/single_controller/test_setup.py index a8a180cba2a..9988ef296f6 100644 --- a/tests/unit/single_controller/test_setup.py +++ b/tests/unit/single_controller/test_setup.py @@ -1990,6 +1990,10 @@ def _spinup_gym(**_): assert factory_kwargs["generation_backend"] == "megatron" assert factory_kwargs["colocated"] is colocated assert factory_kwargs["inference_cluster"] is inference_cluster + assert ( + factory_kwargs["refit_timeout_s"] + == mc.async_rl.generation_fleet_health.refit_timeout_s + ) if gym: assert actor_args.env_handles["nemo_gym"] is fake_gym_actor assert metrics.nemo_gym_init_time_s is not None diff --git a/tests/unit/weight_sync/test_checkpoint_engine_weight_synchronizer.py b/tests/unit/weight_sync/test_checkpoint_engine_weight_synchronizer.py index 732cc5044c6..cd7cd212501 100644 --- a/tests/unit/weight_sync/test_checkpoint_engine_weight_synchronizer.py +++ b/tests/unit/weight_sync/test_checkpoint_engine_weight_synchronizer.py @@ -53,6 +53,7 @@ def _mock_generation(**overrides): gen.update_weights_via_ipc_zmq.return_value = [MagicMock()] gen.update_weights_from_collective.return_value = [MagicMock()] gen.init_collective.return_value = [MagicMock()] + gen.get_refit_payload_mode.return_value = "hf_export" for k, v in overrides.items(): setattr(gen, k, v) return gen @@ -227,7 +228,9 @@ def test_sync_weights_runs_checkpoint_engine_lifecycle(self, mock_ray): sync.sync_weights(kv_scales={"kv": 1.0}) assert not sync.is_stale - sync._policy.prepare_refit_info.assert_called_once() + sync._policy.prepare_refit_info.assert_called_once_with( + refit_payload_mode="hf_export" + ) sync._generation.prepare_refit_info.assert_called_once() assert ( "checkpoint_engine_rpc", diff --git a/tests/unit/weight_sync/test_reshard_rebuild.py b/tests/unit/weight_sync/test_reshard_rebuild.py index b54211070ff..7531687f0ce 100644 --- a/tests/unit/weight_sync/test_reshard_rebuild.py +++ b/tests/unit/weight_sync/test_reshard_rebuild.py @@ -81,7 +81,7 @@ def _reshard(dp_size=4, workers_per_shard=1, dead_shards=(), train_world_size=8) # so a stand-in has to answer. Omitting it is how the reshard rebuild went on # passing its tests while silently defaulting to "nemo". get_collective_sender_spec=lambda: SimpleNamespace(nccl_peer="nemo"), - get_refit_payload_mode=lambda: "bridge_export", + get_refit_payload_mode=lambda: "hf_export", ) for name in ( "rebuild_collective", diff --git a/tests/unit/weight_sync/test_weight_synchronizer.py b/tests/unit/weight_sync/test_weight_synchronizer.py index 02386985090..7afc9f7586f 100644 --- a/tests/unit/weight_sync/test_weight_synchronizer.py +++ b/tests/unit/weight_sync/test_weight_synchronizer.py @@ -82,7 +82,7 @@ def _mock_generation(**overrides): gen.worker_group.workers = [MagicMock()] gen.get_collective_sender_spec.return_value = CollectiveSenderSpec() gen.get_inference_world_size.return_value = None - gen.get_refit_payload_mode.return_value = "bridge_export" + gen.get_refit_payload_mode.return_value = "hf_export" for k, v in overrides.items(): setattr(gen, k, v) return gen @@ -194,7 +194,9 @@ def test_init_communicator(self): sync = IPCWeightSynchronizer(policy, gen) sync.init_communicator() - policy.prepare_refit_info.assert_called_once() + policy.prepare_refit_info.assert_called_once_with( + refit_payload_mode="hf_export" + ) gen.prepare_refit_info.assert_called_once() @patch("nemo_rl.weight_sync.ipc_weight_synchronizer.ray") @@ -394,7 +396,9 @@ def test_init_communicator(self, mock_ray): sync = SGLangColocatedWeightSynchronizer(policy, gen) sync.init_communicator() - policy.prepare_refit_info.assert_called_once() + policy.prepare_refit_info.assert_called_once_with( + refit_payload_mode="hf_export" + ) gen.prepare_refit_info.assert_called_once() def test_phase_restoration_on_transfer_failure(self, mock_ray): @@ -580,7 +584,7 @@ def test_init_communicator_sets_up_collective(self, mock_ray): sync.init_communicator() policy.prepare_refit_info.assert_called_once_with( - refit_payload_mode="bridge_export" + refit_payload_mode="hf_export" ) gen.prepare_refit_info.assert_called_once() policy.init_collective.assert_called_once_with( @@ -687,7 +691,7 @@ def test_init_communicator_ships_wire_safe_refit_info(self, mock_ray): {"tp_size": 4, "ep_size": 1, "etp_size": 4, "pp_size": 1}, 2, 4, - refit_payload_mode="bridge_export", + refit_payload_mode="hf_export", ) gen.prepare_nccl_reshard_refit_info.assert_called_once() (shipped,), _ = gen.prepare_nccl_reshard_refit_info.call_args @@ -763,6 +767,15 @@ def test_m2n_refit_delegates_transfer_and_keeps_megatron_lifecycle( colocated=False, train_cluster=_mock_cluster(), inference_cluster=_mock_cluster(), + refit_timeout_s=17.0, + ) + + mock_m2n_cls.assert_called_once_with( + policy=policy, + generation=gen, + train_cluster=sync._train_cluster, + inference_cluster=sync._inference_cluster, + refit_timeout_s=17.0, ) sync.init_communicator() @@ -805,13 +818,13 @@ def test_non_colocated_sync_sequence(self, mock_ray, refit_backend): == refit_backend ) gen.init_collective.assert_called_once() - assert gen.init_collective.call_args.kwargs["refit_backend"] == refit_backend + assert "refit_backend" not in gen.init_collective.call_args.kwargs assert sync.sync_weights() == {} gen.suspend_for_refit.assert_called_once() policy.offload_before_refit.assert_called_once() policy.swap_weights_via_reshard.assert_called_once_with(is_source=True) - gen.update_weights_from_collective.assert_called_once() + gen.update_weights_from_collective.assert_called_once_with(refit_timeout_s=None) gen.resume_after_refit.assert_called_once() # prepare called for the weights phase and then the kv_cache phase tags = [c.kwargs.get("tags") for c in gen.prepare_for_generation.call_args_list] @@ -820,6 +833,17 @@ def test_non_colocated_sync_sequence(self, mock_ray, refit_backend): policy.preinit_nvshmem.assert_not_called() assert not sync.is_stale + def test_native_refit_rejects_unenforceable_deadline_during_setup(self): + with pytest.raises(NotImplementedError, match="native MCore refit cannot"): + MegatronWeightSynchronizer( + _mock_megatron_policy(), + _mock_megatron_generation(), + colocated=False, + train_cluster=_mock_cluster(), + inference_cluster=_mock_cluster(), + refit_timeout_s=30.0, + ) + def test_non_colocated_requires_clusters(self): with pytest.raises(ValueError): MegatronWeightSynchronizer( diff --git a/tools/refit_verifier.py b/tools/refit_verifier.py index 30ec0835612..892cd33a3f8 100644 --- a/tools/refit_verifier.py +++ b/tools/refit_verifier.py @@ -828,7 +828,9 @@ def main_vllm(): generation_data = prepare_input_data(args.prompt, tokenizer) # prepare refit info - state_dict_info = policy.prepare_refit_info() + state_dict_info = policy.prepare_refit_info( + refit_payload_mode=vllm_inference_policy.get_refit_payload_mode() + ) vllm_inference_policy.prepare_refit_info(state_dict_info) # Perform model refitting From 7e5ddaed43a131bd9845d24363556529de66ce21 Mon Sep 17 00:00:00 2001 From: William Dykas Date: Tue, 8 Sep 2026 15:43:02 -0700 Subject: [PATCH 3/8] fix(megatron): repair native refit CI failures Signed-off-by: William Dykas --- ...mni-30ba3b-circle-count-1n4g-megatron_generation.v1.yaml | 4 ++++ ...otron-omni-30ba3b-clevr-8n4g-megatron_generation.v1.yaml | 4 ++++ nemo_rl/models/policy/workers/megatron_policy_worker.py | 6 ++++++ .../grpo_megatron_generation_colocated_reshard.sh | 4 +++- 4 files changed, 17 insertions(+), 1 deletion(-) diff --git a/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-circle-count-1n4g-megatron_generation.v1.yaml b/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-circle-count-1n4g-megatron_generation.v1.yaml index 8ba4341daa8..2bcda842d2c 100644 --- a/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-circle-count-1n4g-megatron_generation.v1.yaml +++ b/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 diff --git a/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-8n4g-megatron_generation.v1.yaml b/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-8n4g-megatron_generation.v1.yaml index 1a67593604e..f9a2fa45116 100644 --- a/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-8n4g-megatron_generation.v1.yaml +++ b/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-8n4g-megatron_generation.v1.yaml @@ -9,6 +9,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 diff --git a/nemo_rl/models/policy/workers/megatron_policy_worker.py b/nemo_rl/models/policy/workers/megatron_policy_worker.py index 3580bd34100..c6202dea9db 100644 --- a/nemo_rl/models/policy/workers/megatron_policy_worker.py +++ b/nemo_rl/models/policy/workers/megatron_policy_worker.py @@ -3744,6 +3744,12 @@ def _build_colocated_inference_model(self, config: PolicyConfig) -> None: inference_mcfg = plan.inference_megatron_cfg inference_config = {**config, "megatron_cfg": inference_mcfg} + if inference_mcfg.get("transformer_impl") == "inference_optimized": + # The provider snapshot was taken from the training model, whose + # generic GPT spec uses Transformer Engine. Select MCore inference + # linears before constructing the dedicated generation model. + _configure_inference_optimized_layer_spec(plan.initial_model_provider) + print( "[colocated-reshard] building dedicated inference model " f"(inference TP={inference_mcfg['tensor_model_parallel_size']} " diff --git a/tests/functional/grpo_megatron_generation_colocated_reshard.sh b/tests/functional/grpo_megatron_generation_colocated_reshard.sh index c2d5ad2163d..286579e8e28 100755 --- a/tests/functional/grpo_megatron_generation_colocated_reshard.sh +++ b/tests/functional/grpo_megatron_generation_colocated_reshard.sh @@ -21,6 +21,7 @@ mkdir -p $EXP_DIR $LOG_DIR # prepare_for_generation must build a dedicated inference model and swap # weights across differing layouts. A wrong-weights reshard shows up as # generation/training logprob disagreement, hence the mult_prob_error gate. +# Compare raw generation and training logprobs without sampling filters. cd $PROJECT_ROOT uv run coverage run -a --data-file=$PROJECT_ROOT/tests/.coverage --source=$PROJECT_ROOT/nemo_rl \ $PROJECT_ROOT/examples/run_grpo.py \ @@ -34,7 +35,8 @@ uv run coverage run -a --data-file=$PROJECT_ROOT/tests/.coverage --source=$PROJE policy.megatron_cfg.tensor_model_parallel_size=2 \ policy.generation.backend=megatron \ policy.generation.refit_transport=mcore \ - policy.generation.top_k=1 \ + policy.generation.top_k=null \ + policy.generation.top_p=1.0 \ policy.generation.mcore_generation_config.logprobs_mode=raw_logprobs \ ++policy.generation.mcore_generation_config.transformer_impl=inference_optimized \ ++policy.generation.mcore_generation_config.tensor_model_parallel_size=1 \ From ea842d7c4df6e1ea03173065ad4cf0db328ecb2d Mon Sep 17 00:00:00 2001 From: William Dykas Date: Tue, 8 Sep 2026 18:46:29 -0700 Subject: [PATCH 4/8] fix(megatron): address final refit review comments Signed-off-by: William Dykas --- ...eneration-colocated-reshard-async-gym.yaml | 2 + ...olocated-single-controller-async-lag4.yaml | 6 +++ .../generation/megatron/megatron_worker.py | 16 ++++--- .../models/generation/vllm/vllm_backend.py | 3 +- .../policy/workers/megatron_policy_worker.py | 4 +- nemo_rl/weight_sync/nccl_reshard_utils.py | 11 ++++- ...o_megatron_generation_non_colocated_gym.sh | 1 + .../models/policy/test_megatron_worker.py | 45 ++++++++++++++++++- .../weight_sync/test_nccl_reshard_utils.py | 8 ++++ 9 files changed, 87 insertions(+), 9 deletions(-) diff --git a/examples/configs/recipes/llm/grpo-nanov3-30BA3B-2n8g-megatron_generation-colocated-reshard-async-gym.yaml b/examples/configs/recipes/llm/grpo-nanov3-30BA3B-2n8g-megatron_generation-colocated-reshard-async-gym.yaml index ead10f761f9..87fb4170e75 100644 --- a/examples/configs/recipes/llm/grpo-nanov3-30BA3B-2n8g-megatron_generation-colocated-reshard-async-gym.yaml +++ b/examples/configs/recipes/llm/grpo-nanov3-30BA3B-2n8g-megatron_generation-colocated-reshard-async-gym.yaml @@ -23,7 +23,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: diff --git a/examples/configs/recipes/llm/grpo-qwen2.5-math-1.5b-instruct-1n4g-megatron_generation-noncolocated-single-controller-async-lag4.yaml b/examples/configs/recipes/llm/grpo-qwen2.5-math-1.5b-instruct-1n4g-megatron_generation-noncolocated-single-controller-async-lag4.yaml index dc46d0b0be8..4f26f3c9a96 100644 --- a/examples/configs/recipes/llm/grpo-qwen2.5-math-1.5b-instruct-1n4g-megatron_generation-noncolocated-single-controller-async-lag4.yaml +++ b/examples/configs/recipes/llm/grpo-qwen2.5-math-1.5b-instruct-1n4g-megatron_generation-noncolocated-single-controller-async-lag4.yaml @@ -12,5 +12,11 @@ async_rl: policy: generation: backend: megatron + refit_transport: mcore mcore_generation_config: + refit_backend: nccl kv_cache_management_mode: recompute +async_rl: + generation_fleet_health: + # Native MCore refit has no worker-side abort hook for a deadline. + refit_timeout_s: null diff --git a/nemo_rl/models/generation/megatron/megatron_worker.py b/nemo_rl/models/generation/megatron/megatron_worker.py index d453a7c61ba..653bc734d89 100644 --- a/nemo_rl/models/generation/megatron/megatron_worker.py +++ b/nemo_rl/models/generation/megatron/megatron_worker.py @@ -1426,7 +1426,7 @@ def _commit_megatron_bulk_refit_piece( ) if not piece.task.is_mxfp8: assert piece.destination is not None - piece.destination.copy_(tensor) + piece.spec.select(piece.destination).copy_(tensor) return pending = self._generation_m2n_pending.setdefault(piece.task.target_id, {}) @@ -1454,7 +1454,7 @@ def add_piece(task: _MegatronRefitTask, spec: Any) -> None: raise RuntimeError( f"Duplicate Megatron M-to-N target for {spec.name!r}." ) - destination = None if task.is_mxfp8 else spec.select(task.destination) + destination = None if task.is_mxfp8 else task.destination # MXFP8 destinations stage through BF16: the wire payload is logical # BF16 by design, and MXFP8Tensor.dtype is optional metadata that is # None until the destination's first successful update. Pin the @@ -1515,6 +1515,14 @@ def post(ctx: RefitCtx) -> None: return LocalParamSpec(base=None, pre=pre, post=post) + def direct_spec(piece: _MegatronBulkRefitPiece) -> LocalParamSpec: + assert piece.destination is not None + + def pre(base: torch.Tensor) -> RefitCtx: + return RefitCtx(buf=piece.spec.select(base)) + + return LocalParamSpec(base=piece.destination, pre=pre) + specs = {} for layer_name in refit_info["layer_names"]: for param_info in refit_info["per_layer_params"][layer_name]: @@ -1541,9 +1549,7 @@ def post(ctx: RefitCtx) -> None: f"No local Megatron destination maps to M-to-N weight {name!r}." ) specs[name] = ( - staged_spec(piece) - if piece.task.is_mxfp8 - else LocalParamSpec(base=piece.destination) + staged_spec(piece) if piece.task.is_mxfp8 else direct_spec(piece) ) return HFToLocalParamMap(specs=specs) diff --git a/nemo_rl/models/generation/vllm/vllm_backend.py b/nemo_rl/models/generation/vllm/vllm_backend.py index 09024a798c4..dc426d381c4 100644 --- a/nemo_rl/models/generation/vllm/vllm_backend.py +++ b/nemo_rl/models/generation/vllm/vllm_backend.py @@ -38,6 +38,7 @@ _STR_TO_DTYPE, HFToLocalParamMap, LocalParamSpec, + RefitBuilderInterface, RefitCtx, _extract_layer_prefix, ) @@ -271,7 +272,7 @@ def _read_mtp_layer_weights_from_checkpoint( return weights -class VllmInternalWorkerExtension: +class VllmInternalWorkerExtension(RefitBuilderInterface): # Per-PP-stage refit groups, None until init_nccl_reshard_comm_group builds them. # Declared rather than sprung into existence so a rebuild can release the previous # ones without probing, matching AbstractPolicyWorker.model_update_group. None and diff --git a/nemo_rl/models/policy/workers/megatron_policy_worker.py b/nemo_rl/models/policy/workers/megatron_policy_worker.py index c6202dea9db..570816e5099 100644 --- a/nemo_rl/models/policy/workers/megatron_policy_worker.py +++ b/nemo_rl/models/policy/workers/megatron_policy_worker.py @@ -132,6 +132,7 @@ _INDIVIDUAL_EXPERT_RE, HFToLocalParamMap, LocalParamSpec, + RefitBuilderInterface, RefitCtx, _extract_layer_name, _extract_layer_prefix, @@ -420,6 +421,7 @@ class MegatronPolicyWorkerImpl( PolicyCheckpointEngineMixin, AbstractPolicyWorker, ColocatablePolicyInterface, + RefitBuilderInterface, ): # Tests and extension classes that bypass __init__ retain the historical # training/source behavior unless they explicitly select destination. @@ -659,7 +661,7 @@ def __init__( ): # Bridge's generic GPT provider defaults to its training layer spec. # Dedicated inference workers need MCore inference linears instead. - # A model-specific provider already handles this and returns False. + # A model-specific provider already handles this itself. _configure_inference_optimized_layer_spec(self.megatron_cfg.model) self.dtype = runtime_config.dtype self.optimizer_cpu_offload = runtime_config.optimizer_cpu_offload diff --git a/nemo_rl/weight_sync/nccl_reshard_utils.py b/nemo_rl/weight_sync/nccl_reshard_utils.py index e7cd9bcd4ea..bde2c2ec996 100644 --- a/nemo_rl/weight_sync/nccl_reshard_utils.py +++ b/nemo_rl/weight_sync/nccl_reshard_utils.py @@ -815,10 +815,19 @@ def check_nccl_reshard_refit_support(master_config: Any) -> None: # Megatron ETP. The vLLM backend shards experts by index across # its TP ranks, so its EP is either 1 (TP-sharded experts) or equal to TP # (EP-sharded). PP is not yet supported gen-side. - if generation.get("backend") == "vllm": + if backend == "vllm": gen_tp = vllm_cfg.get("tensor_parallel_size", 1) gen_ep = vllm_cfg.get("expert_parallel_size", 1) gen_pp = vllm_cfg.get("pipeline_parallel_size", 1) + # Megatron-source ETP has unit coverage but has not been fully tested + # end to end with a vLLM destination. Keep it disabled until it has. + train_etp = megatron_cfg.get("expert_tensor_parallel_size", 1) + if train_etp != 1: + violations.append( + "policy.megatron_cfg.expert_tensor_parallel_size must be 1 " + "for a vLLM destination with nccl_reshard refit " + f"(got {train_etp})." + ) if gen_ep != 1 and gen_ep != gen_tp: violations.append( "policy.generation.vllm_cfg.expert_parallel_size must be 1 or " diff --git a/tests/functional/grpo_megatron_generation_non_colocated_gym.sh b/tests/functional/grpo_megatron_generation_non_colocated_gym.sh index 50a8edf635f..a9db4478fe7 100755 --- a/tests/functional/grpo_megatron_generation_non_colocated_gym.sh +++ b/tests/functional/grpo_megatron_generation_non_colocated_gym.sh @@ -64,6 +64,7 @@ uv run coverage run -a --data-file=$PROJECT_ROOT/tests/.coverage --source=$PROJE policy.megatron_cfg.context_parallel_size=1 \ policy.megatron_cfg.sequence_parallel=false \ policy.generation.backend=megatron \ + +policy.generation.refit_transport=mcore \ policy.generation.mcore_generation_config.expose_http_server=true \ policy.max_total_sequence_length=512 \ policy.generation.max_new_tokens=128 \ diff --git a/tests/unit/models/policy/test_megatron_worker.py b/tests/unit/models/policy/test_megatron_worker.py index f10d3662703..6dd4b2cebff 100644 --- a/tests/unit/models/policy/test_megatron_worker.py +++ b/tests/unit/models/policy/test_megatron_worker.py @@ -611,7 +611,9 @@ def test_megatron_m2n_unstacks_grouped_experts_into_local_weights() -> None: for expert in range(2): gate_name = f"model.layers.0.mlp.experts.{expert}.gate_proj.weight" up_name = f"model.layers.0.mlp.experts.{expert}.up_proj.weight" - destination = torch.zeros((4, 2), dtype=torch.bfloat16) + destination = torch.nn.Parameter( + torch.zeros((4, 2), dtype=torch.bfloat16), requires_grad=False + ) destinations.append(destination) tasks.append( _make_refit_task( @@ -638,6 +640,10 @@ def test_megatron_m2n_unstacks_grouped_experts_into_local_weights() -> None: worker = object.__new__(MegatronGenerationRefitMixin) param_map = worker._build_destination_hf_to_local_param_map(refit_info, tasks) + orphaned_storage = [destination.data for destination in destinations] + for destination in destinations: + destination.data = torch.full_like(destination, -1) + gate_spec = param_map.get(grouped_gate) gate_ctx = gate_spec.pre(None) gate_ctx.buf[0].fill_(1) @@ -653,6 +659,43 @@ def test_megatron_m2n_unstacks_grouped_experts_into_local_weights() -> None: assert torch.equal(destinations[0][2:], torch.full_like(destinations[0][2:], 3)) assert torch.equal(destinations[1][:2], torch.full_like(destinations[1][:2], 2)) assert torch.equal(destinations[1][2:], torch.full_like(destinations[1][2:], 4)) + assert all(torch.count_nonzero(storage).item() == 0 for storage in orphaned_storage) + + +def test_megatron_m2n_resolves_direct_destination_after_parameter_rebind() -> None: + from nemo_rl.models.generation.megatron.megatron_worker import ( + MegatronGenerationRefitMixin, + ) + + gate_name = "model.layers.0.mlp.gate_proj.weight" + up_name = "model.layers.0.mlp.up_proj.weight" + destination = torch.nn.Parameter( + torch.zeros((4, 2), dtype=torch.bfloat16), requires_grad=False + ) + task = _make_refit_task( + param_name="decoder.layers.0.mlp.linear_fc1.weight", + destination=destination, + dependencies=(gate_name, up_name), + local_specs=((gate_name, "gate"), (up_name, "up")), + ) + refit_info = { + "layer_names": ["model.layers.0"], + "per_layer_params": {"model.layers.0": [{"name": gate_name}]}, + } + worker = object.__new__(MegatronGenerationRefitMixin) + param_map = worker._build_destination_hf_to_local_param_map(refit_info, [task]) + gate_spec = param_map.get(gate_name) + assert gate_spec is not None + assert gate_spec.pre is not None + + orphaned_storage = destination.data + destination.data = torch.full_like(destination, -1) + gate_ctx = gate_spec.pre(gate_spec.base) + gate_ctx.buf.fill_(7) + + assert torch.equal(destination[:2], torch.full_like(destination[:2], 7)) + assert torch.equal(destination[2:], torch.full_like(destination[2:], -1)) + assert torch.count_nonzero(orphaned_storage).item() == 0 @pytest.mark.parametrize( diff --git a/tests/unit/weight_sync/test_nccl_reshard_utils.py b/tests/unit/weight_sync/test_nccl_reshard_utils.py index 5cb659505e4..547da849cf2 100644 --- a/tests/unit/weight_sync/test_nccl_reshard_utils.py +++ b/tests/unit/weight_sync/test_nccl_reshard_utils.py @@ -66,6 +66,14 @@ def test_check_nccl_reshard_refit_support_accepts_valid_config() -> None: check_nccl_reshard_refit_support(_valid_nccl_reshard_config()) +def test_check_nccl_reshard_refit_support_rejects_megatron_etp_with_vllm() -> None: + config = _valid_nccl_reshard_config() + config.policy["megatron_cfg"]["expert_tensor_parallel_size"] = 2 + + with pytest.raises(ValueError, match="expert_tensor_parallel_size must be 1"): + check_nccl_reshard_refit_support(config) + + def test_check_nccl_reshard_refit_support_rejects_reload_api() -> None: config = _valid_nccl_reshard_config() config.policy["generation"]["vllm_cfg"]["refit_with_reload_api"] = True From b95991f946f0bb894700d04c90adac0d28393dd3 Mon Sep 17 00:00:00 2001 From: William Dykas Date: Wed, 9 Sep 2026 04:05:30 -0700 Subject: [PATCH 5/8] fix(megatron): repair refit functional configs Signed-off-by: William Dykas --- .../functional/grpo_megatron_generation_colocated_reshard.sh | 4 ---- .../grpo_megatron_generation_colocated_reshard_async_gym.sh | 2 ++ ...tron_generation_colocated_reshard_gym_single_controller.sh | 2 ++ ...nemotron_omni_gym_video_megatron_single_controller_1n2g.sh | 1 + 4 files changed, 5 insertions(+), 4 deletions(-) diff --git a/tests/functional/grpo_megatron_generation_colocated_reshard.sh b/tests/functional/grpo_megatron_generation_colocated_reshard.sh index 286579e8e28..63ece7fa4ce 100755 --- a/tests/functional/grpo_megatron_generation_colocated_reshard.sh +++ b/tests/functional/grpo_megatron_generation_colocated_reshard.sh @@ -21,7 +21,6 @@ mkdir -p $EXP_DIR $LOG_DIR # prepare_for_generation must build a dedicated inference model and swap # weights across differing layouts. A wrong-weights reshard shows up as # generation/training logprob disagreement, hence the mult_prob_error gate. -# Compare raw generation and training logprobs without sampling filters. cd $PROJECT_ROOT uv run coverage run -a --data-file=$PROJECT_ROOT/tests/.coverage --source=$PROJECT_ROOT/nemo_rl \ $PROJECT_ROOT/examples/run_grpo.py \ @@ -35,9 +34,6 @@ uv run coverage run -a --data-file=$PROJECT_ROOT/tests/.coverage --source=$PROJE policy.megatron_cfg.tensor_model_parallel_size=2 \ policy.generation.backend=megatron \ policy.generation.refit_transport=mcore \ - policy.generation.top_k=null \ - policy.generation.top_p=1.0 \ - policy.generation.mcore_generation_config.logprobs_mode=raw_logprobs \ ++policy.generation.mcore_generation_config.transformer_impl=inference_optimized \ ++policy.generation.mcore_generation_config.tensor_model_parallel_size=1 \ policy.generation.mcore_generation_config.refit_backend=nccl \ diff --git a/tests/functional/grpo_megatron_generation_colocated_reshard_async_gym.sh b/tests/functional/grpo_megatron_generation_colocated_reshard_async_gym.sh index 0412b08fe9f..66811e4aecc 100755 --- a/tests/functional/grpo_megatron_generation_colocated_reshard_async_gym.sh +++ b/tests/functional/grpo_megatron_generation_colocated_reshard_async_gym.sh @@ -64,6 +64,7 @@ uv run coverage run -a --data-file=$PROJECT_ROOT/tests/.coverage --source=$PROJE policy.megatron_cfg.context_parallel_size=1 \ policy.megatron_cfg.sequence_parallel=false \ policy.generation.backend=megatron \ + policy.generation.refit_transport=mcore \ policy.generation.mcore_generation_config.expose_http_server=true \ policy.generation.mcore_generation_config.enable_prefix_caching=true \ policy.max_total_sequence_length=512 \ @@ -71,6 +72,7 @@ uv run coverage run -a --data-file=$PROJECT_ROOT/tests/.coverage --source=$PROJE policy.generation.colocated.enabled=true \ ++policy.generation.mcore_generation_config.transformer_impl=inference_optimized \ ++policy.generation.mcore_generation_config.tensor_model_parallel_size=1 \ + policy.generation.mcore_generation_config.refit_backend=nccl \ grpo.num_prompts_per_step=4 \ grpo.num_generations_per_prompt=2 \ grpo.max_num_steps=10 \ diff --git a/tests/functional/grpo_megatron_generation_colocated_reshard_gym_single_controller.sh b/tests/functional/grpo_megatron_generation_colocated_reshard_gym_single_controller.sh index 23f65fd94c1..470301bdf6f 100755 --- a/tests/functional/grpo_megatron_generation_colocated_reshard_gym_single_controller.sh +++ b/tests/functional/grpo_megatron_generation_colocated_reshard_gym_single_controller.sh @@ -69,9 +69,11 @@ uv run coverage run -a --data-file=$PROJECT_ROOT/tests/.coverage --source=$PROJE policy.megatron_cfg.context_parallel_size=1 \ policy.megatron_cfg.sequence_parallel=false \ policy.generation.backend=megatron \ + policy.generation.refit_transport=mcore \ policy.generation.mcore_generation_config.expose_http_server=true \ ++policy.generation.mcore_generation_config.transformer_impl=inference_optimized \ ++policy.generation.mcore_generation_config.tensor_model_parallel_size=1 \ + policy.generation.mcore_generation_config.refit_backend=nccl \ policy.generation.max_new_tokens=128 \ policy.max_total_sequence_length=512 \ policy.generation.colocated.enabled=true \ diff --git a/tests/functional/nemotron_omni_gym_video_megatron_single_controller_1n2g.sh b/tests/functional/nemotron_omni_gym_video_megatron_single_controller_1n2g.sh index 86d24273d8f..8f76cf9c7ce 100755 --- a/tests/functional/nemotron_omni_gym_video_megatron_single_controller_1n2g.sh +++ b/tests/functional/nemotron_omni_gym_video_megatron_single_controller_1n2g.sh @@ -104,6 +104,7 @@ uv run --no-sync python examples/run_grpo_single_controller.py \ ++policy.megatron_cfg.optimizer.exp_avg_sq_dtype=bfloat16 \ ++policy.megatron_cfg.optimizer.store_param_remainders=false \ policy.generation.backend=megatron \ + policy.generation.refit_transport=mcore \ ++policy.generation.bad_words=null \ policy.generation.colocated.enabled=false \ policy.generation.colocated.resources.num_nodes=1 \ From a8d2bcd61b4ee75ef5e1f270bf63934a44cec52e Mon Sep 17 00:00:00 2001 From: William Dykas Date: Wed, 9 Sep 2026 05:21:39 -0700 Subject: [PATCH 6/8] fix(megatron): repair refit validation and recipes Signed-off-by: William Dykas --- ...ration-noncolocated-single-controller-async-lag4.yaml | 7 +++---- ...ration-noncolocated-single-controller-async-lag4.yaml | 7 +++---- nemo_rl/weight_sync/nccl_reshard_utils.py | 9 +++++++++ 3 files changed, 15 insertions(+), 8 deletions(-) diff --git a/examples/configs/recipes/llm/grpo-qwen2.5-math-1.5b-instruct-1n4g-megatron_generation-noncolocated-single-controller-async-lag4.yaml b/examples/configs/recipes/llm/grpo-qwen2.5-math-1.5b-instruct-1n4g-megatron_generation-noncolocated-single-controller-async-lag4.yaml index 4f26f3c9a96..f1e7f368bcf 100644 --- a/examples/configs/recipes/llm/grpo-qwen2.5-math-1.5b-instruct-1n4g-megatron_generation-noncolocated-single-controller-async-lag4.yaml +++ b/examples/configs/recipes/llm/grpo-qwen2.5-math-1.5b-instruct-1n4g-megatron_generation-noncolocated-single-controller-async-lag4.yaml @@ -9,6 +9,9 @@ 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 @@ -16,7 +19,3 @@ policy: mcore_generation_config: refit_backend: nccl kv_cache_management_mode: recompute -async_rl: - generation_fleet_health: - # Native MCore refit has no worker-side abort hook for a deadline. - refit_timeout_s: null diff --git a/examples/configs/recipes/llm/grpo-qwen2.5-math-1.5b-instruct-1n8g-megatron_generation-noncolocated-single-controller-async-lag4.yaml b/examples/configs/recipes/llm/grpo-qwen2.5-math-1.5b-instruct-1n8g-megatron_generation-noncolocated-single-controller-async-lag4.yaml index 6df01ce1aeb..b2777c2baac 100644 --- a/examples/configs/recipes/llm/grpo-qwen2.5-math-1.5b-instruct-1n8g-megatron_generation-noncolocated-single-controller-async-lag4.yaml +++ b/examples/configs/recipes/llm/grpo-qwen2.5-math-1.5b-instruct-1n8g-megatron_generation-noncolocated-single-controller-async-lag4.yaml @@ -9,6 +9,9 @@ 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 @@ -16,7 +19,3 @@ policy: mcore_generation_config: refit_backend: nccl kv_cache_management_mode: recompute -async_rl: - generation_fleet_health: - # Native MCore refit has no worker-side abort hook for a deadline. - refit_timeout_s: null diff --git a/nemo_rl/weight_sync/nccl_reshard_utils.py b/nemo_rl/weight_sync/nccl_reshard_utils.py index bde2c2ec996..11e14f2e6b0 100644 --- a/nemo_rl/weight_sync/nccl_reshard_utils.py +++ b/nemo_rl/weight_sync/nccl_reshard_utils.py @@ -652,6 +652,15 @@ def check_nccl_reshard_refit_support(master_config: Any) -> None: ) if megatron_enabled: + etp = megatron_cfg.get("expert_tensor_parallel_size", 1) + + # ETP with vLLM generation has not been fully tested yet. + if backend == "vllm" and etp not in (1, None): + violations.append( + "policy.megatron_cfg.expert_tensor_parallel_size must be 1 " + f"for vLLM generation until ETP is fully tested (got {etp})." + ) + # PP-layout knobs that _build_layer_to_pp_stage doesn't yet handle. if megatron_cfg.get("pipeline_model_parallel_layout") is not None: violations.append( From efff98bf7b2797f5081935be5db9bbd50c0aa07a Mon Sep 17 00:00:00 2001 From: wdykas Date: Thu, 10 Sep 2026 10:38:57 -0700 Subject: [PATCH 7/8] test(megatron): initialize OPD state in refit fixture Signed-off-by: wdykas --- tests/unit/models/policy/test_megatron_worker.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/unit/models/policy/test_megatron_worker.py b/tests/unit/models/policy/test_megatron_worker.py index 441f43dfc40..c6c80e6ea86 100644 --- a/tests/unit/models/policy/test_megatron_worker.py +++ b/tests/unit/models/policy/test_megatron_worker.py @@ -1646,6 +1646,7 @@ def test_megatron_prepare_for_logprobs_restores_mxfp8_shared_buffer(monkeypatch) ) worker = object.__new__(MegatronPolicyWorkerImpl) + _disable_opd_full(worker) model = _FakeTrainableModel() move_calls = [] From 4194bf5b808434151716c46ddf2cb940ca3754ac Mon Sep 17 00:00:00 2001 From: wdykas Date: Thu, 10 Sep 2026 15:14:24 -0700 Subject: [PATCH 8/8] test(megatron): compare raw sampling logprobs Signed-off-by: wdykas --- tests/functional/grpo_megatron_generation_topp_topk.sh | 3 +++ 1 file changed, 3 insertions(+) diff --git a/tests/functional/grpo_megatron_generation_topp_topk.sh b/tests/functional/grpo_megatron_generation_topp_topk.sh index 7eaafaf640e..c91b3b205f1 100644 --- a/tests/functional/grpo_megatron_generation_topp_topk.sh +++ b/tests/functional/grpo_megatron_generation_topp_topk.sh @@ -20,6 +20,8 @@ mkdir -p $EXP_DIR $LOG_DIR # Megatron generation with non-default sampling (temperature / top-p / top-k), exercising # the SamplingParams path (the megatron analog of grpo_topp_topk.sh). # Using Qwen2.5-0.5B instead of Qwen3-0.6B because the latter is not supported by Megatron yet +# Request pre-sampling log-probs for parity with policy recomputation; processed log-probs +# include temperature/top-p/top-k renormalization. cd $PROJECT_ROOT uv run coverage run -a --data-file=$PROJECT_ROOT/tests/.coverage --source=$PROJECT_ROOT/nemo_rl \ $PROJECT_ROOT/examples/run_grpo.py \ @@ -33,6 +35,7 @@ uv run coverage run -a --data-file=$PROJECT_ROOT/tests/.coverage --source=$PROJE policy.generation.backend=megatron \ policy.generation.refit_transport=mcore \ policy.generation.mcore_generation_config.refit_backend=nccl \ + policy.generation.mcore_generation_config.logprobs_mode=raw_logprobs \ policy.generation.temperature=0.8 \ policy.generation.top_p=0.9 \ policy.generation.top_k=50 \