diff --git a/docs/guides/models/qwen/qwen3-5.md b/docs/guides/models/qwen/qwen3-5.md index 8b2896aa34c..150442b0321 100644 --- a/docs/guides/models/qwen/qwen3-5.md +++ b/docs/guides/models/qwen/qwen3-5.md @@ -64,6 +64,7 @@ authoritative settings. | Qwen3.5-9B-Base | LLM | GRPO | Megatron | 1n8g | [`grpo-qwen3.5-9b-1n8g-megatron-fp8.yaml`](../../../../examples/configs/recipes/llm/grpo-qwen3.5-9b-1n8g-megatron-fp8.yaml) | | Qwen3.5-35B-A3B-Base | LLM | GRPO | Megatron | 2n8g | [`grpo-qwen3.5-35ba3b-2n8g-megatron-ep16tp2cp2.yaml`](../../../../examples/configs/recipes/llm/grpo-qwen3.5-35ba3b-2n8g-megatron-ep16tp2cp2.yaml) | | Qwen3.5-35B-A3B-Base | LLM | GRPO | Megatron | 2n8g | [`grpo-qwen3.5-35ba3b-2n8g-megatron-ep16tp2-fp8.yaml`](../../../../examples/configs/recipes/llm/grpo-qwen3.5-35ba3b-2n8g-megatron-ep16tp2-fp8.yaml) | +| Qwen3.5-35B-A3B-Base | LLM | GRPO | Megatron | 4n4g | [`grpo-qwen3.5-35ba3b-4n4g-megatron-ep16tp2-mxfp8-trtllm.yaml`](../../../../examples/configs/recipes/llm/grpo-qwen3.5-35ba3b-4n4g-megatron-ep16tp2-mxfp8-trtllm.yaml) | | Qwen3.5-35B-A3B-Base | LLM | GRPO | Megatron | 6n4g | [`grpo-qwen3.5-35ba3b-6n4g-async-1off-bf16-trtllm.yaml`](../../../../examples/configs/recipes/llm/grpo-qwen3.5-35ba3b-6n4g-async-1off-bf16-trtllm.yaml) | | Qwen3.5-35B-A3B-Base | LLM | GRPO | AutoModel | 2n8g | [`grpo-qwen3.5-35ba3b-2n8g-automodel-ep16.yaml`](../../../../examples/configs/recipes/llm/grpo-qwen3.5-35ba3b-2n8g-automodel-ep16.yaml) | | Qwen3.5-35B-A3B-Base | LLM | GRPO | AutoModel | 4n8g | [`grpo-qwen3.5-35ba3b-dapo-4n8g-automodel.yaml`](../../../../examples/configs/recipes/llm/grpo-qwen3.5-35ba3b-dapo-4n8g-automodel.yaml) | @@ -185,6 +186,12 @@ and General FP8 background lives in [FP8 Quantization](../../../fp8.md); the notes below are the Qwen3.5-specific sharp edges. +For BF16 Megatron training with grouped-MoE MXFP8 rollouts and the FlashInfer +TRTLLM MoE backend, use +[`grpo-qwen3.5-35ba3b-4n4g-megatron-ep16tp2-mxfp8-trtllm.yaml`](../../../../examples/configs/recipes/llm/grpo-qwen3.5-35ba3b-4n4g-megatron-ep16tp2-mxfp8-trtllm.yaml). +The recurring 20-step GB200 test checks reward, finite loss, and rollout-versus- +training log-prob parity after each refit. + ### Block Layout and Shape Constraints - Weights are quantized in `[128, 128]` blocks (e4m3) with one fp32 @@ -247,8 +254,8 @@ FP8 noise. ### Limitations -- MXFP8 is not yet supported; tracked in - [#3694](https://github.com/NVIDIA-NeMo/RL/issues/3694). +- Dense Qwen3.5 MXFP8 rollout has not yet completed recurring end-to-end + validation. ## `flash-linear-attention` Performance diff --git a/docs/guides/refit.md b/docs/guides/refit.md index 368c3c123c5..da3b724fe2a 100644 --- a/docs/guides/refit.md +++ b/docs/guides/refit.md @@ -52,7 +52,7 @@ for those cases. |---|---|---|---| | Colocated CUDA IPC | vLLM or SGLang | DTensor or Megatron | Uses the generation backend's standard loader. | | NCCL | vLLM or Megatron | DTensor or Megatron | Uses the standard full-weight loader. | -| NCCL + reload API | vLLM | DTensor or Megatron | Default non-colocated NCCL only. Rejects colocated refit, `nccl_reshard`, sparse delta, NIXL/checkpoint-engine transports, ModelOpt quantization (`real_quant` or `quant_cfg`), Eagle/MTP trainer-refit draft weights, and grouped-MoE MXFP8 slabs. | +| NCCL + reload API | vLLM | DTensor or Megatron | Default non-colocated NCCL only. Rejects colocated refit, `nccl_reshard`, sparse delta, NIXL/checkpoint-engine transports, ModelOpt quantization (`real_quant` or `quant_cfg`), and Eagle/MTP trainer-refit draft weights. | | MCore native refit | Megatron | Megatron | Supports MCore's configured copy-service backend. | | SGLang NCCL weight-update group | SGLang | Megatron | Trainer rank 0 broadcasts finalized HF tensors to the engine leaders. | | NCCL reshard | vLLM or Megatron | Megatron | Supports BF16 and the documented FP8/MXFP8 combinations; see the [NCCL Reshard Refit design](../design-docs/nccl-reshard-refit.md). | @@ -63,9 +63,7 @@ for those cases. Non-colocated SGLang generation is supported with a Megatron policy and `refit_transport: null`; it creates its own NCCL weight-update group. The NIXL restrictions are on the generation backend; both Megatron and DTensor policy -workers can send weights. vLLM refit also rejects grouped MoE MXFP8 expert -slabs such as `mlp.experts.gate_up_proj` and `mlp.experts.down_proj`. Sparse -delta is currently limited to GRPO. NIXL is +workers can send weights. Sparse delta is currently limited to GRPO. NIXL is initialized by the GRPO and distillation setup paths; PPO currently requires colocated generation. diff --git a/examples/configs/recipes/llm/grpo-qwen3.5-35ba3b-4n4g-megatron-ep16tp2-mxfp8-trtllm.yaml b/examples/configs/recipes/llm/grpo-qwen3.5-35ba3b-4n4g-megatron-ep16tp2-mxfp8-trtllm.yaml new file mode 100644 index 00000000000..2aa00af3c83 --- /dev/null +++ b/examples/configs/recipes/llm/grpo-qwen3.5-35ba3b-4n4g-megatron-ep16tp2-mxfp8-trtllm.yaml @@ -0,0 +1,38 @@ +defaults: grpo-qwen3.5-35ba3b-2n8g-megatron-ep16tp2cp2.yaml +checkpointing: + checkpoint_dir: results/grpo-qwen3.5-35ba3b-4n4g-megatron-ep16tp2-mxfp8-trtllm +loss_fn: + use_importance_sampling_correction: true + truncated_importance_sampling_type: tis + truncated_importance_sampling_ratio: 2 +policy: + make_sequence_length_divisible_by: 16 + megatron_cfg: + context_parallel_size: 1 + moe_router_dtype: fp32 + freeze_config: + freeze_vision_model: true + freeze_vision_projection: true + freeze_language_model: false + generation: + vllm_cfg: + tensor_parallel_size: 2 + expert_parallel_size: 2 + gpu_memory_utilization: 0.5 + precision: fp8 + is_mx: true + enforce_eager: false + quantization_ignore_patterns: + - language_model.model.layers.*.linear_attn.* + - visual.* + - lm_head + vllm_kwargs: + moe_backend: flashinfer_trtllm + expert_placement_strategy: linear +logger: + wandb: + name: grpo-qwen3.5-35ba3b-4n4g-megatron-ep16tp2-mxfp8-trtllm +cluster: + gpus_per_node: 4 + num_nodes: 4 + segment_size: 4 diff --git a/nemo_rl/models/generation/vllm/quantization/fp8.py b/nemo_rl/models/generation/vllm/quantization/fp8.py index c254c92e09d..c779b5fb854 100644 --- a/nemo_rl/models/generation/vllm/quantization/fp8.py +++ b/nemo_rl/models/generation/vllm/quantization/fp8.py @@ -570,9 +570,8 @@ def get_quantized_weight_iterator( for k, v in weights: # Grouped MoE experts arrive as fused slabs without a ``.weight`` suffix # (so `_is_fp8_weight` would skip them) and vLLM's grouped loader cannot - # load their per-block scales. Expand them into the per-expert FP8 (w13, w2 -> w1, w2, and w3) - # layout, then reshape to 2D [num_experts, out_features, in_features] -> [num_experts*out_features, in_features] - # so the block scales can be quantized and routed correctly. + # load their scales. Expand them into the per-expert projection layout so + # both values and scales route through the standard expert mapping. if is_grouped_moe_expert_weight_name(k): # Quantize only if vLLM built this layer's experts as FP8. Experts # covered by ``ignored_layers`` (num_{first,last}_layers_in_bf16 / @@ -582,7 +581,12 @@ def get_quantized_weight_iterator( # bf16 slab through instead; vLLM's fused expert mapping loads it # directly, same as a bf16 refit. if _is_fp8_grouped_moe_expert(k, model): - yield from _expand_grouped_moe_expert_to_fp8(k, v) + if global_fp8_config.is_mx: + yield from _expand_grouped_moe_expert_to_mxfp8( + k, v, refit_with_reload_api=refit_with_reload_api + ) + else: + yield from _expand_grouped_moe_expert_to_fp8(k, v) else: yield k, v continue @@ -763,6 +767,24 @@ def _quantize_grouped_experts_blockwise(grouped_moe_expert): return weight_fp8, scale_inv +def _split_grouped_moe_shards( + key: str, weight: torch.Tensor +) -> tuple[str, tuple[tuple[str, torch.Tensor], ...]]: + """Split one grouped expert slab into vLLM projection shards.""" + base, proj = key.rsplit(".", 1) + if proj == "gate_up_proj": + intermediate = weight.shape[1] // 2 + shards = ( + ("gate_proj", weight[:, :intermediate, :]), + ("up_proj", weight[:, intermediate:, :]), + ) + elif proj == "down_proj": + shards = (("down_proj", weight),) + else: + raise ValueError(f"Unsupported grouped MoE projection {proj!r} in {key!r}") + return base, shards + + def _expand_grouped_moe_expert_to_fp8(key, weight): """Expand a grouped Qwen3.5 MoE expert slab into per-expert FP8 weights. @@ -788,15 +810,7 @@ def _expand_grouped_moe_expert_to_fp8(key, weight): A list of ``(name, tensor)`` pairs: for every expert, the FP8 weight and its ``_scale_inv`` for each unfused projection. """ - base, proj = key.rsplit(".", 1) - if proj == "gate_up_proj": - intermediate = weight.shape[1] // 2 - shards = ( - ("gate_proj", weight[:, :intermediate, :]), - ("up_proj", weight[:, intermediate:, :]), - ) - else: - shards = (("down_proj", weight),) + base, shards = _split_grouped_moe_shards(key, weight) entries = [] # gate/up are dim-1 slices; feed the views directly — per-expert rows stay @@ -811,6 +825,23 @@ def _expand_grouped_moe_expert_to_fp8(key, weight): return entries +def _expand_grouped_moe_expert_to_mxfp8( + key: str, weight: torch.Tensor, *, refit_with_reload_api: bool +) -> list[tuple[str, torch.Tensor]]: + """Expand a grouped Qwen3.5 MoE slab into per-expert MXFP8 entries.""" + base, shards = _split_grouped_moe_shards(key, weight) + + entries = [] + scale_suffix = "_scale" if refit_with_reload_api else "_scale_from_checkpoint" + for shard_name, grouped_moe_expert in shards: + for expert_id, expert_weight in enumerate(grouped_moe_expert): + value, scale = quantize_mxfp8_weight(expert_weight.contiguous()) + name = f"{base}.{expert_id}.{shard_name}.weight" + entries.append((name, value)) + entries.append((name + scale_suffix, scale)) + return entries + + # Ref: https://github.com/vllm-project/vllm/blob/275de34170654274616082721348b7edd9741d32/vllm/model_executor/layers/quantization/utils/fp8_utils.py#L1175 # Patches this method to not create new torch.nn.Parameter for layer weights # to maintain weight loaders. @@ -1169,8 +1200,6 @@ def process_weights_after_loading_moe(self, layer) -> None: from vllm.model_executor.layers.quantization.fp8 import make_fp8_moe_kernel assert self.experts_cls is not None - # vLLM 0.28 dropped the `layer` kwarg (0.25 forwarded it only to the - # FlashInfer TRTLLM experts); routing tables still come from the layer. self.moe_kernel = make_fp8_moe_kernel( moe_quant_config=self.moe_quant_config, moe_config=self.moe, diff --git a/nemo_rl/models/generation/vllm/quantization/mxfp8_utils.py b/nemo_rl/models/generation/vllm/quantization/mxfp8_utils.py index 07f29e481b8..6e28e58d27b 100644 --- a/nemo_rl/models/generation/vllm/quantization/mxfp8_utils.py +++ b/nemo_rl/models/generation/vllm/quantization/mxfp8_utils.py @@ -120,6 +120,13 @@ def pad_w13_intermediate( ) intermediate_size = input_tensor.shape[1] // 2 + if intermediate_size == padded_intermediate_size: + return input_tensor + if intermediate_size > padded_intermediate_size: + raise ValueError( + "Cannot pad gated MXFP8 W13 intermediate dimension from " + f"{intermediate_size} to {padded_intermediate_size}." + ) sharded = input_tensor.reshape( input_tensor.shape[0], 2, diff --git a/nemo_rl/models/generation/vllm/utils.py b/nemo_rl/models/generation/vllm/utils.py index 84624afad27..6e564c08d42 100644 --- a/nemo_rl/models/generation/vllm/utils.py +++ b/nemo_rl/models/generation/vllm/utils.py @@ -51,9 +51,6 @@ # The expert-id range vs carry dtype is model-constant, so it is verified on the # first non-empty routed-experts tensor per process and skipped afterwards. G_ROUTED_EXPERTS_RANGE_CHECKED = False -GROUPED_MOE_MXFP8_REFIT_ERROR = ( - "MXFP8 refit does not support grouped MoE expert weights." -) _GROUPED_MOE_EXPERT_WEIGHT_SUFFIXES = ( "mlp.experts.gate_up_proj", "mlp.experts.down_proj", @@ -130,19 +127,6 @@ def is_grouped_moe_expert_weight_name(name: str) -> bool: return name.endswith(_GROUPED_MOE_EXPERT_WEIGHT_SUFFIXES) -def assert_refit_unsupported_grouped_moe_params( - config: VllmConfig, state_dict_info: dict[str, Any] -) -> None: - """Reject grouped MoE MXFP8 state-dict params before refit starts.""" - vllm_cfg = config["vllm_cfg"] - if ( - vllm_cfg.get("precision") == "fp8" - and vllm_cfg.get("is_mx") - and any(is_grouped_moe_expert_weight_name(name) for name in state_dict_info) - ): - raise AssertionError(GROUPED_MOE_MXFP8_REFIT_ERROR) - - def _as_routed_experts_tensor( value: Any, *, device: torch.device, dtype: torch.dtype ) -> torch.Tensor: diff --git a/nemo_rl/models/generation/vllm/vllm_generation.py b/nemo_rl/models/generation/vllm/vllm_generation.py index 250542d1b77..84120111cdf 100644 --- a/nemo_rl/models/generation/vllm/vllm_generation.py +++ b/nemo_rl/models/generation/vllm/vllm_generation.py @@ -52,7 +52,6 @@ from nemo_rl.models.generation.vllm.metric_names import BATCH_DURATION_KEY from nemo_rl.models.generation.vllm.utils import ( aggregate_spec_decode_counters, - assert_refit_unsupported_grouped_moe_params, assert_reload_refit_config_supported, compute_engine_step_metrics, compute_spec_decode_metrics, @@ -1378,8 +1377,6 @@ def shutdown(self) -> bool: def prepare_refit_info(self, state_dict_info: dict[str, Any]) -> None: """Prepare the info for refit.""" - assert_refit_unsupported_grouped_moe_params(self.cfg, state_dict_info) - # Choose the appropriate method based on async_engine setting method_name = ( "prepare_refit_info_async" diff --git a/tests/test_suites/llm/grpo-qwen3.5-35ba3b-4n4g-megatron-ep16tp2-mxfp8-trtllm.sh b/tests/test_suites/llm/grpo-qwen3.5-35ba3b-4n4g-megatron-ep16tp2-mxfp8-trtllm.sh new file mode 100755 index 00000000000..68ea4b6bcbe --- /dev/null +++ b/tests/test_suites/llm/grpo-qwen3.5-35ba3b-4n4g-megatron-ep16tp2-mxfp8-trtllm.sh @@ -0,0 +1,43 @@ +#!/bin/bash +SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd) +source "$SCRIPT_DIR/common.env" + +# ===== BEGIN CONFIG ===== +NUM_NODES=4 +GPUS_PER_NODE=4 +SEGMENT_SIZE=4 +STEPS_PER_RUN=20 +MAX_STEPS=20 +NUM_RUNS=$(( (MAX_STEPS + STEPS_PER_RUN - 1) / STEPS_PER_RUN )) +NUM_MINUTES=240 +# ===== END CONFIG ===== + +exit_if_max_steps_reached + +cd "$PROJECT_ROOT" +uv run examples/run_grpo.py \ + --config "$CONFIG_PATH" \ + grpo.max_num_steps=$MAX_STEPS \ + 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 tests/json_dump_tb_logs.py "$LOG_DIR" --output_path "$JSON_METRICS" + +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" \ + 'mean(data["train/gen_kl_error"]) < 0.003' \ + 'mean(data["train/reward"]) > 0.45' \ + 'all_finite(data["train/loss"])' + # The 4n4g GB200 verification run measured a 0.002345 mean + # gen_kl_error and 0.565 mean reward over 20 steps; all losses were finite. + + rm -rf "$CKPT_DIR" +fi diff --git a/tests/test_suites/nightly_gb200.txt b/tests/test_suites/nightly_gb200.txt index af5ec8d9804..7a334a36ef5 100644 --- a/tests/test_suites/nightly_gb200.txt +++ b/tests/test_suites/nightly_gb200.txt @@ -20,6 +20,9 @@ tests/test_suites/llm/grpo-nanov3-30BA3B-4n4g-megatron_generation-noncolocated-m tests/test_suites/llm/grpo-nanov3-30BA3B-4n4g-megatron_generation-noncolocated-mxfp8-rollouts-packed-refit.sh tests/test_suites/llm/grpo-nanov3-30BA3B-4n4g-megatron_generation-noncolocated-mxfp8-rollouts-flashinfer.sh +# Qwen3.5 grouped-MoE MXFP8 rollout refit correctness +tests/test_suites/llm/grpo-qwen3.5-35ba3b-4n4g-megatron-ep16tp2-mxfp8-trtllm.sh + # BF16 FlashInfer TRTLLM + NCCL reshard numerical correctness tests/test_suites/llm/grpo-qwen3.5-35ba3b-6n4g-async-1off-bf16-trtllm.sh diff --git a/tests/unit/models/generation/test_vllm_backend.py b/tests/unit/models/generation/test_vllm_backend.py index 6d7e55fefe7..6842089a691 100644 --- a/tests/unit/models/generation/test_vllm_backend.py +++ b/tests/unit/models/generation/test_vllm_backend.py @@ -1631,7 +1631,7 @@ def test_worker_prepare_refit_info_forwards_state_dict_info(): @pytest.mark.vllm @pytest.mark.parametrize("refit_with_reload_api", [False, True]) -def test_generation_prepare_refit_info_rejects_mxfp8_grouped_moe( +def test_generation_prepare_refit_info_allows_mxfp8_grouped_moe( monkeypatch, refit_with_reload_api, ): @@ -1652,14 +1652,14 @@ def test_generation_prepare_refit_info_rejects_mxfp8_grouped_moe( generation.worker_group = SimpleNamespace(workers=[leader]) generation.dp_size = 1 generation._refit_membership = None - monkeypatch.setattr(vllm_generation.ray, "get", MagicMock()) + monkeypatch.setattr(vllm_generation.ray, "get", MagicMock(return_value=[None])) - with pytest.raises(AssertionError, match="MXFP8 refit does not support"): - generation.prepare_refit_info( - {"model.layers.0.mlp.experts.gate_up_proj": object()} - ) + state_dict_info = {"model.layers.0.mlp.experts.gate_up_proj": object()} + assert generation.prepare_refit_info(state_dict_info) is None - leader.prepare_refit_info.remote.assert_not_called() + leader.prepare_refit_info.remote.assert_called_once_with( + state_dict_info=state_dict_info + ) leader.prepare_refit_info_async.remote.assert_not_called() diff --git a/tests/unit/models/generation/test_vllm_fp8_quantization.py b/tests/unit/models/generation/test_vllm_fp8_quantization.py index d255e18b3d1..de90aa3c28e 100644 --- a/tests/unit/models/generation/test_vllm_fp8_quantization.py +++ b/tests/unit/models/generation/test_vllm_fp8_quantization.py @@ -2666,6 +2666,118 @@ def test_load_weights_expands_grouped_experts_for_fp8_layers( _assert_dequant_close(weight, scale, source[eid]) +@GROUPED_EXPERT_KEY_SHAPES +@pytest.mark.parametrize( + "refit_with_reload_api,scale_suffix", + [(False, "_scale_from_checkpoint"), (True, "_scale")], + ids=["legacy", "reload-api"], +) +def test_load_weights_expands_grouped_experts_for_mxfp8( + fp8_module, + monkeypatch, + layers_prefix, + wrap_language_model, + refit_with_reload_api, + scale_suffix, +): + """MXFP8 grouped slabs emit per-expert values and E8M0 scales.""" + import torch + + fp8 = fp8_module + fp8.global_fp8_config = types.SimpleNamespace( + use_weight_pow2_scale=False, is_mx=True + ) + model = _grouped_expert_model( + fp8, monkeypatch, torch.float8_e4m3fn, wrap_language_model + ) + + from vllm.model_executor.layers.quantization.utils import mxfp8_utils + + quantized_inputs = [] + + def fake_mxfp8_quantize(weight): + quantized_inputs.append(weight.clone()) + rows, cols = weight.shape + value = weight.to(torch.float8_e4m3fn) + scale = torch.full( + (rows, cols // 32, 1), + 0, + dtype=torch.uint8, + ) + return value, scale + + monkeypatch.setattr(mxfp8_utils, "mxfp8_e4m3_quantize", fake_mxfp8_quantize) + + intermediate, hidden = 32, 64 + gate_up = ( + torch.arange(2 * 2 * intermediate * hidden, dtype=torch.float32) + .remainder(128) + .reshape(2, 2 * intermediate, hidden) + ) + down = ( + torch.arange(2 * hidden * intermediate, dtype=torch.float32) + .remainder(128) + .reshape(2, hidden, intermediate) + ) + loaded = list( + fp8.get_quantized_weight_iterator( + [ + ( + f"{layers_prefix}.0.mlp.experts.gate_up_proj", + gate_up.to(torch.bfloat16), + ), + ( + f"{layers_prefix}.0.mlp.experts.down_proj", + down.to(torch.bfloat16), + ), + ], + types.SimpleNamespace(model=model), + refit_with_reload_api=refit_with_reload_api, + ) + ) + + base = f"{layers_prefix}.0.mlp.experts" + expected_names = [ + f"{base}.{eid}.{proj}.weight{suffix}" + for proj in ("gate_proj", "up_proj") + for eid in (0, 1) + for suffix in ("", scale_suffix) + ] + [ + f"{base}.{eid}.down_proj.weight{suffix}" + for eid in (0, 1) + for suffix in ("", scale_suffix) + ] + assert [name for name, _ in loaded] == expected_names + + entries = dict(loaded) + source_shards = { + "gate_proj": gate_up[:, :intermediate, :], + "up_proj": gate_up[:, intermediate:, :], + "down_proj": down, + } + quantize_call = 0 + for proj, source in source_shards.items(): + for expert_id in (0, 1): + quantize_call += 1 + name = f"{base}.{expert_id}.{proj}.weight" + value = entries[name] + scale = entries[name + scale_suffix] + assert value.dtype == torch.float8_e4m3fn + # The refit path receives bf16 slabs, so expectations must follow + # the same fp32 -> bf16 -> e4m3 rounding chain as production. + expected = source[expert_id].to(torch.bfloat16).to(torch.float8_e4m3fn) + assert torch.equal(value.float(), expected.float()) + assert scale.dtype == torch.uint8 + assert scale.shape == (*value.shape[:-1], value.shape[-1] // 32) + assert torch.all(scale == 1) + assert torch.equal( + quantized_inputs[quantize_call - 1], + source[expert_id].to(torch.bfloat16), + ) + + assert len(quantized_inputs) == 6 + + def _deepseek_v4_lookup_model(): """A DeepSeek V4-shaped model with vLLM 0.29's ambiguous packed mapping.""" from vllm.model_executor.models.utils import WeightsMapper diff --git a/tests/unit/models/generation/test_vllm_mxfp8_utils.py b/tests/unit/models/generation/test_vllm_mxfp8_utils.py index 8d84f7e2613..be99da86f85 100644 --- a/tests/unit/models/generation/test_vllm_mxfp8_utils.py +++ b/tests/unit/models/generation/test_vllm_mxfp8_utils.py @@ -198,3 +198,17 @@ def test_pad_w13_intermediate_preserves_gate_halves(is_gated: bool) -> None: else: torch.testing.assert_close(padded[:, :2], tensor) assert torch.all(padded[:, 2:] == -1) + + +def test_pad_w13_intermediate_reuses_matching_storage() -> None: + pad_w13_intermediate = _load_mxfp8_utils().pad_w13_intermediate + tensor = torch.arange(12).reshape(1, 4, 3) + + assert pad_w13_intermediate(tensor, 2, is_gated=True) is tensor + + +def test_pad_w13_intermediate_rejects_smaller_target() -> None: + pad_w13_intermediate = _load_mxfp8_utils().pad_w13_intermediate + + with pytest.raises(ValueError, match="Cannot pad gated MXFP8 W13"): + pad_w13_intermediate(torch.ones(1, 4, 3), 1, is_gated=True) diff --git a/tests/unit/models/generation/test_vllm_qwen35_mxfp8_recipe.py b/tests/unit/models/generation/test_vllm_qwen35_mxfp8_recipe.py new file mode 100644 index 00000000000..50d06aadd09 --- /dev/null +++ b/tests/unit/models/generation/test_vllm_qwen35_mxfp8_recipe.py @@ -0,0 +1,107 @@ +# 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. + +from pathlib import Path +from typing import Any + +import pytest +from omegaconf import OmegaConf + +from nemo_rl.utils.config import ( + load_config_with_inheritance, + register_omegaconf_resolvers, +) + +PROJECT_ROOT = Path(__file__).resolve().parents[4] +RECIPE_NAME = "grpo-qwen3.5-35ba3b-4n4g-megatron-ep16tp2-mxfp8-trtllm.yaml" +SCRIPT_NAME = "grpo-qwen3.5-35ba3b-4n4g-megatron-ep16tp2-mxfp8-trtllm.sh" + + +def _load_recipe() -> dict[str, Any]: + register_omegaconf_resolvers() + recipe_path = PROJECT_ROOT / "examples/configs/recipes/llm" / RECIPE_NAME + recipe = OmegaConf.to_container( + load_config_with_inheritance(recipe_path), resolve=True + ) + assert isinstance(recipe, dict) + return recipe + + +def test_qwen35_mxfp8_recipe_uses_sync_full_loader_and_trtllm() -> None: + recipe = _load_recipe() + generation = recipe["policy"]["generation"] + vllm_cfg = generation["vllm_cfg"] + + assert recipe["cluster"]["gpus_per_node"] == 4 + assert recipe["cluster"]["num_nodes"] == 4 + assert recipe["cluster"]["segment_size"] == 4 + assert generation.get("refit_transport") is None + assert generation["colocated"]["enabled"] is True + assert vllm_cfg["precision"] == "fp8" + assert vllm_cfg["is_mx"] is True + assert vllm_cfg["enforce_eager"] is False + assert vllm_cfg["expert_parallel_size"] == 2 + assert generation["vllm_kwargs"] == { + "moe_backend": "flashinfer_trtllm", + "expert_placement_strategy": "linear", + } + + +def test_qwen35_mxfp8_recipe_keeps_gdn_and_vision_in_bf16() -> None: + recipe = _load_recipe() + patterns = recipe["policy"]["generation"]["vllm_cfg"][ + "quantization_ignore_patterns" + ] + + assert patterns == [ + "language_model.model.layers.*.linear_attn.*", + "visual.*", + "lm_head", + ] + + +def test_qwen35_mxfp8_test_script_requests_one_topology_segment() -> None: + script_path = PROJECT_ROOT / "tests/test_suites/llm" / SCRIPT_NAME + script = script_path.read_text(encoding="utf-8") + + assert "SEGMENT_SIZE=4" in script + + +@pytest.mark.vllm +def test_qwen35_mxfp8_patterns_match_vllm_module_names() -> None: + pytest.importorskip("vllm") + + from vllm.model_executor.layers.quantization.modelopt import ModelOptMxFp8Config + + recipe = _load_recipe() + patterns = recipe["policy"]["generation"]["vllm_cfg"][ + "quantization_ignore_patterns" + ] + modelopt_config = ModelOptMxFp8Config.from_config( + { + "quant_method": "modelopt", + "quant_algo": "MXFP8", + "ignore": patterns, + "ignored_layers": ["lm_head"], + } + ) + + assert modelopt_config.is_layer_excluded( + "language_model.model.layers.0.linear_attn.in_proj" + ) + assert modelopt_config.is_layer_excluded("visual.blocks.0.mlp.linear_fc2") + assert modelopt_config.is_layer_excluded("lm_head") + assert not modelopt_config.is_layer_excluded( + "language_model.model.layers.0.mlp.experts" + )