From 2d1bdc2db5577ad3c7ea5062b9a149bafd3cd822 Mon Sep 17 00:00:00 2001 From: Ziang Li Date: Thu, 30 Jul 2026 20:56:40 -0700 Subject: [PATCH 1/2] Add NVFP4 QAT launch mode --- docs/advanced/fp8-low-precision.md | 30 ++++++++++++++++++++++++++---- scripts/run_qwen3_30b_a3b.py | 16 ++++++++++++++-- 2 files changed, 40 insertions(+), 6 deletions(-) diff --git a/docs/advanced/fp8-low-precision.md b/docs/advanced/fp8-low-precision.md index ac35116be81..3e768fc59d5 100644 --- a/docs/advanced/fp8-low-precision.md +++ b/docs/advanced/fp8-low-precision.md @@ -33,7 +33,7 @@ forward precision. ✅ = supported; ✗ = not supported. | **BF16** | ✅ baseline | ✗ | ✗ | ✗ | | **FP8 block-wise** | ✅ | ✅ Hopper + Blackwell | ✗ | ✗ | | **MXFP8** | ✅ | ✗ | ✅ Blackwell | ✗ | -| **NVFP4** | ✗ | ✗ | ✗ | 🚧 coming soon | +| **NVFP4** | ✅ with precision drift | ✗ | ✗ | ✅ experimental | Two rules enforced in the reference script (`scripts/run_qwen3_30b_a3b.py`): @@ -41,6 +41,8 @@ Two rules enforced in the reference script * `--rollout-mxfp8` and `--rollout-fp8` are mutually exclusive. * `--train-mxfp8` requires `--rollout-mxfp8` (no MXFP8-train + FP8-rollout combo). +* `--train-nvfp4-qat` requires `--rollout-nvfp4` and is mutually exclusive + with native `--train-nvfp4`. ## Unified training recipe @@ -179,14 +181,35 @@ NVFP4 is FP4 E2M1 with 1D block scaling (group size 16) and a two-level scale reference. Today only **MoE expert GEMMs** are quantized; dense layers stay in their original precision. -The full unified NVFP4 recipe is in development. +The reference script exposes two experimental training modes: + +* `--train-nvfp4-qat` keeps BF16 master weights and GEMMs, but applies the same + TE 1×16 NVFP4 quantize/dequantize operation to routed-expert FC1/FC2 weights + on every forward. Gradients pass through with a straight-through estimator. +* `--train-nvfp4` uses native TE NVFP4 GEMMs and is a separate recipe; do not + combine it with fake QAT. + +Run the fake-QAT path together with NVFP4 rollout: + +```bash +python scripts/run_qwen3_30b_a3b.py \ + --hardware B200 \ + --rollout-nvfp4 \ + --train-nvfp4-qat +``` + +The script sets `OPEN_TRAINING_NVFP4_FAKE_QAT_FLAG=1` for Megatron. It also +forwards the driver's `NVTE_*` and `FLASHINFER_*` variables, so conversion, +training, and rollout can share the same 4over6 scope, E4M3 bound, error mode, +and fast-math settings. The quantized checkpoint does not encode those choices; +set both backends consistently and restore the same environment when resuming. ## Hardware support | GPU | BF16 | FP8 block-wise | MXFP8 | NVFP4 | |---|---|---|---|---| | NVIDIA H100 / H200 | ✅ | ✅ | ✗ | ✗ | -| NVIDIA B200 / B300 / GB200 / GB300 | ✅ | ✅ | ✅ | 🚧 in development | +| NVIDIA B200 / B300 / GB200 / GB300 | ✅ | ✅ | ✅ | ✅ experimental | | NVIDIA A100 | ✅ | ✗ | ✗ | ✗ | | AMD MI300X / MI325 / MI350 / MI355X | ✅ | ✗ | ✗ | ✗ | @@ -197,4 +220,3 @@ The full unified NVFP4 recipe is in development. * AMD hardware today. * Bring-up of a new model architecture, where clean BF16 numerics simplify debugging. - diff --git a/scripts/run_qwen3_30b_a3b.py b/scripts/run_qwen3_30b_a3b.py index 1023164278c..1823d655f62 100644 --- a/scripts/run_qwen3_30b_a3b.py +++ b/scripts/run_qwen3_30b_a3b.py @@ -31,6 +31,7 @@ class ScriptArgs(U.ExecuteTrainConfig): train_fp8: bool = False train_mxfp8: bool = False train_nvfp4: bool = False + train_nvfp4_qat: bool = False enable_megatron_bridge: bool = False enable_mis: bool = False # TODO improve, should be able to override more easily @@ -50,9 +51,18 @@ def __post_init__(self): sum((self.rollout_fp8, self.rollout_mxfp8, self.rollout_int4, self.rollout_nvfp4)) <= 1 ), "only one rollout precision mode can be enabled" assert ( - sum((self.train_fp8, self.train_mxfp8, self.train_nvfp4)) <= 1 + sum((self.train_fp8, self.train_mxfp8, self.train_nvfp4, self.train_nvfp4_qat)) <= 1 ), "only one train precision mode can be enabled" - if any((self.rollout_mxfp8, self.rollout_nvfp4, self.train_mxfp8, self.train_nvfp4)): + assert not self.train_nvfp4_qat or self.rollout_nvfp4, "NVFP4 QAT requires NVFP4 rollout" + if any( + ( + self.rollout_mxfp8, + self.rollout_nvfp4, + self.train_mxfp8, + self.train_nvfp4, + self.train_nvfp4_qat, + ) + ): assert self.hardware in ("B200", "B300", "GB200", "GB300"), "mxfp8 and nvfp4 only support Blackwell GPUs" @@ -217,6 +227,8 @@ def execute(args: ScriptArgs): "OPEN_TRAINING_INT4_FAKE_QAT_FLAG": "1", "OPEN_TRAINING_INT4_GROUP_SIZE": "128", } + elif args.train_nvfp4_qat: + misc_env_vars["OPEN_TRAINING_NVFP4_FAKE_QAT_FLAG"] = "1" if args.train_fp8 or args.train_mxfp8: match args.hardware: From eef435f843d84b6b9b6ba6a60b74b99c8a36387b Mon Sep 17 00:00:00 2001 From: Ziang Li Date: Thu, 30 Jul 2026 21:08:56 -0700 Subject: [PATCH 2/2] Remove NVFP4 QAT documentation changes --- docs/advanced/fp8-low-precision.md | 30 ++++-------------------------- 1 file changed, 4 insertions(+), 26 deletions(-) diff --git a/docs/advanced/fp8-low-precision.md b/docs/advanced/fp8-low-precision.md index 3e768fc59d5..ac35116be81 100644 --- a/docs/advanced/fp8-low-precision.md +++ b/docs/advanced/fp8-low-precision.md @@ -33,7 +33,7 @@ forward precision. ✅ = supported; ✗ = not supported. | **BF16** | ✅ baseline | ✗ | ✗ | ✗ | | **FP8 block-wise** | ✅ | ✅ Hopper + Blackwell | ✗ | ✗ | | **MXFP8** | ✅ | ✗ | ✅ Blackwell | ✗ | -| **NVFP4** | ✅ with precision drift | ✗ | ✗ | ✅ experimental | +| **NVFP4** | ✗ | ✗ | ✗ | 🚧 coming soon | Two rules enforced in the reference script (`scripts/run_qwen3_30b_a3b.py`): @@ -41,8 +41,6 @@ Two rules enforced in the reference script * `--rollout-mxfp8` and `--rollout-fp8` are mutually exclusive. * `--train-mxfp8` requires `--rollout-mxfp8` (no MXFP8-train + FP8-rollout combo). -* `--train-nvfp4-qat` requires `--rollout-nvfp4` and is mutually exclusive - with native `--train-nvfp4`. ## Unified training recipe @@ -181,35 +179,14 @@ NVFP4 is FP4 E2M1 with 1D block scaling (group size 16) and a two-level scale reference. Today only **MoE expert GEMMs** are quantized; dense layers stay in their original precision. -The reference script exposes two experimental training modes: - -* `--train-nvfp4-qat` keeps BF16 master weights and GEMMs, but applies the same - TE 1×16 NVFP4 quantize/dequantize operation to routed-expert FC1/FC2 weights - on every forward. Gradients pass through with a straight-through estimator. -* `--train-nvfp4` uses native TE NVFP4 GEMMs and is a separate recipe; do not - combine it with fake QAT. - -Run the fake-QAT path together with NVFP4 rollout: - -```bash -python scripts/run_qwen3_30b_a3b.py \ - --hardware B200 \ - --rollout-nvfp4 \ - --train-nvfp4-qat -``` - -The script sets `OPEN_TRAINING_NVFP4_FAKE_QAT_FLAG=1` for Megatron. It also -forwards the driver's `NVTE_*` and `FLASHINFER_*` variables, so conversion, -training, and rollout can share the same 4over6 scope, E4M3 bound, error mode, -and fast-math settings. The quantized checkpoint does not encode those choices; -set both backends consistently and restore the same environment when resuming. +The full unified NVFP4 recipe is in development. ## Hardware support | GPU | BF16 | FP8 block-wise | MXFP8 | NVFP4 | |---|---|---|---|---| | NVIDIA H100 / H200 | ✅ | ✅ | ✗ | ✗ | -| NVIDIA B200 / B300 / GB200 / GB300 | ✅ | ✅ | ✅ | ✅ experimental | +| NVIDIA B200 / B300 / GB200 / GB300 | ✅ | ✅ | ✅ | 🚧 in development | | NVIDIA A100 | ✅ | ✗ | ✗ | ✗ | | AMD MI300X / MI325 / MI350 / MI355X | ✅ | ✗ | ✗ | ✗ | @@ -220,3 +197,4 @@ set both backends consistently and restore the same environment when resuming. * AMD hardware today. * Bring-up of a new model architecture, where clean BF16 numerics simplify debugging. +