Skip to content

[WIP] Add experimental NVFP4 fake-QAT RL recipe #2020

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
zianglih wants to merge 2 commits into
base: main
Choose a base branch
from
Draft

[WIP] Add experimental NVFP4 fake-QAT RL recipe #2020

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
2d1bdc2
Add NVFP4 QAT launch mode
zianglih Jul 31, 2026
eef435f
Remove NVFP4 QAT documentation changes
zianglih Jul 31, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
16 changes: 14 additions & 2 deletions scripts/run_qwen3_30b_a3b.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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"


Expand Down Expand Up @@ -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:
Expand Down