From 411d072a2e7caad083ee858de8a20caf8b2043f0 Mon Sep 17 00:00:00 2001
From: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
Date: Wed, 2 Sep 2026 22:08:56 +0530
Subject: [PATCH] Speed up megatron_bridge example tests by ~6x on a single GPU
 (#2296)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

### 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>
---
 .github/workflows/example_tests.yml           |   3 +-
 .../examples/megatron_example_runner.py       | 333 ++++++++++++++++++
 tests/_test_utils/examples/run_command.py     |  48 ++-
 tests/examples/megatron_bridge/conftest.py    |  57 +++
 4 files changed, 437 insertions(+), 4 deletions(-)
 create mode 100644 tests/_test_utils/examples/megatron_example_runner.py
 create mode 100644 tests/examples/megatron_bridge/conftest.py

diff --git a/.github/workflows/example_tests.yml b/.github/workflows/example_tests.yml
index 14cf7247e..68d9a0c1e 100644
--- a/.github/workflows/example_tests.yml
+++ b/.github/workflows/example_tests.yml
@@ -155,7 +155,8 @@ jobs:
     with:
       docker_image: "nvcr.io/nvidia/nemo:26.08"
       example: megatron_bridge
-      timeout_minutes: 75
+      # PR runs single-GPU (~11 min); the nightly runs every test multi-GPU (~58 min).
+      timeout_minutes: ${{ startsWith(github.ref, 'refs/heads/pull-request/') && 30 || 75 }}
       pip_install_extras: "[hf,puzzletron,dev-test]"
       runner: ${{ startsWith(github.ref, 'refs/heads/pull-request/') && 'linux-amd64-gpu-rtxpro6000-latest-1' || 'linux-amd64-gpu-rtxpro6000-latest-2' }}
       allow_failure: ${{ contains(format(',{0},', vars.ALLOW_FAILURE_EXAMPLE_TESTS), ',megatron_bridge,') }}
diff --git a/tests/_test_utils/examples/megatron_example_runner.py b/tests/_test_utils/examples/megatron_example_runner.py
new file mode 100644
index 000000000..ac974a5a0
--- /dev/null
+++ b/tests/_test_utils/examples/megatron_example_runner.py
@@ -0,0 +1,333 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Run ``examples/megatron_bridge`` scripts without a ``torchrun`` launch per step.
+
+A launch spends ~25s importing torch/megatron/modelopt before doing any work, and one example test
+runs several. Single-rank commands drive the script's ``main()`` directly in the pytest process,
+which is where that cost disappears (the suite goes from ~21min to ~4min on one GPU). Multi-rank
+commands drive ``torchrun`` itself in-process, which only saves the launcher's own imports -- the
+pytest process cannot be more than one rank -- but keeps torchrun's fresh worker processes.
+
+The script's own ``get_args()`` still runs, so CLI flags and recipe strings stay covered. Not
+covered: the ``torchrun`` invocation and the ``__main__`` block (``dist.setup()``, ``dist.abort()``
+and ``dist.cleanup()``).
+
+Megatron and torch.distributed.run are imported lazily on purpose: this module is imported at
+collection time, and importing ``megatron.bridge`` initialises CUDA in the pytest process, which
+would hold a context on device 0 for the whole session even when every step runs under torchrun.
+"""
+
+import contextlib
+import gc
+import importlib
+import importlib.util
+import logging
+import os
+import signal
+import sys
+import tempfile
+from pathlib import Path
+from unittest.mock import patch
+
+import torch
+from _test_utils.examples.run_command import MODELOPT_ROOT
+from _test_utils.torch.distributed.utils import get_free_port
+
+import modelopt.torch.utils.distributed as dist
+
+_SINK = None  # never closed: a handler built inside a capture keeps a reference to it
+
+
+def _sink():
+    """One long-lived stream on the real stdout, whose fd is redirected per step.
+
+    A per-step stream closed on exit would leave a permanently-closed file behind for anything that
+    captured it -- e.g. the common module-level ``logging.StreamHandler(sys.stdout)``, where
+    ``sys.stdout`` is this sink while a step runs and the handler outlives it in ``sys.modules``.
+    """
+    global _SINK
+    if _SINK is None:
+        _SINK = os.fdopen(os.dup(1), "w", buffering=1, errors="replace")
+    return _SINK
+
+
+# torchrun's PContext installs handlers for these and never restores them.
+_LAUNCHER_SIGNALS = (signal.SIGTERM, signal.SIGINT, signal.SIGHUP, signal.SIGQUIT)
+
+
+@contextlib.contextmanager
+def _capture_output():
+    """Capture everything a step emits, in the order it happened.
+
+    Three sources have to end up in the same place: ``print`` from the script, records from loggers
+    Megatron-Bridge uses for lines tests assert on, and fd-level writes from ``torchrun`` workers
+    and C extensions. Under pytest's fd capture ``sys.stdout``/``sys.stderr`` are objects over
+    pytest's own temp file rather than fds 1/2, so both a redirect and an fd swap are needed.
+
+    Buffered, not streamed: the caller prints the result, which pytest shows for a failing test.
+    Echoing it live would not reach the CI log anyway -- without ``-s``, fd 1 is already pytest's
+    own capture file. Piping it to stream for real cost the 1-GPU suite 3m43 -> 9m06 in
+    backpressure, so surviving a job-level SIGKILL would need a file outside the process.
+    """
+    holder = [""]
+    with tempfile.TemporaryFile(mode="w+", encoding="utf-8", errors="replace") as tmp:
+        handler = logging.StreamHandler()  # bound to the swapped fd 2 below
+        handler.setFormatter(logging.Formatter("%(message)s"))
+        root = logging.getLogger()
+        with logging._lock:  # any getLogger() from a background thread mutates loggerDict
+            others = [
+                lg for lg in root.manager.loggerDict.values() if isinstance(lg, logging.Logger)
+            ]
+        # Lower levels only for loggers that own their output: a propagating logger still needs its
+        # own level lowered, or it drops the record before root sees it. The boost is load-bearing:
+        # without it test_distill_validate_only cannot see "skipping training ...". It does mean the
+        # capture is a superset of what a torchrun user sees, so these assertions check that the
+        # step did the right work, not that its logging is user-visible.
+        levels = {lg: lg.level for lg in [root, *others] if lg.handlers or not lg.propagate}
+        # Attach the handler only where the record cannot reach root, or it is written twice.
+        handled = [root, *(lg for lg in others if not lg.propagate)]
+        sys.stdout.flush()
+        sys.stderr.flush()
+        sink = _sink()  # before the swap, so it holds the real stdout
+        saved_out, saved_err, saved_sink = os.dup(1), os.dup(2), os.dup(sink.fileno())
+        try:
+            # Inside the try: a raise between here and the restore would otherwise leave fds 1/2
+            # pointing at a temp file the ``with`` then closes, silencing the rest of the session.
+            os.dup2(tmp.fileno(), 1)
+            os.dup2(tmp.fileno(), 2)
+            sink.flush()
+            os.dup2(tmp.fileno(), sink.fileno())  # move its fd rather than replacing the object
+            handler.setStream(sink)
+            for lg in levels:
+                if lg.level > logging.INFO:
+                    lg.setLevel(logging.INFO)
+            for lg in handled:
+                lg.addHandler(handler)
+            with contextlib.redirect_stdout(sink), contextlib.redirect_stderr(sink):
+                yield holder
+        finally:
+            for lg in handled:
+                lg.removeHandler(handler)
+            for lg, level in levels.items():
+                lg.setLevel(level)
+            sink.flush()
+            os.dup2(saved_sink, sink.fileno())
+            os.close(saved_sink)
+            sys.stdout.flush()
+            sys.stderr.flush()
+            os.dup2(saved_out, 1)
+            os.dup2(saved_err, 2)
+            os.close(saved_out)
+            os.close(saved_err)
+            tmp.seek(0)
+            holder[0] = tmp.read()
+
+
+@contextlib.contextmanager
+def _preserved_signal_handlers():
+    """Put back handlers a step installs; with a subprocess launcher, process exit did this.
+
+    Both paths need it: ``PContext.start()`` installs its own, and a Megatron training loop run
+    in-process is if anything likelier to (graceful exit, checkpoint-on-signal). Losing them breaks
+    pytest's Ctrl-C and CI cancellation for the rest of the session.
+    """
+    saved = {s: signal.getsignal(s) for s in _LAUNCHER_SIGNALS}
+    try:
+        yield
+    finally:
+        for sig, handler in saved.items():
+            if handler is not None:  # installed from C; not restorable via signal.signal()
+                signal.signal(sig, handler)
+
+
+def reset_megatron_global_state() -> None:
+    """Drop the global state a finished single-rank step leaves behind.
+
+    Steps that run here share one interpreter, so anything global has to be put back or it leaks
+    into the next test. Failure-tolerant on purpose: this also runs after a failed step, where the
+    model may be half-built, and it must not mask the error that got us here.
+    """
+    with contextlib.suppress(Exception):
+        from megatron.core import parallel_state
+
+        parallel_state.destroy_model_parallel()
+    with contextlib.suppress(Exception):
+        # A separate singleton from the parallel state, initialised by the training path.
+        from megatron.core.rerun_state_machine import destroy_rerun_state_machine
+
+        destroy_rerun_state_machine()
+    with contextlib.suppress(Exception):
+        # Collect first: empty_cache() only returns blocks nothing references, and a finished step
+        # can leave its model reachable from the imported module. Without this a later test ran
+        # against a fragmented allocator ~9x slower.
+        gc.collect()
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+
+
+def _load_example_module(script: str, example_path: str):
+    """Import an example script under a namespaced module name.
+
+    ``import_module("quantize")`` would consult ``sys.modules`` first and could pick up an
+    unrelated top-level module of the same name, and would leave the example cached under a
+    generic name.
+    """
+    example_dir = str(MODELOPT_ROOT / "examples" / example_path)
+    name = f"_modelopt_example_{example_path.replace('/', '_')}_{Path(script).stem}"
+    if name in sys.modules:
+        return sys.modules[name]
+    spec = importlib.util.spec_from_file_location(name, Path(example_dir) / Path(script).name)
+    module = importlib.util.module_from_spec(spec)
+    sys.modules[name] = module
+    sys.path.insert(0, example_dir)  # the script's own sibling imports
+    try:
+        spec.loader.exec_module(module)
+    except BaseException:
+        del sys.modules[name]
+        raise
+    finally:
+        sys.path.remove(example_dir)
+    return module
+
+
+def _require_drivable(script: str, example_path: str) -> None:
+    """Check the script exposes the ``get_args()`` + ``main()`` shape this runner drives."""
+    module = _load_example_module(script, example_path)
+    if not callable(getattr(module, "get_args", None)) or not callable(
+        getattr(module, "main", None)
+    ):
+        raise AssertionError(f"{script} must define get_args() and main(args)")
+
+
+def requested_world_size(cmd_parts: list[str]) -> int | None:
+    """World size a ``torchrun`` command asks for, or ``None`` if it is not a plain integer."""
+    parts = [str(p) for p in cmd_parts]
+    for i, part in enumerate(parts):
+        if part.startswith("--nproc_per_node="):
+            value = part.split("=", 1)[1]
+        elif part == "--nproc_per_node" and i + 1 < len(parts):
+            value = parts[i + 1]  # torchrun accepts the space-separated form too
+        else:
+            continue
+        return int(value) if value.isdigit() else None  # "gpu"/"auto"
+    return None
+
+
+def run_example_in_process(cmd_parts: list[str], example_path: str) -> str:
+    """Drive a single-rank step's real ``get_args()`` + ``main()`` here. Returns its stdout."""
+    # Set unconditionally: the per-test environment restore drops these while the process group
+    # stays initialised, so a later step would otherwise run with a live group and no launch vars.
+    os.environ.setdefault("MASTER_ADDR", "127.0.0.1")
+    os.environ["RANK"] = "0"
+    os.environ["WORLD_SIZE"] = "1"
+    os.environ["LOCAL_RANK"] = "0"
+    if not dist.is_initialized():
+        # Not setdefault: another suite may have left a stale MASTER_PORT in the environment, and
+        # binding an occupied port blocks in rendezvous until timeout instead of failing fast.
+        os.environ["MASTER_PORT"] = str(get_free_port())
+        dist.setup()  # also sets the CUDA device
+
+    script = next(p for p in cmd_parts if str(p).endswith(".py"))
+    argv = [str(p) for p in cmd_parts[cmd_parts.index(script) :]]
+    example_dir = str(MODELOPT_ROOT / "examples" / example_path)
+    cwd = os.getcwd()
+    os.chdir(example_dir)  # match the subprocess path, which runs with the example dir as cwd
+    native = [""]  # bound up front: readable even if the capture fails to start
+    try:
+        module = _load_example_module(str(script), example_path)
+        reset_megatron_global_state()  # a previous step left its own parallel state behind
+        try:
+            # get_args() inside the guard: argparse calls parser.error() -> SystemExit on a flag
+            # that has drifted from the example's parser, which is the failure this suite most
+            # wants to report clearly rather than raise bare out of the runner. argv stays patched
+            # across main() as well, which the subprocess path had.
+            with (
+                _preserved_signal_handlers(),
+                _capture_output() as native,
+                patch.object(sys, "argv", argv),
+            ):
+                args = module.get_args()
+                module.main(args)
+        except SystemExit as e:
+            if e.code not in (0, None):
+                raise RuntimeError(f"{script} exited with code {e.code}") from e
+        return native[0]
+    except BaseException as e:
+        # The caller matches transient-HuggingFace markers on this alongside the traceback.
+        e.captured_output = native[0]
+        raise
+    finally:
+        os.chdir(cwd)
+        print(native[0])  # keep the step's output in the test log
+
+
+def run_torchrun_in_process(cmd_parts: list[str], example_path: str) -> str:
+    """Drive ``torchrun`` itself in-process, letting it spawn fresh workers as usual.
+
+    Used for multi-rank steps. Only the launcher's imports are saved, but the workers are fresh
+    processes, so none of the global state a reused process would inherit applies. Borrowed from
+    Megatron-Bridge's own functional tests.
+    """
+    from torch.distributed.run import main as torchrun_main
+
+    reset_megatron_global_state()  # release device 0 before the workers claim it
+    example_dir = MODELOPT_ROOT / "examples" / example_path
+    script = next(p for p in cmd_parts if str(p).endswith(".py"))
+    argv = [str(p) for p in cmd_parts]
+    argv[argv.index(str(script))] = str(example_dir / Path(script).name)
+    if not any(a.startswith(("--master_port", "--master-port")) for a in argv):
+        # The rendezvous store now lives in a long-lived process, so do not rely on torchrun's
+        # fixed default port being free by the time the next step starts.
+        argv.insert(1, f"--master_port={get_free_port()}")
+    cwd = os.getcwd()
+    os.chdir(example_dir)
+    cap = [""]
+    try:
+        with (
+            _preserved_signal_handlers(),
+            _capture_output() as cap,
+            patch.object(sys, "argv", argv),
+        ):
+            torchrun_main()
+        return cap[0]
+    except BaseException as e:
+        # ChildFailedError only carries the worker's traceback when the entrypoint is decorated
+        # with @record, which the examples are not -- so the captured output is all the caller has.
+        e.captured_output = cap[0]
+        raise
+    finally:
+        os.chdir(cwd)
+        print(cap[0])  # keep the step's output in the test log
+
+
+def run_example_step(cmd_parts: list[str], example_path: str) -> str:
+    """Run an example step without shelling out.
+
+    Every step must be ``torchrun --nproc_per_node=<int> <script>.py``, with the script exposing
+    ``get_args()`` + ``main()``. Anything else raises: falling back to a subprocess would still
+    pass, only ~6x slower, so a step that breaks the convention has to fail instead of quietly
+    costing the suite its speed-up.
+    """
+    script = next((str(p) for p in cmd_parts if str(p).endswith(".py")), None)
+    if script is None:
+        raise AssertionError(f"step must invoke a .py script: {cmd_parts}")
+    world_size = requested_world_size(cmd_parts)
+    if world_size is None:
+        raise AssertionError(f"--nproc_per_node must be a plain integer: {cmd_parts}")
+    if world_size > 1:
+        # torchrun imports the script in fresh children, so get_args()/main() need not exist here.
+        return run_torchrun_in_process(cmd_parts, example_path)
+    _require_drivable(script, example_path)
+    return run_example_in_process(cmd_parts, example_path)
diff --git a/tests/_test_utils/examples/run_command.py b/tests/_test_utils/examples/run_command.py
index 329aaf727..ce7a0248c 100644
--- a/tests/_test_utils/examples/run_command.py
+++ b/tests/_test_utils/examples/run_command.py
@@ -20,6 +20,7 @@ import signal
 import subprocess
 import threading
 import time
+import traceback
 import warnings
 from pathlib import Path
 
@@ -136,6 +137,17 @@ def _run_capturing(cmd_parts: list[str], cwd: Path, env: dict[str, str]) -> tupl
     return returncode, "".join(chunks)
 
 
+# Set by a suite's conftest to run a command without spawning a subprocess (see
+# ``_test_utils.examples.megatron_example_runner``). Unset means always use a subprocess.
+_in_process_runner = None
+
+
+def set_in_process_runner(runner) -> None:
+    """Register ``runner(cmd_parts, example_path) -> str``; it must not fall back to a subprocess."""
+    global _in_process_runner
+    _in_process_runner = runner
+
+
 def run_example_command(
     cmd_parts: list[str],
     example_path: str,
@@ -146,18 +158,48 @@ def run_example_command(
 ) -> str | None:
     """Run an example command, retrying transient HuggingFace access errors."""
     print(f"[{example_path}] Running command: {cmd_parts}")
-    env = env or os.environ.copy()
+    in_process = _in_process_runner if env is None else None
+    if _in_process_runner is not None and env is not None:
+        # The in-process runner uses the ambient environment, so a caller-supplied env would be
+        # silently dropped. Fall back to a subprocess rather than run with the wrong environment.
+        warnings.warn(f"[{example_path}] env= given; running this step as a subprocess")
+    env = os.environ.copy() if env is None else env
     cwd = MODELOPT_ROOT / "examples" / example_path
 
     for attempt in range(hf_max_retries + 1):
         if setup_free_port:
+            # Subprocess steps only: an in-process runner picks its own free port, since env is a
+            # copy it never sees.
             env["MASTER_PORT"] = str(get_free_port())  # fresh port per attempt
-        returncode, output = _run_capturing(cmd_parts, cwd, env)
+        if in_process is not None:
+            # Inside the loop so in-process steps get the same transient-HuggingFace retries;
+            # these tests do hit the Hub (e.g. calib_dataset_name="cnn_dailymail").
+            try:
+                result = in_process(cmd_parts, example_path)
+            except Exception as e:
+                # Re-raise unless it looks transient, so a real failure keeps its traceback
+                # instead of being flattened into CalledProcessError.
+                # format_exc() walks __cause__/__context__: most markers are exception type
+                # names that only ever appear in a formatted traceback, never in str(e) -- a Hub
+                # ConnectionError typically surfaces wrapped in a DatasetGenerationError. The
+                # runner also attaches what the step printed, for launcher-level failures.
+                captured = getattr(e, "captured_output", "")
+                text = f"{traceback.format_exc()}\n{captured}"
+                if attempt == hf_max_retries or not any(
+                    marker in text for marker in _HF_TRANSIENT_MARKERS
+                ):
+                    raise
+                returncode, output = 1, text
+            else:
+                assert result is not None, "an in-process runner must return output, not fall back"
+                return result
+        else:
+            returncode, output = _run_capturing(cmd_parts, cwd, env)
         if returncode == 0:
             return output
         transient = any(marker in output for marker in _HF_TRANSIENT_MARKERS)
         if not transient or attempt == hf_max_retries:
-            raise subprocess.CalledProcessError(returncode, cmd_parts)
+            raise subprocess.CalledProcessError(returncode, cmd_parts, output=output)
         warnings.warn(
             f"[{example_path}] transient HuggingFace access error; retrying in "
             f"{hf_retry_delay_s}s (attempt {attempt + 1}/{hf_max_retries})"
diff --git a/tests/examples/megatron_bridge/conftest.py b/tests/examples/megatron_bridge/conftest.py
new file mode 100644
index 000000000..1093fb7b2
--- /dev/null
+++ b/tests/examples/megatron_bridge/conftest.py
@@ -0,0 +1,57 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Run example steps without a ``torchrun`` launch; see ``_test_utils.examples.megatron_example_runner``."""
+
+import os
+
+import pytest
+from _test_utils.examples.megatron_example_runner import (
+    reset_megatron_global_state,
+    run_example_step,
+)
+from _test_utils.examples.run_command import set_in_process_runner
+
+
+@pytest.fixture(autouse=True)
+def _fast_example_runner():
+    """Run example steps without shelling out to ``torchrun``.
+
+    Per test, not per session: the hook is a module-global in ``run_command``, and this runner
+    raises rather than falling back, so leaving it installed would break any other example suite
+    collected later in the same session (``pytest tests/examples``).
+    """
+    set_in_process_runner(run_example_step)
+    try:
+        yield
+    finally:
+        set_in_process_runner(None)
+
+
+@pytest.fixture(autouse=True)
+def _isolate_megatron_global_state():
+    """Reset shared state around every test so a failure cannot cascade into the next one.
+
+    In-process steps share the interpreter. Besides Megatron's singletons, Transformer-Engine
+    records its chosen attention backend in ``NVTE_*``, which failed a Mamba hybrid that ran after
+    an attention model -- so the environment is restored wholesale rather than by naming variables.
+    """
+    env_before = os.environ.copy()
+    reset_megatron_global_state()
+    try:
+        yield
+    finally:
+        reset_megatron_global_state()
+        os.environ.clear()
+        os.environ.update(env_before)