Skip to content

[OMNIML-6180] Run only the fsdp2 llm_qat training cases on PRs #2925

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

Merged
shengliangxu merged 1 commit into from
Oct 22, 2026
+12 −3
Merged

[OMNIML-6180] Run only the fsdp2 llm_qat training cases on PRs #2925

Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
15 changes: 12 additions & 3 deletions tests/examples/llm_qat/test_llm_qat.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,7 @@


import json
import os

import pytest
import torch
Expand All @@ -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/")
Comment thread
shengliangxu marked this conversation as resolved.
_NIGHTLY_ONLY = pytest.mark.skipif(

@kevalmorabia97 kevalmorabia97 Oct 22, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Can we use pytest.mark.manual marker directly?

@shengliangxu shengliangxu Oct 22, 2026

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Not directly, for now. The example workflow never passes --run-manual (only gpu_tests.yml and unit_tests.yml do, on non-PR runs), so a bare pytest.mark.manual would skip these three cases on the nightly too, and nothing would run them. Making it work means the example runner passes --run-manual off-PR, which would also start the two manual cases in tests/examples/speculative_decoding/test_eagle.py ("skip redundant test, too slow") on every nightly, and local runs would skip the three cases unless --run-manual is given. I can do that in a follow-up (the runner passes --run-manual off-PR, the examples use the marker, and we decide what to do with the two eagle cases). Here the GITHUB_REF check matches the existing precedent in test_torch_quant_to_onnx.py (#2853).

_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",
Expand Down Expand Up @@ -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"
Expand Down Expand Up @@ -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"
Expand Down
Loading