diff --git a/3rdparty/Megatron-Bridge-workspace/Megatron-Bridge b/3rdparty/Megatron-Bridge-workspace/Megatron-Bridge index 554c7b93242..a056408d29a 160000 --- a/3rdparty/Megatron-Bridge-workspace/Megatron-Bridge +++ b/3rdparty/Megatron-Bridge-workspace/Megatron-Bridge @@ -1 +1 @@ -Subproject commit 554c7b9324225aa863eee52e8b8fdde7abced2b1 +Subproject commit a056408d29ad1070fb926512991075998c9e023f diff --git a/docs/design-docs/modelopt-real-quant-architecture.md b/docs/design-docs/modelopt-real-quant-architecture.md new file mode 100644 index 00000000000..c45331c95cc --- /dev/null +++ b/docs/design-docs/modelopt-real-quant-architecture.md @@ -0,0 +1,159 @@ +# ModelOpt Real-Quant Refit Architecture + +NeMo RL supports deployment-style NVFP4 rollout generation while a Megatron +policy is trained with ModelOpt quantization-aware training. During each +policy refit, Megatron-Bridge exports packed quantized tensors and NeMo RL +loads them into a vLLM generation model without rebuilding the model or its +CUDA graphs. + +## Relationship to the overall NeMo RL design + +Real-quant refit extends the standard NeMo RL policy-generation workflow. The +algorithm controller continues to coordinate independent policy and generation +workers through the existing interfaces. ModelOpt changes the representation +used by the policy, while this design changes only how a policy update is +exported and installed in the vLLM generation worker. + +This design builds on: + +- [Design and Philosophy](design-and-philosophy.md), which defines NeMo RL's + controller, worker, isolation, and communication model; +- [Generation Interface](generation.md), which defines generation backends and + their weight-update lifecycle; and +- [Quantization-Aware RL](../guides/quantization-aware-rl.md), which documents + the user workflow, configuration, and supported recipes. + +The real-quant path preserves those abstractions: algorithms still call the +same policy and generation interfaces, and non-quantized and fake-quantized +weight updates continue to use their existing paths. + +## Component responsibilities + +| Component | Responsibility | +|---|---| +| ModelOpt | Quantization configuration, calibration, QAT state, NVFP4 packing, and scale derivation | +| Megatron-Bridge | Megatron-to-Hugging-Face conversion, distributed TP/PP/EP handling, and NVFP4 export | +| NeMo RL | Mode validation, refit scheduling, named-tensor transport, and vLLM reload orchestration | +| vLLM | Checkpoint loading, runtime-layout conversion, kernel selection, stable tensor placement, and KV-cache scale processing | + +NeMo RL delegates quantization math and runtime-kernel conversion to ModelOpt, +Megatron-Bridge, and vLLM. Its vLLM integration is limited to the +format-specific adapters required to connect their public interfaces. + +```mermaid +flowchart LR + A[RL algorithm controller] --> B[ModelOpt QAT policy worker] + B --> C[Megatron-Bridge NVFP4 export] + C --> D[NeMo RL named-tensor transport] + D --> E[vLLM generation worker] + E --> F[Rollouts] + F --> A +``` + +## Quantization modes + +Both modes use block-16 E2M1 NVFP4 weights with E4M3 block scales. + +| Mode | Deployment algorithm | Weights | Activations | +|---|---|---|---| +| W4A4 | `NVFP4` | NVFP4 | NVFP4 with per-projection input scales | +| W4A16 | `W4A16_NVFP4` | NVFP4 | Native model dtype | + +The policy and generation workers must use quantization recipes that resolve +to the same mode. Unsupported or mismatched formats fail during setup. + +Example routed-expert recipes are provided at: + +- [`examples/modelopt/quant_configs/nvfp4_experts.yaml`](../../examples/modelopt/quant_configs/nvfp4_experts.yaml) for W4A4; +- [`examples/modelopt/quant_configs/nvfp4_experts_weightonly.yaml`](../../examples/modelopt/quant_configs/nvfp4_experts_weightonly.yaml) for W4A16. + +## Refit lifecycle + +Every real-quant refit follows the same lifecycle: + +1. NeMo RL validates the policy and generation quantization modes. +2. Megatron-Bridge exports named NVFP4 weights and scale tensors. +3. NeMo RL starts vLLM's layerwise reload lifecycle for the affected + quantized modules. +4. The normal vLLM model loader consumes the exported tensors. +5. vLLM performs post-load conversion and copies the converted tensors into + the existing runtime storage. +6. NeMo RL finalizes the reload and synchronizes the device before the + transport buffer can be reused. + +This lifecycle is shared by collective and CUDA IPC transports. vLLM retains +ownership of runtime tensor layouts and preserves tensor addresses referenced +by CUDA graphs. + +## vLLM compatibility adapters + +NeMo RL registers ModelOpt NVFP4 extensions through vLLM's quantization +registry. The adapters cover: + +- W4A16 dense and fused-MoE execution; +- fused-MoE input-scale loading for W4A4; +- rank-local padding required by W4A16 MoE kernels; and +- preservation of runtime kernel references during repeated refits. + +Native vLLM processing remains responsible for checkpoint-layout restoration, +post-load conversion, kernel construction, and stable-storage copy-back. The +adapters do not define a separate reload implementation. + +## Fused-MoE transport + +Megatron-Bridge exports fused expert projections as expert-batched W13 and W2 +tensors. Each family contains the packed weight, block scale, and global scale. +W4A4 additionally includes one input scale for each projection. + +The receiving adapter validates that each tensor family is complete and maps +it to the checkpoint names accepted by the vLLM model loader. Gated experts +use two W13 shards; non-gated experts use one. + +Fused-MoE refit currently requires every vLLM rank to own the full expert set. +Megatron expert parallelism remains supported because Megatron-Bridge gathers +the exported payload before the vLLM refit. + +## CUDA IPC buffer lifetime + +Layerwise loading may temporarily retain a view of an incoming tensor until a +layer has received its complete payload. NeMo RL ensures that no retained view +aliases a reusable CUDA IPC staging buffer before acknowledging that buffer. +The final acknowledgment is sent only after vLLM finalization and device +synchronization. + +## KV-cache behavior + +W4A4 and W4A16 weight refit do not enable KV-cache quantization. KV-cache +precision remains controlled by vLLM's `kv_cache_dtype` configuration. + +When FP8 KV cache is selected, vLLM owns its scale creation, loading, and +post-load processing. The ModelOpt real-quant adapter does not replace the +vLLM KV-cache processing method or process an attention layer a second time. + +## Configuration + +Set the same quantization recipe for the policy and generation worker, and +enable real-quant generation: + +```yaml +policy: + quant_cfg: examples/modelopt/quant_configs/nvfp4_experts.yaml + + generation: + quant_cfg: examples/modelopt/quant_configs/nvfp4_experts.yaml + real_quant: true +``` + +Layer selection belongs in a purpose-specific quantization recipe. Existing +shared configs should not be changed to carry experiment-specific exclusions. +Accuracy-driven exclusions should be supported by reproducible sensitivity or +AutoQuant results; exclusions required by an unsupported tensor or operator +contract should be documented in the corresponding recipe. + +## Current limitations + +- Real-quant rollout generation requires vLLM. +- Policy export currently uses the Megatron policy path and Megatron-Bridge. +- Supported real-quant formats are dynamic block-16 NVFP4 W4A4 and W4A16. +- Fused-MoE vLLM expert parallelism is not supported during refit. +- Model support is recipe-specific and requires end-to-end validation. diff --git a/docs/guides/quantization-aware-rl.md b/docs/guides/quantization-aware-rl.md index 88c49c2683f..a8a95e0c6b1 100644 --- a/docs/guides/quantization-aware-rl.md +++ b/docs/guides/quantization-aware-rl.md @@ -1,6 +1,6 @@ # Quantization-Aware RL (QARL) -Quantization-Aware RL (QARL) integrates [NVIDIA Model Optimizer (ModelOpt)](https://github.com/NVIDIA/Model-Optimizer) into the NeMo RL training loop, enabling quantization-aware training and generation for both GRPO and on-policy distillation workflows. QARL automatically quantizes a standard model at initialization, maintains quantizer state (amax values) throughout training, and transfers quantized state to vLLM during weight refit. By default, vLLM generation uses fake-quantized modules. For NVFP4 W4A16 rollout experiments, NeMo RL can instead stream packed real-quant ModelOpt NVFP4 weights into vLLM. +Quantization-Aware RL (QARL) integrates [NVIDIA Model Optimizer (ModelOpt)](https://github.com/NVIDIA/Model-Optimizer) into the NeMo RL training loop, enabling quantization-aware training and generation for both GRPO and on-policy distillation workflows. QARL automatically quantizes a standard model at initialization, maintains quantizer state (amax values) throughout training, and transfers quantized state to vLLM during weight refit. By default, vLLM generation uses fake-quantized modules. For NVFP4 W4A4 and W4A16 rollout experiments, NeMo RL can instead stream packed real-quant ModelOpt NVFP4 weights and scales into vLLM. ## Overview @@ -9,9 +9,9 @@ In a standard NeMo RL loop, model weights are trained in full precision and refi There are two vLLM rollout modes: - **Fake-quant rollout**: vLLM receives folded full-precision weights and runs fake-quantized layers. This is the default when `policy.generation.quant_cfg` is set. -- **Real-quant rollout**: vLLM is initialized with ModelOpt NVFP4 kernels and receives packed NVFP4 weights plus scale tensors during every refit. Enable this with `policy.generation.real_quant: true`. +- **Real-quant rollout**: vLLM is initialized with ModelOpt NVFP4 kernels and receives packed NVFP4 weights plus scale tensors during every refit. Enable this with `policy.generation.real_quant: true`. W4A16 keeps activations in their native dtype; W4A4 additionally streams calibrated activation scales and uses an activation-quantizing vLLM kernel. The Megatron policy worker exports the payload through Megatron-Bridge. -See [Verified Configurations](#verified-configurations) for the workflow + recipe combinations that have been empirically validated, and [Supported Quantization Formats](#supported-quantization-formats) for the full set of available formats. W4A4 (`NVFP4_DEFAULT_CFG`) converges for on-policy distillation but has been observed to have convergence issues on GRPO; W4A16 (NVFP4 weights, native-dtype activations) works for GRPO. +See [Verified Configurations](#verified-configurations) for the workflow + recipe combinations that have been empirically validated, and [Supported Quantization Formats](#supported-quantization-formats) for the full set of available formats. Results are recipe- and model-specific: the generic W4A4 `NVFP4_DEFAULT_CFG` has known GRPO convergence issues, while the routed-expert Qwen3 W4A4 real-quant recipe below completed the documented single-seed campaign. ## Verified Configurations @@ -25,8 +25,10 @@ The following workflow + quantization recipe combinations have been validated en | QA-Distillation | W4A4 | `examples/modelopt/quant_configs/nano3_nvfp4_default.yaml` | ✅ Converges | `examples/modelopt/qa_distillation_nano3_megatron.yaml` | | QA-GRPO | W4A16 | `NVFP4_MLP_WEIGHT_ONLY_CFG` | ✅ Smoke tested on MoE | `examples/modelopt/qa_grpo_qwen3_30ba3b_megatron.yaml` | | QA-GRPO real quantization rollout | W4A16 | `examples/modelopt/quant_configs/nvfp4_a16_mlp_only.yaml` with `policy.generation.real_quant: true` | ✅ Converges | `examples/configs/recipes/llm/grpo-qwen3-8b-base-dapo-2n8g-long-megatron-qa-nvfp4-w4a16.yaml` | +| QA-GRPO real quantization rollout | W4A16 | `examples/modelopt/quant_configs/nano3_nvfp4_weightonly.yaml` with `policy.generation.real_quant: true` and the model-specific `policy.generation.real_quant_ignore` list in the example | ✅ Converges tested on hybrid MoE/Mamba | `examples/configs/recipes/llm/grpo-nanov3-30ba3b-4n4g-megatron-qa-nvfp4-w4a16-real.yaml` | +| QA-GRPO real quantization rollout | W4A4 | `examples/modelopt/quant_configs/nvfp4_experts.yaml` with `policy.generation.real_quant: true` | ✅ Completed one 300-step Qwen3-30B-A3B MoE run | `examples/configs/recipes/llm/grpo-qwen3-30ba3b-4n4g-megatron-qa-nvfp4-w4a4-real.yaml` | -The `nvfp4_a16.yaml` custom YAML enables NVFP4 e2m1 weight quantization (with dynamic e4m3 micro-block scales) and leaves activations unquantized; weights are still exercised through both Megatron training and vLLM generation. The `nvfp4_a16_mlp_only.yaml` recipe restricts W4A16 to MLP weights for real-quant rollout. +The `nvfp4_a16.yaml` custom YAML enables NVFP4 e2m1 weight quantization (with dynamic e4m3 micro-block scales) and leaves activations unquantized; weights are still exercised through both Megatron training and vLLM generation. The `nvfp4_a16_mlp_only.yaml` recipe restricts W4A16 to MLP weights for real-quant rollout. The Nano3 `nano3_nvfp4_weightonly.yaml` recipe applies the same W4A16 weight-only format to the supported MLP/MoE weights while keeping Nano3-sensitive Mamba, attention, gate/router, shared-expert, norm, and selected layer paths in BF16 through the model-specific `real_quant_ignore` list in the example config. ## ModelOpt Layer Spec Toggle @@ -87,11 +89,11 @@ sbatch \ ray.sub ``` -## Real-Quant NVFP4 Rollout (W4A16) +## Real-Quant NVFP4 Rollout (W4A4 and W4A16) -Real-quant rollout is intended for checking the deployment-style vLLM path during RL, not only the fake-quant training path. With `policy.generation.real_quant: true`, the Megatron policy worker exports ModelOpt QAT weights as packed NVFP4 tensors during refit, and the vLLM worker loads them into ModelOpt NVFP4 layers. This exercises vLLM's real FP4 kernel path during rollout while the policy training worker remains a QAT model. +Real-quant rollout is intended for checking the deployment-style vLLM path during RL, not only the fake-quant training path. With `policy.generation.real_quant: true`, the Megatron policy worker exports ModelOpt QAT weights through Megatron-Bridge as packed NVFP4 tensors during refit, and the vLLM worker loads them into ModelOpt NVFP4 layers. This exercises vLLM's real FP4 kernel path during rollout while the policy training worker remains a QAT model. -This path is validated for W4A16. +The W4A16 recipes below are validated configurations. W4A4 uses the same refit path plus calibrated per-layer or per-expert activation scales. The Megatron Qwen3 MoE recipe below completed one end-to-end 300-step run. Dense models can use the default real-quant ignore profile. MoE and hybrid models should use a model-specific ignore profile so unsupported or numerically sensitive paths stay in BF16. ### Minimal Configuration @@ -107,12 +109,89 @@ policy: real_quant: true ``` +For routed-expert Qwen3 MoE W4A4, use the block-16 E2M1 recipe in both sections: + +```yaml +policy: + quant_cfg: examples/modelopt/quant_configs/nvfp4_experts.yaml + + generation: + backend: vllm + quant_cfg: examples/modelopt/quant_configs/nvfp4_experts.yaml + real_quant: true +``` + +NeMo RL derives W4A4 versus W4A16 from the effective ModelOpt quantizer formats and rejects a policy/generation mode mismatch. W4A4 requires a native activation-quantizing NVFP4 backend; the weight-only Marlin path is reserved for W4A16. Fused-MoE real-quant refits currently require every vLLM rank to own the full expert set. + +For Nano3 W4A16 real-quant rollout, use the Nano3 weight-only recipe and an explicit model-specific ignore list: + +```yaml +policy: + quant_cfg: examples/modelopt/quant_configs/nano3_nvfp4_weightonly.yaml + + generation: + backend: vllm + quant_cfg: examples/modelopt/quant_configs/nano3_nvfp4_weightonly.yaml + real_quant: true + real_quant_ignore: + - lm_head + - '*output_layer*' + - '*mlp.gate' + - '*router*' + - '*block_sparse_moe.gate*' + - '*self_attention*' + - '*self_attn*' + - '*proj_out.*' + - '*.gate.*' + - '*mlp.shared_expert_gate.*' + - '*linear_attn.conv1d*' + - '*mixer.conv1d*' + - '*.mixer.in_proj*' + - '*.mixer.out_proj*' + - '*.shared_expert.*' + - '*.shared_experts.*' + - '*.norm.*' + - '*.q_proj*' + - '*.k_proj*' + - '*.v_proj*' + - '*.o_proj*' + - '*.qkv_proj*' + - '*.linear_proj*' + - '*.linear_qkv*' + - '*.layers.4.*' + - '*.layers.11.*' + - '*.layers.18.*' + - '*.layers.25.*' + - '*.layers.32.*' + - '*.layers.41.*' + vllm_cfg: + gpu_memory_utilization: 0.35 + enable_prefix_caching: false +``` + The ready-to-run 2-node DAPO long-context recipe is: ```text examples/configs/recipes/llm/grpo-qwen3-8b-base-dapo-2n8g-long-megatron-qa-nvfp4-w4a16.yaml ``` +The ready-to-run Nano3 4-node x 4-GPU smoke recipe is: + +```text +examples/configs/recipes/llm/grpo-nanov3-30ba3b-4n4g-megatron-qa-nvfp4-w4a16-real.yaml +``` + +The Qwen3-30B-A3B W4A4 real-quant recipe is: + +```text +examples/configs/recipes/llm/grpo-qwen3-30ba3b-4n4g-megatron-qa-nvfp4-w4a4-real.yaml +``` + +This recipe contains the 300-step, 256-example-validation campaign settings. +The GB200 nightly driver overrides it to a two-step, 32-example smoke test. +Both paths require the standalone ModelOpt-enabled Megatron-Bridge checkout; +the embedded NeMo RL checkout does not contain the grouped-MoE W4A4 exporter. + For a BF16 baseline, copy the recipe, remove `policy.quant_cfg`, `policy.generation.quant_cfg`, and `policy.generation.real_quant`, and use distinct checkpoint and log directories. @@ -127,11 +206,27 @@ uv run --extra mcore --extra modelopt --extra vllm \ --config examples/configs/recipes/llm/grpo-qwen3-8b-base-dapo-2n8g-long-megatron-qa-nvfp4-w4a16.yaml ``` -For Slurm, wrap the same command in `ray.sub` as shown in [Running QA-GRPO](#running-qa-grpo). Keep the W4A16 and BF16 runs separate and use distinct checkpoint directories. +For Nano3: + +```bash +uv run --extra mcore --extra modelopt --extra vllm \ + examples/run_grpo.py \ + --config examples/configs/recipes/llm/grpo-nanov3-30ba3b-4n4g-megatron-qa-nvfp4-w4a16-real.yaml +``` + +For Qwen3 MoE W4A4: + +```bash +uv run --extra mcore --extra modelopt --extra vllm \ + examples/run_grpo.py \ + --config examples/configs/recipes/llm/grpo-qwen3-30ba3b-4n4g-megatron-qa-nvfp4-w4a4-real.yaml +``` + +For Slurm, wrap the same command in `ray.sub` as shown in [Running QA-GRPO](#running-qa-grpo). Keep each quantized and BF16 comparison arm separate and use distinct checkpoint directories. ### Checkpoints and Fresh Starts -Real-quant rollout is sensitive to stale Megatron conversion checkpoints. For a clean first-step comparison, move aside or remove both the training checkpoint and the converted Megatron checkpoint before launching: +For a clean first-step comparison, use a new, empty `checkpointing.checkpoint_dir`; NeMo RL resumes automatically from the highest `step_*` directory it finds. The Megatron policy path also uses a converted startup checkpoint, so move aside both the training checkpoint and the converted Megatron checkpoint before launching: ```bash # Training checkpoint: matches `checkpointing.checkpoint_dir` in your config. @@ -144,7 +239,7 @@ MEGATRON_CKPT_ROOT="${NRL_MEGATRON_CHECKPOINT_DIR:-${HF_HOME:-$HOME/.cache/huggi mv "$MEGATRON_CKPT_ROOT/" "$MEGATRON_CKPT_ROOT/.old" ``` -If `NRL_MEGATRON_CHECKPOINT_DIR` is set, clear the subdirectory used by the run. On first startup, the log should show that iteration 0 was saved or loaded from a freshly generated conversion checkpoint. +If `NRL_MEGATRON_CHECKPOINT_DIR` is set, move aside the subdirectory used by the run. On first startup, the log should show that iteration 0 was saved or loaded from a freshly generated conversion checkpoint. For long runs on queues with short wall times, enable periodic checkpointing and submit dependency jobs with `afterany` so the next job can resume from the checkpoint written by the previous job. @@ -153,11 +248,14 @@ For long runs on queues with short wall times, enable periodic checkpointing and A healthy W4A16 real-rollout run should include these lines or equivalent vLLM logs: ```text -quantization=modelopt_fp4 +quantization=modelopt +Detected ModelOpt NVFP4 checkpoint Using NvFp4LinearBackend.MARLIN for NVFP4 GEMM MegatronQuantPolicyWorker[rank=0]: Packed ... groups of tensors ``` +A W4A4 run should show the same ModelOpt checkpoint/refit signals but must not select `NvFp4LinearBackend.MARLIN`; NeMo RL rejects that weight-only backend because it would leave activations unquantized. + It should not include: ```text @@ -172,10 +270,11 @@ For an initial sanity check, compare the first `Generation KL Error` with the BF | Symptom | Likely Cause | Action | |---|---|---| -| vLLM does not log `quantization=modelopt_fp4` | `policy.generation.real_quant` is not set or generation is not using vLLM | Check the YAML under `policy.generation` | +| vLLM does not log `quantization=modelopt` | `policy.generation.real_quant` is not set or generation is not using vLLM | Check the YAML under `policy.generation` | | `Using rollout logprobs` appears | The run is bypassing policy/reference logprob computation | Do not use rollout logprobs for real-quant validation | -| First-step W4A16 `Generation KL Error` is much higher than BF16 | Stale converted Megatron checkpoint or refit/export mismatch | Clear checkpoints and rerun; confirm packed tensors are streamed | -| `negative scales` warning appears | Invalid or stale NVFP4 scale tensors reached vLLM | Clear checkpoints and verify `nvfp4_a16_mlp_only.yaml` is used for both policy and generation | +| First-step W4A16 `Generation KL Error` is much higher than BF16 | Stale resume state, a stale converted Megatron checkpoint on the Megatron path, or a refit/export mismatch | Use a fresh training checkpoint directory; on Megatron also move aside the converted startup checkpoint; confirm packed tensors are streamed | +| `negative scales` warning appears | Invalid or stale NVFP4 scale tensors reached vLLM | Use a fresh checkpoint directory and verify the same supported NVFP4 recipe is used for both policy and generation | +| Nano3 first-step KL is high while dense W4A16 is healthy | Nano3-sensitive paths were quantized or the vLLM ignore set does not match the policy recipe | Use `nano3_nvfp4_weightonly.yaml` for policy and generation, and copy the explicit `policy.generation.real_quant_ignore` list from the Nano3 example recipe | | CUDA invalid argument during refit or generation | vLLM consumed malformed packed tensors or stale IPC state | Restart from a fresh job and inspect the first real-quant refit logs | ## Quantization-Aware Distillation (On-Policy QAD) @@ -227,8 +326,8 @@ Generation-specific parameters are added under `policy.generation`: | Parameter | Description | |---|---| | `quant_cfg` | Quantization config used by the vLLM generation worker. For QARL, this should normally match `policy.quant_cfg`. | -| `real_quant` | When `true`, vLLM uses ModelOpt NVFP4 real kernels and receives packed quantized weights during refit. When unset or `false`, vLLM uses fake-quantized generation. | -| `real_quant_ignore` | Optional list of vLLM parameter name patterns that should stay in native dtype during real-quant rollout. If omitted, NeMo RL uses the default ModelOpt NVFP4 ignore set for sensitive layers such as attention and output heads. | +| `real_quant` | When `true`, vLLM uses ModelOpt NVFP4 real kernels and receives packed quantized weights during refit. The effective `quant_cfg` selects W4A4 or W4A16; unsupported and mismatched formats fail during setup. When unset or `false`, vLLM uses fake-quantized generation. | +| `real_quant_ignore` | Optional list of vLLM parameter name patterns that should stay in native dtype during real-quant rollout. If omitted, NeMo RL uses the default ModelOpt NVFP4 ignore set for sensitive layers such as attention and output heads. For Nano3 hybrid MoE/Mamba W4A16 real-quant rollout, use the model-specific list shown in the Nano3 example recipe. | ## Megatron Checkpoint Directory @@ -245,16 +344,16 @@ QARL (via ModelOpt) and NeMo RL's built-in [FP8 training](../fp8.md) (via Transf - **TransformerEngine FP8** focuses on **speeding up pre-training and fine-tuning** using real quantization. It replaces linear layers with FP8-native implementations that compute directly in reduced precision for throughput gains. - **ModelOpt QARL** focuses on **recovering accuracy under quantization** using quantization-aware training. The policy forward pass uses quantized weights and, depending on the recipe, quantized activations while the backward pass uses full-precision gradients, so the model learns to be robust to quantization error. vLLM generation can run fake-quantized layers for W4A8/W4A16 recipes. - W4A16 experiments can also use real ModelOpt NVFP4 kernels. + W4A4 and W4A16 experiments can also use real ModelOpt NVFP4 kernels. ## Supported Quantization Formats -- **Weight quantization**: per-tensor, per-channel, and block-wise formats are all supported. In fake-quant rollout, weights are pre-folded on the policy (Megatron) side before transfer to vLLM. In W4A16 real-quant rollout, weights are packed as NVFP4 tensors and streamed with their scale tensors. -- **Input (activation) quantization**: only per-tensor is supported. The input quantizer amax is synced to vLLM as a per-tensor scalar. +- **Weight quantization**: per-tensor, per-channel, and block-wise formats are supported by fake-quant rollout. Real-quant rollout specifically requires block-16 E2M1 NVFP4 weights with dynamic block scaling, packed on the policy side and streamed with FP8 block scales and global scales. +- **Input (activation) quantization**: fake-quant rollout supports the existing ModelOpt recipes. W4A4 real rollout specifically requires block-16 E2M1 NVFP4 inputs and streams one calibrated global input scale per dense projection or per expert projection. -## Exporting Checkpoints +## Exporting Megatron Checkpoints -After quantization-aware training, the Megatron checkpoint contains BF16 weights alongside quantization metadata (amax values, scales). To export a trained checkpoint to a fully quantized HuggingFace format (with real low-precision weights), use the Megatron-Bridge export tool. The exported checkpoint is ready for deployment with inference engines like vLLM or TensorRT-LLM. +After quantization-aware training, a Megatron checkpoint contains BF16 weights alongside quantization metadata (amax values, scales). To export it to a fully quantized HuggingFace format (with real low-precision weights), use the Megatron-Bridge export tool. The exported checkpoint is ready for deployment with inference engines like vLLM or TensorRT-LLM. From within the NeMo RL container: @@ -281,6 +380,8 @@ uv run --extra mcore --extra modelopt \ - **Generation**: Currently only vLLM is supported for generation. - **DTensor backend**: Quantization support for the DTensor policy worker is not yet implemented. -- **Real-quant rollout**: W4A16 real rollout is supported for dense vLLM ModelOpt NVFP4 layers. -- **Input quantization**: Only per-tensor input (activation) quantization is supported. -- **Model support**: Dense Transformer, MoE (Mixture of Experts), and hybrid MoE/Mamba models are supported on the Megatron policy + vLLM generation path when Megatron-Bridge and ModelOpt support the model architecture and quantization recipe. MoE/Mamba support is currently covered by smoke-tested example configs rather than broad convergence guarantees. +- **Real-quant rollout**: W4A4 and W4A16 are supported for dense and fused-MoE vLLM ModelOpt NVFP4 layers exported from the Megatron policy path. Fused MoE currently requires all experts local to each vLLM rank. Hybrid MoE/Mamba recipes should keep unsupported or sensitive non-MLP paths in BF16 via `real_quant_ignore`. +- **Router Replay (R3)**: R3 is supported on the Megatron policy path. +- **Input quantization**: W4A4 real rollout supports ModelOpt's block-16 E2M1 input format with a global scale per projection; other activation formats remain fake-quant only. +- **Static NVFP4 weights**: The real-quant exporter rejects static weight quantizers because their calibrated per-block amax state cannot be reconstructed after distributed TP/EP gathering. Use the dynamic block-scaling recipes shown above. +- **Model support**: Dense Transformer, MoE (Mixture of Experts), and hybrid MoE/Mamba models are supported on the Megatron policy + vLLM generation path when Megatron-Bridge and ModelOpt support the model architecture and quantization recipe. diff --git a/docs/index.md b/docs/index.md index 4e86e4e47b9..aa3f85e76ce 100644 --- a/docs/index.md +++ b/docs/index.md @@ -196,7 +196,7 @@ Optimize large language models with FP8 quantization for faster training and inf :link-type: doc Run quantization-aware GRPO and distillation using NVIDIA ModelOpt. -Includes NVFP4 W4A16 real rollout. +Includes NVFP4 W4A4 and W4A16 real rollout. ::: :::{grid-item-card} {octicon}`container` Docker Containers @@ -330,6 +330,7 @@ design-docs/training-backends.md design-docs/sequence-packing-and-dynamic-batching.md design-docs/env-vars.md design-docs/nemo-gym-integration.md +design-docs/modelopt-real-quant-architecture.md ``` ```{toctree} diff --git a/examples/configs/recipes/llm/grpo-nanov3-30ba3b-4n4g-megatron-qa-nvfp4-w4a16-real.yaml b/examples/configs/recipes/llm/grpo-nanov3-30ba3b-4n4g-megatron-qa-nvfp4-w4a16-real.yaml new file mode 100644 index 00000000000..d910cda357b --- /dev/null +++ b/examples/configs/recipes/llm/grpo-nanov3-30ba3b-4n4g-megatron-qa-nvfp4-w4a16-real.yaml @@ -0,0 +1,73 @@ +defaults: ../../../../examples/modelopt/qa_grpo_nano3_megatron.yaml +grpo: + max_num_steps: 1 + val_period: 0 +checkpointing: + checkpoint_dir: results/grpo-nanov3-30ba3b-4n4g-megatron-qa-nvfp4-w4a16-real +policy: + model_name: nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16 + tokenizer: + name: ${policy.model_name} + generation: + real_quant: true + # Nano3-specific paths that stay in BF16 during real-quant rollout. + real_quant_ignore: + - lm_head + - '*output_layer*' + - '*mlp.gate' + - '*router*' + - '*block_sparse_moe.gate*' + - '*self_attention*' + - '*self_attn*' + - '*proj_out.*' + - '*.gate.*' + - '*mlp.shared_expert_gate.*' + - '*linear_attn.conv1d*' + - '*mixer.conv1d*' + - '*.mixer.in_proj*' + - '*.mixer.out_proj*' + - '*.shared_expert.*' + - '*.shared_experts.*' + - '*.norm.*' + - '*.q_proj*' + - '*.k_proj*' + - '*.v_proj*' + - '*.o_proj*' + - '*.qkv_proj*' + - '*.linear_proj*' + - '*.linear_qkv*' + - '*.layers.4.*' + - '*.layers.11.*' + - '*.layers.18.*' + - '*.layers.25.*' + - '*.layers.32.*' + - '*.layers.41.*' + vllm_cfg: + # Keep vLLM conservative because the full-parameter Nano3 policy and + # reference model are colocated on these GPUs. With TP=4, this leaves + # refit headroom while providing ample KV cache for the 2K smoke workload. + gpu_memory_utilization: 0.35 + enable_prefix_caching: false + vllm_kwargs: + tokenizer: ${policy.tokenizer.name} +data: + max_input_seq_length: 1024 + train: + dataset_name: DAPOMath17K + default: + prompt_file: null +env: + dapo: + num_workers: 2 + math: + num_workers: 2 + math_verify_impl: dapo_math_verify +logger: + log_dir: logs/grpo-nanov3-30ba3b-4n4g-megatron-qa-nvfp4-w4a16-real + wandb: + name: grpo-nanov3-30ba3b-4n4g-megatron-qa-nvfp4-w4a16-real + tensorboard: + log_dir: tb_logs-grpo-nanov3-30ba3b-4n4g-megatron-qa-nvfp4-w4a16-real +cluster: + gpus_per_node: 4 + num_nodes: 4 diff --git a/examples/configs/recipes/llm/grpo-nemotron3-super-120BA12B-16n4g-megatron-qa-bf16-300step.yaml b/examples/configs/recipes/llm/grpo-nemotron3-super-120BA12B-16n4g-megatron-qa-bf16-300step.yaml new file mode 100644 index 00000000000..88646521124 --- /dev/null +++ b/examples/configs/recipes/llm/grpo-nemotron3-super-120BA12B-16n4g-megatron-qa-bf16-300step.yaml @@ -0,0 +1,77 @@ +defaults: grpo-nemotron3-super-120BA12B-16n8g-megatron.yaml +cluster: + segment_size: 16 + gpus_per_node: 4 +grpo: + max_num_steps: 300 + max_num_epochs: 1000000 + use_leave_one_out_baseline: false + use_dynamic_sampling: true + batch_multiplier: 2 + val_at_start: true + val_at_end: true +loss_fn: + reference_policy_kl_penalty: 0.0 + ratio_clip_max: 0.28 + ratio_clip_c: 10 + use_importance_sampling_correction: true + truncated_importance_sampling_type: tis + truncated_importance_sampling_ratio: 2.0 +checkpointing: + checkpoint_dir: results/super120b-bf16-300step +policy: + tokenizer: + name: ${policy.model_name} + chat_template_kwargs: + enable_thinking: false + logprob_chunk_size: 1024 + max_total_sequence_length: 65536 + refit_buffer_size_gb: 8 + sequence_packing: + train_mb_tokens: 49152 + quant_cfg: null + megatron_cfg: + tensor_model_parallel_size: 2 + context_parallel_size: 2 + mtp_num_layers: 2 + mtp_detach_heads: false + optimizer: + lr: 1.0e-06 + min_lr: 1.0e-06 + weight_decay: 0.1 + adam_beta2: 0.95 + clip_grad: 1.0 + scheduler: + lr_warmup_iters: 0 + lr_warmup_init: 1.0e-06 + start_weight_decay: 0.1 + end_weight_decay: 0.1 + generation: + top_p: 0.95 + max_new_tokens: 32768 + quant_cfg: null + real_quant: false + vllm_cfg: + tensor_parallel_size: 4 + enable_prefix_caching: false + enable_vllm_metrics_logger: false + env_vars: + VLLM_USE_RAY_V2_EXECUTOR_BACKEND: 1 + vllm_kwargs: + # Equal bytes keep KV capacity comparable across different model footprints. + kv_cache_memory_bytes: 12884901888 + mamba_ssm_cache_dtype: float32 + max_num_seqs: 2 + max_num_batched_tokens: 131072 + moe_backend: triton +data: + max_input_seq_length: 1024 + default: + prompt_file: null +logger: + log_dir: logs/super120b-bf16-300step + wandb_enabled: false + wandb: + name: super120b-bf16-300step + tensorboard: + log_dir: tb_logs/super120b-bf16-300step diff --git a/examples/configs/recipes/llm/grpo-nemotron3-super-120BA12B-16n4g-megatron-qa-nvfp4-w4a16-real-300step.yaml b/examples/configs/recipes/llm/grpo-nemotron3-super-120BA12B-16n4g-megatron-qa-nvfp4-w4a16-real-300step.yaml new file mode 100644 index 00000000000..67fd6a66d27 --- /dev/null +++ b/examples/configs/recipes/llm/grpo-nemotron3-super-120BA12B-16n4g-megatron-qa-nvfp4-w4a16-real-300step.yaml @@ -0,0 +1,43 @@ +defaults: grpo-nemotron3-super-120BA12B-16n4g-megatron-qa-bf16-300step.yaml + +checkpointing: + checkpoint_dir: results/super120b-w4a16-real-300step + +policy: + disable_modelopt_layer_spec: true + quant_cfg: examples/modelopt/quant_configs/nvfp4_experts_weightonly.yaml + generation: + quant_cfg: examples/modelopt/quant_configs/nvfp4_experts_weightonly.yaml + real_quant: true + real_quant_ignore: + - lm_head + - '*output_layer*' + - '*mtp*' + - '*router*' + - '*mixer.gate*' + - '*self_attention*' + - '*self_attn*' + - '*.q_proj*' + - '*.k_proj*' + - '*.v_proj*' + - '*.o_proj*' + - '*.qkv_proj*' + - '*.linear_qkv*' + - '*.linear_proj*' + - '*.mixer.in_proj*' + - '*.mixer.out_proj*' + - '*fc1_latent_proj*' + - '*fc2_latent_proj*' + - '*.shared_expert.*' + - '*.shared_experts.*' + vllm_kwargs: + moe_backend: auto + attention_config: + use_trtllm_attention: false + +logger: + log_dir: logs/super120b-w4a16-real-300step + wandb: + name: super120b-w4a16-real-300step + tensorboard: + log_dir: tb_logs/super120b-w4a16-real-300step diff --git a/examples/configs/recipes/llm/grpo-nemotron3-super-120BA12B-16n4g-megatron-qa-nvfp4-w4a4-real-300step.yaml b/examples/configs/recipes/llm/grpo-nemotron3-super-120BA12B-16n4g-megatron-qa-nvfp4-w4a4-real-300step.yaml new file mode 100644 index 00000000000..9dec4582785 --- /dev/null +++ b/examples/configs/recipes/llm/grpo-nemotron3-super-120BA12B-16n4g-megatron-qa-nvfp4-w4a4-real-300step.yaml @@ -0,0 +1,45 @@ +defaults: grpo-nemotron3-super-120BA12B-16n4g-megatron-qa-bf16-300step.yaml + +checkpointing: + checkpoint_dir: results/super120b-w4a4-real-300step + +policy: + disable_modelopt_layer_spec: true + quant_cfg: examples/modelopt/quant_configs/nvfp4_experts.yaml + quant_calib_data: cnn_dailymail + quant_calib_size: 16 + quant_batch_size: 1 + quant_sequence_length: 1024 + generation: + quant_cfg: examples/modelopt/quant_configs/nvfp4_experts.yaml + real_quant: true + real_quant_ignore: + - lm_head + - '*output_layer*' + - '*mtp*' + - '*router*' + - '*mixer.gate*' + - '*self_attention*' + - '*self_attn*' + - '*.q_proj*' + - '*.k_proj*' + - '*.v_proj*' + - '*.o_proj*' + - '*.qkv_proj*' + - '*.linear_qkv*' + - '*.linear_proj*' + - '*.mixer.in_proj*' + - '*.mixer.out_proj*' + - '*fc1_latent_proj*' + - '*fc2_latent_proj*' + - '*.shared_expert.*' + - '*.shared_experts.*' + vllm_kwargs: + moe_backend: auto + +logger: + log_dir: logs/super120b-w4a4-real-300step + wandb: + name: super120b-w4a4-real-300step + tensorboard: + log_dir: tb_logs/super120b-w4a4-real-300step diff --git a/examples/configs/recipes/llm/grpo-qwen3-30ba3b-4n4g-megatron-qa-nvfp4-w4a4-real.yaml b/examples/configs/recipes/llm/grpo-qwen3-30ba3b-4n4g-megatron-qa-nvfp4-w4a4-real.yaml new file mode 100644 index 00000000000..d324a5a47d4 --- /dev/null +++ b/examples/configs/recipes/llm/grpo-qwen3-30ba3b-4n4g-megatron-qa-nvfp4-w4a4-real.yaml @@ -0,0 +1,33 @@ +defaults: ../../../../examples/modelopt/qa_grpo_qwen3_30ba3b_megatron.yaml +cluster: + segment_size: null +grpo: + max_num_steps: 300 + val_at_end: true +checkpointing: + enabled: true + checkpoint_dir: results/grpo-qwen3-30ba3b-4n4g-megatron-qa-nvfp4-w4a4-real + save_period: 30 + keep_top_k: 1 + checkpoint_must_save_by: 00:03:45:00 +policy: + quant_cfg: examples/modelopt/quant_configs/nvfp4_experts.yaml + generation: + quant_cfg: examples/modelopt/quant_configs/nvfp4_experts.yaml + real_quant: true + real_quant_ignore: + - lm_head + - '*output_layer*' + - '*mlp.gate*' + - '*router*' + - '*self_attention*' + - '*self_attn*' + vllm_kwargs: + moe_backend: auto +logger: + log_dir: logs/grpo-qwen3-30ba3b-4n4g-megatron-qa-nvfp4-w4a4-real + wandb: + project: nemo-rl-qarl + name: grpo-qwen3-30ba3b-4n4g-megatron-qa-nvfp4-w4a4-real + tensorboard: + log_dir: tb_logs-grpo-qwen3-30ba3b-4n4g-megatron-qa-nvfp4-w4a4-real diff --git a/examples/modelopt/quant_configs/nvfp4_experts.yaml b/examples/modelopt/quant_configs/nvfp4_experts.yaml new file mode 100644 index 00000000000..412274f60ae --- /dev/null +++ b/examples/modelopt/quant_configs/nvfp4_experts.yaml @@ -0,0 +1,30 @@ +# Routed-expert NVFP4 W4A4 recipe for QARL. +# +# Routed-expert weights and layer-input activations use block-16 E2M1 NVFP4. +# Every non-routed-expert path remains in its native dtype. + +metadata: + recipe_type: ptq + description: Routed-expert NVFP4 W4A4. + +quantize: + algorithm: max + quant_cfg: + - quantizer_name: '*' + enable: false + - quantizer_name: '*.experts.*weight_quantizer' + enable: true + cfg: + block_sizes: + -1: 16 + type: dynamic + scale_bits: e4m3 + num_bits: e2m1 + - quantizer_name: '*.experts.*input_quantizer' + enable: true + cfg: + block_sizes: + -1: 16 + type: dynamic + scale_bits: e4m3 + num_bits: e2m1 diff --git a/examples/modelopt/quant_configs/nvfp4_experts_weightonly.yaml b/examples/modelopt/quant_configs/nvfp4_experts_weightonly.yaml new file mode 100644 index 00000000000..54093ee7f88 --- /dev/null +++ b/examples/modelopt/quant_configs/nvfp4_experts_weightonly.yaml @@ -0,0 +1,23 @@ +# Routed-expert NVFP4 W4A16 recipe for QARL. +# +# This is the weight-only counterpart of ``nvfp4_experts.yaml``: +# routed-expert weights use block-16 E2M1 NVFP4, while activations and every +# non-routed-expert path stay in their native dtype. + +metadata: + recipe_type: ptq + description: Routed-expert NVFP4 W4A16. + +quantize: + algorithm: max + quant_cfg: + - quantizer_name: '*' + enable: false + - quantizer_name: '*.experts.*weight_quantizer' + enable: true + cfg: + block_sizes: + -1: 16 + type: dynamic + scale_bits: e4m3 + num_bits: e2m1 diff --git a/nemo_rl/modelopt/models/generation/vllm_modelopt.py b/nemo_rl/modelopt/models/generation/vllm_modelopt.py new file mode 100644 index 00000000000..b344055d7d5 --- /dev/null +++ b/nemo_rl/modelopt/models/generation/vllm_modelopt.py @@ -0,0 +1,408 @@ +# Copyright 2025 Bytedance Ltd. and/or its affiliates +# 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. + +"""Narrow vLLM extensions for ModelOpt NVFP4 rollout checkpoints. + +vLLM owns checkpoint-layout restoration, layerwise post-load processing, +CUDA-graph-stable tensor placement, and KV-cache scale reload. This module +only supplies the vLLM 0.20 gaps needed here: ModelOpt W4A16 NVFP4 methods, +rank-local Marlin padding, per-projection ModelOpt MoE input-scale loading, +materialization of FlashInfer's global-scale views, and retention of +method-owned MoE kernel references across layerwise reload. +""" + +import copy +from types import MethodType +from typing import Any + +import torch +from torch.nn import Parameter + +NEMO_MODELOPT_W4A4 = "nemo_modelopt_nvfp4" +NEMO_MODELOPT_W4A16 = "nemo_modelopt_w4a16_nvfp4" + +_W4A4_ALGO = "NVFP4" +_W4A16_ALGO = "W4A16_NVFP4" +_registered = False + + +def quantization_method_for_mode(mode: str) -> str: + """Return the registered vLLM quantization method for a rollout mode.""" + if mode == "w4a4": + return NEMO_MODELOPT_W4A4 + if mode == "w4a16": + return NEMO_MODELOPT_W4A16 + raise ValueError(f"Unsupported ModelOpt NVFP4 rollout mode: {mode!r}") + + +def _load_modelopt_moe_input_scale( + moe_layer: Any, + param: torch.nn.Parameter, + loaded_weight: torch.Tensor, + weight_name: str, + shard_id: str, + expert_id: int, + return_success: bool = False, +) -> bool | None: + """Load a ModelOpt input scale without losing the gate/up shard. + + Replaces the ``input_scale`` branch of vLLM v0.20.0's + ``FusedMoE.weight_loader``, whose ``_load_single_value`` writes + ``param.data[expert_id]`` and drops the gate/up (w1/w3) shard index: + https://github.com/vllm-project/vllm/blob/v0.20.0/vllm/model_executor/layers/fused_moe/layer.py#L1025-L1031 + Delete once upstream loads per-projection ModelOpt MoE input scales + correctly. + """ + del weight_name + global_expert_id = expert_id + local_expert_id = moe_layer._map_global_expert_id_to_local_expert_id( + global_expert_id + ) + use_global_scale = bool(getattr(moe_layer.quant_method, "use_global_sf", False)) + if local_expert_id == -1 and not use_global_scale: + return False if return_success else None + + target_expert_id = global_expert_id if use_global_scale else local_expert_id + if shard_id == "w2": + target = param.data[target_expert_id] + elif shard_id in ("w1", "w3"): + shard_index = 0 if shard_id == "w1" else min(1, param.shape[-1] - 1) + target = param.data[target_expert_id, shard_index] + else: + raise ValueError(f"Unexpected ModelOpt MoE shard: {shard_id!r}") + + source = loaded_weight.to(device=target.device, dtype=target.dtype) + target.copy_(source.reshape_as(target)) + return True if return_success else None + + +def _normalized_w4a16_config(config: dict[str, Any]) -> dict[str, Any]: + normalized = copy.deepcopy(config) + quantization = normalized.get("quantization") + target = quantization if isinstance(quantization, dict) else normalized + if str(target.get("quant_algo", "")).upper() != _W4A16_ALGO: + raise ValueError(f"{NEMO_MODELOPT_W4A16} requires quant_algo={_W4A16_ALGO!r}") + # vLLM 0.20 validates known ModelOpt algorithms before dispatching to a + # custom subclass. Normalize only for its parser; class identity selects + # the W4A16 methods below. + target["quant_algo"] = _W4A4_ALGO + return normalized + + +def _canonicalize_nvfp4_scale_(scale: torch.Tensor) -> None: + """Remove the E4M3 sign bit before Marlin's unsigned scale conversion.""" + with torch.no_grad(): + scale.copy_(scale.to(torch.float32).abs().to(scale.dtype)) + + +def _pad_nvfp4_moe_for_marlin( + w13: torch.Tensor, + w13_scale: torch.Tensor, + w2: torch.Tensor, + w2_scale: torch.Tensor, + *, + is_act_and_mul: bool, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, int]: + """Apply rank-local post-load padding required by the Marlin MoE kernel.""" + num_experts = w13.shape[0] + num_shards = 2 if is_act_and_mul else 1 + intermediate_size = w13.shape[1] // num_shards + hidden_size = w13.shape[2] * 2 + if hidden_size % 128 == 0: + tile_size = 64 + elif hidden_size % 64 == 0: + tile_size = 128 + else: + raise ValueError( + f"W4A16 Marlin MoE requires hidden_size divisible by 64, got {hidden_size}" + ) + padded_size = (intermediate_size + tile_size - 1) // tile_size * tile_size + if padded_size == intermediate_size: + return w13, w13_scale, w2, w2_scale, intermediate_size + + def pad_w13(tensor: torch.Tensor) -> torch.Tensor: + tensor = tensor.view( + num_experts, + num_shards, + intermediate_size, + tensor.shape[-1], + ) + tensor = torch.nn.functional.pad( + tensor, + (0, 0, 0, padded_size - intermediate_size), + ) + return tensor.reshape(num_experts, num_shards * padded_size, -1) + + w13 = pad_w13(w13) + w13_scale = pad_w13(w13_scale) + w2 = torch.nn.functional.pad(w2, (0, (padded_size - intermediate_size) // 2)) + w2_scale = torch.nn.functional.pad( + w2_scale, + (0, (padded_size - intermediate_size) // 16), + ) + return w13, w13_scale, w2, w2_scale, padded_size + + +def register_nemo_modelopt_nvfp4() -> None: + """Register NeMo's two ModelOpt NVFP4 configs through vLLM's public API.""" + global _registered + if _registered: + return + + from vllm.model_executor.kernels.linear import ( + MarlinNvFp4LinearKernel, + NvFp4LinearLayerConfig, + ) + from vllm.model_executor.layers.fused_moe.fused_moe_method_base import ( + FusedMoEMethodBase, + ) + from vllm.model_executor.layers.fused_moe.oracle.nvfp4 import ( + NvFp4MoeBackend, + is_global_sf_supported_for_nvfp4_backend, + select_nvfp4_moe_backend, + ) + from vllm.model_executor.layers.linear import ( + register_weight_loader_v2_supported_method, + ) + from vllm.model_executor.layers.quantization import register_quantization_config + from vllm.model_executor.layers.quantization.modelopt import ( + ModelOptNvFp4Config, + ModelOptNvFp4FusedMoE, + ModelOptNvFp4LinearMethod, + ) + from vllm.model_executor.layers.quantization.utils.quant_utils import kNvfp4Static + from vllm.model_executor.utils import replace_parameter + + class NemoModelOptNvFp4FusedMoE(ModelOptNvFp4FusedMoE): + """Native W4A4 MoE plus the vLLM 0.20 input-scale loader fix.""" + + moe_kernel: Any + moe_quant_config: Any + + def create_weights(self, layer: Any, *args: Any, **kwargs: Any) -> None: + super().create_weights(layer, *args, **kwargs) + # Bind to the layer so vLLM's reload metadata sanitizer can remove + # and restore this reference. A partial would retain the whole model. + loader = MethodType(_load_modelopt_moe_input_scale, layer) + layer.w13_input_scale.weight_loader = loader + layer.w2_input_scale.weight_loader = loader + + def process_weights_after_loading(self, layer: Any) -> None: + reload_kernel = self.moe_kernel + reload_quant_config = self.moe_quant_config + super().process_weights_after_loading(layer) + processed_quant_config = self.moe_quant_config + if reload_kernel is None: + # FlashInfer NVFP4 backends in vLLM 0.20 return global activation + # scales as stride-zero expanded views. Materialize them before + # native reload records these Parameters as future copy targets. + layer.w13_input_scale.data = layer.w13_input_scale.data.contiguous() + layer.w2_input_scale.data = layer.w2_input_scale.data.contiguous() + return + + # Native reload copies registered layer tensors into their original + # CUDA-graph-stable storage. The two reciprocal activation scales are + # stored only in FusedMoEQuantConfig, so refresh them explicitly while + # preserving the original config tensor addresses and kernel object. + if reload_quant_config is None or processed_quant_config is None: + raise RuntimeError("W4A4 MoE reload is missing its quant config") + reload_a1_gscale = reload_quant_config.a1_gscale + processed_a1_gscale = processed_quant_config.a1_gscale + reload_a2_gscale = reload_quant_config.a2_gscale + processed_a2_gscale = processed_quant_config.a2_gscale + if ( + reload_a1_gscale is None + or processed_a1_gscale is None + or reload_a2_gscale is None + or processed_a2_gscale is None + ): + raise RuntimeError("W4A4 MoE reload is missing activation scales") + if ( + reload_a1_gscale.shape != processed_a1_gscale.shape + or reload_a2_gscale.shape != processed_a2_gscale.shape + ): + raise RuntimeError("W4A4 MoE activation-scale shape changed on reload") + reload_a1_gscale.copy_(processed_a1_gscale) + reload_a2_gscale.copy_(processed_a2_gscale) + self.moe_kernel = reload_kernel + self.moe_quant_config = reload_quant_config + + class NemoModelOptNvFp4Config(ModelOptNvFp4Config): + FusedMoEMethodCls = NemoModelOptNvFp4FusedMoE + + def get_name(self) -> str: + return NEMO_MODELOPT_W4A4 + + @classmethod + def override_quantization_method( + cls, + hf_quant_cfg: dict[str, Any], + user_quant: str | None, + hf_config: Any = None, + ) -> str | None: + del hf_config + if ( + user_quant == NEMO_MODELOPT_W4A4 + and cls._extract_modelopt_quant_algo(hf_quant_cfg) == _W4A4_ALGO + ): + return NEMO_MODELOPT_W4A4 + return None + + @register_weight_loader_v2_supported_method + class NemoModelOptW4A16LinearMethod(ModelOptNvFp4LinearMethod): + """ModelOpt NVFP4 weights with BF16/FP16 Marlin activations.""" + + def __init__(self, quant_config: object) -> None: + self.quant_config = quant_config + self.marlin_input_dtype = None + self.kernel = MarlinNvFp4LinearKernel(NvFp4LinearLayerConfig()) + + def create_weights(self, layer: Any, *args: Any, **kwargs: Any) -> None: + super().create_weights(layer, *args, **kwargs) + del layer.input_scale + + # Adapted from vLLM v0.20.0 ModelOptNvFp4LinearMethod + # .process_weights_after_loading/.apply with input-scale/alpha handling + # removed for weight-only W4A16: + # https://github.com/vllm-project/vllm/blob/v0.20.0/vllm/model_executor/layers/quantization/modelopt.py#L1169-L1208 + # Re-sync on vLLM bumps; delete if upstream gains a native W4A16 + # NVFP4 method. + def process_weights_after_loading(self, layer: Any) -> None: + layer.weight_global_scale = Parameter( + layer.weight_scale_2.max().to(torch.float32), + requires_grad=False, + ) + del layer.weight_scale_2 + _canonicalize_nvfp4_scale_(layer.weight_scale) + self.kernel.process_weights_after_loading(layer) + + def apply( + self, + layer: Any, + x: torch.Tensor, + bias: torch.Tensor | None = None, + ) -> torch.Tensor: + return self.kernel.apply_weights(layer=layer, x=x, bias=bias) + + class NemoModelOptW4A16FusedMoE(ModelOptNvFp4FusedMoE): + """ModelOpt W4A16 MoE using vLLM's NVFP4 Marlin implementation.""" + + moe_kernel: Any + moe_quant_config: Any + + def __init__(self, quant_config: object, moe_config: object) -> None: + # Duplicates vLLM v0.20.0 ModelOptNvFp4FusedMoE.__init__ except + # activation_key=None (weight-only); the base hard-wires + # kNvfp4Dynamic and offers no hook: + # https://github.com/vllm-project/vllm/blob/v0.20.0/vllm/model_executor/layers/quantization/modelopt.py#L1218-L1234 + # Intentionally calls FusedMoEMethodBase.__init__ to skip the + # parent __init__; do not replace it with super().__init__(). + # Re-sync on vLLM bumps. + FusedMoEMethodBase.__init__(self, moe_config) + self.quant_config = quant_config + self.nvfp4_backend, self.experts_cls = select_nvfp4_moe_backend( + config=self.moe, + weight_key=kNvfp4Static, + activation_key=None, + ) + self.use_global_sf = is_global_sf_supported_for_nvfp4_backend( + self.nvfp4_backend + ) + + def create_weights( + self, + layer: Any, + num_experts: int, + hidden_size: int, + intermediate_size_per_partition: int, + params_dtype: torch.dtype, + **extra_weight_attrs: Any, + ) -> None: + super().create_weights( + layer, + num_experts, + hidden_size, + intermediate_size_per_partition, + params_dtype, + **extra_weight_attrs, + ) + del layer.w13_input_scale + del layer.w2_input_scale + + def process_weights_after_loading(self, layer: Any) -> None: + reload_kernel = self.moe_kernel + reload_quant_config = self.moe_quant_config + original_intermediate_size = ( + layer.moe_config.intermediate_size_per_partition + ) + if self.nvfp4_backend == NvFp4MoeBackend.MARLIN: + w13, w13_scale, w2, w2_scale, padded_size = _pad_nvfp4_moe_for_marlin( + layer.w13_weight, + layer.w13_weight_scale, + layer.w2_weight, + layer.w2_weight_scale, + is_act_and_mul=self.moe.is_act_and_mul, + ) + replace_parameter(layer, "w13_weight", w13) + replace_parameter(layer, "w13_weight_scale", w13_scale) + replace_parameter(layer, "w2_weight", w2) + replace_parameter(layer, "w2_weight_scale", w2_scale) + _canonicalize_nvfp4_scale_(layer.w13_weight_scale) + _canonicalize_nvfp4_scale_(layer.w2_weight_scale) + layer.moe_config.intermediate_size_per_partition = padded_size + # W4A16 checkpoint metadata deliberately omits activation scales so + # layerwise reload never waits for tensors that do not exist. The + # native Marlin converter accepts None and removes these attributes. + layer.w13_input_scale = None + layer.w2_input_scale = None + try: + super().process_weights_after_loading(layer) + finally: + layer.moe_config.intermediate_size_per_partition = ( + original_intermediate_size + ) + if reload_kernel is not None: + self.moe_kernel = reload_kernel + self.moe_quant_config = reload_quant_config + + class NemoModelOptW4A16Config(ModelOptNvFp4Config): + LinearMethodCls = NemoModelOptW4A16LinearMethod + FusedMoEMethodCls = NemoModelOptW4A16FusedMoE + + def get_name(self) -> str: + return NEMO_MODELOPT_W4A16 + + @classmethod + def override_quantization_method( + cls, + hf_quant_cfg: dict[str, Any], + user_quant: str | None, + hf_config: Any = None, + ) -> str | None: + del hf_config + if ( + user_quant == NEMO_MODELOPT_W4A16 + and cls._extract_modelopt_quant_algo(hf_quant_cfg) == _W4A16_ALGO + ): + return NEMO_MODELOPT_W4A16 + return None + + @classmethod + def from_config(cls, config: dict[str, Any]) -> Any: + return super().from_config(_normalized_w4a16_config(config)) + + register_quantization_config(NEMO_MODELOPT_W4A4)(NemoModelOptNvFp4Config) + register_quantization_config(NEMO_MODELOPT_W4A16)(NemoModelOptW4A16Config) + _registered = True diff --git a/nemo_rl/modelopt/models/generation/vllm_modelopt_patch.py b/nemo_rl/modelopt/models/generation/vllm_modelopt_patch.py deleted file mode 100644 index 95e0e28aad7..00000000000 --- a/nemo_rl/modelopt/models/generation/vllm_modelopt_patch.py +++ /dev/null @@ -1,284 +0,0 @@ -# Copyright 2025 Bytedance Ltd. and/or its affiliates -# 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. - -"""vLLM ModelOpt NVFP4 patches for dense rollout weight reloads.""" - -import torch -from torch.nn import Parameter - -_DENSE_HF_PARAMS = ("weight", "weight_scale", "weight_scale_2") -_MODELOPT_W4A16_QUANT_MODES = frozenset({"w4a16_nvfp4", "nvfp4_w4a16"}) -_MODELOPT_W4A16_ATTR = "_nrl_weight_only_w4a16" -_ORIGINAL_NVFP4_CONFIG_FROM_CONFIG_ATTR = "_nrl_original_from_config" -_ORIGINAL_LINEAR_APPLY_ATTR = "_nrl_original_apply" - - -def _unwrap_vllm_model(model: torch.nn.Module) -> torch.nn.Module: - return model.model if hasattr(model, "model") else model - - -def _canonicalize_nvfp4_weight_scale(layer: torch.nn.Module) -> None: - weight_scale = layer.weight_scale - scale = weight_scale.data.to(torch.float32).abs().to(weight_scale.dtype) - weight_scale.data.copy_(scale) - - -def _requests_w4a16_modelopt_config(config: dict) -> bool: - quant_mode = config.get("quant_mode") - if ( - isinstance(quant_mode, str) - and quant_mode.lower() in _MODELOPT_W4A16_QUANT_MODES - ): - return True - if config.get("weight_only") is True: - return True - - nested = config.get("quantization") - return isinstance(nested, dict) and _requests_w4a16_modelopt_config(nested) - - -def _is_w4a16_modelopt_quant_config(quant_config) -> bool: - return bool(getattr(quant_config, _MODELOPT_W4A16_ATTR, False)) - - -def _modelopt_nvfp4_config_from_config(cls, *args, **kwargs): - original_from_config = getattr(cls, _ORIGINAL_NVFP4_CONFIG_FROM_CONFIG_ATTR) - quant_config = original_from_config(*args, **kwargs) - - original_config = kwargs.get("original_config") - if isinstance(original_config, dict) and _requests_w4a16_modelopt_config( - original_config - ): - setattr(quant_config, _MODELOPT_W4A16_ATTR, True) - - return quant_config - - -def _convert_nvfp4_linear_kernel_format(quant_method, layer: torch.nn.Module) -> None: - kernel = getattr(quant_method, "kernel", None) - if kernel is not None: - kernel.process_weights_after_loading(layer) - return - - from vllm.model_executor.layers.quantization.utils.nvfp4_utils import ( - convert_to_nvfp4_linear_kernel_format, - ) - - convert_to_nvfp4_linear_kernel_format(quant_method.backend, layer) - - -def _convert_w4a16_linear_kernel_format(layer: torch.nn.Module) -> None: - from vllm.model_executor.layers.quantization.utils.marlin_utils_fp4 import ( - prepare_fp4_layer_for_marlin, - ) - - prepare_fp4_layer_for_marlin(layer) - - -def _capture_modelopt_dense_param_reload_meta(layer: torch.nn.Module) -> None: - if not hasattr(layer, "_nrl_modelopt_param_meta"): - layer._nrl_modelopt_param_meta = {} - layer._nrl_modelopt_weight_loaders = {} - elif not hasattr(layer, "_nrl_modelopt_weight_loaders"): - layer._nrl_modelopt_weight_loaders = {} - - for param_name in _DENSE_HF_PARAMS: - if param_name in layer._nrl_modelopt_param_meta: - continue - param = getattr(layer, param_name) - meta = { - "shape": tuple(param.shape), - "dtype": param.dtype, - "device": str(param.device), - "param_class": type(param), - } - if hasattr(param, "_input_dim"): - meta["input_dim"] = param._input_dim - if hasattr(param, "_output_dim"): - meta["output_dim"] = param._output_dim - layer._nrl_modelopt_param_meta[param_name] = meta - if hasattr(param, "weight_loader"): - layer._nrl_modelopt_weight_loaders[param_name] = param.weight_loader - - -def _modelopt_dense_process_w4a16_weights(self, layer: torch.nn.Module) -> None: - """Convert dense ModelOpt NVFP4 W4A16 weights for Marlin weight-only GEMM.""" - _capture_modelopt_dense_param_reload_meta(layer) - - weight_global_scale = layer.weight_scale_2.max().to(torch.float32) - layer.weight_global_scale = Parameter(weight_global_scale, requires_grad=False) - delattr(layer, "weight_scale_2") - - for attr in ( - "input_scale", - "input_global_scale", - "alpha", - "input_global_scale_inv", - ): - if hasattr(layer, attr): - delattr(layer, attr) - - _canonicalize_nvfp4_weight_scale(layer) - _convert_w4a16_linear_kernel_format(layer) - - -def _modelopt_dense_process_weights(self, layer: torch.nn.Module) -> None: - """Convert dense ModelOpt NVFP4 weights after initial load or refit.""" - if _is_w4a16_modelopt_quant_config(getattr(self, "quant_config", None)): - _modelopt_dense_process_w4a16_weights(self, layer) - return - - _capture_modelopt_dense_param_reload_meta(layer) - - input_global_scale = torch.ones( - (), - dtype=torch.float32, - device=layer.weight.device, - ) - layer.input_global_scale = Parameter(input_global_scale, requires_grad=False) - - weight_global_scale = layer.weight_scale_2.max().to(torch.float32) - layer.weight_global_scale = Parameter(weight_global_scale, requires_grad=False) - delattr(layer, "weight_scale_2") - - layer.alpha = Parameter( - layer.input_global_scale * layer.weight_global_scale, - requires_grad=False, - ) - layer.input_global_scale_inv = Parameter( - (1.0 / layer.input_global_scale).to(torch.float32), - requires_grad=False, - ) - - _canonicalize_nvfp4_weight_scale(layer) - _convert_nvfp4_linear_kernel_format(self, layer) - - -def _modelopt_dense_apply( - self, - layer: torch.nn.Module, - x: torch.Tensor, - bias: torch.Tensor | None = None, -) -> torch.Tensor: - if _is_w4a16_modelopt_quant_config(getattr(self, "quant_config", None)): - from vllm.model_executor.layers.quantization.utils.marlin_utils_fp4 import ( - apply_fp4_marlin_linear, - ) - - return apply_fp4_marlin_linear( - input=x, - weight=layer.weight, - weight_scale=layer.weight_scale, - weight_global_scale=layer.weight_global_scale, - workspace=layer.workspace, - size_n=layer.output_size_per_partition, - size_k=layer.input_size_per_partition, - bias=bias, - ) - - original_apply = getattr(type(self), _ORIGINAL_LINEAR_APPLY_ATTR, None) - if original_apply is not None: - return original_apply(self, layer, x, bias) - - return self.kernel.apply_weights(layer=layer, x=x, bias=bias) - - -def prepare_modelopt_for_weight_reload(model, device=None) -> None: - """Prepare a dense ModelOpt-vLLM model for one weight reload cycle.""" - inner_model = _unwrap_vllm_model(model) - for module in inner_model.modules(): - layer_meta = getattr(module, "_nrl_modelopt_param_meta", None) - if layer_meta is None: - continue - for param_name, meta in layer_meta.items(): - param = getattr(module, param_name, None) - weight_loader = module._nrl_modelopt_weight_loaders.get(param_name) - param_class = meta["param_class"] - if ( - param is None - or tuple(param.shape) != tuple(meta["shape"]) - or param.dtype != meta["dtype"] - or ( - weight_loader is not None - and ( - not isinstance(param, param_class) - or not hasattr(param, "weight_loader") - ) - ) - ): - data = torch.empty( - meta["shape"], - dtype=meta["dtype"], - device=device or meta["device"], - ) - if param_class is not Parameter and weight_loader is not None: - kwargs = {"data": data, "weight_loader": weight_loader} - if "input_dim" in meta: - kwargs["input_dim"] = meta["input_dim"] - if "output_dim" in meta: - kwargs["output_dim"] = meta["output_dim"] - replacement = param_class(**kwargs) - else: - replacement = Parameter(data, requires_grad=False) - if weight_loader is not None: - replacement.weight_loader = weight_loader - setattr(module, param_name, replacement) - - -def modelopt_process_weights_after_loading(model) -> None: - """Run vLLM ModelOpt post-load processing for dense quantized layers.""" - actual_model = _unwrap_vllm_model(model) - - for module in actual_model.modules(): - quant_method = getattr(module, "quant_method", None) - if quant_method.__class__.__name__ == "ModelOptNvFp4LinearMethod": - quant_method.process_weights_after_loading(module) - - -_patched = False - - -def apply_modelopt_nvfp4_patches() -> None: - """Patch vLLM's dense ModelOpt NVFP4 method for rollout refits.""" - global _patched - - if _patched: - return - - from vllm.model_executor.layers.quantization.modelopt import ( - ModelOptNvFp4Config, - ModelOptNvFp4LinearMethod, - ) - - if not hasattr(ModelOptNvFp4Config, _ORIGINAL_NVFP4_CONFIG_FROM_CONFIG_ATTR): - setattr( - ModelOptNvFp4Config, - _ORIGINAL_NVFP4_CONFIG_FROM_CONFIG_ATTR, - ModelOptNvFp4Config._from_config, - ) - ModelOptNvFp4Config._from_config = classmethod(_modelopt_nvfp4_config_from_config) - - if not hasattr(ModelOptNvFp4LinearMethod, _ORIGINAL_LINEAR_APPLY_ATTR): - setattr( - ModelOptNvFp4LinearMethod, - _ORIGINAL_LINEAR_APPLY_ATTR, - ModelOptNvFp4LinearMethod.apply, - ) - ModelOptNvFp4LinearMethod.process_weights_after_loading = ( - _modelopt_dense_process_weights - ) - ModelOptNvFp4LinearMethod.apply = _modelopt_dense_apply - - _patched = True diff --git a/nemo_rl/modelopt/models/generation/vllm_quant_backend.py b/nemo_rl/modelopt/models/generation/vllm_quant_backend.py index 7298f56094b..de22810dabe 100644 --- a/nemo_rl/modelopt/models/generation/vllm_quant_backend.py +++ b/nemo_rl/modelopt/models/generation/vllm_quant_backend.py @@ -12,38 +12,526 @@ # See the License for the specific language governing permissions and # limitations under the License. -import gc import os import types +from collections.abc import Iterator from contextlib import ExitStack, contextmanager +from typing import Any import torch import vllm # noqa: F401 +import zmq from modelopt.torch.quantization.nn.modules.tensor_quantizer import TensorQuantizer from nemo_rl.modelopt.utils import ( - iter_quant_ignore_name_candidates, + MODELOPT_REAL_QUANT_ZMQ_TIMEOUT_MS, matches_quant_ignore_pattern, ) -from nemo_rl.models.generation.vllm.vllm_backend import VllmInternalWorkerExtension -from nemo_rl.models.policy.utils import ( - IPCProtocol, - calculate_aligned_size, - rebuild_cuda_tensor_from_ipc, +from nemo_rl.models.generation.vllm.vllm_backend import ( + IPCWeightManifestError, + VllmInternalWorkerExtension, + WeightUpdateFinalizer, + WeightUpdateTransport, ) +_FUSED_MODELOPT_MOE_SUFFIXES = { + ".experts.w13_weight": "w13_weight", + ".experts.w13_weight_scale": "w13_weight_scale", + ".experts.w13_weight_scale_2": "w13_weight_scale_2", + ".experts.w2_weight": "down_proj.weight", + ".experts.w2_weight_scale": "down_proj.weight_scale", + ".experts.w2_weight_scale_2": "down_proj.weight_scale_2", + ".experts.w13_input_scale": "w13_input_scale", + ".experts.w2_input_scale": "w2_input_scale", +} + + +def _match_fused_modelopt_moe_weight(name: str) -> tuple[str, str] | None: + return next( + ( + (suffix, target) + for suffix, target in _FUSED_MODELOPT_MOE_SUFFIXES.items() + if name.endswith(suffix) + ), + None, + ) + + +def _w13_num_shards_from_state_dict_info( + state_dict_info: dict[str, Any], + *, + require_input_scales: bool = False, +) -> dict[str, int]: + """Validate complete fused-MoE families and resolve their W13 layout.""" + num_shards_by_prefix: dict[str, int] = {} + input_shards_by_prefix: dict[str, int] = {} + targets_by_prefix: dict[str, set[str]] = {} + for name, (shape, _dtype) in state_dict_info.items(): + matched = _match_fused_modelopt_moe_weight(name) + if matched is None: + continue + suffix, target = matched + prefix = name[: -len(suffix)] + if target.startswith("down_proj."): + target = "w2_" + target.removeprefix("down_proj.") + targets_by_prefix.setdefault(prefix, set()).add(target) + if target == "w13_input_scale": + if len(shape) == 1: + input_shards = 1 + elif len(shape) == 2 and shape[1] in {1, 2}: + input_shards = shape[1] + else: + raise ValueError( + f"Expected one or two W13 input scales per expert for {name}, " + f"got {tuple(shape)}" + ) + input_shards_by_prefix[prefix] = input_shards + if target != "w13_weight_scale_2": + continue + if len(shape) == 1: + num_shards = 1 + elif len(shape) == 2 and shape[1] in {1, 2}: + num_shards = shape[1] + else: + raise ValueError( + f"Expected one or two W13 global scales per expert for {name}, " + f"got {tuple(shape)}" + ) + num_shards_by_prefix[prefix] = num_shards + + required_targets = { + "w13_weight", + "w13_weight_scale", + "w13_weight_scale_2", + "w2_weight", + "w2_weight_scale", + "w2_weight_scale_2", + } + if require_input_scales: + required_targets.update({"w13_input_scale", "w2_input_scale"}) + for prefix, targets in targets_by_prefix.items(): + missing = required_targets - targets + if missing: + raise RuntimeError( + f"Incomplete ModelOpt MoE export family for {prefix}: " + f"missing {sorted(missing)}" + ) + if set(num_shards_by_prefix) != set(targets_by_prefix): + missing = set(targets_by_prefix) - set(num_shards_by_prefix) + raise RuntimeError( + "ModelOpt MoE export families are missing W13 global scales: " + f"{sorted(missing)}" + ) + if require_input_scales: + mismatched = { + prefix + for prefix, num_shards in num_shards_by_prefix.items() + if input_shards_by_prefix.get(prefix) != num_shards + } + if mismatched: + raise RuntimeError( + "ModelOpt MoE W13 input/global scale layouts disagree for: " + f"{sorted(mismatched)}" + ) + return num_shards_by_prefix + + +def _batch_fused_modelopt_moe_weights( + weights: list[tuple[str, torch.Tensor]], + *, + w13_num_shards_by_prefix: dict[str, int], +) -> list[tuple[str, torch.Tensor]]: + """Map fused ModelOpt payloads to vLLM per-projection checkpoint names. + + Large expert weights and block scales stay batched so vLLM can + tensor-parallel-shard the full ``[E, ...]`` tensor at once. Its scalar + loader still requires an expert id, so only the tiny per-expert global + scales are exposed as scalar views. + """ + batched: list[tuple[str, torch.Tensor]] = [] + for name, tensor in weights: + matched = _match_fused_modelopt_moe_weight(name) + if matched is None: + batched.append((name, tensor)) + continue + + suffix, target = matched + prefix = name[: -len(suffix)] + if tensor.ndim == 0: + raise ValueError( + f"Fused ModelOpt MoE tensor must have an expert dimension: {name}" + ) + + if target in {"w13_weight", "w13_weight_scale"}: + target_suffix = "weight" if target == "w13_weight" else "weight_scale" + if w13_num_shards_by_prefix.get(prefix) == 1: + batched.append( + ( + f"{prefix}.experts.0.up_proj.{target_suffix}", + tensor, + ) + ) + continue + if tensor.ndim < 2 or tensor.shape[1] % 2 != 0: + raise ValueError( + f"Expected fused gate/up tensor with an even projection " + f"dimension for {name}, got {tuple(tensor.shape)}" + ) + gate, up = tensor.chunk(2, dim=1) + batched.extend( + ( + f"{prefix}.experts.0.{projection}.{target_suffix}", + shard, + ) + for projection, shard in ( + ("gate_proj", gate), + ("up_proj", up), + ) + ) + continue + + if target == "w13_input_scale": + if tensor.ndim == 1: + tensor = tensor[:, None] + if tensor.ndim != 2 or tensor.shape[1] not in {1, 2}: + raise ValueError( + f"Expected one or two W13 input scales per expert for {name}, " + f"got {tuple(tensor.shape)}" + ) + if tensor.shape[1] == 1: + batched.extend( + ( + f"{prefix}.experts.{expert_id}.up_proj.input_scale", + expert_scale[0], + ) + for expert_id, expert_scale in enumerate(tensor.unbind(0)) + ) + continue + for expert_id, expert_scale in enumerate(tensor.unbind(0)): + batched.append( + ( + f"{prefix}.experts.{expert_id}.gate_proj.input_scale", + expert_scale[0], + ) + ) + batched.append( + ( + f"{prefix}.experts.{expert_id}.up_proj.input_scale", + expert_scale[1], + ) + ) + continue + + if target == "w2_input_scale": + if tensor.ndim == 2 and tensor.shape[1] == 1: + tensor = tensor[:, 0] + if tensor.ndim != 1: + raise ValueError( + f"Expected one down-projection input scale per expert for " + f"{name}, got {tuple(tensor.shape)}" + ) + batched.extend( + (f"{prefix}.experts.{expert_id}.down_proj.input_scale", scale) + for expert_id, scale in enumerate(tensor.unbind(0)) + ) + continue + + if target == "w13_weight_scale_2": + if tensor.ndim == 1: + tensor = tensor[:, None] + if tensor.ndim != 2 or tensor.shape[1] not in {1, 2}: + raise ValueError( + f"Expected one or two W13 global scales per expert for {name}, " + f"got {tuple(tensor.shape)}" + ) + if tensor.shape[1] == 1: + batched.extend( + ( + f"{prefix}.experts.{expert_id}.up_proj.weight_scale_2", + expert_scale[0], + ) + for expert_id, expert_scale in enumerate(tensor.unbind(0)) + ) + continue + for expert_id, expert_scale in enumerate(tensor.unbind(0)): + batched.append( + ( + f"{prefix}.experts.{expert_id}.gate_proj.weight_scale_2", + expert_scale[0], + ) + ) + batched.append( + ( + f"{prefix}.experts.{expert_id}.up_proj.weight_scale_2", + expert_scale[1], + ) + ) + continue + + if not target.endswith("weight_scale_2"): + batched.append((f"{prefix}.experts.0.{target}", tensor)) + continue + + if tensor.ndim == 1: + expert_scales = tensor + elif tensor.ndim == 2 and tensor.shape[1] == 1: + expert_scales = tensor[:, 0] + else: + raise ValueError( + f"Expected one global scale per expert for {name}, got " + f"shape {tuple(tensor.shape)}" + ) + + batched.extend( + (f"{prefix}.experts.{expert_id}.{target}", expert_scale) + for expert_id, expert_scale in enumerate(expert_scales.unbind(0)) + ) + + return batched + + +def _detach_pending_layerwise_weights( + reload_roots: tuple[torch.nn.Module, ...], + source_storage_ptrs: set[int], +) -> None: + """Own deferred weights before a transport buffer may be reused. + + Completed layers have already released their buffered arguments, so this + clones only tensors from a layer split across transport batches. Only the + cached layerwise-reload subgraphs are inspected. + """ + if not source_storage_ptrs: + return + from vllm.model_executor.model_loader.reload.layerwise import get_layerwise_info + + for reload_root in reload_roots: + for module in reload_root.modules(): + info = get_layerwise_info(module) + for _, arguments in info.loaded_weights: + loaded_weight = arguments.arguments.get("loaded_weight") + if not isinstance(loaded_weight, torch.Tensor): + continue + if loaded_weight.untyped_storage().data_ptr() in source_storage_ptrs: + arguments.arguments["loaded_weight"] = loaded_weight.clone() + + +def _iter_modelopt_quant_modules( + model: torch.nn.Module, +) -> list[tuple[str, torch.nn.Module]]: + """Return modules whose runtime layout is owned by vLLM ModelOpt methods.""" + from vllm.model_executor.layers.quantization.modelopt import ( + ModelOptNvFp4FusedMoE, + ModelOptNvFp4LinearMethod, + ) + + method_types = (ModelOptNvFp4FusedMoE, ModelOptNvFp4LinearMethod) + return [ + (module_name, module) + for module_name, module in model.named_modules() + if isinstance(getattr(module, "quant_method", None), method_types) + ] + + +def _modelopt_layerwise_reload_roots( + model: torch.nn.Module, + *, + include_fp8_kv_cache: bool, +) -> list[torch.nn.Module]: + """Select disjoint roots that require vLLM's native reload lifecycle. + + Ordinary parameters are already updated in place by vLLM's checkpoint + loaders. Restricting layerwise reconstruction to ModelOpt runtime layouts + and attention scale owners avoids materializing unrelated non-persistent + buffers. In vLLM 0.20, whole-model reconstruction can otherwise break a + derived buffer that aliases a child parameter (for example Nemotron-H's + ``conv_weights`` view of ``conv1d.weight``). + """ + from vllm.model_executor.layers.attention import Attention, MLAAttention + from vllm.model_executor.layers.quantization.kv_cache import BaseKVCacheMethod + + modelopt_modules = {module for _, module in _iter_modelopt_quant_modules(model)} + attention_types = (Attention, MLAAttention) + quant_roots: list[torch.nn.Module] = [] + attention_roots: list[torch.nn.Module] = [] + visited: set[torch.nn.Module] = set() + + def collect(module: torch.nn.Module) -> None: + if module in visited: + return + visited.add(module) + if ( + include_fp8_kv_cache + and isinstance(module, attention_types) + and isinstance(getattr(module, "quant_method", None), BaseKVCacheMethod) + and "fp8" in str(getattr(module, "kv_cache_dtype", "auto")).lower() + ): + attention_roots.append(module) + return + if module in modelopt_modules: + quant_roots.append(module) + return + for child in module.children(): + collect(child) + + collect(model) + # Match vLLM's ordering contract: process quantized modules before the + # attention owners that finalize KV-cache scales. + return quant_roots + attention_roots + + +def _require_complete_modelopt_layerwise_reload(model: torch.nn.Module) -> None: + """Reject ModelOpt layers that vLLM would otherwise finalize partially.""" + candidates = _iter_modelopt_quant_modules(model) + + if not candidates: + return + + from vllm.model_executor.model_loader.reload.layerwise import get_layerwise_info + + incomplete = [] + for module_name, module in candidates: + info = get_layerwise_info(module) + if info.load_numel_total is None: + # A completed layer is processed and reset immediately by vLLM. + continue + if info.load_numel == info.load_numel_total: + continue + buffered = sorted({name for name, _ in info.loaded_weights}) + incomplete.append( + f"{module_name or ''}: {info.load_numel}/" + f"{info.load_numel_total} elements, buffered={buffered}" + ) + + if incomplete: + details = "; ".join(incomplete[:8]) + suffix = "; ..." if len(incomplete) > 8 else "" + raise RuntimeError( + "ModelOpt layerwise reload is incomplete for " + f"{len(incomplete)} layer(s): {details}{suffix}" + ) + + if os.environ.get("VLLM_MODELOPT_REAL_QUANT", "0") == "1": - from nemo_rl.modelopt.models.generation.vllm_modelopt_patch import ( - apply_modelopt_nvfp4_patches, + from nemo_rl.modelopt.models.generation.vllm_modelopt import ( + register_nemo_modelopt_nvfp4, ) - apply_modelopt_nvfp4_patches() + register_nemo_modelopt_nvfp4() class VllmQuantInternalWorkerExtension(VllmInternalWorkerExtension): + _nrl_w13_num_shards_by_prefix: dict[str, int] + _nrl_modelopt_reload_roots: tuple[torch.nn.Module, ...] | None = None + + def maybe_init_zmq(self) -> None: + """Use a longer timeout only for ModelOpt real-quant refits.""" + super().maybe_init_zmq() + if self._is_real_quant_model(): + self.zmq_socket.setsockopt(zmq.SNDTIMEO, MODELOPT_REAL_QUANT_ZMQ_TIMEOUT_MS) + self.zmq_socket.setsockopt(zmq.RCVTIMEO, MODELOPT_REAL_QUANT_ZMQ_TIMEOUT_MS) + def _is_real_quant_model(self) -> bool: return os.environ.get("VLLM_MODELOPT_REAL_QUANT", "0") == "1" + def _get_modelopt_reload_roots(self) -> tuple[torch.nn.Module, ...]: + """Return the invariant ModelOpt layerwise-reload subgraphs.""" + if self._nrl_modelopt_reload_roots is None: + self._nrl_modelopt_reload_roots = tuple( + _modelopt_layerwise_reload_roots( + self.model_runner.model, + include_fp8_kv_cache=self._uses_fp8_kv_cache(), + ) + ) + return self._nrl_modelopt_reload_roots + + @contextmanager + def _weight_update_lifecycle( + self, transport: WeightUpdateTransport + ) -> Iterator[WeightUpdateFinalizer]: + """Use vLLM's native layerwise reload lifecycle for real quantization.""" + if not self._is_real_quant_model(): + with super()._weight_update_lifecycle(transport) as finalize: + yield finalize + return + + from vllm.config import set_current_vllm_config + from vllm.model_executor.model_loader.reload import ( + finalize_layerwise_reload, + initialize_layerwise_reload, + ) + + model = self.model_runner.model + reload_roots = self._get_modelopt_reload_roots() + + def finalize() -> None: + try: + with torch.device(self.device): + _require_complete_modelopt_layerwise_reload(model) + for reload_root in reload_roots: + finalize_layerwise_reload(reload_root, self.model_config) + # Fence completion for both collective return and the IPC + # COMPLETE acknowledgment. Data-batch ACKs use the hook below. + torch.accelerator.synchronize() + except Exception as error: + if transport == "ipc": + raise RuntimeError( + f"ModelOpt real-quant refit post-processing failed: {error}" + ) from error + raise + + try: + # Layerwise loading may reconstruct backend CustomOps as soon as a + # layer becomes complete. Keep vLLM's worker config available for + # that online processing as well as deferred finalization. + with set_current_vllm_config(self.model_runner.vllm_config): + with torch.device(self.device): + for reload_root in reload_roots: + initialize_layerwise_reload(reload_root) + yield finalize + except IPCWeightManifestError as error: + raise RuntimeError( + f"ModelOpt real-quant refit rejected: {error}" + ) from error + except Exception as error: + if transport == "collective": + raise RuntimeError( + "ModelOpt real-quant collective refit failed" + ) from error + raise + + def _weight_update_errors_are_fatal(self) -> bool: + return self._is_real_quant_model() + + def _synchronize_before_ipc_data_ack(self) -> None: + """Fence all accelerator streams used by ModelOpt post-load methods.""" + if self._is_real_quant_model(): + torch.accelerator.synchronize() + return + super()._synchronize_before_ipc_data_ack() + + def prepare_refit_info(self, state_dict_info: dict[str, Any]) -> None: + super().prepare_refit_info(state_dict_info) + if not self._is_real_quant_model(): + return + self._get_modelopt_reload_roots() + quant_config = ( + self.model_runner.vllm_config.model_config.hf_config.quantization_config + ) + self._nrl_w13_num_shards_by_prefix = _w13_num_shards_from_state_dict_info( + state_dict_info, + require_input_scales=( + str(quant_config.get("quant_algo", "")).upper() == "NVFP4" + ), + ) + if ( + self._nrl_w13_num_shards_by_prefix + and self.model_runner.vllm_config.parallel_config.enable_expert_parallel + ): + raise RuntimeError( + "Fused ModelOpt MoE refits require all experts local; " + "vLLM expert parallelism is unsupported" + ) + @contextmanager def _patch_named_parameters_to_include_buffers(self, model): """Temporarily patches model.named_parameters() to also yield input_quantizer buffers. @@ -84,43 +572,47 @@ def _load_weights(self, weights): applied during export), so no fold_weight step is needed here. """ if self._is_real_quant_model(): + weights = list(weights) + source_storage_ptrs = { + tensor.untyped_storage().data_ptr() for _, tensor in weights + } quant_config = ( self.model_runner.vllm_config.model_config.hf_config.quantization_config ) ignore_patterns = quant_config.get("ignore", []) or [] - # Built lazily on first use: only the rare ignored, floating-point - # weights (typically just lm_head) need a parameter lookup, so most - # refit chunks skip the full named_parameters() scan entirely. - params = None filtered = [] for name, weight in weights: suffix = name.rsplit(".", 1)[-1] ignored = matches_quant_ignore_pattern(name, ignore_patterns) - if ignored and suffix in {"weight_scale", "weight_scale_2"}: + if ignored and suffix in { + "weight_scale", + "weight_scale_2", + "input_scale", + }: continue - if ignored and suffix == "weight" and weight.is_floating_point(): - if params is None: - params = dict(self.model_runner.model.named_parameters()) - copied = False - for candidate in iter_quant_ignore_name_candidates(name): - param = params.get(candidate) - if param is not None and tuple(param.shape) == tuple( - weight.shape - ): - param.data.copy_( - weight.to(device=param.device, dtype=param.dtype) - ) - copied = True - break - if copied: - continue - filtered.append((name, weight)) - weights = filtered + if any( + _match_fused_modelopt_moe_weight(name) is not None + for name, _ in filtered + ): + weights = _batch_fused_modelopt_moe_weights( + filtered, + w13_num_shards_by_prefix=self._nrl_w13_num_shards_by_prefix, + ) + else: + weights = filtered if not weights: return None - return super()._load_weights(weights) + try: + with torch.device(self.device): + return super()._load_weights(weights) + finally: + with torch.device(self.device): + _detach_pending_layerwise_weights( + self._get_modelopt_reload_roots(), + source_storage_ptrs, + ) with ExitStack() as contexts: for _, child in self.model_runner.model.named_children(): @@ -129,79 +621,6 @@ def _load_weights(self, weights): ) return super()._load_weights(weights) - def update_weights_via_ipc_zmq(self) -> bool: - """Receive and update weights through CUDA IPC.""" - if not self._is_real_quant_model(): - return super().update_weights_via_ipc_zmq() - - from nemo_rl.modelopt.models.generation.vllm_modelopt_patch import ( - modelopt_process_weights_after_loading, - prepare_modelopt_for_weight_reload, - ) - - prepare_modelopt_for_weight_reload(self.model_runner.model, device=self.device) - self.maybe_init_zmq() - while True: - payload = self.zmq_socket.recv_pyobj() - - if payload == IPCProtocol.COMPLETE: - modelopt_process_weights_after_loading(self.model_runner.model) - torch.cuda.synchronize() - self.zmq_socket.send(IPCProtocol.ACK.value.encode()) - break - - ipc_handle, list_keys, used_bytes = payload - buffer = rebuild_cuda_tensor_from_ipc(ipc_handle, self.device.index) - - weights = [] - offset = 0 - for key in list_keys: - shape, dtype = self.state_dict_info[key] - if isinstance(shape, list): - shape = torch.Size(shape) - - size_in_bytes = dtype.itemsize * shape.numel() - weight = ( - buffer[offset : offset + size_in_bytes] - .view(dtype=dtype) - .view(shape) - ) - weights.append((key, weight)) - - offset += calculate_aligned_size(size_in_bytes) - - assert offset == used_bytes, ( - "Offset is not equal to used bytes, usually indicate inaccurate " - "info like keys or cached dtype in state_dict_info" - ) - - self._load_weights(weights) - torch.cuda.synchronize() - - del weights, buffer - self.zmq_socket.send(IPCProtocol.ACK.value.encode()) - - self._maybe_process_fp8_kv_cache() - gc.collect() - torch.cuda.empty_cache() - return True - - def update_weights_from_collective(self) -> bool: - """Receive and update weights through collective communication.""" - if not self._is_real_quant_model(): - return super().update_weights_from_collective() - - from nemo_rl.modelopt.models.generation.vllm_modelopt_patch import ( - modelopt_process_weights_after_loading, - prepare_modelopt_for_weight_reload, - ) - - prepare_modelopt_for_weight_reload(self.model_runner.model, device=self.device) - result = super().update_weights_from_collective() - if result: - modelopt_process_weights_after_loading(self.model_runner.model) - return result - def get_weight_snapshot(self, name: str) -> torch.Tensor: """Return a CPU copy of a named parameter for before/after comparison.""" model = self.model_runner.model diff --git a/nemo_rl/modelopt/models/generation/vllm_quant_worker.py b/nemo_rl/modelopt/models/generation/vllm_quant_worker.py index d7c81d3a340..3f797ec4b7d 100644 --- a/nemo_rl/modelopt/models/generation/vllm_quant_worker.py +++ b/nemo_rl/modelopt/models/generation/vllm_quant_worker.py @@ -30,9 +30,17 @@ _EXTRA_ENV_VARS = ( "VLLM_QUANT_CFG", "VLLM_MODELOPT_REAL_QUANT", + "PYTHONPATH", ) +def _quant_cfg_for_worker_env(quant_cfg: str) -> str: + expanded = os.path.expanduser(quant_cfg) + if os.path.isfile(expanded): + return os.path.abspath(expanded) + return quant_cfg + + def _configure_quant_engine_kwargs( cfg: VllmConfig, llm_kwargs: dict[str, Any], @@ -42,27 +50,39 @@ def _configure_quant_engine_kwargs( ) real_quant = bool(cfg.get("real_quant")) if real_quant: - from nemo_rl.modelopt.models.generation.vllm_modelopt_patch import ( - apply_modelopt_nvfp4_patches, + from nemo_rl.modelopt.models.generation.vllm_modelopt import ( + quantization_method_for_mode, + register_nemo_modelopt_nvfp4, + ) + from nemo_rl.modelopt.utils import ( + build_vllm_modelopt_nvfp4_config, + resolve_nvfp4_real_quant_mode, ) - from nemo_rl.modelopt.utils import build_vllm_modelopt_nvfp4_config - apply_modelopt_nvfp4_patches() + quant_cfg = cfg.get("quant_cfg") + if not quant_cfg: + raise ValueError("NVFP4 real quantization requires a non-empty quant_cfg.") + mode = resolve_nvfp4_real_quant_mode(quant_cfg) + register_nemo_modelopt_nvfp4() + os.environ.pop("VLLM_QUANT_CFG", None) os.environ["VLLM_MODELOPT_REAL_QUANT"] = "1" hf_overrides = llm_kwargs.setdefault("hf_overrides", {}) hf_overrides["quantization_config"] = build_vllm_modelopt_nvfp4_config( + mode=mode, ignore=cfg.get("real_quant_ignore"), ) - llm_kwargs["quantization"] = "modelopt" + llm_kwargs["quantization"] = quantization_method_for_mode(mode) else: llm_kwargs["worker_cls"] = ( "nemo_rl.modelopt.models.generation.vllm_quant_patch.FakeQuantWorker" ) # Expert fakequant needs a decomposed MoE path; explicit user config still wins. llm_kwargs.setdefault("moe_backend", "triton") + os.environ.pop("VLLM_MODELOPT_REAL_QUANT", None) + os.environ.pop("VLLM_QUANT_CFG", None) if cfg["quant_cfg"]: - os.environ["VLLM_QUANT_CFG"] = cfg["quant_cfg"] + os.environ["VLLM_QUANT_CFG"] = _quant_cfg_for_worker_env(cfg["quant_cfg"]) @ray.remote( diff --git a/nemo_rl/modelopt/models/policy/workers/megatron_quant_policy_worker.py b/nemo_rl/modelopt/models/policy/workers/megatron_quant_policy_worker.py index 395dc10d187..96b06082595 100644 --- a/nemo_rl/modelopt/models/policy/workers/megatron_quant_policy_worker.py +++ b/nemo_rl/modelopt/models/policy/workers/megatron_quant_policy_worker.py @@ -13,13 +13,17 @@ # limitations under the License. +import hashlib +import json import os +import warnings +from collections.abc import Generator, Mapping from contextlib import contextmanager -from typing import Generator +from pathlib import Path -import modelopt.torch.quantization as mtq import ray import torch +import zmq from megatron.bridge.training.post_training.checkpointing import ( has_modelopt_state, load_modelopt_state, @@ -36,39 +40,120 @@ quantize_model, symlink_pre_quantized_model, ) +from nemo_rl.modelopt.utils import ( + MODELOPT_REAL_QUANT_ZMQ_TIMEOUT_MS, + resolve_nvfp4_real_quant_mode, +) from nemo_rl.models.policy.utils import get_runtime_env_for_policy_worker from nemo_rl.models.policy.workers.megatron_policy_worker import ( MegatronPolicyWorkerImpl, ) -@contextmanager -def _w4a16_modelopt_exporter(): - """Temporarily adapt Bridge's NVFP4 exporter for W4A16 rollout metadata. - - Get this removed when Bridge fixed the logic here. +def _quant_checkpoint_cache_suffix(config: Mapping[str, object]) -> str: + """Build a short suffix for HF->Megatron checkpoints with ModelOpt state.""" + keys = ( + "quant_cfg", + "quant_calib_data", + "quant_calib_size", + "quant_batch_size", + "quant_sequence_length", + "disable_modelopt_layer_spec", + ) + payload = {key: config.get(key) for key in keys} + quant_cfg = payload["quant_cfg"] + path = Path(quant_cfg).expanduser() if isinstance(quant_cfg, str) else None + if path is not None and path.is_file(): + payload["quant_cfg"] = { + "path": path.resolve().as_posix(), + "sha256": hashlib.sha256(path.read_bytes()).hexdigest(), + } + digest = hashlib.sha256( + json.dumps(payload, sort_keys=True, separators=(",", ":")).encode() + ).hexdigest()[:12] + return f"_modelopt_{digest}" + + +def _find_other_quant_checkpoint_caches( + base_pretrained_path: str, + selected_pretrained_path: str, +) -> list[Path]: + """Find valid quantized startup caches other than the selected cache.""" + base_path = Path(base_pretrained_path) + parent = base_path.parent + if not parent.is_dir(): + return [] + + selected_path = Path(selected_pretrained_path) + hashed_prefix = f"{base_path.name}_modelopt_" + legacy_name = f"{base_path.name}_quantized" + caches = [] + for candidate in parent.iterdir(): + if candidate == selected_path or not ( + candidate.name.startswith(hashed_prefix) or candidate.name == legacy_name + ): + continue + iter0_path = candidate / "iter_0000000" + if iter0_path.exists() and has_modelopt_state(iter0_path.as_posix()): + caches.append(candidate) + return sorted(caches) + + +def _warn_if_other_quant_checkpoint_caches( + base_pretrained_path: str, + selected_pretrained_path: str, +) -> None: + """Warn when a different quantization config already has a startup cache.""" + other_caches = _find_other_quant_checkpoint_caches( + base_pretrained_path, + selected_pretrained_path, + ) + if not other_caches: + return + + warnings.warn( + "Found quantized startup checkpoint cache(s) created with a different " + f"quantization configuration: {', '.join(map(str, other_caches))}. " + f"They will not be reused; the selected cache is {selected_pretrained_path}. " + "This startup-cache check does not validate resumed training checkpoints. " + "When changing an experiment configuration, use a new " + "checkpointing.checkpoint_dir.", + UserWarning, + stacklevel=2, + ) + + +def _set_quantization_model_specs(model_config, disable_modelopt_layer_spec: bool): + """Select quantization-compatible specs across Bridge hybrid API versions. + + Recent Megatron-Bridge revisions load Nemotron-H checkpoints as a + ``HybridModelProvider`` and use ``hybrid_stack_spec``. Older revisions use + the deprecated ``mamba_stack_spec`` field. Setting only the latter leaves + the recent provider to infer the local ModelOpt stack when + ``restore_modelopt_state=True``; that stack contains ``SequentialMLP`` and + cannot restore quantizers with both tensor and expert parallelism enabled. """ - from megatron.bridge.models.conversion import modelopt_utils - from modelopt.torch.export.quant_utils import QUANTIZATION_W4A16_NVFP4 - - original_get_exporter = modelopt_utils.get_modelopt_quant_exporter - - def _get_modelopt_quant_exporter(quant_mode: str): - if quant_mode.lower() == "w4a16_nvfp4": - return QUANTIZATION_W4A16_NVFP4, modelopt_utils.quantize_nvfp4_weight - return original_get_exporter(quant_mode) - - modelopt_utils.get_modelopt_quant_exporter = _get_modelopt_quant_exporter - try: - yield - finally: - modelopt_utils.get_modelopt_quant_exporter = original_get_exporter + model_config.transformer_layer_spec = get_quantization_layer_spec( + disable_modelopt_layer_spec + ) + stack_spec = get_quantization_mamba_stack_spec(disable_modelopt_layer_spec) + if hasattr(model_config, "hybrid_stack_spec"): + model_config.hybrid_stack_spec = stack_spec + elif hasattr(model_config, "mamba_stack_spec"): + model_config.mamba_stack_spec = stack_spec @ray.remote( runtime_env=get_runtime_env_for_policy_worker("megatron_quant_policy_worker") ) # pragma: no cover class MegatronQuantPolicyWorker(MegatronPolicyWorkerImpl): + def maybe_init_zmq(self) -> None: + """Use a longer timeout only for ModelOpt real-quant refits.""" + super().maybe_init_zmq() + if self._use_real_quant_refit(): + self.zmq_socket.setsockopt(zmq.SNDTIMEO, MODELOPT_REAL_QUANT_ZMQ_TIMEOUT_MS) + self.zmq_socket.setsockopt(zmq.RCVTIMEO, MODELOPT_REAL_QUANT_ZMQ_TIMEOUT_MS) + def __init__(self, config, *args, **kwargs): """Initialize the MegatronQuantPolicyWorker.""" megatron_cfg = config.get("megatron_cfg", {}) or {} @@ -106,9 +191,6 @@ def __init__(self, config, *args, **kwargs): self.reference_state_dict[name] = item.detach().to( device="cpu", non_blocking=True, copy=True ) - if self.rank == 0: - print(f"Quantized model: {self.model}") - mtq.print_quant_summary(self.model) def _quantize(self, model): """Quantize the model if the model is not quantized yet.""" @@ -131,8 +213,8 @@ def _patch_validate_model_paths(self): """Patch validate_model_paths to handle quantized checkpoint paths. In cases like distillation where the teacher model is the same as the student model, - we need to save an extra quantized checkpoint. This patch checks for modelopt state - and redirects to a _quantized suffix path. It also handles pre-quantized model symlinks. + we need to save an extra quantized checkpoint. This patch routes auto-converted HF + checkpoints to a ModelOpt-specific cache path. It also handles pre-quantized model symlinks. """ if getattr(megatron_policy_worker.validate_model_paths, "_is_patched", False): return @@ -143,11 +225,15 @@ def _validate_model_paths(config): original_validate_model_paths(config) ) + if config.get("pretrained_checkpoint") is not None: + return hf_model_name, pretrained_path, pt_checkpoint_exists + + base_pretrained_path = pretrained_path + pretrained_path += _quant_checkpoint_cache_suffix(config) iter0_path = os.path.join(pretrained_path, "iter_0000000") - if pt_checkpoint_exists and not has_modelopt_state(iter0_path): - pretrained_path += "_quantized" - iter0_path = os.path.join(pretrained_path, "iter_0000000") - pt_checkpoint_exists = os.path.exists(iter0_path) + pt_checkpoint_exists = os.path.exists(iter0_path) and has_modelopt_state( + iter0_path + ) pre_quantized_model_path = os.environ.get( "NRL_PRE_QUANTIZED_MEGATRON_MODEL_PATH" @@ -156,6 +242,12 @@ def _validate_model_paths(config): symlink_pre_quantized_model(pre_quantized_model_path, pretrained_path) pt_checkpoint_exists = True + if not pt_checkpoint_exists and self.rank == 0: + _warn_if_other_quant_checkpoint_caches( + base_pretrained_path, + pretrained_path, + ) + return hf_model_name, pretrained_path, pt_checkpoint_exists _validate_model_paths._is_patched = True @@ -183,13 +275,9 @@ def _setup_model_and_optimizer(policy_cfg, megatron_cfg, *args, **kwargs): "disable_modelopt_layer_spec", False ) megatron_cfg.model.restore_modelopt_state = True - megatron_cfg.model.transformer_layer_spec = get_quantization_layer_spec( - disable_modelopt_layer_spec + _set_quantization_model_specs( + megatron_cfg.model, disable_modelopt_layer_spec ) - if hasattr(megatron_cfg.model, "mamba_stack_spec"): - megatron_cfg.model.mamba_stack_spec = ( - get_quantization_mamba_stack_spec(disable_modelopt_layer_spec) - ) return original_setup_model_and_optimizer( policy_cfg, megatron_cfg, *args, **kwargs @@ -360,31 +448,53 @@ def save_checkpoint(self, *args, **kwargs): return super().save_checkpoint(*args, **kwargs) def _use_real_quant_refit(self) -> bool: - generation_cfg = self.cfg.get("generation") or {} + generation_cfg = self.cfg["generation"] return ( - generation_cfg.get("backend") == "vllm" + generation_cfg["backend"] == "vllm" and generation_cfg.get("quant_cfg") is not None and bool(generation_cfg.get("real_quant")) ) - def _iter_real_quant_refit_params(self, kv_scales=None): + def _get_real_quant_mode(self) -> str: + """Resolve and cross-check the training and rollout quantization modes.""" + cached_mode = getattr(self, "_real_quant_mode", None) + if cached_mode is not None: + return cached_mode + + policy_quant_cfg = self.cfg.get("quant_cfg") + generation_quant_cfg = self.cfg["generation"].get("quant_cfg") + policy_mode = resolve_nvfp4_real_quant_mode(policy_quant_cfg) + generation_mode = resolve_nvfp4_real_quant_mode(generation_quant_cfg) + if policy_mode != generation_mode: + raise ValueError( + "Real-quant refit requires matching policy and generation " + f"quantization modes, got {policy_mode} from {policy_quant_cfg!r} " + f"and {generation_mode} from {generation_quant_cfg!r}." + ) + self._real_quant_mode = policy_mode + return policy_mode + + def _iter_real_quant_refit_params( + self, + kv_scales: dict[str, float] | None = None, + ) -> Generator[tuple[str, torch.Tensor], None, None]: """Export packed NVFP4 weights and scales for real-quant vLLM rollout.""" from nemo_rl.modelopt.utils import DEFAULT_NVFP4_IGNORE - generation_cfg = self.cfg.get("generation") or {} - vllm_cfg = generation_cfg.get("vllm_cfg", {}) + generation_cfg = self.cfg["generation"] + vllm_cfg = generation_cfg["vllm_cfg"] ignore = generation_cfg.get("real_quant_ignore") if ignore is None: ignore = DEFAULT_NVFP4_IGNORE - with _w4a16_modelopt_exporter(): - yield from self.megatron_bridge.export_hf_weights_modelopt( - [self.model], - quant_mode="w4a16_nvfp4", - cpu=True, - show_progress=False, - conversion_tasks=self.refit_conversion_tasks, - ignore_patterns=ignore, - ) + mode = self._get_real_quant_mode() + yield from self.megatron_bridge.export_hf_weights_modelopt( + [self.model], + quant_mode="nvfp4" if mode == "w4a4" else "w4a16_nvfp4", + cpu=True, + show_progress=False, + conversion_tasks=self.refit_conversion_tasks, + ignore_patterns=ignore, + ) if self.draft_model is not None: from nemo_rl.models.megatron.draft import export_eagle_weights_to_hf @@ -392,7 +502,7 @@ def _iter_real_quant_refit_params(self, kv_scales=None): for name, tensor in export_eagle_weights_to_hf(self.draft_model): yield f"draft.{name}", tensor - if not vllm_cfg.get("kv_cache_dtype", "").startswith("fp8"): + if not vllm_cfg["kv_cache_dtype"].startswith("fp8"): return from nemo_rl.models.generation.vllm.quantization.fp8_train_utils import ( @@ -417,7 +527,10 @@ def _iter_real_quant_refit_params(self, kv_scales=None): ) @staticmethod - def _find_weight_quantizer(module, param_weight): + def _find_weight_quantizer( + module: object, + param_weight: object, + ) -> object | None: """Find the enabled weight quantizer that corresponds to ``param_weight``. Uses ModelOpt's ``QuantModule.iter_weights_for_calibration`` to discover @@ -431,6 +544,7 @@ def _find_weight_quantizer(module, param_weight): return None if not isinstance(module, QuantModule): return None + for weight, wq in module.iter_weights_for_calibration(): if ( param_weight is weight diff --git a/nemo_rl/modelopt/models/policy/workers/utils.py b/nemo_rl/modelopt/models/policy/workers/utils.py index d8f9065a854..ef90e9922e5 100644 --- a/nemo_rl/modelopt/models/policy/workers/utils.py +++ b/nemo_rl/modelopt/models/policy/workers/utils.py @@ -25,13 +25,17 @@ modelopt_mamba_stack_spec, transformer_engine_mamba_stack_spec, ) +from megatron.core import parallel_state from megatron.core.post_training.modelopt.gpt.model_specs import get_gpt_modelopt_spec from modelopt.torch.quantization.config import need_calibration from modelopt.torch.utils.dataset_utils import ( create_forward_loop, get_dataset_dataloader, ) -from modelopt.torch.utils.plugins import megatron_prefill +from modelopt.torch.utils.plugins import ( + get_megatron_calibration_forward_loop, + megatron_prefill, +) from torch.utils.data import DataLoader, Dataset from nemo_rl.algorithms.utils import get_tokenizer as _base_get_tokenizer @@ -81,10 +85,63 @@ def __len__(self): def get_forward_loop_func( + *, is_megatron: bool, - calib_dataloader: DataLoader, + tokenizer, + dataset_name: str, + batch_size: int, + num_samples: int, + sample_length: int, + device: torch.device, ): - """Gets the forward loop function for the model.""" + """Build the calibration forward loop for the requested backend and data.""" + if is_megatron and dataset_name == "random": + # The upstream helper owns CP/DP sharding for named datasets. The local + # synthetic loop has no tokenizer/dataset for that helper to partition. + cp_size = parallel_state.get_context_parallel_world_size() + if cp_size > 1: + raise RuntimeError( + "Random ModelOpt Megatron calibration requires " + f"context_parallel_size=1, got {cp_size}; use a named dataset" + ) + + if dataset_name != "random" and is_megatron: + return get_megatron_calibration_forward_loop( + tokenizer, + dataset_name=dataset_name, + batch_size=batch_size, + num_samples=num_samples, + seq_length=sample_length, + device=device, + apply_chat_template=False, + pack=True, + ) + + if dataset_name == "random": + calib_dataloader = DataLoader( + _DictDataset( + { + "input_ids": torch.randint( + 0, + 100, + (num_samples, sample_length), + device=device, + ) + } + ), + batch_size=batch_size, + ) + else: + calib_dataloader = get_dataset_dataloader( + dataset_name=dataset_name, + tokenizer=tokenizer, + batch_size=batch_size, + num_samples=num_samples, + device=device, + include_labels=False, + max_sample_length=sample_length, + ) + if not is_megatron: return create_forward_loop(dataloader=calib_dataloader) @@ -136,31 +193,23 @@ def quantize_model( if hasattr(model, "device") else next(model.parameters()).device ) - if data == "random": - calib_size = 1 - calib_dataloader = DataLoader( - _DictDataset( - {"input_ids": torch.randint(0, 100, (1, 5), device=device)} - ), - batch_size=1, - ) - else: - calib_dataloader = get_dataset_dataloader( - dataset_name=data, - tokenizer=tokenizer, - batch_size=batch_size - if batch_size is not None - else DEFAULT_CALIB_BATCH_SIZE, - num_samples=calib_size, - device=device, - include_labels=False, - max_sample_length=( - max_sample_length - if max_sample_length is not None - else DEFAULT_CALIB_SAMPLE_LENGTH - ), - ) - forward_loop = get_forward_loop_func(is_megatron, calib_dataloader) + calib_batch_size = ( + batch_size if batch_size is not None else DEFAULT_CALIB_BATCH_SIZE + ) + calib_sample_length = ( + max_sample_length + if max_sample_length is not None + else DEFAULT_CALIB_SAMPLE_LENGTH + ) + forward_loop = get_forward_loop_func( + is_megatron=is_megatron, + tokenizer=tokenizer, + dataset_name=data, + batch_size=calib_batch_size, + num_samples=calib_size, + sample_length=calib_sample_length, + device=device, + ) model = mtq.quantize(model, mtq_cfg, forward_loop) mtq.print_quant_summary(model) diff --git a/nemo_rl/modelopt/utils.py b/nemo_rl/modelopt/utils.py index c6771e888b3..ed9443b8e26 100644 --- a/nemo_rl/modelopt/utils.py +++ b/nemo_rl/modelopt/utils.py @@ -12,17 +12,21 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Lightweight quantization config resolver usable by both Megatron and vLLM workers.""" +"""Lightweight ModelOpt helpers shared by Megatron and vLLM workers.""" from __future__ import annotations +from collections.abc import Mapping, Sequence from fnmatch import fnmatchcase -from typing import Any, Iterator +from typing import Any, Iterator, Literal + +MODELOPT_REAL_QUANT_ZMQ_TIMEOUT_MS = 600_000 _QUANT_IGNORE_NAME_SUFFIXES = ( ".weight", ".weight_scale", ".weight_scale_2", + ".input_scale", ) # Layers kept in native dtype by the real-quant vLLM rollout. Shared between the @@ -37,6 +41,9 @@ "*self_attn*", ] +NVFP4RealQuantMode = Literal["w4a4", "w4a16"] +_NVFP4_REAL_QUANT_MODES = frozenset({"w4a4", "w4a16"}) + def _iter_quant_ignore_suffix_variants(name: str) -> Iterator[str]: """Yield ``name`` and, if it ends in a known quant suffix, the stripped form.""" @@ -71,6 +78,7 @@ def matches_quant_ignore_pattern(name: str, patterns: list[str]) -> bool: def build_vllm_modelopt_nvfp4_config( *, + mode: NVFP4RealQuantMode, ignore: list[str] | None = None, ) -> dict[str, Any]: """Build the HuggingFace quantization_config consumed by vLLM ModelOpt NVFP4. @@ -79,28 +87,198 @@ def build_vllm_modelopt_nvfp4_config( ``mtq.quantize``. vLLM expects the deployment/export-side ``quantization_config`` shape instead. """ - return { - "quant_method": "modelopt", - "config_groups": { - "group_0": { - "input_activations": None, - "weights": { - "dynamic": False, - "num_bits": 4, - "type": "float", - "group_size": 16, - }, - "targets": ["Linear"], - } - }, - "ignore": ignore if ignore is not None else list(DEFAULT_NVFP4_IGNORE), - "quant_algo": "NVFP4", - "quant_mode": "w4a16_nvfp4", - "weight_only": True, - "group_size": 16, - "producer": {"name": "modelopt"}, + from modelopt.torch.export.convert_hf_config import ( + convert_hf_quant_config_format, + ) + + if mode not in _NVFP4_REAL_QUANT_MODES: + raise ValueError( + f"Unsupported NVFP4 real-quant mode {mode!r}; expected 'w4a4' or 'w4a16'." + ) + return convert_hf_quant_config_format( + { + "producer": {"name": "modelopt"}, + "quantization": { + "quant_algo": "NVFP4" if mode == "w4a4" else "W4A16_NVFP4", + "group_size": 16, + "exclude_modules": ( + ignore if ignore is not None else list(DEFAULT_NVFP4_IGNORE) + ), + }, + } + ) + + +def _resolve_effective_quantizer_formats( + quant_cfg: Sequence[Mapping[str, Any]], + *, + source: str, +) -> tuple[list[object], list[object]]: + """Resolve enabled weight and input formats from ordered ModelOpt entries.""" + states: dict[str, dict[str, tuple[bool, object | None]]] = { + "weight_quantizer": {}, + "input_quantizer": {}, } + def _updated_state( + current: tuple[bool, object | None], + entry: Mapping[str, Any], + ) -> tuple[bool, object | None]: + enabled, format_cfg = current + if entry.get("cfg") is not None: + format_cfg = entry["cfg"] + enabled = entry["enable"] + return enabled, format_cfg + + for entry in quant_cfg: + pattern = entry["quantizer_name"] + + # Parent-scoped overrides describe exclusions, not a model-wide format. + if entry.get("parent_class") is not None: + continue + + if pattern == "*": + if entry.get("cfg") is not None: + raise ValueError( + f"Real quantization for {source!r} cannot infer weight and " + "activation formats from an enabled catch-all quantizer entry." + ) + for kind_states in states.values(): + for existing_pattern, current in tuple(kind_states.items()): + kind_states[existing_pattern] = _updated_state( + current, + entry, + ) + continue + + matching_kinds = [kind for kind in states if kind in pattern] + if not matching_kinds: + continue + if len(matching_kinds) != 1: + raise ValueError( + f"Quantization config {source!r} has ambiguous quantizer pattern " + f"{pattern!r}." + ) + + kind = matching_kinds[0] + kind_states = states[kind] + if pattern == f"*{kind}": + # A generic selector overrides every previously described subset. + for existing_pattern, current in tuple(kind_states.items()): + kind_states[existing_pattern] = _updated_state( + current, + entry, + ) + kind_states[pattern] = _updated_state( + kind_states.get(pattern, (False, None)), + entry, + ) + + weight_formats = [ + format_cfg + for enabled, format_cfg in states["weight_quantizer"].values() + if enabled + ] + input_formats = [ + format_cfg + for enabled, format_cfg in states["input_quantizer"].values() + if enabled + ] + return weight_formats, input_formats + + +def _is_float_format(value: object, exponent_bits: int, mantissa_bits: int) -> bool: + """Return whether a ModelOpt numeric-format value names the requested float.""" + if isinstance(value, str): + return value.lower() == f"e{exponent_bits}m{mantissa_bits}" + return value == (exponent_bits, mantissa_bits) or value == [ + exponent_bits, + mantissa_bits, + ] + + +def _validate_nvfp4_quantizer_format( + format_cfg: object, + *, + quantizer_name: str, + source: str, +) -> None: + """Validate the block-16 E2M1 format supported by vLLM ModelOpt NVFP4.""" + if not isinstance(format_cfg, Mapping): + raise ValueError( + f"Real quantization for {source!r} requires a single NVFP4 " + f"{quantizer_name} format; got {format_cfg!r}." + ) + + block_sizes = format_cfg.get("block_sizes") + block_size = None + block_type = None + scale_bits = None + if isinstance(block_sizes, Mapping): + block_size = block_sizes.get(-1, block_sizes.get("-1")) + block_type = block_sizes.get("type") + scale_bits = block_sizes.get("scale_bits") + + if not ( + _is_float_format(format_cfg.get("num_bits"), 2, 1) + and block_size == 16 + and block_type == "dynamic" + and _is_float_format(scale_bits, 4, 3) + ): + raise ValueError( + f"Real quantization for {source!r} supports only block-16 NVFP4 " + f"(E2M1 with E4M3 dynamic scales) {quantizer_name}; got " + f"{dict(format_cfg)!r}." + ) + + +def resolve_nvfp4_real_quant_mode(quant_cfg: str) -> NVFP4RealQuantMode: + """Resolve a ModelOpt training config to its supported real-quant mode. + + The mode is derived from all effective weight and input quantizer entries, + including model-specific selectors, rather than from a config name. NVFP4 + weights with no enabled input quantizer resolve to W4A16; block-16 NVFP4 + input quantization resolves to W4A4. FP8, W4A8, sequential, mixed, and other + activation formats fail loudly because vLLM would otherwise run a kernel + with incompatible semantics. + """ + if not isinstance(quant_cfg, str) or not quant_cfg: + raise ValueError("NVFP4 real quantization requires a non-empty quant_cfg.") + + from modelopt.torch.quantization.config import QuantizeConfig + + resolved = resolve_quant_cfg(quant_cfg) + normalized = QuantizeConfig(**resolved) + entries = [ + entry.model_dump(mode="python", exclude_none=True) + for entry in normalized.quant_cfg + ] + + weight_formats, input_formats = _resolve_effective_quantizer_formats( + entries, source=quant_cfg + ) + if not weight_formats: + raise ValueError( + f"Real quantization for {quant_cfg!r} requires enabled NVFP4 weights." + ) + for weight_format in weight_formats: + _validate_nvfp4_quantizer_format( + weight_format, + quantizer_name="weights", + source=quant_cfg, + ) + + if not input_formats: + return "w4a16" + + for input_format in input_formats: + _validate_nvfp4_quantizer_format( + input_format, + quantizer_name="input activations", + source=quant_cfg, + ) + return "w4a4" + def resolve_quant_cfg(quant_cfg: str) -> dict[str, Any]: """Resolve a quantization config string into a dict consumable by ``mtq.quantize``. diff --git a/nemo_rl/models/generation/megatron/megatron_worker.py b/nemo_rl/models/generation/megatron/megatron_worker.py index 71148b06b25..367c96cb1e9 100644 --- a/nemo_rl/models/generation/megatron/megatron_worker.py +++ b/nemo_rl/models/generation/megatron/megatron_worker.py @@ -385,6 +385,12 @@ def prepare_for_generation(self, tags=None, **kwargs) -> None: self.model = self.move_model( self.model, "cuda", move_params=True, move_grads=False ) + # DP inference schedules requests independently, so a forward pre-hook + # cannot safely launch a parameter all-gather from only the rank that + # received work. Gather once across every worker, then keep the hooks + # disabled until the next training step completes. + if self._forward_pre_hook_enabled(): + self._disable_forward_pre_hook_until_next_train_step(param_sync=True) lang_module = unwrap_model(self.model) lang_module.eval() diff --git a/nemo_rl/models/generation/vllm/vllm_backend.py b/nemo_rl/models/generation/vllm/vllm_backend.py index 8344b77330c..7b4310ce9b2 100644 --- a/nemo_rl/models/generation/vllm/vllm_backend.py +++ b/nemo_rl/models/generation/vllm/vllm_backend.py @@ -15,7 +15,9 @@ import re import socket import traceback -from typing import Any +from collections.abc import Callable, Iterable, Iterator, Sequence +from contextlib import contextmanager +from typing import Any, Literal import torch import zmq @@ -39,6 +41,66 @@ ) +WeightUpdateTransport = Literal["ipc", "collective"] +WeightUpdateFinalizer = Callable[[], None] + + +def _format_refit_key_error(label: str, keys: set[str]) -> str: + """Format a bounded refit-key diagnostic.""" + ordered = sorted(keys) + suffix = " ..." if len(ordered) > 8 else "" + return f"{label} ({len(ordered)}): {ordered[:8]}{suffix}" + + +class IPCWeightManifestError(RuntimeError): + """An IPC transfer did not match the prepared state-dict manifest.""" + + +class _IPCWeightManifest: + """Validate an IPC stream against its prepared state-dict manifest.""" + + def __init__(self, expected_keys: Iterable[str]) -> None: + self.expected_keys = set(expected_keys) + self.loaded_keys: set[str] = set() + self.errors: list[str] = [] + + def validate_batch(self, keys: Sequence[str]) -> set[str] | None: + batch_keys: set[str] = set() + duplicate_keys: set[str] = set() + for key in keys: + if key in batch_keys: + duplicate_keys.add(key) + batch_keys.add(key) + duplicate_keys.update(self.loaded_keys & batch_keys) + unexpected_keys = batch_keys - self.expected_keys + if duplicate_keys: + self.errors.append( + _format_refit_key_error("duplicate keys", duplicate_keys) + ) + if unexpected_keys: + self.errors.append( + _format_refit_key_error("unexpected keys", unexpected_keys) + ) + return None if self.errors else batch_keys + + def record_loaded(self, keys: set[str]) -> None: + self.loaded_keys.update(keys) + + def record_load_failure(self, error: Exception) -> None: + message = f"{type(error).__name__}: {error}" + if len(message) > 512: + message = message[:512] + " ..." + self.errors.append(f"weight load failed: {message}") + + def require_complete(self) -> None: + details = list(self.errors) + missing_keys = self.expected_keys - self.loaded_keys + if missing_keys: + details.append(_format_refit_key_error("missing keys", missing_keys)) + if details: + raise IPCWeightManifestError("; ".join(details)) + + def fix_gemma3_vision_weight_name(key: str) -> str: """Re-insert the `vision_model` segment into Gemma3 vision-tower weights. @@ -180,18 +242,16 @@ def prepare_sparse_delta_refit_info( applier = self._get_sparse_delta_applier() return sorted(applier.discover_native_skips(state_dict_info)) + def _uses_fp8_kv_cache(self) -> bool: + """Return whether this worker owns an FP8 KV cache.""" + vllm_config = getattr(self.model_runner, "vllm_config", None) + cache_config = getattr(vllm_config, "cache_config", None) + kv_cache_dtype = getattr(cache_config, "cache_dtype", None) + return kv_cache_dtype is not None and "fp8" in str(kv_cache_dtype).lower() + def _maybe_process_fp8_kv_cache(self) -> None: """Process weights after loading for FP8 KV cache (static scales).""" - use_fp8_kv_cache = False - if hasattr(self.model_runner.vllm_config, "cache_config"): - kv_cache_dtype = getattr( - self.model_runner.vllm_config.cache_config, "cache_dtype", None - ) - use_fp8_kv_cache = ( - kv_cache_dtype is not None and "fp8" in str(kv_cache_dtype).lower() - ) - - if not use_fp8_kv_cache: + if not self._uses_fp8_kv_cache(): return # FP8 KV cache: process KV scales after weight loading @@ -380,6 +440,36 @@ def _get_sparse_delta_applier(self) -> Any: ) return self._sparse_delta_applier + @contextmanager + def _weight_update_lifecycle( + self, transport: WeightUpdateTransport + ) -> Iterator[WeightUpdateFinalizer]: + """Provide setup/finalization around a transport-owned weight update.""" + del transport + from vllm.config import set_current_vllm_config + from vllm.model_executor.model_loader.utils import ( + process_weights_after_loading, + ) + + def finalize() -> None: + with set_current_vllm_config(self.model_runner.vllm_config): + process_weights_after_loading( + self.model_runner.model, self.model_config, self.device + ) + + yield finalize + # Preserve the IPC lifetime boundary: the COMPLETE ACK is sent before + # this optional second pass, just as it was before lifecycle hooks. + self._maybe_process_fp8_kv_cache() + + def _weight_update_errors_are_fatal(self) -> bool: + """Whether transport errors should propagate instead of returning False.""" + return False + + def _synchronize_before_ipc_data_ack(self) -> None: + """Fence work consuming one IPC data batch before its acknowledgment.""" + torch.cuda.current_stream().synchronize() + @wrap_with_nvtx_name("vllm_internal_worker_extension/update_weights_via_ipc_zmq") def update_weights_via_ipc_zmq(self) -> bool: """Receive and update model weights via ZMQ IPC socket. @@ -388,79 +478,90 @@ def update_weights_via_ipc_zmq(self) -> bool: bool: True if weights were successfully updated. """ buffer = None + weight = None weights = None try: self.maybe_init_zmq() - while True: - # Blocking receive with timeout (this is the main operation) - payload = self.zmq_socket.recv_pyobj() - - if payload == IPCProtocol.COMPLETE: - # means the update is done - from vllm.config import set_current_vllm_config - from vllm.model_executor.model_loader.utils import ( - process_weights_after_loading, - ) - - with set_current_vllm_config(self.model_runner.vllm_config): - process_weights_after_loading( - self.model_runner.model, self.model_config, self.device + manifest = _IPCWeightManifest(self.state_dict_info) + with self._weight_update_lifecycle("ipc") as finalize: + while True: + # Blocking receive with timeout (this is the main operation) + payload = self.zmq_socket.recv_pyobj() + + if payload == IPCProtocol.COMPLETE: + # A REP socket must reply even when validation or finalization + # fails, otherwise the sender remains blocked until timeout. + try: + manifest.require_complete() + finalize() + finally: + self.zmq_socket.send(IPCProtocol.ACK.value.encode()) + break + + batch_keys = None + batch_error = None + try: + ipc_handle, list_keys, used_bytes = payload + batch_keys = manifest.validate_batch(list_keys) + if batch_keys is None: + continue + + buffer = rebuild_cuda_tensor_from_ipc( + ipc_handle, self.device.index ) - self.zmq_socket.send(IPCProtocol.ACK.value.encode()) - break - - ipc_handle, list_keys, used_bytes = payload - buffer = rebuild_cuda_tensor_from_ipc(ipc_handle, self.device.index) - - weight = None - weights = [] - offset = 0 - for key in list_keys: - shape, dtype = self.state_dict_info[key] # pyrefly - if isinstance(shape, list): - shape = torch.Size(shape) - - # Get the weight from the buffer - size_in_bytes = dtype.itemsize * shape.numel() - weight = ( - buffer[offset : offset + size_in_bytes] - .view(dtype=dtype) - .view(shape) - ) - weights.append((key, weight)) - - # Move offset to the next weight - aligned_size = calculate_aligned_size(size_in_bytes) - offset += aligned_size - - assert offset == used_bytes, ( - "Offset is not equal to used bytes, usually indicate inaccurate info like keys or cached dtype in state_dict_info" - ) - - # Load weights into the model - self._load_weights(weights) - - torch.cuda.current_stream().synchronize() - - # CRITICAL: Delete views before ACK to prevent corruption. - # 'weights' contains views into IPC shared memory. Even though load_weights() - # copied the data, Python may not garbage collect these view objects immediately. - # If sender reuses the buffer before GC runs, old views would read corrupted data. - # Explicit del ensures immediate cleanup before sending ACK. - del weight, weights, buffer - weight = None - weights = None - buffer = None - self.zmq_socket.send(IPCProtocol.ACK.value.encode()) - - # Process weights after loading for FP8 KV cache - self._maybe_process_fp8_kv_cache() + weights = [] + offset = 0 + for key in list_keys: + shape, dtype = self.state_dict_info[key] # pyrefly + if isinstance(shape, list): + shape = torch.Size(shape) + + size_in_bytes = dtype.itemsize * shape.numel() + weight = ( + buffer[offset : offset + size_in_bytes] + .view(dtype=dtype) + .view(shape) + ) + weights.append((key, weight)) + offset += calculate_aligned_size(size_in_bytes) + + assert offset == used_bytes, ( + "Offset is not equal to used bytes, usually indicate " + "inaccurate info like keys or cached dtype in " + "state_dict_info" + ) + self._load_weights(weights) + except Exception as error: + batch_error = error + finally: + # Synchronize before releasing or ACKing an IPC allocation, + # including when a loader failed after scheduling CUDA work. + if buffer is not None: + try: + self._synchronize_before_ipc_data_ack() + except Exception as error: + if batch_error is None: + batch_error = error + + if batch_error is not None: + manifest.record_load_failure(batch_error) + elif batch_keys is not None: + manifest.record_loaded(batch_keys) + + # Drop every view before ACK permits sender-side reuse. + del weight, weights, buffer + weight = None + weights = None + buffer = None + self.zmq_socket.send(IPCProtocol.ACK.value.encode()) gc.collect() torch.cuda.empty_cache() return True except Exception as e: + if self._weight_update_errors_are_fatal(): + raise print( f"Error in VllmInternalWorkerExtension.update_weights_via_ipc_zmq: {e}.\n" f"{traceback.format_exc()}" @@ -477,29 +578,19 @@ def update_weights_from_collective(self) -> bool: "Please call prepare_refit_info when initializing the worker." ) - load_model_weight_func = self._load_weights - try: - packed_broadcast_consumer( - iterator=iter(self.state_dict_info.items()), - group=self.model_update_group, - src=0, - post_unpack_func=load_model_weight_func, - ) - - # Process weights after loading - from vllm.config import set_current_vllm_config - from vllm.model_executor.model_loader.utils import ( - process_weights_after_loading, - ) - - with set_current_vllm_config(self.model_runner.vllm_config): - process_weights_after_loading( - self.model_runner.model, self.model_config, self.device + with self._weight_update_lifecycle("collective") as finalize: + packed_broadcast_consumer( + iterator=iter(self.state_dict_info.items()), + group=self.model_update_group, + src=0, + post_unpack_func=self._load_weights, ) - self._maybe_process_fp8_kv_cache() + finalize() except Exception as e: + if self._weight_update_errors_are_fatal(): + raise print( f"Error in VllmInternalWorkerExtension.update_weights_from_collective: {e}" ) diff --git a/nemo_rl/models/megatron/community_import.py b/nemo_rl/models/megatron/community_import.py index 2d5f866a68f..631bd3673ed 100644 --- a/nemo_rl/models/megatron/community_import.py +++ b/nemo_rl/models/megatron/community_import.py @@ -154,8 +154,14 @@ def import_model_from_hf_name( ] if transformer_layer_spec is not None: model_provider.transformer_layer_spec = transformer_layer_spec - if mamba_stack_spec is not None and hasattr(model_provider, "mamba_stack_spec"): - model_provider.mamba_stack_spec = mamba_stack_spec + if mamba_stack_spec is not None: + # HybridModelProvider superseded the deprecated Mamba-only field. A + # MambaModelProvider normalizes mamba_stack_spec only in __post_init__, + # so assignments made here must target the canonical field directly. + if hasattr(model_provider, "hybrid_stack_spec"): + model_provider.hybrid_stack_spec = mamba_stack_spec + elif hasattr(model_provider, "mamba_stack_spec"): + model_provider.mamba_stack_spec = mamba_stack_spec model_provider.finalize() from megatron.core import parallel_state diff --git a/nemo_rl/models/megatron/setup.py b/nemo_rl/models/megatron/setup.py index a88cd334114..1403ea9bd9b 100644 --- a/nemo_rl/models/megatron/setup.py +++ b/nemo_rl/models/megatron/setup.py @@ -46,6 +46,7 @@ initialize_megatron, set_jit_fusion_options, ) +from megatron.bridge.training.model_load_save import load_model_config from megatron.bridge.training.optim import setup_optimizer from megatron.bridge.training.setup import ( _create_peft_pre_wrap_hook, @@ -54,7 +55,6 @@ from megatron.bridge.training.state import GlobalState from megatron.bridge.training.tokenizers.tokenizer import build_tokenizer from megatron.bridge.training.utils.pg_utils import get_pg_collection -from megatron.bridge.utils.instantiate_utils import InstantiationMode from megatron.bridge.utils.vocab_utils import calculate_padded_vocab_size from megatron.core import parallel_state from megatron.core.process_groups_config import ProcessGroupCollection @@ -519,9 +519,9 @@ def setup_model_config( _patch_hf_config_double_instantiation() try: - cfg_from_pretrained = ConfigContainer.from_yaml( - pretrained_run_config, mode=InstantiationMode.STRICT - ) + # Enter through Bridge's checkpoint loader so its compatibility + # migrations run before the serialized model config is instantiated. + model_cfg, _ = load_model_config(os.path.dirname(pretrained_run_config)) except Exception as e: # Add helpful context as a note to the exception e.add_note( @@ -536,9 +536,6 @@ def setup_model_config( ) raise - model_cfg = cfg_from_pretrained.model - cfg_from_pretrained.logger = LoggerConfig() - # Apply parallelism settings _apply_parallelism_config(model_cfg, config) @@ -704,7 +701,11 @@ def _apply_moe_config(model_cfg: Any, config: PolicyConfig) -> None: "moe_flex_dispatcher_backend" ] if "moe_hybridep_num_sms" in config["megatron_cfg"]: - model_cfg.moe_hybridep_num_sms = config["megatron_cfg"]["moe_hybridep_num_sms"] + num_sms = config["megatron_cfg"]["moe_hybridep_num_sms"] + if hasattr(TransformerConfig, "moe_flex_dispatcher_num_sms"): + model_cfg.moe_flex_dispatcher_num_sms = num_sms + else: + model_cfg.moe_hybridep_num_sms = num_sms # HybridEP environment variables # These are required by DeepEP's hybrid-ep branch for NVLink domain configuration. diff --git a/nemo_rl/models/policy/utils.py b/nemo_rl/models/policy/utils.py index aaa90531c83..a421d9bda6f 100644 --- a/nemo_rl/models/policy/utils.py +++ b/nemo_rl/models/policy/utils.py @@ -415,6 +415,17 @@ def pack_tensor(buffer, tensor, used_bytes) -> int: buffer_b: torch.Tensor | None = None current_buffer: torch.Tensor | None = None + def release_staging_buffers() -> None: + """Release acyclic IPC buffers without scanning the worker object graph.""" + nonlocal buffer_a, buffer_b, current_buffer + + had_buffers = buffer_a is not None or buffer_b is not None + current_buffer = None + buffer_a = None + buffer_b = None + if had_buffers: + torch.cuda.empty_cache() + used_bytes = 0 param_names = [] await_recv = False @@ -462,8 +473,13 @@ def pack_tensor(buffer, tensor, used_bytes) -> int: if await_recv: zmq_socket.recv() - # Final synchronization and completion signal + # The receiver synchronizes and drops every IPC view before ACKing a + # group, so the final data ACK is the staging buffers' safe lifetime + # boundary. Reclaim them before asking the receiver to run its final + # post-load conversion, which can otherwise retain both large buffers + # for the whole conversion and amplify a single-rank tail. torch.cuda.current_stream().synchronize() + release_staging_buffers() zmq_socket.send_pyobj(IPCProtocol.COMPLETE) zmq_socket.recv() @@ -488,15 +504,10 @@ def pack_tensor(buffer, tensor, used_bytes) -> int: ) from e finally: - # Clean up buffers in finally block to ensure cleanup even on exceptions - if buffer_a is not None: - del buffer_a - if buffer_b is not None: - del buffer_b - - # Force garbage collection and clear CUDA cache - gc.collect() - torch.cuda.empty_cache() + # Tensor references are acyclic and deterministic; a full gc.collect() + # scans the entire model object graph and can become a multi-second + # rank straggler without releasing anything that refcounting cannot. + release_staging_buffers() def rebuild_cuda_tensor_from_ipc( diff --git a/nemo_rl/models/policy/workers/megatron_policy_worker.py b/nemo_rl/models/policy/workers/megatron_policy_worker.py index dea0c079a14..cf0e2621ed3 100644 --- a/nemo_rl/models/policy/workers/megatron_policy_worker.py +++ b/nemo_rl/models/policy/workers/megatron_policy_worker.py @@ -484,10 +484,12 @@ def _forward_pre_hook_enabled(self) -> bool: return False return len(getattr(self.model, "remove_forward_pre_hook_handles", {})) > 0 - def _disable_forward_pre_hook_until_next_train_step(self) -> None: + def _disable_forward_pre_hook_until_next_train_step( + self, *, param_sync: bool = False + ) -> None: assert isinstance(self.model, DistributedDataParallel) if self._forward_pre_hook_enabled(): - self.disable_forward_pre_hook(param_sync=False) + self.disable_forward_pre_hook(param_sync=param_sync) model_config = get_model_config(self.model) self._first_train_step_param_sync_func = model_config.param_sync_func model_config.param_sync_func = None @@ -2005,9 +2007,6 @@ def broadcast_weights_for_collective( post_iter_func=lambda x: x[1], ) - def _use_real_quant_refit(self) -> bool: - return False - def prepare_for_lp_inference(self): self.model = self.move_model(self.model, "cuda", move_grads=False) self.model.eval() @@ -2230,14 +2229,7 @@ def move_model( else: # Ordinary offload case if move_params: - new_state_dict = {} - for name, item in model.state_dict().items(): - if isinstance(item, torch.Tensor): - item = item.detach().to( - device=device, non_blocking=True, copy=True - ) - new_state_dict[name] = item - model.load_state_dict(new_state_dict) + model.to(device=device, non_blocking=True) return model def move_optimizer(self, device: str): diff --git a/nemo_rl/utils/checkpoint.py b/nemo_rl/utils/checkpoint.py index ba92251fbe0..8c8de3129d9 100644 --- a/nemo_rl/utils/checkpoint.py +++ b/nemo_rl/utils/checkpoint.py @@ -46,6 +46,21 @@ PathLike = Union[str, "os.PathLike[Any]"] +def _load_megatron_common_state_dict(iteration_dir: Path) -> dict[str, Any]: + """Load common state from either legacy or current MCore checkpoints.""" + # Keep the optional MCore dependency out of DTensor and Automodel imports. + try: + from megatron.core.dist_checkpointing import load_common_state_dict + except ImportError as error: + raise RuntimeError( + "Megatron-Core is required to inspect optimizer state in the distributed " + f"checkpoint at {iteration_dir}. Install NeMo-RL with the `mcore` extra." + ) from error + + # MCore accepts Path today but deprecates it in favor of str. + return load_common_state_dict(str(iteration_dir)) + + class PretrainedCheckpointConfig(TypedDict): """Configuration for restoring initial weights from a pre-existing Megatron checkpoint. @@ -198,11 +213,16 @@ def get_resume_paths( if optimizer_path.exists(): return weights_path, optimizer_path - # Megatron path - common_pt_path = weights_path / "iter_0000000" / "common.pt" - if common_pt_path.exists(): - common_pt_obj = torch.load(common_pt_path, map_location="cpu") - if "optimizer" in common_pt_obj: + # Megatron path. MCore's public loader supports both legacy checkpoints, + # which store common state in common.pt, and current torch_dist + # checkpoints, which embed it as a common_state ShardedObject. + iteration_dir = weights_path / "iter_0000000" + is_megatron_checkpoint = (iteration_dir / "common.pt").exists() or ( + iteration_dir / "metadata.json" + ).exists() + if is_megatron_checkpoint: + common_state_dict = _load_megatron_common_state_dict(iteration_dir) + if "optimizer" in common_state_dict: # In Megatron, optimizer_path is only a flag to indicate that the optimizer # state is embedded in the weights_path. We will actually load the optimizer # state from the weights_path. diff --git a/pyrefly.toml b/pyrefly.toml index 75315f5d862..55e57fd32a9 100644 --- a/pyrefly.toml +++ b/pyrefly.toml @@ -142,6 +142,7 @@ project-includes = [ "nemo_rl/modelopt/__init__.py", "nemo_rl/modelopt/models/__init__.py", "nemo_rl/modelopt/models/generation/__init__.py", + "nemo_rl/modelopt/models/generation/vllm_modelopt.py", "nemo_rl/modelopt/models/generation/vllm_quant_worker.py", "nemo_rl/modelopt/models/policy/__init__.py", "nemo_rl/modelopt/models/policy/workers/__init__.py", diff --git a/tests/functional/_bridge_to_mlm_helper.py b/tests/functional/_bridge_to_mlm_helper.py index 0d5a4a633a0..bf410d65d39 100644 --- a/tests/functional/_bridge_to_mlm_helper.py +++ b/tests/functional/_bridge_to_mlm_helper.py @@ -16,13 +16,17 @@ The two formats share torch_dist sharded weights. They differ only in metadata: - bridge: iter_*/run_config.yaml + common.pt (no ``args``) + bridge: iter_*/run_config.yaml + common state (no ``args``) MLM: iter_*/common.pt with ``args`` (no run_config.yaml) The conversion is lossless for our purposes: copy/symlink everything, drop ``run_config.yaml``, and inject an ``argparse.Namespace`` into ``common.pt`` populated with the TP/PP fields bridge's MLM-load path consults. +Bridge checkpoints may store common state either in the legacy ``common.pt`` +file or, with current MCore, inside the torch_dist checkpoint. MCore's public +loader handles both layouts. + This script exists to support functional testing of the ``pretrained_checkpoint.format=megatron_lm`` code path without requiring an upstream MLM training run to produce a fixture. @@ -34,6 +38,7 @@ import torch import yaml +from megatron.core import dist_checkpointing def bridge_to_mlm(bridge_iter_dir: str, mlm_iter_dir: str) -> None: @@ -59,11 +64,7 @@ def bridge_to_mlm(bridge_iter_dir: str, mlm_iter_dir: str) -> None: model_cfg: dict[str, Any] = run_config.get("model", {}) or {} ckpt_cfg: dict[str, Any] = run_config.get("checkpoint", {}) or {} - common = torch.load( - os.path.join(bridge_iter_dir, "common.pt"), - map_location="cpu", - weights_only=False, - ) + common = dist_checkpointing.load_common_state_dict(str(bridge_iter_dir)) # Bridge's _extract_megatron_lm_args_from_state_dict reads these via # getattr-with-defaults; only TP/PP must be accurate, the rest are flags # that don't affect a pretrained_checkpoint load. diff --git a/tests/functional/grpo_megatron_mbridge_restore.sh b/tests/functional/grpo_megatron_mbridge_restore.sh index 490db81fca5..15f624ece5a 100755 --- a/tests/functional/grpo_megatron_mbridge_restore.sh +++ b/tests/functional/grpo_megatron_mbridge_restore.sh @@ -144,6 +144,10 @@ run_restore_phase \ echo "[INFO] Phase 3: convert bridge → MLM and restore" MLM_CKPT="$EXP_DIR/mlm_ckpt/iter_0000000" trap "rm -rf $EXP_DIR/mlm_ckpt" EXIT +# Current torch_dist checkpoints store common state through MCore, so the +# converter needs the checked-out Megatron-LM package on its import path. +MEGATRON_LM_SRC="$PROJECT_ROOT/3rdparty/Megatron-Bridge-workspace/Megatron-Bridge/3rdparty/Megatron-LM" +PYTHONPATH="$MEGATRON_LM_SRC:${PYTHONPATH:-}" \ uv run --no-sync python "$SCRIPT_DIR/_bridge_to_mlm_helper.py" \ --bridge-iter-dir "$BRIDGE_CKPT" \ --mlm-iter-dir "$MLM_CKPT" diff --git a/tests/functional/modelopt_quant_rollout.sh b/tests/functional/modelopt_quant_rollout.sh index 857103ad654..91ce14809c1 100644 --- a/tests/functional/modelopt_quant_rollout.sh +++ b/tests/functional/modelopt_quant_rollout.sh @@ -60,7 +60,7 @@ run_quant_rollout_case() { cd "$PROJECT_ROOT" NRL_MEGATRON_CHECKPOINT_DIR="$megatron_cache_dir" \ - uv run --extra modelopt --group test \ + uv run --no-sync --extra modelopt --group test \ coverage run -a --data-file="$PROJECT_ROOT/tests/.coverage" --source="$PROJECT_ROOT/nemo_rl" \ "$PROJECT_ROOT/examples/run_grpo.py" \ --config "$PROJECT_ROOT/examples/modelopt/qa_grpo_math_megatron.yaml" \ @@ -96,9 +96,9 @@ run_quant_rollout_case() { "$@" \ 2>&1 | tee "$run_log" - uv run --extra modelopt --group test tests/json_dump_tb_logs.py "$log_dir" --output_path "$metrics_json" + uv run --no-sync --extra modelopt --group test tests/json_dump_tb_logs.py "$log_dir" --output_path "$metrics_json" - uv run --extra modelopt --group test tests/check_metrics.py "$metrics_json" \ + uv run --no-sync --extra modelopt --group test tests/check_metrics.py "$metrics_json" \ "data[\"train/gen_kl_error\"][\"1\"] < $gen_kl_error_step1_max" \ "max(data[\"train/token_mult_prob_error\"]) < $token_mult_prob_error_max" @@ -109,12 +109,12 @@ run_quant_rollout_case() { assert_not_grep "VLLM_QUANT_CFG" "$run_log" else assert_grep "FakeQuantWorker" "$run_log" - assert_grep "VLLM_QUANT_CFG" "$run_log" + assert_grep "VllmQuantGenerationWorker.*Inserted [1-9][0-9]* quantizers" "$run_log" assert_not_grep "Detected ModelOpt NVFP4 checkpoint" "$run_log" fi } run_quant_rollout_case w4a16_real_quant examples/modelopt/quant_configs/nvfp4_a16_mlp_only.yaml true 0.003 1.05 Qwen/Qwen2.5-0.5B "$@" -run_quant_rollout_case w4a8_fake_quant examples/modelopt/quant_configs/nvfp4_w4a8_fp8.yaml false 0.006 1.06 Qwen/Qwen2.5-0.5B "$@" +run_quant_rollout_case w4a8_fake_quant examples/modelopt/quant_configs/nvfp4_w4a8_fp8.yaml false 0.02 1.15 Qwen/Qwen2.5-0.5B "$@" echo "[PASS] ModelOpt W4A16 real-quant and W4A8 fake-quant rollout functional test" diff --git a/tests/test_suites/llm/common.env b/tests/test_suites/llm/common.env index 427cd4a7de8..c3d363b0999 100644 --- a/tests/test_suites/llm/common.env +++ b/tests/test_suites/llm/common.env @@ -24,6 +24,16 @@ exit_if_max_steps_reached() { echo "[INFO] Steps so far: $STEPS_SO_FAR, running till $MAX_STEPS steps" } +assert_not_grep() { + local pattern=$1 + local file=$2 + local message=$3 + if grep -q -- "$pattern" "$file"; then + echo "[ERROR] $message" + exit 1 + fi +} + # EXP_NAME may be pre-exported by a caller (e.g. a TQ wrapper that delegates # to a base recipe script but wants its own log/ckpt dirs and wandb name). # If unset, fall back to the conventional derivation from $0. diff --git a/tests/test_suites/llm/grpo-nanov3-30ba3b-4n4g-megatron-qa-nvfp4-w4a16-real.sh b/tests/test_suites/llm/grpo-nanov3-30ba3b-4n4g-megatron-qa-nvfp4-w4a16-real.sh new file mode 100755 index 00000000000..ca587a2f417 --- /dev/null +++ b/tests/test_suites/llm/grpo-nanov3-30ba3b-4n4g-megatron-qa-nvfp4-w4a16-real.sh @@ -0,0 +1,57 @@ +#!/bin/bash +# One-step Nano3 W4A16 real-quant rollout check. This validates that Megatron +# exports ModelOpt NVFP4 packed tensors, vLLM loads them through the real +# ModelOpt kernel path, and generation/policy logprobs stay aligned. +SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd) +source "$SCRIPT_DIR/common.env" + +# ===== BEGIN CONFIG ===== +NUM_NODES=4 +GPUS_PER_NODE=4 +STEPS_PER_RUN=1 +MAX_STEPS=1 +NUM_RUNS=$(( (MAX_STEPS + STEPS_PER_RUN - 1) / STEPS_PER_RUN )) # Round up +NUM_MINUTES=180 +SNAPSHOT_MEGATRON_BRIDGE=1 +# ===== END CONFIG ===== + +exit_if_max_steps_reached + +cd "$PROJECT_ROOT" +uv run --no-sync examples/run_grpo.py \ + --config "$CONFIG_PATH" \ + grpo.max_num_steps=$MAX_STEPS \ + cluster.num_nodes=$NUM_NODES \ + cluster.gpus_per_node=$GPUS_PER_NODE \ + logger.log_dir="$LOG_DIR" \ + logger.wandb_enabled=True \ + logger.wandb.project=nemo-rl \ + logger.wandb.name="$EXP_NAME" \ + logger.monitor_gpus=True \ + logger.tensorboard_enabled=True \ + checkpointing.enabled=True \ + checkpointing.checkpoint_dir="$CKPT_DIR" \ + "$@" \ + 2>&1 | tee "$RUN_LOG" + +uv run --no-sync tests/json_dump_tb_logs.py "$LOG_DIR" --output_path "$JSON_METRICS" + +grep -q "VllmQuantInternalWorkerExtension" "$RUN_LOG" +grep -q "Detected ModelOpt NVFP4 checkpoint" "$RUN_LOG" +grep -q "quantization=modelopt" "$RUN_LOG" +assert_not_grep "FakeQuantWorker" "$RUN_LOG" \ + "Real-quant run unexpectedly used FakeQuantWorker" +assert_not_grep "VLLM_QUANT_CFG" "$RUN_LOG" \ + "Real-quant run unexpectedly took the fake-quant VLLM_QUANT_CFG path" + +MAX_RECORDED_STEP=$(jq -r 'if has("train/loss") then (."train/loss" | keys | map(tonumber) | max // 0) else 0 end' "$JSON_METRICS") +if [[ $MAX_RECORDED_STEP -lt $MAX_STEPS ]]; then + echo "[ERROR] Expected train/loss through step $MAX_STEPS, found step $MAX_RECORDED_STEP" + exit 1 +fi + +uv run --no-sync tests/check_metrics.py "$JSON_METRICS" \ + 'data["train/gen_kl_error"]["1"] < 0.003' \ + 'max(data["train/token_mult_prob_error"]) < 1.05' \ + 'data["train/loss"]["1"] > 0.0' \ + 'data["train/num_valid_samples"]["1"] > 0' diff --git a/tests/test_suites/llm/grpo-nemotron3-super-120BA12B-16n4g-megatron-qa-bf16-300step.sh b/tests/test_suites/llm/grpo-nemotron3-super-120BA12B-16n4g-megatron-qa-bf16-300step.sh new file mode 100755 index 00000000000..c759586f13c --- /dev/null +++ b/tests/test_suites/llm/grpo-nemotron3-super-120BA12B-16n4g-megatron-qa-bf16-300step.sh @@ -0,0 +1,50 @@ +#!/bin/bash +# One-step BF16 smoke for the public 300-step Super 120B-A12B recipe. +SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd) +source "$SCRIPT_DIR/common.env" + +# ===== BEGIN CONFIG ===== +NUM_NODES=16 +GPUS_PER_NODE=4 +SEGMENT_SIZE=16 +STEPS_PER_RUN=1 +MAX_STEPS=1 +NUM_RUNS=1 +NUM_MINUTES=240 +# ===== END CONFIG ===== + +exit_if_max_steps_reached + +cd "$PROJECT_ROOT" +uv run --no-sync examples/run_grpo.py \ + --config "$CONFIG_PATH" \ + grpo.max_num_steps=$MAX_STEPS \ + cluster.num_nodes=$NUM_NODES \ + cluster.gpus_per_node=$GPUS_PER_NODE \ + grpo.val_at_start=false \ + grpo.val_at_end=true \ + grpo.max_val_samples=16 \ + grpo.val_batch_size=16 \ + logger.log_dir="$LOG_DIR" \ + logger.wandb_enabled=True \ + logger.wandb.project=nemo-rl \ + logger.wandb.name="$EXP_NAME" \ + logger.monitor_gpus=True \ + logger.tensorboard_enabled=True \ + checkpointing.enabled=false \ + "$@" \ + 2>&1 | tee "$RUN_LOG" + +uv run --no-sync tests/json_dump_tb_logs.py "$LOG_DIR" --output_path "$JSON_METRICS" + +assert_not_grep "VllmQuantInternalWorkerExtension" "$RUN_LOG" \ + "BF16 baseline unexpectedly loaded the quant rollout worker extension" +assert_not_grep "FakeQuantWorker" "$RUN_LOG" \ + "BF16 baseline unexpectedly used FakeQuantWorker" + +uv run --no-sync tests/check_metrics.py "$JSON_METRICS" \ + 'data["train/num_valid_samples"]["1"] >= 64' \ + 'data["train/reward"]["1"] >= 0.4' \ + 'data["train/gen_kl_error"]["1"] < 0.03' \ + 'data["train/token_mult_prob_error"]["1"] < 1.05' \ + 'data["validation/accuracy"]["1"] >= 0.5' diff --git a/tests/test_suites/llm/grpo-nemotron3-super-120BA12B-16n4g-megatron-qa-nvfp4-w4a16-real-300step.sh b/tests/test_suites/llm/grpo-nemotron3-super-120BA12B-16n4g-megatron-qa-nvfp4-w4a16-real-300step.sh new file mode 100755 index 00000000000..bab884781e6 --- /dev/null +++ b/tests/test_suites/llm/grpo-nemotron3-super-120BA12B-16n4g-megatron-qa-nvfp4-w4a16-real-300step.sh @@ -0,0 +1,55 @@ +#!/bin/bash +# One-step W4A16 real-quant smoke for the public 300-step Super 120B-A12B recipe. +SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd) +source "$SCRIPT_DIR/common.env" + +# ===== BEGIN CONFIG ===== +NUM_NODES=16 +GPUS_PER_NODE=4 +SEGMENT_SIZE=16 +STEPS_PER_RUN=1 +MAX_STEPS=1 +NUM_RUNS=1 +NUM_MINUTES=240 +# ===== END CONFIG ===== + +exit_if_max_steps_reached + +cd "$PROJECT_ROOT" +uv run --no-sync examples/run_grpo.py \ + --config "$CONFIG_PATH" \ + grpo.max_num_steps=$MAX_STEPS \ + cluster.num_nodes=$NUM_NODES \ + cluster.gpus_per_node=$GPUS_PER_NODE \ + grpo.val_at_start=false \ + grpo.val_at_end=true \ + grpo.max_val_samples=16 \ + grpo.val_batch_size=16 \ + logger.log_dir="$LOG_DIR" \ + logger.wandb_enabled=True \ + logger.wandb.project=nemo-rl \ + logger.wandb.name="$EXP_NAME" \ + logger.monitor_gpus=True \ + logger.tensorboard_enabled=True \ + checkpointing.enabled=false \ + "$@" \ + 2>&1 | tee "$RUN_LOG" + +uv run --no-sync tests/json_dump_tb_logs.py "$LOG_DIR" --output_path "$JSON_METRICS" + +# Real-quant rollout must go through the ModelOpt NVFP4 vLLM kernel path. +grep -q "VllmQuantInternalWorkerExtension" "$RUN_LOG" +grep -q "Detected ModelOpt NVFP4 checkpoint" "$RUN_LOG" +grep -q "quantization=modelopt" "$RUN_LOG" +grep -q "nvfp4_experts_weightonly.yaml" "$RUN_LOG" +assert_not_grep "FakeQuantWorker" "$RUN_LOG" \ + "Real-quant run unexpectedly used FakeQuantWorker" +assert_not_grep "VLLM_QUANT_CFG" "$RUN_LOG" \ + "Real-quant run unexpectedly took the fake-quant VLLM_QUANT_CFG path" + +uv run --no-sync tests/check_metrics.py "$JSON_METRICS" \ + 'data["train/num_valid_samples"]["1"] >= 64' \ + 'data["train/reward"]["1"] >= 0.4' \ + 'data["train/gen_kl_error"]["1"] < 0.03' \ + 'data["train/token_mult_prob_error"]["1"] < 1.05' \ + 'data["validation/accuracy"]["1"] >= 0.5' diff --git a/tests/test_suites/llm/grpo-nemotron3-super-120BA12B-16n4g-megatron-qa-nvfp4-w4a4-real-300step.sh b/tests/test_suites/llm/grpo-nemotron3-super-120BA12B-16n4g-megatron-qa-nvfp4-w4a4-real-300step.sh new file mode 100755 index 00000000000..1a324e0e197 --- /dev/null +++ b/tests/test_suites/llm/grpo-nemotron3-super-120BA12B-16n4g-megatron-qa-nvfp4-w4a4-real-300step.sh @@ -0,0 +1,57 @@ +#!/bin/bash +# One-step W4A4 real-quant smoke for the public 300-step Super 120B-A12B recipe. +SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd) +source "$SCRIPT_DIR/common.env" + +# ===== BEGIN CONFIG ===== +NUM_NODES=16 +GPUS_PER_NODE=4 +SEGMENT_SIZE=16 +STEPS_PER_RUN=1 +MAX_STEPS=1 +NUM_RUNS=1 +NUM_MINUTES=240 +# ===== END CONFIG ===== + +exit_if_max_steps_reached + +cd "$PROJECT_ROOT" +uv run --no-sync examples/run_grpo.py \ + --config "$CONFIG_PATH" \ + grpo.max_num_steps=$MAX_STEPS \ + cluster.num_nodes=$NUM_NODES \ + cluster.gpus_per_node=$GPUS_PER_NODE \ + grpo.val_at_start=false \ + grpo.val_at_end=true \ + grpo.max_val_samples=16 \ + grpo.val_batch_size=16 \ + logger.log_dir="$LOG_DIR" \ + logger.wandb_enabled=True \ + logger.wandb.project=nemo-rl \ + logger.wandb.name="$EXP_NAME" \ + logger.monitor_gpus=True \ + logger.tensorboard_enabled=True \ + checkpointing.enabled=false \ + "$@" \ + 2>&1 | tee "$RUN_LOG" + +uv run --no-sync tests/json_dump_tb_logs.py "$LOG_DIR" --output_path "$JSON_METRICS" + +# Real-quant rollout must go through the ModelOpt NVFP4 vLLM kernel path. +grep -q "VllmQuantInternalWorkerExtension" "$RUN_LOG" +grep -q "Detected ModelOpt NVFP4 checkpoint" "$RUN_LOG" +grep -q "quantization=modelopt" "$RUN_LOG" +grep -q "examples/modelopt/quant_configs/nvfp4_experts.yaml" "$RUN_LOG" +assert_not_grep "FakeQuantWorker" "$RUN_LOG" \ + "Real-quant run unexpectedly used FakeQuantWorker" +assert_not_grep "VLLM_QUANT_CFG" "$RUN_LOG" \ + "Real-quant run unexpectedly took the fake-quant VLLM_QUANT_CFG path" +assert_not_grep "Using NvFp4LinearBackend.MARLIN" "$RUN_LOG" \ + "W4A4 run unexpectedly fell back to the weight-only Marlin backend" + +uv run --no-sync tests/check_metrics.py "$JSON_METRICS" \ + 'data["train/num_valid_samples"]["1"] >= 64' \ + 'data["train/reward"]["1"] >= 0.4' \ + 'data["train/gen_kl_error"]["1"] < 0.1' \ + 'data["train/token_mult_prob_error"]["1"] < 1.15' \ + 'data["validation/accuracy"]["1"] >= 0.5' diff --git a/tests/test_suites/llm/grpo-qwen3-30ba3b-4n4g-megatron-qa-nvfp4-w4a4-real.sh b/tests/test_suites/llm/grpo-qwen3-30ba3b-4n4g-megatron-qa-nvfp4-w4a4-real.sh new file mode 100755 index 00000000000..c87cbc43b55 --- /dev/null +++ b/tests/test_suites/llm/grpo-qwen3-30ba3b-4n4g-megatron-qa-nvfp4-w4a4-real.sh @@ -0,0 +1,70 @@ +#!/bin/bash +# Two-step GB200 smoke test for Megatron -> vLLM W4A4 real-quant refits. +SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd) +source "$SCRIPT_DIR/common.env" + +# ===== BEGIN CONFIG ===== +NUM_NODES=4 +GPUS_PER_NODE=4 +STEPS_PER_RUN=2 +MAX_STEPS=2 +NUM_RUNS=$(( (MAX_STEPS + STEPS_PER_RUN - 1) / STEPS_PER_RUN )) +NUM_MINUTES=180 +SNAPSHOT_MEGATRON_BRIDGE=1 +# ===== END CONFIG ===== + +exit_if_max_steps_reached + +cd "$PROJECT_ROOT" +uv run --no-sync examples/run_grpo.py \ + --config "$CONFIG_PATH" \ + grpo.max_num_steps=$MAX_STEPS \ + grpo.val_at_start=true \ + grpo.val_at_end=true \ + grpo.max_val_samples=32 \ + grpo.val_batch_size=32 \ + cluster.num_nodes=$NUM_NODES \ + cluster.gpus_per_node=$GPUS_PER_NODE \ + logger.log_dir="$LOG_DIR" \ + logger.wandb_enabled=True \ + logger.wandb.project=nemo-rl \ + logger.wandb.name="$EXP_NAME" \ + logger.monitor_gpus=True \ + logger.tensorboard_enabled=True \ + checkpointing.enabled=True \ + checkpointing.checkpoint_dir="$CKPT_DIR" \ + "$@" \ + 2>&1 | tee "$RUN_LOG" + +uv run --no-sync tests/json_dump_tb_logs.py "$LOG_DIR" --output_path "$JSON_METRICS" + +grep -q "VllmQuantInternalWorkerExtension" "$RUN_LOG" +grep -q "Detected ModelOpt NVFP4 checkpoint" "$RUN_LOG" +grep -q "quantization=modelopt" "$RUN_LOG" +assert_not_grep "FakeQuantWorker" "$RUN_LOG" \ + "Real-quant run unexpectedly used FakeQuantWorker" +assert_not_grep "VLLM_QUANT_CFG" "$RUN_LOG" \ + "Real-quant run unexpectedly took the fake-quant VLLM_QUANT_CFG path" +assert_not_grep "Using NvFp4LinearBackend.MARLIN" "$RUN_LOG" \ + "W4A4 run unexpectedly fell back to the weight-only Marlin backend" + +MAX_RECORDED_STEP=$(jq -r 'if has("train/loss") then (."train/loss" | keys | map(tonumber) | max // 0) else 0 end' "$JSON_METRICS") +if [[ $MAX_RECORDED_STEP -lt $MAX_STEPS ]]; then + echo "[ERROR] Expected train/loss through step $MAX_STEPS, found step $MAX_RECORDED_STEP" + exit 1 +fi + +uv run --no-sync tests/check_metrics.py "$JSON_METRICS" \ + "data[\"train/reward\"][\"$MAX_STEPS\"] >= 0.25" \ + "data[\"validation/accuracy\"][\"$MAX_STEPS\"] >= 0.4" \ + "data[\"train/gen_kl_error\"][\"$MAX_STEPS\"] < 0.03" \ + "data[\"train/js_divergence_error\"][\"$MAX_STEPS\"] < 0.007" \ + "data[\"train/approx_entropy\"][\"$MAX_STEPS\"] < 0.35" + +mapfile -t TRAIN_DATA_FILES < <( + find "$LOG_DIR" -type f -name 'train_data_step*.jsonl' -print | sort -V +) +if [[ ${#TRAIN_DATA_FILES[@]} -ne $MAX_STEPS ]]; then + echo "[ERROR] Expected $MAX_STEPS rollout JSONL files, found ${#TRAIN_DATA_FILES[@]}" + exit 1 +fi diff --git a/tests/test_suites/llm/grpo-qwen3-8b-base-dapo-2n8g-long-megatron-qa-nvfp4-w4a16.sh b/tests/test_suites/llm/grpo-qwen3-8b-base-dapo-2n8g-long-megatron-qa-nvfp4-w4a16.sh index 17bfe9ace30..c6fd070aec1 100755 --- a/tests/test_suites/llm/grpo-qwen3-8b-base-dapo-2n8g-long-megatron-qa-nvfp4-w4a16.sh +++ b/tests/test_suites/llm/grpo-qwen3-8b-base-dapo-2n8g-long-megatron-qa-nvfp4-w4a16.sh @@ -9,6 +9,7 @@ STEPS_PER_RUN=10 MAX_STEPS=10 NUM_RUNS=$(( (MAX_STEPS + STEPS_PER_RUN - 1) / STEPS_PER_RUN )) # Round up NUM_MINUTES=240 +SNAPSHOT_MEGATRON_BRIDGE=1 # ===== END CONFIG ===== exit_if_max_steps_reached @@ -32,18 +33,21 @@ uv run examples/run_grpo.py \ # Convert tensorboard logs to json uv run tests/json_dump_tb_logs.py $LOG_DIR --output_path $JSON_METRICS -# Only run metrics if the target step is reached -if [[ $(jq 'to_entries | .[] | select(.key == "train/loss") | .value | keys | map(tonumber) | max' $JSON_METRICS) -ge $MAX_STEPS ]]; then - uv run tests/check_metrics.py $JSON_METRICS \ - 'median(data["train/token_mult_prob_error"]) < 1.1' \ - 'max(data["train/gen_kl_error"]) < 0.003' \ - 'max(data["train/reward"]) > -0.9' +if ! grep -q "VllmQuantInternalWorkerExtension" "$RUN_LOG"; then echo "ERROR: VllmQuantInternalWorkerExtension not found in real-quant run" >&2; exit 1; fi +if ! grep -q "Detected ModelOpt NVFP4 checkpoint" "$RUN_LOG"; then echo "ERROR: 'Detected ModelOpt NVFP4 checkpoint' not found in real-quant run" >&2; exit 1; fi +if grep -q "FakeQuantWorker" "$RUN_LOG"; then echo "ERROR: FakeQuantWorker unexpectedly present in real-quant run" >&2; exit 1; fi +if grep -q "VLLM_QUANT_CFG" "$RUN_LOG"; then echo "ERROR: VLLM_QUANT_CFG unexpectedly present in real-quant run" >&2; exit 1; fi - if ! grep -q "VllmQuantInternalWorkerExtension" "$RUN_LOG"; then echo "ERROR: VllmQuantInternalWorkerExtension not found in real-quant run" >&2; exit 1; fi - if ! grep -q "Detected ModelOpt NVFP4 checkpoint" "$RUN_LOG"; then echo "ERROR: 'Detected ModelOpt NVFP4 checkpoint' not found in real-quant run" >&2; exit 1; fi - if grep -q "FakeQuantWorker" "$RUN_LOG"; then echo "ERROR: FakeQuantWorker unexpectedly present in real-quant run" >&2; exit 1; fi - if grep -q "VLLM_QUANT_CFG" "$RUN_LOG"; then echo "ERROR: VLLM_QUANT_CFG unexpectedly present in real-quant run" >&2; exit 1; fi - - # Clean up checkpoint directory after successful run to save space. - rm -rf "$CKPT_DIR" +MAX_RECORDED_STEP=$(jq -r 'if has("train/loss") then (."train/loss" | keys | map(tonumber) | max // 0) else 0 end' "$JSON_METRICS") +if [[ $MAX_RECORDED_STEP -lt $MAX_STEPS ]]; then + echo "[ERROR] Expected train/loss through step $MAX_STEPS, found step $MAX_RECORDED_STEP" + exit 1 fi + +uv run tests/check_metrics.py $JSON_METRICS \ + 'median(data["train/token_mult_prob_error"]) < 1.1' \ + 'max(data["train/gen_kl_error"]) < 0.003' \ + 'max(data["train/reward"]) > -0.9' + +# Clean up checkpoint directory after successful run to save space. +rm -rf "$CKPT_DIR" diff --git a/tests/test_suites/nightly_gb200.txt b/tests/test_suites/nightly_gb200.txt index afedb577f59..8ef59da164e 100644 --- a/tests/test_suites/nightly_gb200.txt +++ b/tests/test_suites/nightly_gb200.txt @@ -18,6 +18,10 @@ tests/test_suites/llm/grpo-llama3.2-1b-instruct-1n4g-megatron_generation.sh # Functional moonlight run tests/test_suites/llm/grpo-moonlight-16ba3b-4n4g-megatron.sh +# ModelOpt quant rollout smoke tests +tests/test_suites/llm/grpo-nanov3-30ba3b-4n4g-megatron-qa-nvfp4-w4a16-real.sh +tests/test_suites/llm/grpo-qwen3-30ba3b-4n4g-megatron-qa-nvfp4-w4a4-real.sh + # Functional VLM run tests/test_suites/vlm/vlm_grpo-qwen2.5-vl-3b-instruct-clevr-1n4g-dtensor2tp1.v1.sh tests/test_suites/vlm/vlm_grpo-qwen2.5-vl-3b-instruct-clevr-1n4g-megatrontp1.v1.sh diff --git a/tests/test_suites/release_gb200.txt b/tests/test_suites/release_gb200.txt index f0c346509db..f25470654fc 100644 --- a/tests/test_suites/release_gb200.txt +++ b/tests/test_suites/release_gb200.txt @@ -26,6 +26,11 @@ tests/test_suites/llm/grpo-dapomath17k-dsv3-32n4g-megatron.sh # GPT-OSS tests/test_suites/llm/grpo-gptoss-20b-8n4g-megatron.sh +# Nemotron-3 Super 120B-A12B BF16 / real-quant comparison smokes +tests/test_suites/llm/grpo-nemotron3-super-120BA12B-16n4g-megatron-qa-bf16-300step.sh +tests/test_suites/llm/grpo-nemotron3-super-120BA12B-16n4g-megatron-qa-nvfp4-w4a16-real-300step.sh +tests/test_suites/llm/grpo-nemotron3-super-120BA12B-16n4g-megatron-qa-nvfp4-w4a4-real-300step.sh + ####### # SFT # ####### diff --git a/tests/unit/models/generation/test_vllm_backend.py b/tests/unit/models/generation/test_vllm_backend.py index 2fd83278116..3144efb929b 100644 --- a/tests/unit/models/generation/test_vllm_backend.py +++ b/tests/unit/models/generation/test_vllm_backend.py @@ -153,6 +153,38 @@ def packed_broadcast_consumer(iterator, group, src, post_unpack_func): ] +@pytest.mark.vllm +def test_update_weights_via_ipc_acks_manifest_error_and_returns_false(monkeypatch): + from nemo_rl.models.generation.vllm import vllm_backend + from nemo_rl.models.policy.utils import IPCProtocol + + class FakeSocket: + def __init__(self): + self.sent = [] + + def recv_pyobj(self): + return IPCProtocol.COMPLETE + + def send(self, payload): + self.sent.append(payload) + + ext = vllm_backend.VllmInternalWorkerExtension.__new__( + vllm_backend.VllmInternalWorkerExtension + ) + ext.state_dict_info = {"model.weight": (torch.Size([1]), torch.float32)} + ext.zmq_socket = FakeSocket() + ext.maybe_init_zmq = lambda: None + + @contextlib.contextmanager + def lifecycle(_transport): + yield lambda: pytest.fail("an incomplete transfer must not be finalized") + + ext._weight_update_lifecycle = lifecycle + + assert ext.update_weights_via_ipc_zmq() is False + assert ext.zmq_socket.sent == [IPCProtocol.ACK.value.encode()] + + @pytest.mark.vllm def test_read_mtp_layer_weights_from_checkpoint_filters_and_reads(tmp_path): """Only the requested MTP layer tensors are read, across the shards holding them.""" diff --git a/tests/unit/models/generation/test_vllm_modelopt_real_quant_config.py b/tests/unit/models/generation/test_vllm_modelopt_real_quant_config.py index 207d397755c..f1d75458d6b 100644 --- a/tests/unit/models/generation/test_vllm_modelopt_real_quant_config.py +++ b/tests/unit/models/generation/test_vllm_modelopt_real_quant_config.py @@ -16,31 +16,121 @@ import os import sys import types +import weakref +from contextlib import contextmanager, nullcontext import pytest import torch -from nemo_rl.modelopt.models.generation.vllm_modelopt_patch import ( - _canonicalize_nvfp4_weight_scale, - _convert_nvfp4_linear_kernel_format, - _modelopt_dense_apply, - _modelopt_dense_process_weights, - apply_modelopt_nvfp4_patches, - modelopt_process_weights_after_loading, - prepare_modelopt_for_weight_reload, +import nemo_rl.modelopt.models.generation.vllm_modelopt as vllm_modelopt +import nemo_rl.modelopt.utils as modelopt_utils +from nemo_rl.modelopt.models.generation.vllm_modelopt import ( + NEMO_MODELOPT_W4A4, + NEMO_MODELOPT_W4A16, + _pad_nvfp4_moe_for_marlin, + quantization_method_for_mode, + register_nemo_modelopt_nvfp4, ) from nemo_rl.modelopt.utils import ( build_vllm_modelopt_nvfp4_config, iter_quant_ignore_name_candidates, matches_quant_ignore_pattern, + resolve_nvfp4_real_quant_mode, resolve_quant_cfg, ) +@pytest.fixture(autouse=True) +def _install_optional_modelopt_config_api(monkeypatch): + """Provide ModelOpt's config APIs when the optional dependency is absent.""" + try: + import modelopt.torch.export.convert_hf_config # noqa: F401 + import modelopt.torch.quantization.config # noqa: F401 + + return + except ImportError: + pass + + module_names = ( + "modelopt", + "modelopt.recipe", + "modelopt.torch", + "modelopt.torch.export", + "modelopt.torch.quantization", + ) + for module_name in module_names: + module = types.ModuleType(module_name) + module.__path__ = [] + monkeypatch.setitem(sys.modules, module_name, module) + + def missing_recipe(config_name): + raise FileNotFoundError(config_name) + + sys.modules["modelopt.recipe"].load_config = missing_recipe + + convert_module = types.ModuleType("modelopt.torch.export.convert_hf_config") + + def convert_hf_quant_config_format(config): + quantization = config["quantization"] + algo = quantization["quant_algo"] + group = { + "weights": { + "dynamic": False, + "num_bits": 4, + "type": "float", + "group_size": quantization["group_size"], + }, + "targets": ["Linear"], + } + if algo == "NVFP4": + group["input_activations"] = dict(group["weights"]) + return { + "config_groups": {"group_0": group}, + "ignore": quantization["exclude_modules"], + "quant_algo": algo, + "producer": config["producer"], + "quant_method": "modelopt", + } + + convert_module.convert_hf_quant_config_format = convert_hf_quant_config_format + monkeypatch.setitem( + sys.modules, + "modelopt.torch.export.convert_hf_config", + convert_module, + ) + + config_module = types.ModuleType("modelopt.torch.quantization.config") + + class QuantizerCfgEntry: + def __init__(self, entry): + self.entry = {"enable": True, **entry} + + def model_dump(self, **kwargs): + del kwargs + return { + key: value for key, value in self.entry.items() if value is not None + } + + class QuantizeConfig: + def __init__(self, quant_cfg, **kwargs): + del kwargs + self.quant_cfg = [QuantizerCfgEntry(entry) for entry in quant_cfg] + + config_module.QuantizeConfig = QuantizeConfig + monkeypatch.setitem( + sys.modules, + "modelopt.torch.quantization.config", + config_module, + ) + + def _import_vllm_quant_backend(monkeypatch): """Import the NeMo-RL backend without requiring the vLLM C extension.""" monkeypatch.delenv("VLLM_MODELOPT_REAL_QUANT", raising=False) - monkeypatch.setitem(sys.modules, "vllm", types.ModuleType("vllm")) + vllm_module = types.ModuleType("vllm") + vllm_module.__path__ = [] + monkeypatch.setitem(sys.modules, "vllm", vllm_module) + _install_fake_vllm_reload(monkeypatch) _install_fake_modelopt_tensor_quantizer(monkeypatch) sys.modules.pop("nemo_rl.modelopt.models.generation.vllm_quant_backend", None) sys.modules.pop("nemo_rl.models.generation.vllm.vllm_backend", None) @@ -52,6 +142,347 @@ def _import_vllm_quant_backend(monkeypatch): pytest.skip(f"could not import vLLM quant backend: {exc}") +def _base_vllm_backend(): + return sys.modules["nemo_rl.models.generation.vllm.vllm_backend"] + + +def _install_fake_vllm_reload(monkeypatch): + """Install the public vLLM layerwise-reload API used by real-quant refits.""" + module_names = ( + "vllm.model_executor", + "vllm.model_executor.layers", + "vllm.model_executor.layers.quantization", + "vllm.model_executor.model_loader", + ) + for module_name in module_names: + module = types.ModuleType(module_name) + module.__path__ = [] + monkeypatch.setitem(sys.modules, module_name, module) + + config_module = types.ModuleType("vllm.config") + config_module.current = None + + @contextmanager + def set_current_vllm_config(config): + previous = config_module.current + config_module.current = config + try: + yield + finally: + config_module.current = previous + + def get_current_vllm_config(): + if config_module.current is None: + raise AssertionError("Current vLLM config is not set") + return config_module.current + + config_module.set_current_vllm_config = set_current_vllm_config + config_module.get_current_vllm_config = get_current_vllm_config + monkeypatch.setitem(sys.modules, "vllm.config", config_module) + + reload_module = types.ModuleType("vllm.model_executor.model_loader.reload") + reload_module.__path__ = [] + reload_module.initialize_layerwise_reload = lambda model: None + reload_module.finalize_layerwise_reload = lambda model, model_config: None + layerwise_module = types.ModuleType( + "vllm.model_executor.model_loader.reload.layerwise" + ) + layerwise_module.get_layerwise_info = lambda module: types.SimpleNamespace( + loaded_weights=[], + load_numel=0, + load_numel_total=None, + ) + monkeypatch.setitem( + sys.modules, + "vllm.model_executor.model_loader.reload", + reload_module, + ) + monkeypatch.setitem( + sys.modules, + "vllm.model_executor.model_loader.reload.layerwise", + layerwise_module, + ) + modelopt_module = types.ModuleType( + "vllm.model_executor.layers.quantization.modelopt" + ) + modelopt_module.ModelOptNvFp4FusedMoE = type("ModelOptNvFp4FusedMoE", (), {}) + modelopt_module.ModelOptNvFp4LinearMethod = type( + "ModelOptNvFp4LinearMethod", (), {} + ) + monkeypatch.setitem( + sys.modules, + "vllm.model_executor.layers.quantization.modelopt", + modelopt_module, + ) + attention_module = types.ModuleType("vllm.model_executor.layers.attention") + attention_module.Attention = type("Attention", (torch.nn.Module,), {}) + attention_module.MLAAttention = type("MLAAttention", (torch.nn.Module,), {}) + monkeypatch.setitem( + sys.modules, + "vllm.model_executor.layers.attention", + attention_module, + ) + kv_cache_module = types.ModuleType( + "vllm.model_executor.layers.quantization.kv_cache" + ) + kv_cache_module.BaseKVCacheMethod = type("BaseKVCacheMethod", (), {}) + monkeypatch.setitem( + sys.modules, + "vllm.model_executor.layers.quantization.kv_cache", + kv_cache_module, + ) + return reload_module + + +def _install_fake_registered_vllm_modelopt(monkeypatch): + """Install the public vLLM registration surface used by vllm_modelopt.""" + module_names = ( + "vllm", + "vllm.model_executor", + "vllm.model_executor.kernels", + "vllm.model_executor.layers", + "vllm.model_executor.layers.fused_moe", + "vllm.model_executor.layers.fused_moe.oracle", + "vllm.model_executor.layers.quantization", + "vllm.model_executor.layers.quantization.utils", + ) + for module_name in module_names: + module = types.ModuleType(module_name) + module.__path__ = [] + monkeypatch.setitem(sys.modules, module_name, module) + + registry = {} + events = [] + + quantization_module = sys.modules["vllm.model_executor.layers.quantization"] + + def register_quantization_config(name): + def register(config_cls): + registry[name] = config_cls + return config_cls + + return register + + quantization_module.register_quantization_config = register_quantization_config + + weight_loader_v2_supported = [] + linear_module = types.ModuleType("vllm.model_executor.layers.linear") + + def register_weight_loader_v2_supported_method(method_cls: type) -> type: + weight_loader_v2_supported.append(method_cls.__name__) + return method_cls + + linear_module.register_weight_loader_v2_supported_method = ( + register_weight_loader_v2_supported_method + ) + monkeypatch.setitem( + sys.modules, + "vllm.model_executor.layers.linear", + linear_module, + ) + + class FakeModelOptNvFp4LinearMethod: + def __init__(self, quant_config): + self.quant_config = quant_config + + def create_weights(self, layer, *args, **kwargs): + del args, kwargs + if not hasattr(layer, "input_scale"): + layer.input_scale = torch.nn.Parameter(torch.ones(1)) + + class FakeModelOptNvFp4FusedMoE: + def __init__(self, quant_config, moe_config): + self.quant_config = quant_config + self.moe = moe_config + self.moe_kernel = None + self.moe_quant_config = None + + def create_weights(self, layer, *args, **kwargs): + events.append(("native_create_weights", layer, args, kwargs)) + num_experts = args[0] if args else 1 + if not hasattr(layer, "w13_input_scale"): + layer.w13_input_scale = torch.nn.Parameter(torch.zeros(num_experts, 2)) + if not hasattr(layer, "w2_input_scale"): + layer.w2_input_scale = torch.nn.Parameter(torch.zeros(num_experts)) + + def get_fused_moe_quant_config(self, layer): + del layer + return object() + + def process_weights_after_loading(self, layer): + events.append( + ( + "native_process_moe", + getattr( + getattr(layer, "moe_config", None), + "intermediate_size_per_partition", + None, + ), + ) + ) + self.moe_kernel = types.SimpleNamespace(source="native", layer=layer) + w13_input_scale = getattr(layer, "w13_input_scale", None) + w2_input_scale = getattr(layer, "w2_input_scale", None) + if isinstance(w13_input_scale, torch.Tensor) and isinstance( + w2_input_scale, torch.Tensor + ): + self.moe_quant_config = types.SimpleNamespace( + source="native", + a1_gscale=1.0 / w13_input_scale, + a2_gscale=1.0 / w2_input_scale, + ) + else: + self.moe_quant_config = types.SimpleNamespace(source="native") + + class FakeModelOptNvFp4Config: + LinearMethodCls = FakeModelOptNvFp4LinearMethod + FusedMoEMethodCls = FakeModelOptNvFp4FusedMoE + + def __init__(self, group_size=16): + self.group_size = group_size + + @classmethod + def from_config(cls, config): + target = config.get("quantization", config) + instance = cls(group_size=target.get("group_size", 16)) + instance.parsed_config = config + return instance + + @classmethod + def _extract_modelopt_quant_algo(cls, config): + del cls + target = config.get("quantization", config) + return str(target.get("quant_algo", "")).upper() + + modelopt_module = types.ModuleType( + "vllm.model_executor.layers.quantization.modelopt" + ) + modelopt_module.ModelOptNvFp4Config = FakeModelOptNvFp4Config + modelopt_module.ModelOptNvFp4LinearMethod = FakeModelOptNvFp4LinearMethod + modelopt_module.ModelOptNvFp4FusedMoE = FakeModelOptNvFp4FusedMoE + monkeypatch.setitem( + sys.modules, + "vllm.model_executor.layers.quantization.modelopt", + modelopt_module, + ) + + class FakeFusedMoEMethodBase: + def __init__(self, moe_config): + self.moe = moe_config + self.moe_kernel = None + self.moe_quant_config = None + + fused_method_module = types.ModuleType( + "vllm.model_executor.layers.fused_moe.fused_moe_method_base" + ) + fused_method_module.FusedMoEMethodBase = FakeFusedMoEMethodBase + monkeypatch.setitem( + sys.modules, + "vllm.model_executor.layers.fused_moe.fused_moe_method_base", + fused_method_module, + ) + + class FakeNvFp4LinearLayerConfig: + pass + + class FakeMarlinNvFp4LinearKernel: + def __init__(self, config): + self.config = config + + def process_weights_after_loading(self, layer): + events.append(("process_marlin_kernel", layer)) + + def apply_weights(self, **kwargs): + events.append(("apply_marlin_kernel", kwargs)) + return "output" + + linear_kernel_module = types.ModuleType("vllm.model_executor.kernels.linear") + linear_kernel_module.MarlinNvFp4LinearKernel = FakeMarlinNvFp4LinearKernel + linear_kernel_module.NvFp4LinearLayerConfig = FakeNvFp4LinearLayerConfig + monkeypatch.setitem( + sys.modules, + "vllm.model_executor.kernels.linear", + linear_kernel_module, + ) + + class FakeMarlinExperts: + pass + + oracle_module = types.ModuleType( + "vllm.model_executor.layers.fused_moe.oracle.nvfp4" + ) + oracle_module.NvFp4MoeBackend = types.SimpleNamespace(MARLIN="marlin") + + def convert_to_nvfp4_moe_kernel_format(**kwargs): + events.append( + ( + "convert_moe", + kwargs["layer"].moe_config.intermediate_size_per_partition, + ) + ) + return tuple( + kwargs[name] + for name in ( + "w13", + "w13_scale", + "w13_scale_2", + "a13_scale", + "w2", + "w2_scale", + "w2_scale_2", + "a2_scale", + ) + ) + + oracle_module.convert_to_nvfp4_moe_kernel_format = ( + convert_to_nvfp4_moe_kernel_format + ) + oracle_module.is_global_sf_supported_for_nvfp4_backend = lambda backend: False + oracle_module.select_nvfp4_moe_backend = lambda **kwargs: ( + oracle_module.NvFp4MoeBackend.MARLIN, + FakeMarlinExperts, + ) + oracle_module.make_nvfp4_moe_kernel = lambda **kwargs: events.append( + ("make_moe_kernel", kwargs) + ) or types.SimpleNamespace( + fused_experts=types.SimpleNamespace( + process_weights_after_loading=lambda layer: events.append( + ("process_moe", layer) + ) + ) + ) + monkeypatch.setitem( + sys.modules, + "vllm.model_executor.layers.fused_moe.oracle.nvfp4", + oracle_module, + ) + + quant_utils_module = types.ModuleType( + "vllm.model_executor.layers.quantization.utils.quant_utils" + ) + quant_utils_module.kNvfp4Static = object() + monkeypatch.setitem( + sys.modules, + "vllm.model_executor.layers.quantization.utils.quant_utils", + quant_utils_module, + ) + + utils_module = types.ModuleType("vllm.model_executor.utils") + + def replace_parameter(layer, name, value): + if name in layer._parameters: + del layer._parameters[name] + setattr(layer, name, torch.nn.Parameter(value, requires_grad=False)) + + utils_module.replace_parameter = replace_parameter + monkeypatch.setitem(sys.modules, "vllm.model_executor.utils", utils_module) + return types.SimpleNamespace( + registry=registry, + events=events, + weight_loader_v2_supported=weight_loader_v2_supported, + ) + + def _install_fake_modelopt_tensor_quantizer(monkeypatch): """Install the minimal ModelOpt module hierarchy needed by vLLM backend import.""" module_names = [ @@ -94,16 +525,66 @@ class FakeTensorQuantizer(torch.nn.Module): ].tensor_quantizer = tensor_quantizer_module -def test_w4a16_real_quant_config_keeps_weight_only_default(): - cfg = build_vllm_modelopt_nvfp4_config() +def _make_real_quant_extension(backend, model, ignore): + extension = object.__new__(backend.VllmQuantInternalWorkerExtension) + extension.device = torch.device("cpu") + extension._nrl_w13_num_shards_by_prefix = {} + extension.model_runner = types.SimpleNamespace( + model=model, + vllm_config=types.SimpleNamespace( + parallel_config=types.SimpleNamespace(enable_expert_parallel=False), + model_config=types.SimpleNamespace( + hf_config=types.SimpleNamespace(quantization_config={"ignore": ignore}) + ), + ), + ) + return extension + + +def _patch_real_quant_load(monkeypatch, backend, forwarded=None): + monkeypatch.setattr( + backend.VllmQuantInternalWorkerExtension, + "_is_real_quant_model", + lambda self: True, + ) + if forwarded is not None: + monkeypatch.setattr( + backend.VllmInternalWorkerExtension, + "_load_weights", + lambda self, weights: forwarded.extend(weights) or "loaded", + ) + + +def _mark_as_modelopt_layer(model): + modelopt_module = sys.modules["vllm.model_executor.layers.quantization.modelopt"] + model.quant_method = modelopt_module.ModelOptNvFp4LinearMethod() + return model + + +def test_base_ipc_data_ack_fence_synchronizes_current_stream_once(monkeypatch): + _import_vllm_quant_backend(monkeypatch) + backend = _base_vllm_backend() + extension = object.__new__(backend.VllmInternalWorkerExtension) + calls = [] + stream = types.SimpleNamespace(synchronize=lambda: calls.append("sync")) + monkeypatch.setattr( + backend.torch.cuda, + "current_stream", + lambda: calls.append("current_stream") or stream, + ) + + extension._synchronize_before_ipc_data_ack() + + assert calls == ["current_stream", "sync"] + + +def test_w4a16_real_quant_config_is_weight_only(): + cfg = build_vllm_modelopt_nvfp4_config(mode="w4a16") group = cfg["config_groups"]["group_0"] assert cfg["quant_method"] == "modelopt" - assert cfg["quant_algo"] == "NVFP4" - assert cfg["quant_mode"] == "w4a16_nvfp4" - assert cfg["weight_only"] is True - assert cfg["group_size"] == 16 - assert group["input_activations"] is None + assert cfg["quant_algo"] == "W4A16_NVFP4" + assert "input_activations" not in group assert group["weights"] == { "dynamic": False, "num_bits": 4, @@ -121,14 +602,51 @@ def test_w4a16_real_quant_config_keeps_weight_only_default(): ] +def test_w4a4_real_quant_config_has_static_input_activations(): + cfg = build_vllm_modelopt_nvfp4_config(mode="w4a4") + + group = cfg["config_groups"]["group_0"] + assert cfg["quant_method"] == "modelopt" + assert cfg["quant_algo"] == "NVFP4" + assert group["input_activations"] == { + "dynamic": False, + "num_bits": 4, + "type": "float", + "group_size": 16, + } + + +def test_real_quant_config_rejects_unsupported_mode(): + with pytest.raises(ValueError, match="expected 'w4a4' or 'w4a16'"): + build_vllm_modelopt_nvfp4_config(mode="w4a8") # type: ignore[arg-type] + + +@pytest.mark.parametrize( + ("mode", "method"), + [("w4a4", NEMO_MODELOPT_W4A4), ("w4a16", NEMO_MODELOPT_W4A16)], +) +def test_quantization_method_for_mode_uses_registered_names(mode, method): + assert quantization_method_for_mode(mode) == method + + +def test_quantization_method_for_mode_rejects_unknown_mode(): + with pytest.raises(ValueError, match="Unsupported ModelOpt NVFP4 rollout mode"): + quantization_method_for_mode("w4a8") + + def test_real_quant_config_allows_explicit_ignore_override(): - cfg = build_vllm_modelopt_nvfp4_config(ignore=["lm_head"]) + ignore = ["lm_head", "*.mixer.in_proj*"] + cfg = build_vllm_modelopt_nvfp4_config(mode="w4a16", ignore=ignore) - assert cfg["ignore"] == ["lm_head"] + assert cfg["ignore"] == ignore + assert matches_quant_ignore_pattern( + "model.layers.0.mixer.in_proj.weight", + cfg["ignore"], + ) def test_default_ignore_patterns_match_expected_layers(): - ignore_patterns = build_vllm_modelopt_nvfp4_config()["ignore"] + ignore_patterns = build_vllm_modelopt_nvfp4_config(mode="w4a16")["ignore"] assert matches_quant_ignore_pattern( "model.layers.0.self_attn.o_proj.weight", ignore_patterns @@ -144,6 +662,9 @@ def test_default_ignore_patterns_match_expected_layers(): assert matches_quant_ignore_pattern( "model.layers.0.mlp.gate.weight_scale", ignore_patterns ) + assert matches_quant_ignore_pattern( + "model.layers.0.mlp.gate.input_scale", ignore_patterns + ) assert not matches_quant_ignore_pattern( "model.layers.0.mlp.experts.0.w1.weight", ignore_patterns ) @@ -164,18 +685,28 @@ def test_quant_ignore_name_candidates_include_model_prefix_and_base_names(): "lm_head.weight_scale", "lm_head", ] + assert list(iter_quant_ignore_name_candidates("model.lm_head.input_scale")) == [ + "model.lm_head.input_scale", + "model.lm_head", + "lm_head.input_scale", + "lm_head", + ] -def test_configure_quant_engine_kwargs_for_fake_quant(monkeypatch): +def test_configure_quant_engine_kwargs_for_fake_quant(monkeypatch, tmp_path): worker_mod = pytest.importorskip( "nemo_rl.modelopt.models.generation.vllm_quant_worker" ) monkeypatch.delenv("VLLM_QUANT_CFG", raising=False) monkeypatch.delenv("VLLM_MODELOPT_REAL_QUANT", raising=False) + quant_cfg = "quant.yaml" + (tmp_path / quant_cfg).touch() + monkeypatch.chdir(tmp_path) + llm_kwargs = {} worker_mod._configure_quant_engine_kwargs( - {"quant_cfg": "examples/modelopt/quant_configs/nvfp4_w4a8_fp8.yaml"}, + {"quant_cfg": quant_cfg}, llm_kwargs, ) @@ -185,24 +716,32 @@ def test_configure_quant_engine_kwargs_for_fake_quant(monkeypatch): assert llm_kwargs["worker_extension_cls"] == ( "nemo_rl.modelopt.models.generation.vllm_quant_backend.VllmQuantInternalWorkerExtension" ) - assert os.environ["VLLM_QUANT_CFG"] == ( - "examples/modelopt/quant_configs/nvfp4_w4a8_fp8.yaml" - ) + assert os.environ["VLLM_QUANT_CFG"] == os.path.abspath(quant_cfg) assert "quantization" not in llm_kwargs -def test_configure_quant_engine_kwargs_for_real_quant(monkeypatch): +def test_quant_worker_forwards_snapshot_pythonpath_to_inner_vllm_workers(): worker_mod = pytest.importorskip( "nemo_rl.modelopt.models.generation.vllm_quant_worker" ) - patch_mod = pytest.importorskip( - "nemo_rl.modelopt.models.generation.vllm_modelopt_patch" + + assert "PYTHONPATH" in worker_mod._EXTRA_ENV_VARS + + +def test_configure_quant_engine_kwargs_for_real_quant(monkeypatch): + worker_mod = pytest.importorskip( + "nemo_rl.modelopt.models.generation.vllm_quant_worker" ) monkeypatch.delenv("VLLM_QUANT_CFG", raising=False) monkeypatch.delenv("VLLM_MODELOPT_REAL_QUANT", raising=False) - patch_calls = [] + registration_calls = [] + monkeypatch.setattr( + vllm_modelopt, + "register_nemo_modelopt_nvfp4", + lambda: registration_calls.append(True), + ) monkeypatch.setattr( - patch_mod, "apply_modelopt_nvfp4_patches", lambda: patch_calls.append(True) + modelopt_utils, "resolve_nvfp4_real_quant_mode", lambda _: "w4a16" ) llm_kwargs = {} @@ -215,25 +754,53 @@ def test_configure_quant_engine_kwargs_for_real_quant(monkeypatch): llm_kwargs, ) - assert patch_calls == [True] + assert registration_calls == [True] assert os.environ["VLLM_MODELOPT_REAL_QUANT"] == "1" assert "VLLM_QUANT_CFG" not in os.environ assert "worker_cls" not in llm_kwargs - assert llm_kwargs["quantization"] == "modelopt" + assert llm_kwargs["quantization"] == NEMO_MODELOPT_W4A16 assert llm_kwargs["hf_overrides"]["quantization_config"] == ( - build_vllm_modelopt_nvfp4_config(ignore=["lm_head"]) + build_vllm_modelopt_nvfp4_config(mode="w4a16", ignore=["lm_head"]) ) -def test_configure_quant_engine_kwargs_preserves_hf_overrides(monkeypatch): +@pytest.mark.parametrize("mode", ["w4a4", "w4a16"]) +@pytest.mark.parametrize("kv_cache_dtype", ["auto", "fp8"]) +def test_configure_real_quant_preserves_kv_cache_dtype( + monkeypatch, + mode, + kv_cache_dtype, +): worker_mod = pytest.importorskip( "nemo_rl.modelopt.models.generation.vllm_quant_worker" ) - patch_mod = pytest.importorskip( - "nemo_rl.modelopt.models.generation.vllm_modelopt_patch" + monkeypatch.setattr(vllm_modelopt, "register_nemo_modelopt_nvfp4", lambda: None) + monkeypatch.setattr( + modelopt_utils, + "resolve_nvfp4_real_quant_mode", + lambda _: mode, + ) + + llm_kwargs = {"kv_cache_dtype": kv_cache_dtype} + worker_mod._configure_quant_engine_kwargs( + {"quant_cfg": "NVFP4_EXPERTS_ONLY_CFG", "real_quant": True}, + llm_kwargs, + ) + + assert llm_kwargs["kv_cache_dtype"] == kv_cache_dtype + assert llm_kwargs["quantization"] == quantization_method_for_mode(mode) + assert "kv_cache" not in llm_kwargs["hf_overrides"]["quantization_config"] + + +def test_configure_quant_engine_kwargs_preserves_hf_overrides(monkeypatch): + worker_mod = pytest.importorskip( + "nemo_rl.modelopt.models.generation.vllm_quant_worker" ) monkeypatch.delenv("VLLM_MODELOPT_REAL_QUANT", raising=False) - monkeypatch.setattr(patch_mod, "apply_modelopt_nvfp4_patches", lambda: None) + monkeypatch.setattr(vllm_modelopt, "register_nemo_modelopt_nvfp4", lambda: None) + monkeypatch.setattr( + modelopt_utils, "resolve_nvfp4_real_quant_mode", lambda _: "w4a16" + ) llm_kwargs = {"hf_overrides": {"trust_remote_code": True}} worker_mod._configure_quant_engine_kwargs( @@ -360,28 +927,183 @@ def test_vllm_modelopt_backend_imports_without_gpt_oss_helper(monkeypatch): _import_vllm_quant_backend(monkeypatch) -def test_vllm_modelopt_backend_applies_real_quant_patch_on_import(monkeypatch): - patch_mod = pytest.importorskip( - "nemo_rl.modelopt.models.generation.vllm_modelopt_patch" +def test_real_quant_backend_uses_modelopt_refit_timeout(monkeypatch): + backend = _import_vllm_quant_backend(monkeypatch) + events = [] + + class FakeSocket: + def setsockopt(self, option, value): + events.append(("setsockopt", option, value)) + + def connect(self, address): + events.append(("connect", address)) + + class FakeContext: + def socket(self, socket_type): + events.append(("socket", socket_type)) + return FakeSocket() + + extension = object.__new__(backend.VllmQuantInternalWorkerExtension) + extension.get_zmq_address = lambda: "ipc:///tmp/modelopt-test.sock" + monkeypatch.setattr(backend.zmq, "Context", FakeContext) + monkeypatch.setattr( + backend.VllmQuantInternalWorkerExtension, + "_is_real_quant_model", + lambda _self: True, ) + + extension.maybe_init_zmq() + + assert events[0] == ("socket", backend.zmq.REP) + assert ("setsockopt", backend.zmq.LINGER, 0) in events + assert ("connect", "ipc:///tmp/modelopt-test.sock") in events + assert events[-2:] == [ + ( + "setsockopt", + backend.zmq.SNDTIMEO, + modelopt_utils.MODELOPT_REAL_QUANT_ZMQ_TIMEOUT_MS, + ), + ( + "setsockopt", + backend.zmq.RCVTIMEO, + modelopt_utils.MODELOPT_REAL_QUANT_ZMQ_TIMEOUT_MS, + ), + ] + + +def test_vllm_modelopt_backend_registers_real_quant_configs_on_import(monkeypatch): calls = [] monkeypatch.setenv("VLLM_MODELOPT_REAL_QUANT", "1") monkeypatch.setitem(sys.modules, "vllm", types.ModuleType("vllm")) _install_fake_modelopt_tensor_quantizer(monkeypatch) monkeypatch.setattr( - patch_mod, - "apply_modelopt_nvfp4_patches", - lambda: calls.append("patched"), + vllm_modelopt, + "register_nemo_modelopt_nvfp4", + lambda: calls.append("registered"), ) sys.modules.pop("nemo_rl.modelopt.models.generation.vllm_quant_backend", None) importlib.import_module("nemo_rl.modelopt.models.generation.vllm_quant_backend") - assert calls == ["patched"] + assert calls == ["registered"] -def test_real_quant_load_weights_copies_ignored_float_weights(monkeypatch): +def test_modelopt_moe_manifest_requires_complete_w4a4_family(monkeypatch): + backend = _import_vllm_quant_backend(monkeypatch) + prefix = "model.layers.0.mixer" + state_dict_info = { + f"{prefix}.experts.w13_weight": ((2, 4, 3), torch.uint8), + f"{prefix}.experts.w13_weight_scale": ((2, 4, 1), torch.uint8), + f"{prefix}.experts.w13_weight_scale_2": ((2, 2), torch.float32), + f"{prefix}.experts.w13_input_scale": ((2, 2), torch.float32), + f"{prefix}.experts.w2_weight": ((2, 3, 4), torch.uint8), + f"{prefix}.experts.w2_weight_scale": ((2, 1, 4), torch.uint8), + f"{prefix}.experts.w2_weight_scale_2": ((2,), torch.float32), + f"{prefix}.experts.w2_input_scale": ((2,), torch.float32), + } + + assert backend._w13_num_shards_from_state_dict_info( + state_dict_info, require_input_scales=True + ) == {prefix: 2} + + legacy_state_dict_info = dict(state_dict_info) + legacy_state_dict_info[f"{prefix}.experts.w13_weight_scale_2"] = ( + (2,), + torch.float32, + ) + legacy_state_dict_info[f"{prefix}.experts.w13_input_scale"] = ( + (2,), + torch.float32, + ) + assert backend._w13_num_shards_from_state_dict_info( + legacy_state_dict_info, require_input_scales=True + ) == {prefix: 1} + + mismatched_state_dict_info = dict(state_dict_info) + mismatched_state_dict_info[f"{prefix}.experts.w13_input_scale"] = ( + (2, 1), + torch.float32, + ) + with pytest.raises(RuntimeError, match="input/global scale layouts disagree"): + backend._w13_num_shards_from_state_dict_info( + mismatched_state_dict_info, require_input_scales=True + ) + + del state_dict_info[f"{prefix}.experts.w2_input_scale"] + with pytest.raises(RuntimeError, match="missing.*w2_input_scale"): + backend._w13_num_shards_from_state_dict_info( + state_dict_info, require_input_scales=True + ) + + +def test_real_quant_load_weights_batches_full_experts_and_expands_global_scales( + monkeypatch, +): + backend = _import_vllm_quant_backend(monkeypatch) + + class ModelOptNvFp4FusedMoE: + quant_config = types.SimpleNamespace(get_name=lambda: NEMO_MODELOPT_W4A16) + + def make_model(expert_map): + model = torch.nn.Module() + model.moe = torch.nn.Module() + model.moe.quant_method = ModelOptNvFp4FusedMoE() + model.moe._expert_map = expert_map + model.moe.local_num_experts = 2 if expert_map is None else 1 + model.moe.global_num_experts = 2 + # ModelOpt assigns the same quant config to attention's FP8 KV method; + # this must not be mistaken for expert parallelism. + model.attention = torch.nn.Module() + model.attention.quant_method = ModelOptNvFp4FusedMoE() + return model + + prefix = "model.layers.0.mlp" + w13_weight = torch.arange(24).reshape(2, 4, 3) + w13_scale_2 = torch.tensor([[1.0], [2.0]]) + state_dict_info = { + f"{prefix}.experts.w13_weight": ((2, 4, 3), torch.uint8), + f"{prefix}.experts.w13_weight_scale": ((2, 4, 1), torch.uint8), + f"{prefix}.experts.w13_weight_scale_2": ((2, 1), torch.float32), + f"{prefix}.experts.w2_weight": ((2, 3, 2), torch.uint8), + f"{prefix}.experts.w2_weight_scale": ((2, 3, 1), torch.uint8), + f"{prefix}.experts.w2_weight_scale_2": ((2,), torch.float32), + } + + batched_forwarded = [] + extension = _make_real_quant_extension(backend, make_model(None), []) + extension.prepare_refit_info(state_dict_info) + extension._nrl_w13_num_shards_by_prefix = {prefix: 1} + _patch_real_quant_load(monkeypatch, backend, batched_forwarded) + assert ( + extension._load_weights( + [ + (f"{prefix}.experts.w13_weight", w13_weight), + (f"{prefix}.experts.w13_weight_scale_2", w13_scale_2), + ] + ) + == "loaded" + ) + assert [name for name, _ in batched_forwarded] == [ + f"{prefix}.experts.0.up_proj.weight", + f"{prefix}.experts.0.up_proj.weight_scale_2", + f"{prefix}.experts.1.up_proj.weight_scale_2", + ] + assert batched_forwarded[0][1] is w13_weight + torch.testing.assert_close(batched_forwarded[1][1], w13_scale_2[0, 0]) + torch.testing.assert_close(batched_forwarded[2][1], w13_scale_2[1, 0]) + + extension = _make_real_quant_extension( + backend, + make_model(torch.tensor([0, -1])), + [], + ) + extension.model_runner.vllm_config.parallel_config.enable_expert_parallel = True + with pytest.raises(RuntimeError, match="all experts local"): + extension.prepare_refit_info(state_dict_info) + + +def test_real_quant_load_weights_forwards_ignored_float_weights(monkeypatch): backend = _import_vllm_quant_backend(monkeypatch) class TinyModel(torch.nn.Module): @@ -391,34 +1113,9 @@ def __init__(self): self.keep = torch.nn.Linear(2, 2, bias=False) model = TinyModel() - extension = object.__new__(backend.VllmQuantInternalWorkerExtension) - extension.model_runner = types.SimpleNamespace( - model=model, - vllm_config=types.SimpleNamespace( - model_config=types.SimpleNamespace( - hf_config=types.SimpleNamespace( - quantization_config={"ignore": ["lm_head"]} - ) - ) - ), - ) - forwarded = [] - - def fake_base_load_weights(self, weights): - forwarded.extend(weights) - return "loaded" - - monkeypatch.setattr( - backend.VllmQuantInternalWorkerExtension, - "_is_real_quant_model", - lambda self: True, - ) - monkeypatch.setattr( - backend.VllmInternalWorkerExtension, - "_load_weights", - fake_base_load_weights, - ) + extension = _make_real_quant_extension(backend, model, ["lm_head"]) + _patch_real_quant_load(monkeypatch, backend, forwarded) ignored_weight = torch.full_like(model.lm_head.weight, 7.0) kept_weight = torch.full_like(model.keep.weight, 3.0) @@ -434,81 +1131,278 @@ def fake_base_load_weights(self, weights): == "loaded" ) - torch.testing.assert_close(model.lm_head.weight, ignored_weight) - assert [name for name, _ in forwarded] == ["keep.weight"] - torch.testing.assert_close(forwarded[0][1], kept_weight) + assert [name for name, _ in forwarded] == ["lm_head.weight", "keep.weight"] + torch.testing.assert_close(forwarded[0][1], ignored_weight) + torch.testing.assert_close(forwarded[1][1], kept_weight) -def test_real_quant_load_weights_returns_when_only_ignored_weights(monkeypatch): +def test_real_quant_load_weights_returns_when_only_ignored_scales(monkeypatch): backend = _import_vllm_quant_backend(monkeypatch) model = torch.nn.Module() model.lm_head = torch.nn.Linear(2, 2, bias=False) - extension = object.__new__(backend.VllmQuantInternalWorkerExtension) - extension.model_runner = types.SimpleNamespace( - model=model, - vllm_config=types.SimpleNamespace( - model_config=types.SimpleNamespace( - hf_config=types.SimpleNamespace( - quantization_config={"ignore": ["lm_head"]} - ) - ) - ), - ) - monkeypatch.setattr( - backend.VllmQuantInternalWorkerExtension, - "_is_real_quant_model", - lambda self: True, - ) + extension = _make_real_quant_extension(backend, model, ["lm_head"]) + _patch_real_quant_load(monkeypatch, backend) assert ( extension._load_weights( [ - ("lm_head.weight", torch.full_like(model.lm_head.weight, 1.5)), ("lm_head.weight_scale", torch.ones(1)), ("lm_head.weight_scale_2", torch.ones(1)), ] ) is None ) - torch.testing.assert_close( - model.lm_head.weight, - torch.full_like(model.lm_head.weight, 1.5), - ) -def test_real_quant_load_weights_forwards_ignored_shape_mismatch(monkeypatch): +def test_real_quant_load_weights_forwards_ignored_weights_to_vllm_loader(monkeypatch): backend = _import_vllm_quant_backend(monkeypatch) model = torch.nn.Module() model.lm_head = torch.nn.Linear(2, 2, bias=False) - extension = object.__new__(backend.VllmQuantInternalWorkerExtension) - extension.model_runner = types.SimpleNamespace( - model=model, - vllm_config=types.SimpleNamespace( - model_config=types.SimpleNamespace( - hf_config=types.SimpleNamespace( - quantization_config={"ignore": ["lm_head"]} - ) - ) - ), - ) forwarded = [] + extension = _make_real_quant_extension(backend, model, ["lm_head"]) + _patch_real_quant_load(monkeypatch, backend, forwarded) + + mismatched = torch.ones(1, dtype=model.lm_head.weight.dtype) + + assert extension._load_weights([("lm_head.weight", mismatched)]) == "loaded" + assert forwarded == [("lm_head.weight", mismatched)] + + +def test_real_quant_load_weights_detaches_pending_layerwise_views(monkeypatch): + backend = _import_vllm_quant_backend(monkeypatch) + layerwise_mod = sys.modules["vllm.model_executor.model_loader.reload.layerwise"] + model = torch.nn.Module() + model.reload_root = torch.nn.Linear(2, 2, bias=False) + model.unrelated = torch.nn.Linear(2, 2, bias=False) + extension = _make_real_quant_extension(backend, model, []) + extension._nrl_modelopt_reload_roots = (model.reload_root,) + _patch_real_quant_load(monkeypatch, backend, []) + + source = torch.arange(4, dtype=torch.float32) + incoming = source.view(2, 2) + bound_arguments = types.SimpleNamespace(arguments={"loaded_weight": incoming}) + layerwise_info = types.SimpleNamespace(loaded_weights=[("weight", bound_arguments)]) + inspected = [] + + def get_layerwise_info(module): + inspected.append(module) + if module is model.reload_root: + return layerwise_info + return types.SimpleNamespace(loaded_weights=[]) + + monkeypatch.setattr( + layerwise_mod, + "get_layerwise_info", + get_layerwise_info, + ) + + assert extension._load_weights([("reload_root.weight", incoming)]) == "loaded" + + detached = bound_arguments.arguments["loaded_weight"] + assert detached.untyped_storage().data_ptr() != source.untyped_storage().data_ptr() + torch.testing.assert_close(detached, incoming) + source.zero_() + torch.testing.assert_close(detached, torch.arange(4).view(2, 2).float()) + assert inspected == [model.reload_root] + + +def test_real_quant_pre_ack_fence_is_device_wide_and_load_does_not_fence( + monkeypatch, +): + backend = _import_vllm_quant_backend(monkeypatch) + model = torch.nn.Linear(1, 1) + extension = _make_real_quant_extension(backend, model, []) + extension._nrl_modelopt_reload_roots = (model,) + extension.device = types.SimpleNamespace(type="cuda") + events = [] monkeypatch.setattr( backend.VllmQuantInternalWorkerExtension, "_is_real_quant_model", - lambda self: True, + lambda _self: True, ) monkeypatch.setattr( backend.VllmInternalWorkerExtension, "_load_weights", - lambda self, weights: forwarded.extend(weights) or "loaded", + lambda _self, _weights: events.append("load") or "loaded", + ) + monkeypatch.setattr( + backend, + "_detach_pending_layerwise_weights", + lambda _roots, _storage_ptrs: events.append("detach"), + ) + monkeypatch.setattr(backend.torch, "device", lambda _device: nullcontext()) + monkeypatch.setattr( + backend.torch.accelerator, + "synchronize", + lambda: events.append("sync"), + ) + monkeypatch.setattr( + backend.torch.cuda, + "current_stream", + lambda: pytest.fail("real quant must use one device-wide IPC ACK fence"), ) - mismatched = torch.ones(1, dtype=model.lm_head.weight.dtype) + assert extension._load_weights([("weight", torch.ones(1))]) == "loaded" + assert events == ["load", "detach"] - assert extension._load_weights([("lm_head.weight", mismatched)]) == "loaded" - assert forwarded == [("lm_head.weight", mismatched)] + extension._synchronize_before_ipc_data_ack() + assert events == ["load", "detach", "sync"] + + +@pytest.mark.parametrize("load_numel", [0, 10]) +def test_real_quant_rejects_incomplete_modelopt_layerwise_reload( + monkeypatch, load_numel +): + backend = _import_vllm_quant_backend(monkeypatch) + layerwise_mod = sys.modules["vllm.model_executor.model_loader.reload.layerwise"] + + modelopt_module = types.ModuleType( + "vllm.model_executor.layers.quantization.modelopt" + ) + modelopt_base = type("ModelOptNvFp4FusedMoE", (), {}) + modelopt_module.ModelOptNvFp4FusedMoE = modelopt_base + modelopt_module.ModelOptNvFp4LinearMethod = type( + "ModelOptNvFp4LinearMethod", (), {} + ) + monkeypatch.setitem( + sys.modules, + "vllm.model_executor.layers.quantization.modelopt", + modelopt_module, + ) + experts = torch.nn.Module() + experts.quant_method = type("NemoModelOptNvFp4FusedMoE", (modelopt_base,), {})() + model = torch.nn.Module() + model.experts = experts + info = types.SimpleNamespace( + load_numel=load_numel, + load_numel_total=12, + loaded_weights=[("w13_weight", object())] if load_numel else [], + ) + monkeypatch.setattr(layerwise_mod, "get_layerwise_info", lambda _module: info) + + with pytest.raises( + RuntimeError, + match=rf"experts: {load_numel}/12 elements", + ): + backend._require_complete_modelopt_layerwise_reload(model) + + +def test_real_quant_accepts_processed_modelopt_layerwise_reload(monkeypatch): + backend = _import_vllm_quant_backend(monkeypatch) + layerwise_mod = sys.modules["vllm.model_executor.model_loader.reload.layerwise"] + + modelopt_module = types.ModuleType( + "vllm.model_executor.layers.quantization.modelopt" + ) + modelopt_module.ModelOptNvFp4FusedMoE = type("ModelOptNvFp4FusedMoE", (), {}) + modelopt_base = type("ModelOptNvFp4LinearMethod", (), {}) + modelopt_module.ModelOptNvFp4LinearMethod = modelopt_base + monkeypatch.setitem( + sys.modules, + "vllm.model_executor.layers.quantization.modelopt", + modelopt_module, + ) + linear = torch.nn.Module() + linear.quant_method = type("NemoModelOptW4A16LinearMethod", (modelopt_base,), {})() + model = torch.nn.Module() + model.linear = linear + info = types.SimpleNamespace( + load_numel=0, + load_numel_total=None, + loaded_weights=[], + ) + monkeypatch.setattr(layerwise_mod, "get_layerwise_info", lambda _module: info) + + backend._require_complete_modelopt_layerwise_reload(model) + + +def test_real_quant_scopes_native_reload_away_from_mamba_alias_buffers( + monkeypatch, +): + backend = _import_vllm_quant_backend(monkeypatch) + modelopt_module = sys.modules["vllm.model_executor.layers.quantization.modelopt"] + attention_module = sys.modules["vllm.model_executor.layers.attention"] + kv_cache_module = sys.modules["vllm.model_executor.layers.quantization.kv_cache"] + + class MambaMixer(torch.nn.Module): + def __init__(self): + super().__init__() + self.conv1d = torch.nn.Linear(3, 2, bias=False) + self.register_buffer( + "conv_weights", + self.conv1d.weight.detach().view(-1), + persistent=False, + ) + + class KVAttention(attention_module.Attention): + def __init__(self): + super().__init__() + self.quant_method = kv_cache_module.BaseKVCacheMethod() + self.kv_cache_dtype = "fp8" + self.projection = _mark_as_modelopt_layer(torch.nn.Linear(1, 1)) + + model = torch.nn.Module() + model.mamba = MambaMixer() + model.experts = _mark_as_modelopt_layer(torch.nn.Linear(1, 1)) + model.attention = KVAttention() + + assert backend._modelopt_layerwise_reload_roots( + model, + include_fp8_kv_cache=False, + ) == [model.experts, model.attention.projection] + assert backend._modelopt_layerwise_reload_roots( + model, + include_fp8_kv_cache=True, + ) == [model.experts, model.attention] + + model.attention.kv_cache_dtype = "auto" + assert backend._modelopt_layerwise_reload_roots( + model, + include_fp8_kv_cache=True, + ) == [model.experts, model.attention.projection] + model.attention.kv_cache_dtype = "fp8" + + model.shared = torch.nn.Module() + model.shared.experts = model.experts + assert backend._modelopt_layerwise_reload_roots( + model, + include_fp8_kv_cache=True, + ) == [model.experts, model.attention] + + for roots in ( + backend._modelopt_layerwise_reload_roots(model, include_fp8_kv_cache=False), + backend._modelopt_layerwise_reload_roots(model, include_fp8_kv_cache=True), + ): + assert model.mamba not in roots + assert model.mamba.conv1d not in roots + + +def test_real_quant_caches_scoped_reload_roots(monkeypatch): + backend = _import_vllm_quant_backend(monkeypatch) + model = torch.nn.Linear(1, 1) + extension = _make_real_quant_extension(backend, model, []) + extension._nrl_modelopt_reload_roots = None + selected_roots = [model] + calls = [] + + def select_modelopt_roots(model_arg, *, include_fp8_kv_cache): + calls.append((model_arg, include_fp8_kv_cache)) + return selected_roots + + monkeypatch.setattr( + backend, + "_modelopt_layerwise_reload_roots", + select_modelopt_roots, + ) + + first = extension._get_modelopt_reload_roots() + second = extension._get_modelopt_reload_roots() + + assert first is second + assert first == (model,) + assert calls == [(model, False)] def test_fake_quant_load_weights_exposes_input_quantizer_buffers(monkeypatch): @@ -553,15 +1447,24 @@ def fake_base_load_weights(self, weights): torch.testing.assert_close(child.input_quantizer_amax, torch.tensor([3.0])) -def test_real_quant_collective_reload_runs_modelopt_hooks(monkeypatch): +def test_real_quant_reload_keeps_vllm_config_active_during_layerwise_processing( + monkeypatch, +): backend = _import_vllm_quant_backend(monkeypatch) - patch_mod = pytest.importorskip( - "nemo_rl.modelopt.models.generation.vllm_modelopt_patch" - ) + config_mod = sys.modules["vllm.config"] + reload_mod = sys.modules["vllm.model_executor.model_loader.reload"] - model = torch.nn.Linear(1, 1) + model = _mark_as_modelopt_layer(torch.nn.Linear(1, 1)) + vllm_config = object() + model_config = object() extension = object.__new__(backend.VllmQuantInternalWorkerExtension) - extension.model_runner = types.SimpleNamespace(model=model) + extension.model_runner = types.SimpleNamespace( + model=model, + vllm_config=vllm_config, + ) + extension.model_config = model_config + extension.device = torch.device("cpu") + extension._nrl_modelopt_reload_roots = (model,) calls = [] monkeypatch.setattr( @@ -570,39 +1473,106 @@ def test_real_quant_collective_reload_runs_modelopt_hooks(monkeypatch): lambda self: True, ) monkeypatch.setattr( - patch_mod, - "prepare_modelopt_for_weight_reload", - lambda model_arg, device: calls.append(("prepare", model_arg, device)), + reload_mod, + "initialize_layerwise_reload", + lambda root: calls.append(("initialize", root)), ) + + def finalize(root, config): + assert config_mod.get_current_vllm_config() is vllm_config + calls.append(("finalize", root, config)) + + monkeypatch.setattr(reload_mod, "finalize_layerwise_reload", finalize) monkeypatch.setattr( - patch_mod, - "modelopt_process_weights_after_loading", - lambda model_arg: calls.append(("process", model_arg)), + backend.torch.accelerator, + "synchronize", + lambda: calls.append("sync"), ) + + with extension._weight_update_lifecycle("collective") as finish: + # FlashInferExperts performs this lookup when online layer processing + # reconstructs its kernel during the yielded weight-load phase. + assert config_mod.get_current_vllm_config() is vllm_config + calls.append("load") + finish() + + assert config_mod.current is None + assert calls == [ + ("initialize", model), + "load", + ("finalize", model, model_config), + "sync", + ] + + +def test_real_quant_collective_reload_uses_vllm_layerwise_lifecycle(monkeypatch): + backend = _import_vllm_quant_backend(monkeypatch) + base_backend = _base_vllm_backend() + reload_mod = sys.modules["vllm.model_executor.model_loader.reload"] + + model = _mark_as_modelopt_layer(torch.nn.Linear(1, 1)) + model_config = object() + extension = object.__new__(backend.VllmQuantInternalWorkerExtension) + extension.model_runner = types.SimpleNamespace( + model=model, + vllm_config=object(), + ) + extension.model_config = model_config + extension.device = torch.device("cpu") + extension.state_dict_info = {} + extension.model_update_group = object() + calls = [] + monkeypatch.setattr( - backend.VllmInternalWorkerExtension, - "update_weights_from_collective", + backend.VllmQuantInternalWorkerExtension, + "_is_real_quant_model", lambda self: True, ) - extension.device = torch.device("cpu") + monkeypatch.setattr( + reload_mod, + "initialize_layerwise_reload", + lambda model_arg: calls.append(("initialize", model_arg)), + ) + monkeypatch.setattr( + base_backend, + "packed_broadcast_consumer", + lambda **kwargs: calls.append(("consume", kwargs["post_unpack_func"].__name__)), + ) + monkeypatch.setattr( + reload_mod, + "finalize_layerwise_reload", + lambda model_arg, config_arg: calls.append(("finalize", model_arg, config_arg)), + ) + monkeypatch.setattr( + backend.torch.accelerator, + "synchronize", + lambda: calls.append("sync"), + ) assert extension.update_weights_from_collective() is True assert calls == [ - ("prepare", model, torch.device("cpu")), - ("process", model), + ("initialize", model), + ("consume", "_load_weights"), + ("finalize", model, model_config), + "sync", ] -def test_real_quant_collective_reload_skips_processing_when_base_fails(monkeypatch): +def test_real_quant_collective_reload_raises_on_failure(monkeypatch): backend = _import_vllm_quant_backend(monkeypatch) - patch_mod = pytest.importorskip( - "nemo_rl.modelopt.models.generation.vllm_modelopt_patch" - ) + base_backend = _base_vllm_backend() + reload_mod = sys.modules["vllm.model_executor.model_loader.reload"] - model = torch.nn.Linear(1, 1) + model = _mark_as_modelopt_layer(torch.nn.Linear(1, 1)) extension = object.__new__(backend.VllmQuantInternalWorkerExtension) - extension.model_runner = types.SimpleNamespace(model=model) + extension.model_runner = types.SimpleNamespace( + model=model, + vllm_config=object(), + ) + extension.model_config = object() extension.device = torch.device("cpu") + extension.state_dict_info = {} + extension.model_update_group = object() calls = [] monkeypatch.setattr( @@ -611,23 +1581,26 @@ def test_real_quant_collective_reload_skips_processing_when_base_fails(monkeypat lambda self: True, ) monkeypatch.setattr( - patch_mod, - "prepare_modelopt_for_weight_reload", - lambda model_arg, device: calls.append(("prepare", model_arg, device)), - ) - monkeypatch.setattr( - patch_mod, - "modelopt_process_weights_after_loading", - lambda model_arg: calls.append(("process", model_arg)), + reload_mod, + "initialize_layerwise_reload", + lambda model_arg: calls.append(("initialize", model_arg)), ) + + def _raise_consume(**kwargs): + raise ValueError("broadcast boom") + + monkeypatch.setattr(base_backend, "packed_broadcast_consumer", _raise_consume) monkeypatch.setattr( - backend.VllmInternalWorkerExtension, - "update_weights_from_collective", - lambda self: False, + reload_mod, + "finalize_layerwise_reload", + lambda _model, _model_config: pytest.fail( + "a failed transfer must not be finalized" + ), ) - assert extension.update_weights_from_collective() is False - assert calls == [("prepare", model, torch.device("cpu"))] + with pytest.raises(RuntimeError, match="collective refit failed"): + extension.update_weights_from_collective() + assert calls == [("initialize", model)] def test_non_real_quant_collective_reload_delegates(monkeypatch): @@ -648,11 +1621,11 @@ def test_non_real_quant_collective_reload_delegates(monkeypatch): assert extension.update_weights_from_collective() == "delegated" -def test_real_quant_ipc_complete_processes_modelopt_and_acks(monkeypatch): +def test_real_quant_ipc_complete_finalizes_vllm_layerwise_reload_and_acks( + monkeypatch, +): backend = _import_vllm_quant_backend(monkeypatch) - patch_mod = pytest.importorskip( - "nemo_rl.modelopt.models.generation.vllm_modelopt_patch" - ) + reload_mod = sys.modules["vllm.model_executor.model_loader.reload"] from nemo_rl.models.policy.utils import IPCProtocol class FakeSocket: @@ -665,14 +1638,19 @@ def recv_pyobj(self): def send(self, payload): self.sent.append(payload) - model = torch.nn.Linear(1, 1) + model = _mark_as_modelopt_layer(torch.nn.Linear(1, 1)) + model_config = object() socket = FakeSocket() extension = object.__new__(backend.VllmQuantInternalWorkerExtension) - extension.model_runner = types.SimpleNamespace(model=model) + extension.model_runner = types.SimpleNamespace( + model=model, + vllm_config=object(), + ) + extension.model_config = model_config extension.device = torch.device("cpu") extension.zmq_socket = socket + extension.state_dict_info = {} extension.maybe_init_zmq = lambda: None - extension._maybe_process_fp8_kv_cache = lambda: None calls = [] monkeypatch.setattr( @@ -681,48 +1659,188 @@ def send(self, payload): lambda self: True, ) monkeypatch.setattr( - patch_mod, - "prepare_modelopt_for_weight_reload", - lambda model_arg, device: calls.append(("prepare", model_arg, device)), + reload_mod, + "initialize_layerwise_reload", + lambda model_arg: calls.append(("initialize", model_arg)), + ) + monkeypatch.setattr( + reload_mod, + "finalize_layerwise_reload", + lambda model_arg, config_arg: calls.append(("finalize", model_arg, config_arg)), ) monkeypatch.setattr( - patch_mod, - "modelopt_process_weights_after_loading", - lambda model_arg: calls.append(("process", model_arg)), + backend.torch.accelerator, + "synchronize", + lambda: calls.append("sync"), ) - monkeypatch.setattr(backend.torch.cuda, "synchronize", lambda: calls.append("sync")) monkeypatch.setattr( backend.torch.cuda, "empty_cache", lambda: calls.append("empty") ) assert extension.update_weights_via_ipc_zmq() is True assert calls == [ - ("prepare", model, torch.device("cpu")), - ("process", model), + ("initialize", model), + ("finalize", model, model_config), "sync", "empty", ] assert socket.sent == [IPCProtocol.ACK.value.encode()] -def test_real_quant_ipc_payload_loads_weights_and_handles_gpt_oss(monkeypatch): +def test_real_quant_ipc_finalize_failure_acks_complete(monkeypatch): + backend = _import_vllm_quant_backend(monkeypatch) + reload_mod = sys.modules["vllm.model_executor.model_loader.reload"] + from nemo_rl.models.policy.utils import IPCProtocol + + socket = types.SimpleNamespace( + recv_pyobj=lambda: IPCProtocol.COMPLETE, + sent=[], + ) + socket.send = socket.sent.append + extension = object.__new__(backend.VllmQuantInternalWorkerExtension) + extension.model_runner = types.SimpleNamespace( + model=_mark_as_modelopt_layer(torch.nn.Linear(1, 1)), + vllm_config=object(), + ) + extension.model_config = object() + extension.device = torch.device("cpu") + extension.zmq_socket = socket + extension.state_dict_info = {} + extension.maybe_init_zmq = lambda: None + monkeypatch.setattr( + backend.VllmQuantInternalWorkerExtension, + "_is_real_quant_model", + lambda _self: True, + ) + + def fail_finalize(_model, _model_config): + raise RuntimeError("bad scales") + + monkeypatch.setattr( + reload_mod, + "finalize_layerwise_reload", + fail_finalize, + ) + + with pytest.raises( + RuntimeError, match="ModelOpt real-quant refit post-processing failed" + ): + extension.update_weights_via_ipc_zmq() + assert socket.sent == [IPCProtocol.ACK.value.encode()] + + +@pytest.mark.parametrize( + ("payload_groups", "state_dict_info", "error"), + [ + ( + [["decoder.weight"]], + { + "decoder.weight": ([1], torch.float32), + "decoder.bias": ([1], torch.float32), + }, + "missing keys", + ), + ( + [["decoder.weight"], ["decoder.weight"]], + {"decoder.weight": ([1], torch.float32)}, + "duplicate keys", + ), + ( + [["decoder.weight", "decoder.weight"]], + {"decoder.weight": ([1], torch.float32)}, + "duplicate keys", + ), + ( + [["unexpected.weight"]], + {"decoder.weight": ([1], torch.float32)}, + "unexpected keys", + ), + ], +) +def test_real_quant_ipc_rejects_invalid_key_manifest( + monkeypatch, payload_groups, state_dict_info, error +): backend = _import_vllm_quant_backend(monkeypatch) - patch_mod = pytest.importorskip( - "nemo_rl.modelopt.models.generation.vllm_modelopt_patch" + base_backend = _base_vllm_backend() + reload_mod = sys.modules["vllm.model_executor.model_loader.reload"] + from nemo_rl.models.policy.utils import IPCProtocol + + payload_buffer = torch.tensor([1.0], dtype=torch.float32).view(torch.uint8) + used_bytes = base_backend.calculate_aligned_size(payload_buffer.numel()) + payloads = [ + ("ipc-handle", keys, used_bytes * len(keys)) for keys in payload_groups + ] + [IPCProtocol.COMPLETE] + + class FakeSocket: + def __init__(self): + self.payloads = iter(payloads) + self.sent = [] + + def recv_pyobj(self): + return next(self.payloads) + + def send(self, payload): + self.sent.append(payload) + + extension = object.__new__(backend.VllmQuantInternalWorkerExtension) + extension.model_runner = types.SimpleNamespace( + model=torch.nn.Linear(1, 1), + vllm_config=object(), + ) + extension.model_config = object() + extension.device = torch.device("cuda:0") + extension.zmq_socket = FakeSocket() + extension.state_dict_info = state_dict_info + extension.maybe_init_zmq = lambda: None + extension._load_weights = lambda _weights: None + monkeypatch.setattr( + backend.VllmQuantInternalWorkerExtension, + "_is_real_quant_model", + lambda _self: True, + ) + monkeypatch.setattr( + reload_mod, + "finalize_layerwise_reload", + lambda _model, _model_config: pytest.fail( + "an invalid refit must not be finalized" + ), + ) + monkeypatch.setattr( + base_backend, + "rebuild_cuda_tensor_from_ipc", + lambda _ipc_handle, _device_index: payload_buffer, + ) + monkeypatch.setattr( + base_backend.torch.cuda, + "current_stream", + lambda: types.SimpleNamespace(synchronize=lambda: None), ) + monkeypatch.setattr(backend.torch.accelerator, "synchronize", lambda: None) + + with pytest.raises(RuntimeError, match=error): + extension.update_weights_via_ipc_zmq() + assert extension.zmq_socket.sent == [IPCProtocol.ACK.value.encode()] * len(payloads) + + +def test_real_quant_ipc_payload_loads_weights_and_handles_gpt_oss(monkeypatch): + backend = _import_vllm_quant_backend(monkeypatch) + base_backend = _base_vllm_backend() + reload_mod = sys.modules["vllm.model_executor.model_loader.reload"] from nemo_rl.models.policy.utils import IPCProtocol payload_weight = torch.tensor([1.0, 2.0], dtype=torch.float32) payload_buffer = payload_weight.view(torch.uint8) - used_bytes = backend.calculate_aligned_size(payload_weight.nbytes) + used_bytes = base_backend.calculate_aligned_size(payload_weight.nbytes) loaded = [] calls = [] + view_refs = [] class FakeSocket: def __init__(self): self.payloads = iter( [ ("ipc-handle", ["decoder.weight"], used_bytes), + ("ipc-handle", ["decoder.bias"], used_bytes), IPCProtocol.COMPLETE, ] ) @@ -732,9 +1850,14 @@ def recv_pyobj(self): return next(self.payloads) def send(self, payload): + if len(self.sent) < 2: + assert view_refs + assert all(view_ref() is None for view_ref in view_refs) + calls.append("views_released") self.sent.append(payload) - model = torch.nn.Linear(1, 1) + model = _mark_as_modelopt_layer(torch.nn.Linear(1, 1)) + model_config = object() extension = object.__new__(backend.VllmQuantInternalWorkerExtension) extension.model_runner = types.SimpleNamespace( model=model, @@ -742,12 +1865,21 @@ def send(self, payload): model_config=types.SimpleNamespace(architectures=["GptOssForCausalLM"]) ), ) + extension.model_config = model_config extension.device = torch.device("cuda:0") extension.zmq_socket = FakeSocket() - extension.state_dict_info = {"decoder.weight": ([2], torch.float32)} + extension.state_dict_info = { + "decoder.weight": ([2], torch.float32), + "decoder.bias": ([2], torch.float32), + } extension.maybe_init_zmq = lambda: None - extension._maybe_process_fp8_kv_cache = lambda: calls.append("kv") - extension._load_weights = lambda weights: loaded.extend(weights) + + def load_weights(weights): + for name, weight in weights: + view_refs.append(weakref.ref(weight)) + loaded.append((name, weight.clone())) + + extension._load_weights = load_weights monkeypatch.setattr( backend.VllmQuantInternalWorkerExtension, @@ -755,40 +1887,53 @@ def send(self, payload): lambda self: True, ) monkeypatch.setattr( - patch_mod, - "prepare_modelopt_for_weight_reload", - lambda model_arg, device: calls.append(("prepare", model_arg, device)), + reload_mod, + "initialize_layerwise_reload", + lambda model_arg: calls.append(("initialize", model_arg)), ) monkeypatch.setattr( - patch_mod, - "modelopt_process_weights_after_loading", - lambda model_arg: calls.append(("process", model_arg)), + reload_mod, + "finalize_layerwise_reload", + lambda model_arg, config_arg: calls.append(("finalize", model_arg, config_arg)), ) monkeypatch.setattr( - backend, + base_backend, "rebuild_cuda_tensor_from_ipc", lambda ipc_handle, device_index: payload_buffer, ) - monkeypatch.setattr(backend.torch.cuda, "synchronize", lambda: calls.append("sync")) + monkeypatch.setattr( + base_backend.torch.cuda, + "current_stream", + lambda: pytest.fail("real quant must not use a current-stream IPC ACK fence"), + ) + monkeypatch.setattr( + backend.torch.accelerator, + "synchronize", + lambda: calls.append("sync"), + ) monkeypatch.setattr( backend.torch.cuda, "empty_cache", lambda: calls.append("empty") ) - monkeypatch.setattr(backend.gc, "collect", lambda: calls.append("gc")) + monkeypatch.setattr(base_backend.gc, "collect", lambda: calls.append("gc")) assert extension.update_weights_via_ipc_zmq() is True assert extension.zmq_socket.sent == [ IPCProtocol.ACK.value.encode(), IPCProtocol.ACK.value.encode(), + IPCProtocol.ACK.value.encode(), ] - assert loaded[0][0] == "decoder.weight" - torch.testing.assert_close(loaded[0][1], payload_weight) + assert [name for name, _ in loaded] == ["decoder.weight", "decoder.bias"] + for _, loaded_weight in loaded: + torch.testing.assert_close(loaded_weight, payload_weight) assert calls == [ - ("prepare", model, torch.device("cuda:0")), + ("initialize", model), "sync", - ("process", model), + "views_released", + "sync", + "views_released", + ("finalize", model, model_config), "sync", - "kv", "gc", "empty", ] @@ -854,6 +1999,167 @@ def __init__(self, enabled, amax): } +def _nvfp4_source_format() -> dict: + return { + "num_bits": "e2m1", + "block_sizes": { + -1: 16, + "type": "dynamic", + "scale_bits": "e4m3", + }, + } + + +def test_resolve_nvfp4_real_quant_mode_detects_model_specific_w4a16(monkeypatch): + resolved = { + "quant_cfg": [ + {"quantizer_name": "*", "enable": False}, + { + "quantizer_name": "*mixer.experts.*weight_quantizer", + "cfg": _nvfp4_source_format(), + }, + { + "quantizer_name": "*mlp.experts*weight_quantizer", + "cfg": _nvfp4_source_format(), + }, + ], + "algorithm": "max", + } + monkeypatch.setattr(modelopt_utils, "resolve_quant_cfg", lambda _: resolved) + + assert resolve_nvfp4_real_quant_mode("custom-nvfp4-config") == "w4a16" + + +def test_resolve_nvfp4_real_quant_mode_detects_w4a4(monkeypatch): + resolved = { + "quant_cfg": [ + {"quantizer_name": "*", "enable": False}, + { + "quantizer_name": "*mlp.experts*weight_quantizer", + "cfg": _nvfp4_source_format(), + }, + { + "quantizer_name": "*mlp.experts*input_quantizer", + "cfg": _nvfp4_source_format(), + }, + { + "quantizer_name": "*mlp.experts*input_quantizer", + "parent_class": "nn.LeakyReLU", + "enable": False, + }, + ], + "algorithm": "max", + } + monkeypatch.setattr(modelopt_utils, "resolve_quant_cfg", lambda _: resolved) + + assert resolve_nvfp4_real_quant_mode("not-named-after-the-format") == "w4a4" + + +def test_resolve_nvfp4_real_quant_mode_honors_late_generic_disable(monkeypatch): + resolved = { + "quant_cfg": [ + {"quantizer_name": "*", "enable": False}, + { + "quantizer_name": "*weight_quantizer", + "cfg": _nvfp4_source_format(), + }, + { + "quantizer_name": "*mlp.experts*input_quantizer", + "cfg": _nvfp4_source_format(), + }, + {"quantizer_name": "*input_quantizer", "enable": False}, + ], + "algorithm": "max", + } + monkeypatch.setattr(modelopt_utils, "resolve_quant_cfg", lambda _: resolved) + + assert resolve_nvfp4_real_quant_mode("disabled-input") == "w4a16" + + +@pytest.mark.parametrize( + ("weight_format", "input_format", "error_match"), + [ + ( + {"num_bits": "e4m3", "axis": None}, + {"num_bits": "e4m3", "axis": None}, + "only block-16 NVFP4.*weights", + ), + ( + _nvfp4_source_format(), + {"num_bits": "e4m3", "axis": None}, + "only block-16 NVFP4.*input activations", + ), + ( + _nvfp4_source_format(), + [_nvfp4_source_format()], + "single NVFP4 input activations format", + ), + ( + _nvfp4_source_format(), + { + "num_bits": "e2m1", + "block_sizes": { + -1: 32, + "type": "dynamic", + "scale_bits": "e4m3", + }, + }, + "only block-16 NVFP4.*input activations", + ), + ], + ids=["fp8", "w4a8", "sequential-activation", "unsupported-nvfp4-block"], +) +def test_resolve_nvfp4_real_quant_mode_rejects_unsupported_formats( + monkeypatch, + weight_format, + input_format, + error_match, +): + resolved = { + "quant_cfg": [ + {"quantizer_name": "*", "enable": False}, + {"quantizer_name": "*weight_quantizer", "cfg": weight_format}, + {"quantizer_name": "*input_quantizer", "cfg": input_format}, + ], + "algorithm": "max", + } + monkeypatch.setattr(modelopt_utils, "resolve_quant_cfg", lambda _: resolved) + + with pytest.raises(ValueError, match=error_match): + resolve_nvfp4_real_quant_mode("unsupported-real-quant-config") + + +def test_resolve_nvfp4_real_quant_mode_rejects_mixed_activation_formats( + monkeypatch, +): + resolved = { + "quant_cfg": [ + {"quantizer_name": "*", "enable": False}, + { + "quantizer_name": "*mixer.experts.*weight_quantizer", + "cfg": _nvfp4_source_format(), + }, + { + "quantizer_name": "*mlp.experts*weight_quantizer", + "cfg": _nvfp4_source_format(), + }, + { + "quantizer_name": "*mixer.experts.*input_quantizer", + "cfg": _nvfp4_source_format(), + }, + { + "quantizer_name": "*mlp.experts*input_quantizer", + "cfg": {"num_bits": "e4m3", "axis": None}, + }, + ], + "algorithm": "max", + } + monkeypatch.setattr(modelopt_utils, "resolve_quant_cfg", lambda _: resolved) + + with pytest.raises(ValueError, match="only block-16 NVFP4.*input activations"): + resolve_nvfp4_real_quant_mode("mixed-input-formats") + + def test_resolve_quant_cfg_passes_relative_names_to_modelopt(monkeypatch): modelopt_recipe = pytest.importorskip("modelopt.recipe") captured = {} @@ -935,335 +2241,385 @@ def test_resolve_quant_cfg_rejects_recipe_without_quant_cfg(monkeypatch): resolve_quant_cfg("missing-quant-cfg") -def test_vllm_reload_canonicalizes_nvfp4_scales_before_kernel_conversion(): - layer = torch.nn.Module() - layer.weight_scale = torch.nn.Parameter( - torch.tensor([[1.0, -2.0], [-0.5, 4.0]]), - requires_grad=False, - ) - - _canonicalize_nvfp4_weight_scale(layer) - - torch.testing.assert_close( - layer.weight_scale, - torch.tensor([[1.0, 2.0], [0.5, 4.0]]), - ) - - -def test_prepare_modelopt_for_weight_reload_restores_deleted_dense_params(): - layer = torch.nn.Module() - layer.weight = torch.nn.Parameter(torch.ones(2, 2), requires_grad=False) - layer.weight_scale = torch.nn.Parameter(torch.ones(2, 1), requires_grad=False) - layer._nrl_modelopt_param_meta = { - "weight": { - "shape": (2, 2), - "dtype": torch.float32, - "device": "cpu", - "param_class": torch.nn.Parameter, - }, - "weight_scale_2": { - "shape": (1,), - "dtype": torch.float32, - "device": "cpu", - "param_class": torch.nn.Parameter, - }, - } - layer._nrl_modelopt_weight_loaders = {} - model = torch.nn.Module() - model.layer = layer - - prepare_modelopt_for_weight_reload(model, device=torch.device("cpu")) - - assert hasattr(layer, "weight_scale_2") - assert tuple(layer.weight_scale_2.shape) == (1,) - assert layer.weight_scale_2.dtype == torch.float32 - assert layer.weight_scale_2.device.type == "cpu" - +def test_register_nemo_modelopt_nvfp4_uses_public_vllm_registry(monkeypatch): + fake_vllm = _install_fake_registered_vllm_modelopt(monkeypatch) + monkeypatch.setattr(vllm_modelopt, "_registered", False) -def test_prepare_modelopt_for_weight_reload_restores_loader_class_when_shape_matches(): - class FakeModelWeightParameter(torch.nn.Parameter): - def __new__(cls, data, **kwargs): - return super().__new__(cls, data=data, requires_grad=False) + register_nemo_modelopt_nvfp4() - def __init__(self, data, weight_loader, input_dim=1, output_dim=0): - self.weight_loader = weight_loader - self._input_dim = input_dim - self._output_dim = output_dim - - def fake_merged_loader(param, loaded_weight, shard_id): - shard_size = loaded_weight.shape[0] - shard_offset = shard_id * shard_size - param.data.narrow(0, shard_offset, shard_size).copy_(loaded_weight) - - layer = torch.nn.Module() - layer.weight = torch.nn.Parameter(torch.zeros(4, 2), requires_grad=False) - layer._nrl_modelopt_param_meta = { - "weight": { - "shape": (4, 2), - "dtype": torch.float32, - "device": "cpu", - "param_class": FakeModelWeightParameter, - "input_dim": 1, - "output_dim": 0, - }, - } - layer._nrl_modelopt_weight_loaders = {"weight": fake_merged_loader} - model = torch.nn.Module() - model.layer = layer - - prepare_modelopt_for_weight_reload(model, device=torch.device("cpu")) - - assert isinstance(layer.weight, FakeModelWeightParameter) - assert layer.weight.weight_loader is fake_merged_loader - layer.weight.data.zero_() - layer.weight.weight_loader(layer.weight, torch.ones(2, 2), 1) - torch.testing.assert_close(layer.weight[:2], torch.zeros(2, 2)) - torch.testing.assert_close(layer.weight[2:], torch.ones(2, 2)) - - -def test_prepare_modelopt_for_weight_reload_restores_plain_parameter_loader(): - def fake_loader(param, loaded_weight): - param.data.copy_(loaded_weight) - - layer = torch.nn.Module() - layer.weight = torch.nn.Parameter(torch.zeros(2, 2), requires_grad=False) - layer._nrl_modelopt_param_meta = { - "weight": { - "shape": (2, 2), - "dtype": torch.float32, - "device": "cpu", - "param_class": torch.nn.Parameter, - }, + assert set(fake_vllm.registry) == { + NEMO_MODELOPT_W4A4, + NEMO_MODELOPT_W4A16, } - layer._nrl_modelopt_weight_loaders = {"weight": fake_loader} - model = torch.nn.Module() - model.layer = layer - - prepare_modelopt_for_weight_reload(model, device=torch.device("cpu")) - - assert isinstance(layer.weight, torch.nn.Parameter) - assert layer.weight.weight_loader is fake_loader - layer.weight.weight_loader(layer.weight, torch.ones(2, 2)) - torch.testing.assert_close(layer.weight, torch.ones(2, 2)) + w4a4_config = fake_vllm.registry[NEMO_MODELOPT_W4A4]() + assert w4a4_config.get_name() == NEMO_MODELOPT_W4A4 + source_config = {"quant_algo": "W4A16_NVFP4", "group_size": 16} + w4a16_config = fake_vllm.registry[NEMO_MODELOPT_W4A16].from_config(source_config) + assert source_config["quant_algo"] == "W4A16_NVFP4" + assert w4a16_config.parsed_config["quant_algo"] == "NVFP4" + assert w4a16_config.get_name() == NEMO_MODELOPT_W4A16 -def test_modelopt_process_weights_after_loading_runs_dense_quant_method(): - calls = [] - - class ModelOptNvFp4LinearMethod: - def process_weights_after_loading(self, layer): - calls.append(layer) - - model = torch.nn.Module() - model.layer = torch.nn.Module() - model.layer.quant_method = ModelOptNvFp4LinearMethod() - - modelopt_process_weights_after_loading(model) + with pytest.raises(ValueError, match="requires quant_algo='W4A16_NVFP4'"): + fake_vllm.registry[NEMO_MODELOPT_W4A16].from_config({"quant_algo": "NVFP4"}) - assert calls == [model.layer] +def test_registered_configs_select_only_the_exact_custom_override(monkeypatch): + fake_vllm = _install_fake_registered_vllm_modelopt(monkeypatch) + monkeypatch.setattr(vllm_modelopt, "_registered", False) + register_nemo_modelopt_nvfp4() + w4a4_config_cls = fake_vllm.registry[NEMO_MODELOPT_W4A4] + w4a16_config_cls = fake_vllm.registry[NEMO_MODELOPT_W4A16] -def test_apply_modelopt_nvfp4_patches_updates_vllm_method(monkeypatch): - patch_mod = pytest.importorskip( - "nemo_rl.modelopt.models.generation.vllm_modelopt_patch" + assert ( + w4a4_config_cls.override_quantization_method( + {"quant_algo": "NVFP4"}, NEMO_MODELOPT_W4A4 + ) + == NEMO_MODELOPT_W4A4 ) - module_names = [ - "vllm", - "vllm.model_executor", - "vllm.model_executor.layers", - "vllm.model_executor.layers.quantization", - ] - for module_name in module_names: - monkeypatch.setitem(sys.modules, module_name, types.ModuleType(module_name)) - modelopt_module = types.ModuleType( - "vllm.model_executor.layers.quantization.modelopt" + assert ( + w4a16_config_cls.override_quantization_method( + {"quantization": {"quant_algo": "W4A16_NVFP4"}}, + NEMO_MODELOPT_W4A16, + ) + == NEMO_MODELOPT_W4A16 ) - - class FakeModelOptNvFp4Config: - @classmethod - def _from_config(cls, **kwargs): - return types.SimpleNamespace() - - def fake_linear_apply(self, layer, x, bias=None): - pass - - class FakeModelOptNvFp4LinearMethod: - apply = fake_linear_apply - process_weights_after_loading = None - - modelopt_module.ModelOptNvFp4Config = FakeModelOptNvFp4Config - modelopt_module.ModelOptNvFp4LinearMethod = FakeModelOptNvFp4LinearMethod - monkeypatch.setitem( - sys.modules, - "vllm.model_executor.layers.quantization.modelopt", - modelopt_module, + assert ( + w4a4_config_cls.override_quantization_method( + {"quant_algo": "NVFP4"}, NEMO_MODELOPT_W4A16 + ) + is None ) - monkeypatch.setattr(patch_mod, "_patched", False) - - apply_modelopt_nvfp4_patches() - apply_modelopt_nvfp4_patches() - - cfg = FakeModelOptNvFp4Config._from_config( - quant_method="NVFP4", - kv_cache_quant_method=None, - exclude_modules=[], - original_config=build_vllm_modelopt_nvfp4_config(), - group_size=16, + assert ( + w4a4_config_cls.override_quantization_method( + {"quant_algo": "W4A16_NVFP4"}, NEMO_MODELOPT_W4A4 + ) + is None ) - - assert getattr(cfg, "_nrl_weight_only_w4a16") is True - assert FakeModelOptNvFp4LinearMethod._nrl_original_apply is fake_linear_apply assert ( - FakeModelOptNvFp4LinearMethod.process_weights_after_loading - is _modelopt_dense_process_weights + w4a16_config_cls.override_quantization_method( + {"quant_algo": "W4A16_NVFP4"}, "modelopt" + ) + is None ) - assert FakeModelOptNvFp4LinearMethod.apply is _modelopt_dense_apply - assert patch_mod._patched is True -def test_convert_nvfp4_linear_kernel_format_uses_vllm_fallback(monkeypatch): - calls = [] - module_names = [ - "vllm", - "vllm.model_executor", - "vllm.model_executor.layers", - "vllm.model_executor.layers.quantization", - "vllm.model_executor.layers.quantization.utils", +def test_registered_w4a16_dense_method_supports_weight_loader_v2(monkeypatch): + fake_vllm = _install_fake_registered_vllm_modelopt(monkeypatch) + monkeypatch.setattr(vllm_modelopt, "_registered", False) + + register_nemo_modelopt_nvfp4() + + w4a16_config_cls = fake_vllm.registry[NEMO_MODELOPT_W4A16] + assert fake_vllm.weight_loader_v2_supported == [ + w4a16_config_cls.LinearMethodCls.__name__ ] - for module_name in module_names: - monkeypatch.setitem(sys.modules, module_name, types.ModuleType(module_name)) - nvfp4_utils = types.ModuleType( - "vllm.model_executor.layers.quantization.utils.nvfp4_utils" - ) - def fake_convert(backend, layer): - calls.append((backend, layer)) - nvfp4_utils.convert_to_nvfp4_linear_kernel_format = fake_convert - monkeypatch.setitem( - sys.modules, - "vllm.model_executor.layers.quantization.utils.nvfp4_utils", - nvfp4_utils, - ) - layer = torch.nn.Module() - quant_method = types.SimpleNamespace(backend="backend") +def test_registered_w4a4_moe_loader_is_sanitizer_compatible(monkeypatch): + fake_vllm = _install_fake_registered_vllm_modelopt(monkeypatch) + monkeypatch.setattr(vllm_modelopt, "_registered", False) + register_nemo_modelopt_nvfp4() - _convert_nvfp4_linear_kernel_format(quant_method, layer) + config = fake_vllm.registry[NEMO_MODELOPT_W4A4]() + quant_method = config.FusedMoEMethodCls(config, object()) - assert calls == [("backend", layer)] + class FakeMoeLayer(torch.nn.Module): + def __init__(self): + super().__init__() + self.quant_method = quant_method + self.w13_input_scale = torch.nn.Parameter(torch.zeros(2, 2)) + self.w2_input_scale = torch.nn.Parameter(torch.zeros(2)) + + def _map_global_expert_id_to_local_expert_id(self, expert_id): + return expert_id + + layer = FakeMoeLayer() + quant_method.create_weights(layer) + + w13_loader = layer.w13_input_scale.weight_loader + assert isinstance(w13_loader, types.MethodType) + assert w13_loader.__self__ is layer + + layer_ref_sentinel = object() + layer.w13_input_scale.weight_loader = w13_loader.__func__.__get__( + layer_ref_sentinel + ) + assert layer.w13_input_scale.weight_loader.__self__ is layer_ref_sentinel + layer.w13_input_scale.weight_loader = ( + layer.w13_input_scale.weight_loader.__func__.__get__(layer) + ) + w13_loader = layer.w13_input_scale.weight_loader + assert w13_loader.__self__ is layer + + assert w13_loader( + layer.w13_input_scale, + torch.tensor(1.0), + "gate.input_scale", + "w1", + 0, + True, + ) + assert w13_loader( + layer.w13_input_scale, + torch.tensor(2.0), + "up.input_scale", + "w3", + 0, + True, + ) + assert layer.w2_input_scale.weight_loader( + layer.w2_input_scale, + torch.tensor(3.0), + "down.input_scale", + "w2", + 1, + True, + ) + + torch.testing.assert_close(layer.w13_input_scale[0], torch.tensor([1.0, 2.0])) + torch.testing.assert_close(layer.w2_input_scale, torch.tensor([0.0, 3.0])) + + +def test_registered_w4a4_moe_materializes_initial_input_scales(monkeypatch): + fake_vllm = _install_fake_registered_vllm_modelopt(monkeypatch) + monkeypatch.setattr(vllm_modelopt, "_registered", False) + register_nemo_modelopt_nvfp4() + config = fake_vllm.registry[NEMO_MODELOPT_W4A4]() + quant_method = config.FusedMoEMethodCls(config, object()) + layer = torch.nn.Module() + w13_input_scale = torch.nn.Parameter( + torch.tensor([2.0]).expand(4), requires_grad=False + ) + w2_input_scale = torch.nn.Parameter( + torch.tensor([3.0]).expand(4), requires_grad=False + ) + layer.register_parameter("w13_input_scale", w13_input_scale) + layer.register_parameter("w2_input_scale", w2_input_scale) + + quant_method.process_weights_after_loading(layer) + + assert layer.w13_input_scale is w13_input_scale + assert layer.w2_input_scale is w2_input_scale + assert layer.w13_input_scale.is_contiguous() + assert layer.w2_input_scale.is_contiguous() + torch.testing.assert_close(layer.w13_input_scale, torch.full((4,), 2.0)) + torch.testing.assert_close(layer.w2_input_scale, torch.full((4,), 3.0)) + with torch.no_grad(): + layer.w13_input_scale.copy_(torch.arange(4, dtype=torch.float32)) + layer.w2_input_scale.copy_(torch.arange(4, dtype=torch.float32)) + + +def test_registered_w4a4_moe_refreshes_stable_activation_scales(monkeypatch): + fake_vllm = _install_fake_registered_vllm_modelopt(monkeypatch) + monkeypatch.setattr(vllm_modelopt, "_registered", False) + register_nemo_modelopt_nvfp4() + config = fake_vllm.registry[NEMO_MODELOPT_W4A4]() + quant_method = config.FusedMoEMethodCls(config, object()) + + original_kernel = object() + original_a1_gscale = torch.full((4,), 1.0) + original_a2_gscale = torch.full((4,), 0.5) + original_quant_config = types.SimpleNamespace( + a1_gscale=original_a1_gscale, + a2_gscale=original_a2_gscale, + ) + quant_method.moe_kernel = original_kernel + quant_method.moe_quant_config = original_quant_config + a1_data_ptr = original_a1_gscale.data_ptr() + a2_data_ptr = original_a2_gscale.data_ptr() -def test_modelopt_dense_process_uses_vllm_kernel_api(): layer = torch.nn.Module() - layer.weight = torch.nn.Parameter(torch.ones(2, 1), requires_grad=False) - layer.weight._input_dim = 1 - layer.weight._output_dim = 0 - layer.weight.weight_loader = lambda param, loaded_weight: None - layer.weight_scale = torch.nn.Parameter( - torch.tensor([[1.0, -2.0], [-0.5, 4.0]]), - requires_grad=False, + layer.register_parameter( + "w13_input_scale", + torch.nn.Parameter(torch.full((4,), 4.0), requires_grad=False), + ) + layer.register_parameter( + "w2_input_scale", + torch.nn.Parameter(torch.full((4,), 5.0), requires_grad=False), ) - layer.weight_scale_2 = torch.nn.Parameter(torch.tensor([2.0]), requires_grad=False) - - calls = [] - - class FakeKernel: - def process_weights_after_loading(self, processed_layer): - calls.append(processed_layer) - torch.testing.assert_close( - processed_layer.weight_scale, - torch.tensor([[1.0, 2.0], [0.5, 4.0]]), - ) - - quant_method = types.SimpleNamespace(kernel=FakeKernel()) - - _modelopt_dense_process_weights(quant_method, layer) - assert calls == [layer] - assert not hasattr(layer, "weight_scale_2") - torch.testing.assert_close(layer.input_global_scale, torch.tensor(1.0)) - torch.testing.assert_close(layer.weight_global_scale, torch.tensor(2.0)) - torch.testing.assert_close(layer.alpha, torch.tensor(2.0)) - torch.testing.assert_close(layer.input_global_scale_inv, torch.tensor(1.0)) - assert set(layer._nrl_modelopt_param_meta) == { - "weight", - "weight_scale", - "weight_scale_2", - } - assert layer._nrl_modelopt_param_meta["weight"]["input_dim"] == 1 - assert layer._nrl_modelopt_param_meta["weight"]["output_dim"] == 0 - assert "weight" in layer._nrl_modelopt_weight_loaders + quant_method.process_weights_after_loading(layer) + assert quant_method.moe_kernel is original_kernel + assert quant_method.moe_quant_config is original_quant_config + assert original_quant_config.a1_gscale.data_ptr() == a1_data_ptr + assert original_quant_config.a2_gscale.data_ptr() == a2_data_ptr + torch.testing.assert_close(original_quant_config.a1_gscale, torch.full((4,), 0.25)) + torch.testing.assert_close(original_quant_config.a2_gscale, torch.full((4,), 0.2)) -def test_modelopt_dense_process_w4a16_uses_marlin_weight_only(monkeypatch): - module_names = [ - "vllm", - "vllm.model_executor", - "vllm.model_executor.layers", - "vllm.model_executor.layers.quantization", - "vllm.model_executor.layers.quantization.utils", - ] - for module_name in module_names: - monkeypatch.setitem(sys.modules, module_name, types.ModuleType(module_name)) - marlin_utils = types.ModuleType( - "vllm.model_executor.layers.quantization.utils.marlin_utils_fp4" +def test_registered_w4a16_dense_method_uses_marlin_weight_only(monkeypatch): + fake_vllm = _install_fake_registered_vllm_modelopt(monkeypatch) + monkeypatch.setattr(vllm_modelopt, "_registered", False) + register_nemo_modelopt_nvfp4() + config = fake_vllm.registry[NEMO_MODELOPT_W4A16].from_config( + {"quant_algo": "W4A16_NVFP4", "group_size": 16} ) - calls = [] - - def fake_prepare(layer): - calls.append(("prepare", layer)) - layer.workspace = torch.empty(1) + quant_method = config.LinearMethodCls(config) - def fake_apply(**kwargs): - calls.append(("apply", kwargs)) - return "out" - - marlin_utils.prepare_fp4_layer_for_marlin = fake_prepare - marlin_utils.apply_fp4_marlin_linear = fake_apply - monkeypatch.setitem( - sys.modules, - "vllm.model_executor.layers.quantization.utils.marlin_utils_fp4", - marlin_utils, - ) + created_layer = torch.nn.Module() + quant_method.create_weights(created_layer) + assert not hasattr(created_layer, "input_scale") layer = torch.nn.Module() layer.weight = torch.nn.Parameter(torch.ones(2, 1), requires_grad=False) layer.weight_scale = torch.nn.Parameter( - torch.tensor([[1.0, -2.0], [-0.5, 4.0]]), + torch.tensor([[-1.0, 2.0], [0.5, -4.0]]), requires_grad=False, ) - layer.weight_scale_2 = torch.nn.Parameter(torch.tensor([2.0]), requires_grad=False) - layer.input_scale = torch.nn.Parameter(torch.tensor([3.0]), requires_grad=False) - layer.input_global_scale = torch.nn.Parameter( - torch.tensor([4.0]), - requires_grad=False, - ) - layer.alpha = torch.nn.Parameter(torch.tensor([5.0]), requires_grad=False) - layer.input_global_scale_inv = torch.nn.Parameter( - torch.tensor([6.0]), + layer.weight_scale_2 = torch.nn.Parameter( + torch.tensor([2.0, 3.0]), requires_grad=False, ) layer.output_size_per_partition = 2 layer.input_size_per_partition = 2 - quant_method = types.SimpleNamespace( - quant_config=types.SimpleNamespace(_nrl_weight_only_w4a16=True) - ) - _modelopt_dense_process_weights(quant_method, layer) - result = _modelopt_dense_apply(quant_method, layer, torch.ones(1, 2)) + quant_method.process_weights_after_loading(layer) + output = quant_method.apply(layer, torch.ones(1, 2)) - assert result == "out" - assert calls[0] == ("prepare", layer) - assert calls[1][0] == "apply" - assert not hasattr(layer, "input_scale") - assert not hasattr(layer, "input_global_scale") - assert not hasattr(layer, "alpha") - assert not hasattr(layer, "input_global_scale_inv") + assert output == "output" assert not hasattr(layer, "weight_scale_2") - torch.testing.assert_close(layer.weight_global_scale, torch.tensor(2.0)) torch.testing.assert_close( layer.weight_scale, torch.tensor([[1.0, 2.0], [0.5, 4.0]]), ) - assert calls[1][1]["weight_global_scale"] is layer.weight_global_scale + torch.testing.assert_close(layer.weight_global_scale, torch.tensor(3.0)) + assert fake_vllm.events[0] == ("process_marlin_kernel", layer) + event_name, kernel_args = fake_vllm.events[1] + assert event_name == "apply_marlin_kernel" + assert kernel_args["layer"] is layer + torch.testing.assert_close(kernel_args["x"], torch.ones(1, 2)) + assert kernel_args["bias"] is None + + +@pytest.mark.parametrize( + ("is_act_and_mul", "packed_hidden_size", "expected_padded_size"), + [ + (False, 64, 192), + (True, 32, 256), + ], +) +def test_pad_nvfp4_moe_for_marlin_uses_hidden_size_tile_alignment( + is_act_and_mul, + packed_hidden_size, + expected_padded_size, +): + num_shards = 2 if is_act_and_mul else 1 + intermediate_size = 144 + w13 = torch.ones( + 1, + num_shards * intermediate_size, + packed_hidden_size, + ) + w13_scale = torch.ones(1, num_shards * intermediate_size, 2) + w2 = torch.ones(1, 2, intermediate_size // 2) + w2_scale = torch.ones(1, 2, intermediate_size // 16) + + padded_w13, padded_w13_scale, padded_w2, padded_w2_scale, padded_size = ( + _pad_nvfp4_moe_for_marlin( + w13, + w13_scale, + w2, + w2_scale, + is_act_and_mul=is_act_and_mul, + ) + ) + + assert padded_size == expected_padded_size + assert padded_w13.shape == ( + 1, + num_shards * expected_padded_size, + packed_hidden_size, + ) + assert padded_w13_scale.shape == (1, num_shards * expected_padded_size, 2) + assert padded_w2.shape == (1, 2, expected_padded_size // 2) + assert padded_w2_scale.shape == (1, 2, expected_padded_size // 16) + + padded_w13 = padded_w13.view( + 1, num_shards, expected_padded_size, packed_hidden_size + ) + padded_w13_scale = padded_w13_scale.view(1, num_shards, expected_padded_size, 2) + assert torch.all(padded_w13[:, :, :intermediate_size] == 1) + assert torch.count_nonzero(padded_w13[:, :, intermediate_size:]) == 0 + assert torch.all(padded_w13_scale[:, :, :intermediate_size] == 1) + assert torch.count_nonzero(padded_w13_scale[:, :, intermediate_size:]) == 0 + assert torch.all(padded_w2[..., : intermediate_size // 2] == 1) + assert torch.count_nonzero(padded_w2[..., intermediate_size // 2 :]) == 0 + assert torch.all(padded_w2_scale[..., : intermediate_size // 16] == 1) + assert torch.count_nonzero(padded_w2_scale[..., intermediate_size // 16 :]) == 0 + + +def test_registered_w4a16_moe_create_weights_keeps_checkpoint_layout(monkeypatch): + fake_vllm = _install_fake_registered_vllm_modelopt(monkeypatch) + monkeypatch.setattr(vllm_modelopt, "_registered", False) + register_nemo_modelopt_nvfp4() + config = fake_vllm.registry[NEMO_MODELOPT_W4A16].from_config( + {"quant_algo": "W4A16_NVFP4", "group_size": 16} + ) + quant_method = config.FusedMoEMethodCls( + config, + types.SimpleNamespace(is_act_and_mul=False), + ) + layer = torch.nn.Module() + + quant_method.create_weights( + layer, + num_experts=2, + hidden_size=4096, + intermediate_size_per_partition=672, + params_dtype=torch.bfloat16, + ) + + assert not hasattr(layer, "w13_input_scale") + assert not hasattr(layer, "w2_input_scale") + assert fake_vllm.events == [ + ( + "native_create_weights", + layer, + (2, 4096, 672, torch.bfloat16), + {}, + ) + ] + + +def test_registered_w4a16_moe_preserves_kernel_during_reload(monkeypatch): + fake_vllm = _install_fake_registered_vllm_modelopt(monkeypatch) + monkeypatch.setattr(vllm_modelopt, "_registered", False) + register_nemo_modelopt_nvfp4() + config = fake_vllm.registry[NEMO_MODELOPT_W4A16].from_config( + {"quant_algo": "W4A16_NVFP4", "group_size": 16} + ) + quant_method = config.FusedMoEMethodCls( + config, + types.SimpleNamespace(is_act_and_mul=False), + ) + original_kernel = object() + original_quant_config = object() + quant_method.moe_kernel = original_kernel + quant_method.moe_quant_config = original_quant_config + + layer = torch.nn.Module() + layer.w13_weight = torch.nn.Parameter(torch.ones(1, 80, 32)) + layer.w13_weight_scale = torch.nn.Parameter(-torch.ones(1, 80, 2)) + layer.w13_weight_scale_2 = torch.nn.Parameter(torch.ones(1, 1)) + layer.w2_weight = torch.nn.Parameter(torch.ones(1, 2, 40)) + layer.w2_weight_scale = torch.nn.Parameter(-torch.ones(1, 2, 5)) + layer.w2_weight_scale_2 = torch.nn.Parameter(torch.ones(1)) + layer.moe_config = types.SimpleNamespace(intermediate_size_per_partition=80) + layer.shared_experts = None + layer._maybe_init_expert_routing_tables = lambda: None + + quant_method.process_weights_after_loading(layer) + + assert quant_method.moe_kernel is original_kernel + assert quant_method.moe_quant_config is original_quant_config + assert layer.moe_config.intermediate_size_per_partition == 80 + assert layer.w13_weight.shape == (1, 128, 32) + assert layer.w13_weight_scale.shape == (1, 128, 2) + assert layer.w2_weight.shape == (1, 2, 64) + assert layer.w2_weight_scale.shape == (1, 2, 8) + assert torch.all(layer.w13_weight_scale >= 0) + assert torch.all(layer.w2_weight_scale >= 0) + assert fake_vllm.events == [("native_process_moe", 128)] diff --git a/tests/unit/models/megatron/test_megatron_setup.py b/tests/unit/models/megatron/test_megatron_setup.py index 02c634e3fe9..53f25fbbf4a 100644 --- a/tests/unit/models/megatron/test_megatron_setup.py +++ b/tests/unit/models/megatron/test_megatron_setup.py @@ -707,7 +707,7 @@ def test_hybridep_env_vars_auto_set_with_warning(self, monkeypatch): _apply_moe_config(model_cfg, config) assert model_cfg.moe_flex_dispatcher_backend == "hybridep" - assert model_cfg.moe_hybridep_num_sms == 32 + assert model_cfg.moe_flex_dispatcher_num_sms == 32 # min(ep_size=8, 64) == 8 assert os.environ["NUM_OF_HYBRID_EP_RANKS_PER_NVLINK_DOMAIN"] == "8" # int(ep_size=8 > 4) == 1 @@ -719,6 +719,25 @@ def test_hybridep_env_vars_auto_set_with_warning(self, monkeypatch): ) assert any("USE_MNNVL not configured" in m for m in warn_messages) + def test_hybridep_num_sms_supports_old_mcore(self, monkeypatch): + """The existing recipe key still targets legacy MCore releases.""" + from nemo_rl.models.megatron.setup import _apply_moe_config + + monkeypatch.setenv("NUM_OF_HYBRID_EP_RANKS_PER_NVLINK_DOMAIN", "8") + monkeypatch.setenv("USE_MNNVL", "1") + model_cfg = SimpleNamespace() + config = self._base_moe_cfg( + expert_model_parallel_size=8, + moe_flex_dispatcher_backend="hybridep", + moe_hybridep_num_sms=32, + ) + + with patch("nemo_rl.models.megatron.setup.TransformerConfig", new=object): + _apply_moe_config(model_cfg, config) + + assert not hasattr(model_cfg, "moe_flex_dispatcher_num_sms") + assert model_cfg.moe_hybridep_num_sms == 32 + def test_hybridep_env_vars_from_explicit_config(self, monkeypatch): """Explicit hybridep_* config keys override defaults without warnings.""" from nemo_rl.models.megatron.setup import _apply_moe_config @@ -1991,11 +2010,12 @@ def test_megatron_bridge_with_hf_config_overrides_warns(self, tmp_path, request) "megatron_cfg": {}, } - mock_cfg = MagicMock() - mock_cfg.model = self._make_model_cfg_mock() + mock_model_cfg = self._make_model_cfg_mock() - with patch("nemo_rl.models.megatron.setup.ConfigContainer") as mock_cc: - mock_cc.from_yaml.return_value = mock_cfg + with patch( + "nemo_rl.models.megatron.setup.load_model_config", + return_value=(mock_model_cfg, None), + ) as mock_load_model_config: with pytest.warns( UserWarning, match="hf_config_overrides is set but will be ignored" ): @@ -2007,6 +2027,8 @@ def test_megatron_bridge_with_hf_config_overrides_warns(self, tmp_path, request) pretrained_path=str(tmp_path), ) + mock_load_model_config.assert_called_once_with(str(tmp_path)) + def test_megatron_bridge_without_hf_config_overrides_no_warning( self, tmp_path, request ): @@ -2027,11 +2049,12 @@ def test_megatron_bridge_without_hf_config_overrides_no_warning( "megatron_cfg": {}, } - mock_cfg = MagicMock() - mock_cfg.model = self._make_model_cfg_mock() + mock_model_cfg = self._make_model_cfg_mock() - with patch("nemo_rl.models.megatron.setup.ConfigContainer") as mock_cc: - mock_cc.from_yaml.return_value = mock_cfg + with patch( + "nemo_rl.models.megatron.setup.load_model_config", + return_value=(mock_model_cfg, None), + ) as mock_load_model_config: with _warnings.catch_warnings(): _warnings.simplefilter("error", UserWarning) # Should not raise @@ -2043,6 +2066,40 @@ def test_megatron_bridge_without_hf_config_overrides_no_warning( pretrained_path=str(tmp_path), ) + mock_load_model_config.assert_called_once_with(str(tmp_path)) + + def test_hf_conversion_loads_model_config_from_iteration_dir( + self, tmp_path, request + ): + """Converted HF caches enter through Bridge's compatibility loader.""" + from nemo_rl.models.megatron.setup import setup_model_config + + self._apply_patches(request) + + iteration_dir = tmp_path / "iter_0000000" + iteration_dir.mkdir() + (iteration_dir / "run_config.yaml").touch() + mock_model_cfg = self._make_model_cfg_mock() + + config = { + "pretrained_checkpoint": None, + "megatron_cfg": {}, + } + + with patch( + "nemo_rl.models.megatron.setup.load_model_config", + return_value=(mock_model_cfg, None), + ) as mock_load_model_config: + setup_model_config( + config, + rank=0, + dtype=torch.bfloat16, + hf_model_name="test-model", + pretrained_path=str(tmp_path), + ) + + mock_load_model_config.assert_called_once_with(str(iteration_dir)) + @pytest.mark.mcore class TestHandleModelImport: diff --git a/tests/unit/models/policy/test_megatron_quant_worker.py b/tests/unit/models/policy/test_megatron_quant_worker.py index 07a8d0b9ff3..b82b4855cbd 100644 --- a/tests/unit/models/policy/test_megatron_quant_worker.py +++ b/tests/unit/models/policy/test_megatron_quant_worker.py @@ -85,13 +85,14 @@ def _make_real_quant_worker(): worker_cls = MegatronQuantPolicyWorker.__ray_metadata__.modified_class worker = object.__new__(worker_cls) worker.cfg = { + "quant_cfg": "examples/modelopt/quant_configs/nvfp4_a16_mlp_only.yaml", "generation": { "backend": "vllm", "quant_cfg": "examples/modelopt/quant_configs/nvfp4_a16_mlp_only.yaml", "real_quant": True, "real_quant_ignore": ["lm_head"], - "vllm_cfg": {}, - } + "vllm_cfg": {"kv_cache_dtype": "auto"}, + }, } worker.model = object() worker.draft_model = None @@ -101,6 +102,44 @@ def _make_real_quant_worker(): return worker +@requires_weight_folding +def test_modelopt_policy_worker_uses_real_quant_refit_timeout(monkeypatch): + from nemo_rl.modelopt.models.policy.workers import megatron_quant_policy_worker + + events = [] + + class FakeSocket: + def setsockopt(self, option, value): + events.append(("setsockopt", option, value)) + + def bind(self, address): + events.append(("bind", address)) + + class FakeContext: + def socket(self, socket_type): + events.append(("socket", socket_type)) + return FakeSocket() + + worker_cls = MegatronQuantPolicyWorker.__ray_metadata__.modified_class + worker = object.__new__(worker_cls) + worker._use_real_quant_refit = lambda: True + worker.get_zmq_address = lambda: "ipc:///tmp/modelopt-test.sock" + monkeypatch.setattr(megatron_quant_policy_worker.zmq, "Context", FakeContext) + + worker.maybe_init_zmq() + + timeout = megatron_quant_policy_worker.MODELOPT_REAL_QUANT_ZMQ_TIMEOUT_MS + assert events == [ + ("socket", megatron_quant_policy_worker.zmq.REQ), + ("setsockopt", megatron_quant_policy_worker.zmq.SNDTIMEO, 120_000), + ("setsockopt", megatron_quant_policy_worker.zmq.RCVTIMEO, 120_000), + ("setsockopt", megatron_quant_policy_worker.zmq.LINGER, 0), + ("bind", "ipc:///tmp/modelopt-test.sock"), + ("setsockopt", megatron_quant_policy_worker.zmq.SNDTIMEO, timeout), + ("setsockopt", megatron_quant_policy_worker.zmq.RCVTIMEO, timeout), + ] + + def create_quant_megatron_test_config(model_name, tp=1, pp=1, precision="float32"): """Wrap the base Megatron test config with quantization fields.""" config = create_megatron_test_config( @@ -143,6 +182,63 @@ def test_modelopt_layer_spec_config_selects_layer_specs(): assert get_quantization_mamba_stack_spec(False) is modelopt_mamba_stack_spec +@requires_weight_folding +def test_quantization_model_specs_support_hybrid_and_legacy_mamba_providers(): + from nemo_rl.modelopt.models.policy.workers.megatron_quant_policy_worker import ( + _set_quantization_model_specs, + ) + from nemo_rl.modelopt.models.policy.workers.utils import ( + get_quantization_layer_spec, + get_quantization_mamba_stack_spec, + ) + + hybrid_config = SimpleNamespace(hybrid_stack_spec=None) + _set_quantization_model_specs(hybrid_config, True) + assert hybrid_config.transformer_layer_spec is get_quantization_layer_spec(True) + assert hybrid_config.hybrid_stack_spec is get_quantization_mamba_stack_spec(True) + + legacy_config = SimpleNamespace(mamba_stack_spec=None) + _set_quantization_model_specs(legacy_config, False) + assert ( + legacy_config.transformer_layer_spec.func + is get_quantization_layer_spec(False).func + ) + assert legacy_config.mamba_stack_spec is get_quantization_mamba_stack_spec(False) + + +@requires_weight_folding +def test_warns_when_other_quantized_startup_caches_exist(tmp_path, monkeypatch): + from nemo_rl.modelopt.models.policy.workers import megatron_quant_policy_worker + + base_path = tmp_path / "model" + selected_path = tmp_path / "model_modelopt_selected" + old_hashed_path = tmp_path / "model_modelopt_old" + legacy_path = tmp_path / "model_quantized" + invalid_path = tmp_path / "model_modelopt_invalid" + for cache_path in (old_hashed_path, legacy_path, invalid_path): + (cache_path / "iter_0000000").mkdir(parents=True) + + monkeypatch.setattr( + megatron_quant_policy_worker, + "has_modelopt_state", + lambda path: "invalid" not in path, + ) + + with pytest.warns( + UserWarning, + match=r"checkpointing\.checkpoint_dir", + ) as warning_records: + megatron_quant_policy_worker._warn_if_other_quant_checkpoint_caches( + base_path.as_posix(), + selected_path.as_posix(), + ) + + message = str(warning_records[0].message) + assert old_hashed_path.as_posix() in message + assert legacy_path.as_posix() in message + assert invalid_path.as_posix() not in message + + @requires_weight_folding def test_real_quant_refit_detection_requires_vllm_quant_cfg_and_flag(): worker = _make_real_quant_worker() @@ -177,6 +273,28 @@ def test_iter_real_quant_refit_params_uses_megatron_bridge_export(): assert kwargs["ignore_patterns"] == ["lm_head"] +@requires_weight_folding +def test_iter_real_quant_refit_params_exports_w4a4_mode(): + worker = _make_real_quant_worker() + quant_cfg = "NVFP4_EXPERTS_ONLY_CFG" + worker.cfg["quant_cfg"] = quant_cfg + worker.cfg["generation"]["quant_cfg"] = quant_cfg + + list(worker._iter_real_quant_refit_params()) + + _, kwargs = worker.megatron_bridge.calls[0] + assert kwargs["quant_mode"] == "nvfp4" + + +@requires_weight_folding +def test_iter_real_quant_refit_params_rejects_policy_generation_mode_mismatch(): + worker = _make_real_quant_worker() + worker.cfg["generation"]["quant_cfg"] = "NVFP4_EXPERTS_ONLY_CFG" + + with pytest.raises(ValueError, match="matching policy and generation"): + list(worker._iter_real_quant_refit_params()) + + @requires_weight_folding def test_iter_params_with_optional_kv_scales_uses_real_quant_export(monkeypatch): worker = _make_real_quant_worker() @@ -246,6 +364,63 @@ def __init__(self): torch.testing.assert_close(output[1][1], torch.tensor([3.0])) +@requires_weight_folding +def test_folded_quantizer_error_includes_parameter_name(monkeypatch): + from nemo_rl.models.policy.workers.megatron_policy_worker import ( + MegatronPolicyWorkerImpl, + ) + + class FailingQuantizer: + def __call__(self, _weight): + raise ValueError("invalid quantizer state") + + worker_cls = MegatronQuantPolicyWorker.__ray_metadata__.modified_class + worker = object.__new__(worker_cls) + worker.cfg = { + "generation": { + "backend": "vllm", + "quant_cfg": "FP8_DEFAULT_CFG", + "real_quant": False, + } + } + worker.rank = 0 + task = SimpleNamespace( + param_name="decoder.layers.0.mlp.linear_fc2.weight", + global_param_name="decoder.layers.0.mlp.linear_fc2.weight", + param_weight=torch.ones(2, 2), + megatron_module=object(), + mapping=SimpleNamespace(hf_param="model.layers.0.mlp.down_proj.weight"), + ) + worker.refit_conversion_tasks = [task] + + monkeypatch.setattr( + worker, + "_find_weight_quantizer", + lambda *_args: FailingQuantizer(), + ) + + def access_refit_task_weights(self, kv_scales=None): + for refit_task in self.refit_conversion_tasks: + yield refit_task.param_name, refit_task.param_weight + + monkeypatch.setattr( + MegatronPolicyWorkerImpl, + "_iter_params_with_optional_kv_scales", + access_refit_task_weights, + ) + + with pytest.raises(RuntimeError) as exc_info: + list(worker._iter_params_with_optional_kv_scales()) + + assert ( + "Failed to apply weight quantizer for param " + "'decoder.layers.0.mlp.linear_fc2.weight'" + ) in str(exc_info.value) + assert isinstance(exc_info.value.__cause__, ValueError) + assert str(exc_info.value.__cause__) == "invalid quantizer state" + assert worker.refit_conversion_tasks == [task] + + @requires_weight_folding def test_stream_weights_via_ipc_zmq_uses_real_quant_generator_without_move( monkeypatch, diff --git a/tests/unit/models/policy/test_megatron_worker.py b/tests/unit/models/policy/test_megatron_worker.py index 793cc52bd90..9727b50d06b 100644 --- a/tests/unit/models/policy/test_megatron_worker.py +++ b/tests/unit/models/policy/test_megatron_worker.py @@ -53,6 +53,31 @@ def train(self): self.train_called = True +class _ModelWithNonSerializableExtraState(torch.nn.Module): + def __init__(self): + super().__init__() + self.weight = torch.nn.Parameter(torch.ones(1)) + self.register_buffer("scale", torch.ones(1)) + + def get_extra_state(self): + raise AssertionError("moving a module must not serialize its extra state") + + +def test_megatron_move_model_does_not_serialize_extra_state(): + from nemo_rl.models.policy.workers.megatron_policy_worker import ( + MegatronPolicyWorkerImpl, + ) + + worker = object.__new__(MegatronPolicyWorkerImpl) + model = _ModelWithNonSerializableExtraState() + + moved_model = MegatronPolicyWorkerImpl.move_model(worker, model, "cpu") + + assert moved_model is model + assert model.weight.device.type == "cpu" + assert model.scale.device.type == "cpu" + + def test_megatron_prepare_for_training_restores_optimizer(): from nemo_rl.models.policy.workers.megatron_policy_worker import ( MegatronPolicyWorkerImpl, @@ -154,7 +179,13 @@ def test_compute_moe_grad_scale_clamps_zero_valid_tokens(): assert torch.allclose(scale_fn(), torch.tensor(1.0)) -def test_disable_forward_pre_hook_until_next_step_uses_worker_override(): +@pytest.mark.parametrize( + ("kwargs", "expected_param_sync"), + [({}, False), ({"param_sync": True}, True)], +) +def test_disable_forward_pre_hook_until_next_step_uses_worker_override( + kwargs: dict[str, bool], expected_param_sync: bool +) -> None: source_path = ( Path(__file__).parents[4] / "nemo_rl/models/policy/workers/megatron_policy_worker.py" @@ -204,14 +235,55 @@ def disable_forward_pre_hook(self, param_sync=True): param_sync ) - worker._disable_forward_pre_hook_until_next_train_step() + worker._disable_forward_pre_hook_until_next_train_step(**kwargs) - assert disable_calls == [False] + assert disable_calls == [expected_param_sync] assert worker._first_train_step_param_sync_func == "sync" assert model_config.param_sync_func is None assert worker._first_train_step_forward_pre_hook_disabled is True +def test_prepare_for_generation_disables_param_gather_hook_before_wake( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from nemo_rl.models.generation.megatron import megatron_worker + + events = [] + model = SimpleNamespace( + config=SimpleNamespace(flash_decode=True), + eval=lambda: events.append("eval"), + ) + worker = object.__new__(megatron_worker.MegatronGenerationMixin) + worker.cfg = { + "generation": {"mcore_generation_config": {"cuda_graph_impl": "none"}} + } + worker.model = model + worker.is_generation_colocated = True + worker.should_disable_forward_pre_hook = True + worker.move_model = lambda model, device, **kwargs: ( + events.append("move_to_cuda") or model + ) + 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._inference_engine_initialized = True + worker._wake = lambda: events.append("wake_engine") + + monkeypatch.setattr(megatron_worker, "log_gpu_memory", lambda *_: None) + monkeypatch.setattr(megatron_worker, "unwrap_model", lambda model: model) + + worker.prepare_for_generation() + + assert events == [ + "move_to_cuda", + ("disable_hook", True), + "eval", + "wake_engine", + ] + assert model.config.flash_decode is False + + def create_megatron_test_config( model_name: str, tp: int = 1, @@ -1368,6 +1440,9 @@ def test_megatron_checkpoint_save_kill_and_restore( weights_path=weights_path, optimizer_path=optimizer_path, ) + # save_checkpoint() may use MCore's async save path. Complete the + # write before inspecting the checkpoint or terminating its workers. + policy1.finalize_async_save() # Verify checkpoint was created assert os.path.exists(checkpoint_dir), "Checkpoint directory not created" diff --git a/tests/unit/models/policy/test_modelopt_worker_utils.py b/tests/unit/models/policy/test_modelopt_worker_utils.py index 0a40199cca8..41bdc80c534 100644 --- a/tests/unit/models/policy/test_modelopt_worker_utils.py +++ b/tests/unit/models/policy/test_modelopt_worker_utils.py @@ -66,6 +66,9 @@ def _install_optional_dependency_stubs(): ) plugins = _ensure_module("modelopt.torch.utils.plugins") + plugins.get_megatron_calibration_forward_loop = lambda tokenizer, **kwargs: ( + lambda model: None + ) plugins.megatron_prefill = lambda model, input_ids, skip_return_logits: None gpt_provider = _ensure_module("megatron.bridge.models.gpt_provider") @@ -120,11 +123,18 @@ def test_get_tokenizer_applies_modelopt_calibration_defaults(monkeypatch): assert tokenizer.model_max_length == 128 -def test_megatron_forward_loop_prefills_batch_input_ids(monkeypatch): +def test_get_forward_loop_func_prefills_random_megatron_data(monkeypatch): seen = [] - dataloader = DataLoader( - worker_utils._DictDataset({"input_ids": torch.tensor([[1, 2, 3]])}), - batch_size=1, + input_ids = torch.tensor([[1, 2, 3]]) + monkeypatch.setattr( + worker_utils.parallel_state, + "get_context_parallel_world_size", + lambda: 1, + ) + monkeypatch.setattr( + worker_utils.torch, + "randint", + lambda *args, **kwargs: input_ids, ) monkeypatch.setattr( worker_utils, @@ -134,14 +144,43 @@ def test_megatron_forward_loop_prefills_batch_input_ids(monkeypatch): ), ) - loop = worker_utils.get_forward_loop_func(True, dataloader) + loop = worker_utils.get_forward_loop_func( + is_megatron=True, + tokenizer=None, + dataset_name="random", + batch_size=1, + num_samples=1, + sample_length=3, + device=torch.device("cpu"), + ) loop("model") assert seen[0][0] == "model" - torch.testing.assert_close(seen[0][1], torch.tensor([[1, 2, 3]])) + torch.testing.assert_close(seen[0][1], input_ids) assert seen[0][2] is True +def test_get_forward_loop_func_rejects_random_megatron_context_parallelism( + monkeypatch, +): + monkeypatch.setattr( + worker_utils.parallel_state, + "get_context_parallel_world_size", + lambda: 2, + ) + + with pytest.raises(RuntimeError, match="context_parallel_size=1, got 2"): + worker_utils.get_forward_loop_func( + is_megatron=True, + tokenizer="tokenizer", + dataset_name="random", + batch_size=1, + num_samples=1, + sample_length=16, + device=torch.device("cpu"), + ) + + def test_quantize_model_skips_forward_loop_for_weight_only_config(monkeypatch): model = torch.nn.Linear(1, 1) calls = [] @@ -155,12 +194,15 @@ def test_quantize_model_skips_forward_loop_for_weight_only_config(monkeypatch): monkeypatch.setattr( worker_utils.mtq, "quantize", - lambda model_arg, cfg, forward_loop: calls.append( - (model_arg, cfg, forward_loop) - ) - or model_arg, + lambda model_arg, cfg, forward_loop: ( + calls.append((model_arg, cfg, forward_loop)) or model_arg + ), + ) + monkeypatch.setattr( + worker_utils.mtq, + "print_quant_summary", + lambda model: None, ) - monkeypatch.setattr(worker_utils.mtq, "print_quant_summary", lambda model: None) worker_utils.quantize_model( model, @@ -189,7 +231,7 @@ def test_quantize_model_requires_calibration_data(monkeypatch): ) -def test_quantize_model_uses_random_calibration_loop(monkeypatch): +def test_quantize_model_delegates_calibration_loop_building(monkeypatch): model = torch.nn.Linear(1, 1) calls = [] @@ -198,41 +240,56 @@ def test_quantize_model_uses_random_calibration_loop(monkeypatch): monkeypatch.setattr( worker_utils, "get_forward_loop_func", - lambda is_megatron, dataloader: ( - "loop", - is_megatron, - len(dataloader.dataset), - ), + lambda **kwargs: calls.append(("loop", kwargs)) or "forward-loop", ) monkeypatch.setattr( worker_utils.mtq, "quantize", - lambda model_arg, cfg, forward_loop: calls.append(forward_loop) or model_arg, + lambda model_arg, cfg, forward_loop: ( + calls.append(("quantize", model_arg, cfg, forward_loop)) or model_arg + ), + ) + monkeypatch.setattr( + worker_utils.mtq, + "print_quant_summary", + lambda model: None, ) - monkeypatch.setattr(worker_utils.mtq, "print_quant_summary", lambda model: None) worker_utils.quantize_model( model, "activation-cfg", tokenizer=None, calib_size=8, + batch_size=2, + max_sample_length=16, is_megatron=True, data="random", ) - assert calls == [("loop", True, 1)] + assert calls == [ + ( + "loop", + { + "is_megatron": True, + "tokenizer": None, + "dataset_name": "random", + "batch_size": 2, + "num_samples": 8, + "sample_length": 16, + "device": model.weight.device, + }, + ), + ("quantize", model, {}, "forward-loop"), + ] -def test_quantize_model_uses_named_calibration_dataset(monkeypatch): - model = torch.nn.Linear(1, 1) +def test_get_forward_loop_func_uses_named_generic_dataset(monkeypatch): dataset = worker_utils._DictDataset( {"input_ids": torch.ones(2, 3, dtype=torch.long)} ) dataloader = DataLoader(dataset, batch_size=2) calls = [] - monkeypatch.setattr(worker_utils, "resolve_quant_cfg", lambda quant_cfg: {}) - monkeypatch.setattr(worker_utils, "need_calibration", lambda cfg: True) monkeypatch.setattr( worker_utils, "get_dataset_dataloader", @@ -240,46 +297,85 @@ def test_quantize_model_uses_named_calibration_dataset(monkeypatch): ) monkeypatch.setattr( worker_utils, - "get_forward_loop_func", - lambda is_megatron, calib_dataloader: ( - "loop", - is_megatron, - calib_dataloader, + "create_forward_loop", + lambda **kwargs: calls.append(("loop", kwargs)) or "forward-loop", + ) + + result = worker_utils.get_forward_loop_func( + is_megatron=False, + tokenizer="tokenizer", + dataset_name="cnn_dailymail", + batch_size=4, + num_samples=8, + sample_length=16, + device=torch.device("cpu"), + ) + + assert result == "forward-loop" + assert calls == [ + ( + "dataset", + { + "dataset_name": "cnn_dailymail", + "tokenizer": "tokenizer", + "batch_size": 4, + "num_samples": 8, + "device": torch.device("cpu"), + "include_labels": False, + "max_sample_length": 16, + }, ), + ("loop", {"dataloader": dataloader}), + ] + + +def test_get_forward_loop_func_uses_modelopt_megatron_calibration_loop(monkeypatch): + calls = [] + + monkeypatch.setattr( + worker_utils.parallel_state, + "get_context_parallel_world_size", + lambda: 2, ) monkeypatch.setattr( - worker_utils.mtq, - "quantize", - lambda model_arg, cfg, forward_loop: calls.append( - ("quantize", model_arg, cfg, forward_loop) - ) - or model_arg, + worker_utils, + "get_megatron_calibration_forward_loop", + lambda tokenizer, **kwargs: ( + calls.append(("calibration", tokenizer, kwargs)) or "megatron-loop" + ), + ) + monkeypatch.setattr( + worker_utils, + "get_dataset_dataloader", + lambda **kwargs: pytest.fail("Megatron calibration must use ModelOpt's helper"), ) - monkeypatch.setattr(worker_utils.mtq, "print_quant_summary", lambda model: None) - worker_utils.quantize_model( - model, - "activation-cfg", + result = worker_utils.get_forward_loop_func( + is_megatron=True, tokenizer="tokenizer", - calib_size=8, + dataset_name="calibration.jsonl", batch_size=4, - data="cnn_dailymail", - max_sample_length=16, + num_samples=8, + sample_length=16, + device=torch.device("cpu"), ) - assert calls[0] == ( - "dataset", - { - "dataset_name": "cnn_dailymail", - "tokenizer": "tokenizer", - "batch_size": 4, - "num_samples": 8, - "device": model.weight.device, - "include_labels": False, - "max_sample_length": 16, - }, - ) - assert calls[1] == ("quantize", model, {}, ("loop", False, dataloader)) + assert result == "megatron-loop" + assert calls == [ + ( + "calibration", + "tokenizer", + { + "dataset_name": "calibration.jsonl", + "batch_size": 4, + "num_samples": 8, + "seq_length": 16, + "device": torch.device("cpu"), + "apply_chat_template": False, + "pack": True, + }, + ), + ] def test_get_modelopt_checkpoint_dir_env_precedence(monkeypatch): diff --git a/tests/unit/models/policy/test_utils.py b/tests/unit/models/policy/test_utils.py index a847b3b752f..852f8081b64 100644 --- a/tests/unit/models/policy/test_utils.py +++ b/tests/unit/models/policy/test_utils.py @@ -18,6 +18,7 @@ import time import traceback import unittest.mock +import weakref import pytest import torch @@ -136,6 +137,63 @@ def getsockopt(self, _option): return 0 +def test_stream_weights_releases_buffers_before_complete_without_full_gc( + monkeypatch, +): + """The final data ACK is sufficient to reclaim both acyclic IPC buffers.""" + + tensor = torch.ones(4, dtype=torch.float32) + buffer_refs = [] + events = [] + original_empty = torch.empty + + def tracking_empty(*args, **kwargs): + buffer = original_empty(*args, **kwargs) + buffer_refs.append(weakref.ref(buffer)) + return buffer + + def empty_cache(): + events.append("empty_cache") + assert len(buffer_refs) == 2 + assert all(buffer_ref() is None for buffer_ref in buffer_refs) + + class ReleaseAwareSocket(_FakeIpcSocket): + def send_pyobj(self, payload): + if payload == IPCProtocol.COMPLETE: + assert events == ["empty_cache"] + assert all(buffer_ref() is None for buffer_ref in buffer_refs) + super().send_pyobj(payload) + + monkeypatch.setattr(torch, "empty", tracking_empty) + monkeypatch.setattr(torch.cuda, "is_available", lambda: False) + monkeypatch.setattr( + torch.cuda, + "current_stream", + lambda: unittest.mock.Mock(synchronize=lambda: None), + ) + monkeypatch.setattr(torch.cuda, "empty_cache", empty_cache) + monkeypatch.setattr( + "nemo_rl.models.policy.utils.get_handle_from_tensor", + lambda _buffer: ("ipc-handle",), + ) + monkeypatch.setattr( + "nemo_rl.models.policy.utils.gc.collect", + lambda: pytest.fail("IPC buffer cleanup must not scan the full object graph"), + ) + + socket = ReleaseAwareSocket() + stream_weights_via_ipc_zmq_impl( + params_generator=iter([("weight", tensor)]), + buffer_size_bytes=4096, + zmq_socket=socket, + rank=0, + worker_name="test_worker", + ) + + assert events == ["empty_cache"] + assert socket.sent[-1] == IPCProtocol.COMPLETE + + def test_stream_weights_via_ipc_zmq_uses_cuda_buffer_for_cpu_tensors(monkeypatch): """CPU-exported tensors should still be packed into CUDA IPC buffers.""" diff --git a/tests/unit/utils/test_checkpoint.py b/tests/unit/utils/test_checkpoint.py index fca8a1d3ec0..885412c1eb6 100644 --- a/tests/unit/utils/test_checkpoint.py +++ b/tests/unit/utils/test_checkpoint.py @@ -23,6 +23,7 @@ import torch import yaml +import nemo_rl.utils.checkpoint as checkpoint_module from nemo_rl.utils.checkpoint import CheckpointManager @@ -402,6 +403,30 @@ def test_save_optimizer_flag_initialization(checkpoint_config): assert manager.save_optimizer is False +@pytest.mark.parametrize("model_component", ["policy", "value"]) +@pytest.mark.parametrize("common_state_marker", ["common.pt", "metadata.json"]) +def test_get_resume_paths_detects_megatron_optimizer( + checkpoint_dir, monkeypatch, model_component, common_state_marker +): + checkpoint_path = checkpoint_dir / "step_1" + iteration_dir = checkpoint_path / model_component / "weights" / "iter_0000000" + iteration_dir.mkdir(parents=True) + (iteration_dir / common_state_marker).touch() + load_common_state = MagicMock(return_value={"optimizer": {}}) + monkeypatch.setattr( + checkpoint_module, "_load_megatron_common_state_dict", load_common_state + ) + + weights_path, optimizer_path = CheckpointManager.get_resume_paths( + checkpoint_path, + model_component=model_component, + ) + + assert weights_path == checkpoint_path / model_component / "weights" + assert optimizer_path == checkpoint_path / model_component / "optimizer" + load_common_state.assert_called_once_with(iteration_dir) + + @pytest.mark.parametrize("model_component", ["policy", "value"]) def test_get_resume_paths_missing_optimizer( checkpoint_manager, checkpoint_dir, model_component @@ -461,30 +486,71 @@ def test_get_resume_paths_defaults_to_policy(checkpoint_dir): assert optimizer_path == expected_optimizer_path -@pytest.mark.parametrize("model_component", ["policy", "value"]) -def test_get_resume_paths_embedded_megatron_optimizer(checkpoint_dir, model_component): - """MCore uses a nonexistent optimizer path as an embedded-state load flag.""" +def test_get_resume_paths_warns_when_megatron_optimizer_missing( + checkpoint_dir, monkeypatch +): checkpoint_path = checkpoint_dir / "step_1" - expected_weights_path = checkpoint_path / model_component / "weights" - common_pt_path = expected_weights_path / "iter_0000000" / "common.pt" - common_pt_path.parent.mkdir(parents=True) - torch.save( - { - "optimizer": {"state": {}}, - "opt_param_scheduler": {"num_steps": 8}, - }, - common_pt_path, + iteration_dir = checkpoint_path / "policy" / "weights" / "iter_0000000" + iteration_dir.mkdir(parents=True) + (iteration_dir / "metadata.json").touch() + monkeypatch.setattr( + checkpoint_module, + "_load_megatron_common_state_dict", + MagicMock(return_value={"args": {}}), ) - expected_optimizer_path = checkpoint_path / model_component / "optimizer" - assert not expected_optimizer_path.exists() - weights_path, optimizer_path = CheckpointManager.get_resume_paths( - checkpoint_path, - model_component=model_component, + with pytest.warns(UserWarning, match="Optimizer state not found"): + weights_path, optimizer_path = CheckpointManager.get_resume_paths( + checkpoint_path + ) + + assert weights_path == checkpoint_path / "policy" / "weights" + assert optimizer_path is None + + +def test_get_resume_paths_prefers_dtensor_optimizer(checkpoint_dir, monkeypatch): + checkpoint_path = checkpoint_dir / "step_1" + optimizer_path = checkpoint_path / "policy" / "optimizer" + optimizer_path.mkdir(parents=True) + iteration_dir = checkpoint_path / "policy" / "weights" / "iter_0000000" + iteration_dir.mkdir(parents=True) + (iteration_dir / "metadata.json").touch() + load_common_state = MagicMock(side_effect=AssertionError("must not be called")) + monkeypatch.setattr( + checkpoint_module, "_load_megatron_common_state_dict", load_common_state ) - assert weights_path == expected_weights_path - assert optimizer_path == expected_optimizer_path + weights_path, returned_optimizer_path = CheckpointManager.get_resume_paths( + checkpoint_path + ) + + assert weights_path == checkpoint_path / "policy" / "weights" + assert returned_optimizer_path == optimizer_path + load_common_state.assert_not_called() + + +@pytest.mark.parametrize( + "load_error", + [ + RuntimeError("Megatron-Core is required"), + OSError("checkpoint is unreadable"), + ], +) +def test_get_resume_paths_propagates_megatron_load_failure( + checkpoint_dir, monkeypatch, load_error +): + checkpoint_path = checkpoint_dir / "step_1" + iteration_dir = checkpoint_path / "policy" / "weights" / "iter_0000000" + iteration_dir.mkdir(parents=True) + (iteration_dir / "metadata.json").touch() + monkeypatch.setattr( + checkpoint_module, + "_load_megatron_common_state_dict", + MagicMock(side_effect=load_error), + ) + + with pytest.raises(type(load_error), match=str(load_error)): + CheckpointManager.get_resume_paths(checkpoint_path) def test_get_best_checkpoint_path_no_checkpoints(checkpoint_manager, checkpoint_dir): diff --git a/uv.lock b/uv.lock index 0bfea3edf89..ca978d833c3 100644 --- a/uv.lock +++ b/uv.lock @@ -2591,7 +2591,7 @@ requires-dist = [ { name = "flask-restful", marker = "extra == 'mlm'" }, { name = "flask-restful", marker = "extra == 'training'" }, { name = "hypercorn", marker = "extra == 'dev'" }, - { name = "mamba-ssm", marker = "extra == 'ssm'", specifier = "~=2.2" }, + { name = "mamba-ssm", marker = "extra == 'ssm'", git = "https://github.com/state-spaces/mamba.git?rev=0048fbf2e7b2f214dcbe703ea3dec2b9647595e1" }, { name = "megatron-energon", extras = ["av-decode"], marker = "extra == 'dev'", specifier = "~=7.0" }, { name = "multi-storage-client", marker = "extra == 'dev'", specifier = "~=0.50" }, { name = "numpy" }, @@ -2612,7 +2612,7 @@ requires-dist = [ { name = "tiktoken", marker = "extra == 'training'" }, { name = "torch", specifier = ">=2.6.0" }, { name = "tqdm", marker = "extra == 'dev'" }, - { name = "transformer-engine", extras = ["core-cu13", "pytorch"], marker = "extra == 'te'", git = "https://github.com/NVIDIA/TransformerEngine.git?rev=4220403e831d29e93868f7793693ea83f6b8b05b" }, + { name = "transformer-engine", extras = ["core-cu13", "pytorch"], marker = "extra == 'te'", git = "https://github.com/NVIDIA/TransformerEngine.git?rev=b9d690e042b1c4e455214e7dab65d6d3512c05d6" }, { name = "transformers", marker = "extra == 'mlm'" }, { name = "transformers", marker = "extra == 'training'" }, { name = "wandb", marker = "extra == 'mlm'" }, @@ -2654,7 +2654,7 @@ linting = [ ] no-pypi-wheels = [ { name = "emerging-optimizers", git = "https://github.com/NVIDIA-NeMo/Emerging-Optimizers.git?rev=v0.2.0" }, - { name = "flash-mla", git = "https://github.com/deepseek-ai/FlashMLA?rev=9edee0c022cd0938148a18e334203b0aab43aa19" }, + { name = "flash-mla", git = "https://github.com/deepseek-ai/FlashMLA?rev=nv_dev" }, ] test = [ { name = "coverage" },