From 992db8940fdcb08b31ceaba34628c8423776746b Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Wed, 21 Oct 2026 22:03:01 +0000 Subject: [PATCH] Run only the fsdp2 llm_qat training cases on PRs PR runners have one GPU, so the deepspeed and ddp cases of test_qwen3_qat_nvfp4 and the deepspeed case of test_qwen3_qad_nvfp4 exercise the backend wrapper and ModelOpt's patches but not sharding, and each takes 100-109 s on a normal runner (146-157 s on a slow one). On pull requests only the fsdp2 cases run: fsdp2 is the documented default and has the most ModelOpt-specific code (the post-backward patch, the accelerate fix, the sharded save, distillation that requires FSDP2). The other cases keep their ids and run in the scheduled CI, selected with the same GITHUB_REF check the torch_onnx test uses. QLoRA still trains on ddp on PRs. Testing: on a CPU host with run_example_command stubbed, the PR environment runs the two fsdp2 cases and skips the other three with the reason; GITHUB_REF=refs/heads/main and an unset GITHUB_REF run all five. Not run on a GPU here. Signed-off-by: Shengliang Xu --- tests/examples/llm_qat/test_llm_qat.py | 15 ++++++++++++--- 1 file changed, 12 insertions(+), 3 deletions(-) diff --git a/tests/examples/llm_qat/test_llm_qat.py b/tests/examples/llm_qat/test_llm_qat.py index bec89bd1b97..846fdb05eae 100644 --- a/tests/examples/llm_qat/test_llm_qat.py +++ b/tests/examples/llm_qat/test_llm_qat.py @@ -15,6 +15,7 @@ import json +import os import pytest import torch @@ -32,6 +33,14 @@ # Backends that need gradient checkpointing GRADIENT_CHECKPOINTING_BACKENDS = {"ddp", "deepspeed"} +# PR runners have one GPU, so a backend case only exercises the wrapper and ModelOpt's patches, not +# sharding. PRs run fsdp2 (the default, with the most ModelOpt-specific code); the other backends run +# in the scheduled CI. QLoRA still trains on ddp on PRs. +_PR_CI = os.environ.get("GITHUB_REF", "").startswith("refs/heads/pull-request/") +_NIGHTLY_ONLY = pytest.mark.skipif( + _PR_CI, reason="Only fsdp2 runs on PRs; runs in the scheduled CI" +) + # Fast training overrides (short runs with frequent eval) FAST_TRAIN_ARGS = [ "--model_max_length", @@ -152,8 +161,8 @@ def test_dataset_utils_pretokenize(tiny_qwen3_path, tmp_path): @pytest.mark.parametrize("backend", [ "fsdp2", - "deepspeed", - "ddp", + pytest.param("deepspeed", marks=_NIGHTLY_ONLY), + pytest.param("ddp", marks=_NIGHTLY_ONLY), ]) def test_qwen3_qat_nvfp4(nvfp4_ptq_dir, dataset_cache_dir, tmp_path, backend): qat_output_dir = tmp_path / "qat" @@ -227,7 +236,7 @@ def test_qwen3_lora_qat_nvfp4(tiny_qwen3_path, tmp_path): @pytest.mark.parametrize("backend", [ "fsdp2", - "deepspeed", + pytest.param("deepspeed", marks=_NIGHTLY_ONLY), ]) def test_qwen3_qad_nvfp4(tiny_qwen3_path, nvfp4_ptq_dir, dataset_cache_dir, tmp_path, backend): qad_output_dir = tmp_path / "qad"