mirror of
https://github.com/NVIDIA/Model-Optimizer.git
synced 2026-10-02 03:14:52 +08:00
Keval MorabiaandClaude Opus 5
411d072a2e
### What does this PR do? Type of change: Test infrastructure / CI time `tests/examples/megatron_bridge` spends most of its time importing Python, not testing. Each step of a test spawns `torchrun`, and the new process spends **~25s importing torch/megatron/modelopt** before doing any work. A single `test_qad` run pays that **six** times — three steps, plus a spawned child per distributed checkpoint save, because Megatron-Core's async writer uses `mp_mode="spawn"` and spawn re-imports `__main__`. Profiled with phase timers in the example scripts: | | share of `test_qad[qwen3]` | |---|---| | Python imports (6 process launches × ~26s) | **~76%** | | actual compute (`mtq.quantize` 4.2s, model build 0.24s, export 0.06s) | ~5s | `run_example_command` now dispatches each step internally instead of shelling out: - **single-rank steps** run directly in the pytest process, driving the script's own `get_args()` + `main()` — no new interpreter, no re-import; - **multi-rank steps** drive `torch.distributed.run.main()` in-process with patched `sys.argv`, the same pattern Megatron-Bridge uses in its own functional tests. ### Results | suite | before | after | |---|---|---| | `tests/examples/megatron_bridge`, **1 GPU** | **21m27** | **4m03 - 6m17** | | `tests/examples/megatron_bridge`, **2 GPU** | 26m10 | 25m03 | Both figures are on current `main` (17 tests). The 2-GPU number is up from 20m41 before merging #2276, which added a test and made one previously single-rank step multi-rank. The 1-GPU figure is a range, not a best case. Across ten runs on a verified-idle box the suite lands at ~4m most of the time and at ~6m otherwise, always with the same result (17 passed). Per-test durations show the entire spread is one test: `test_qad[qwen3]` runs at ~15s or at ~149s. It is 25s in isolation, 15s after `test_distill.py` and ~25s after `test_prune_minitron.py`, so it needs the full sequence and does not reproduce on demand — three attempts to catch it under instrumentation all landed on fast runs. In those, CUDA state immediately before it is 46 MiB allocated / 68 MiB reserved / 7 segments / 22 MiB inactive-split, and the preceding test's 498/984 MiB is fully reclaimed, so a fragmented allocator is measured *not* to be the cause in the fast path at least. Left documented rather than guessed at: correctness is unaffected across every run, and the worst case sits inside the 30-minute PR budget (CI `Run tests` 902s). The single-GPU path is the big win, and it is the one the per-PR runner uses — that job now finishes in **8 minutes** in CI. Multi-rank steps still launch worker processes that re-import, so the 2-GPU nightly improves far less. This also fixes the timeouts under coverage. With `--cov` (how CI runs it), on the same three tests: in-process **3 passed in 1m15**, subprocess **3 failed on `Timeout (>360.0s)` in 18m57**. ### CI timeout The 2-GPU nightly runs every test multi-GPU and measured **58 minutes against a 60-minute cap** — too close to be reliable. `timeout_minutes` is now ref-conditional, mirroring the `runner` line directly below it: **30 minutes on PRs** (single-GPU, ~8 min) and **75 on nightly**. Keeping the nightly at full multi-GPU coverage is deliberate. Making individual tests single-rank cut it to ~7 minutes, but it gives up the parallel-path coverage that is the whole point of the 2-GPU job, and it surfaced a real fragility: `test_prune_minitron[nemotron_h]` fails with *"No scores collected for importance estimation"* when it runs single-rank after the full distill file. It passes alone and after any single preceding test — multi-rank tests are immune because `torchrun` gives them fresh worker processes. Nightly is the right place to spend the wall-clock. ### What is and isn't covered Each script's real `get_args()` still runs, so CLI flags, defaults and recipe-string resolution stay covered. Not covered for single-rank steps: the `torchrun` invocation itself and the `__main__` block (`dist.setup()` / `dist.abort()`). Multi-rank steps still go through the real launcher. **No test file changes.** The tests still read as "launch this torchrun command" and their assertions are untouched. ### Keeping it that way There is no toggle and no fallback. Every step in this suite must be `torchrun --nproc_per_node=<int> <script>.py` with the script exposing `get_args()` + `main()`, and `run_example_step` raises otherwise — it returns `str`, not `str | None`, so a step cannot quietly become a subprocess. That matters because a silent fallback still *passes*, just ~6x slower, so a new script or test could cost the suite its speed-up with nothing to show for it. Both guards verified by breaking them on purpose, each failing in ~1.4s rather than burning a run: | broken convention | result | |---|---| | `--nproc_per_node=gpu` | `AssertionError: --nproc_per_node must be a plain integer: [...]` | | step invoking `generate_vllm.py` (no `get_args`) | `AssertionError: generate_vllm.py must define get_args() and main(args)` | ### Layout The runner lives in `tests/_test_utils/examples/megatron_example_runner.py`, next to the `run_command.py` it plugs into and the other per-example helpers. It is deliberately not under `tests/_test_utils/torch/megatron/`: both files there import megatron at module top, whereas this one must not, since importing `megatron.bridge` would initialise CUDA in the pytest process and hold a context on device 0 for the whole session. ### Isolation Sharing one interpreter means anything global has to be put back between steps, or one failing test cascades into the next. Each of these was previously cleaned up by `torchrun` simply exiting: - **`NVTE_*`** — Transformer-Engine records its chosen attention backend in the environment, so a Mamba hybrid failed after an attention model ran. The environment is restored wholesale rather than by naming variables. - **Allocator** — `empty_cache()` frees nothing while a finished step's model is still reachable; a later test ran **9x slower** (162s vs 18s) against a fragmented allocator until `gc.collect()` was added first. - **Parallel state and the rerun state machine** — two separate singletons; `destroy_model_parallel()` does not touch the latter. - **Signal handlers** — `PContext.start()` installs its own `SIGTERM`/`SIGINT`/`SIGHUP`/`SIGQUIT` handlers and never restores them. With a subprocess launcher, process exit did that for us; in-process they are saved and put back, or pytest's Ctrl-C and CI cancellation would break for the rest of the session. Verified rather than assumed: injecting a failure mid-test (after a model was built and parallel state left live) gives **1 failed, 2 passed**, with the surviving tests at full speed. ### Coverage Coverage of the exercised code **improves**. In subprocess mode the child imports modelopt as `site-packages/modelopt/...` while pytest measures `modelopt/...`, so the data never merges — which is also why the subprocess report showed exactly double the statement count. | module | subprocess | in-process | |---|---|---| | `unified_export_megatron.py` | 8% | **43%** | | `mcore_custom.py` | 34% | **44%** | ### Testing All in `nvcr.io/nvidia/nemo:26.08` on 2x RTX 6000 Ada, with per-test caps enforced. | run | result | |---|---| | `tests/examples/megatron_bridge`, 1 GPU | 17 passed — 4m03 (6m17 worst of 10 runs) | | same, subprocess baseline | 15 passed, 1 skipped — 21m27 | | `tests/examples/megatron_bridge`, 2 GPU | 17 passed — 25m03 | | CI 1-GPU example job (`megatron / run-test`) | passed — `Run tests` 902s, job ~20m (30m cap) | | CI 2-GPU nightly | passed — `Run tests` 3162s, job 58m (75m cap) | | cascade check (injected mid-test failure) | 1 failed, 2 passed, survivors at full speed | | pre-commit | clean | ### Before your PR is "*Ready for review*" - Is this change backward compatible?: ✅ — test-only; no source or public API changes - If you copied code from any other sources or added a new PIP dependency, did you follow guidance in `CONTRIBUTING.md`: N/A — no new dependencies (`pytest-forked` was evaluated and rejected: `import megatron.bridge` initialises CUDA, and CUDA cannot be re-initialised in a forked child) - Did you write any new necessary tests?: N/A — this changes how existing tests are executed - Did you update [Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?: N/A — internal test infrastructure, not user-facing - Did you get Claude approval on this PR?: ✅ — reviewed by Claude and CodeRabbit, all threads addressed --------- Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> Co-authored-by: Claude Opus 5 <noreply@anthropic.com>