fix(tests): clean up sys.modules pollution in training fixtures

[rob-luke]

What does this PR do ?

Fix two cooperating fixture bugs in tests/unit_tests/training/ that cause order-dependent failures when the affected files are run as a subset (e.g. focused local iteration), without affecting the full-suite CI run.

Changelog

• tests/unit_tests/training/test_timers.py
  • Replace direct sys.modules["torch.distributed"] = dist_stub with monkeypatch.setitem(sys.modules, "torch.distributed", dist_stub) so the stub is reverted at fixture teardown. The previous assignment bypassed monkeypatch's undo log and leaked the stub for the rest of the pytest session, causing later tests that triggered a fresh torchao import to fail with ModuleNotFoundError: No module named 'torch.distributed._tensor'; 'torch.distributed' is not a package.
• tests/unit_tests/training/test_train_ft_mlflow_logging.py
  • Set wandb.__spec__ on the fake wandb stub via importlib.util.spec_from_loader("wandb", loader=None). types.ModuleType("wandb") defaults __spec__ to None, which made accelerate's importlib.util.find_spec("wandb") raise ValueError: wandb.__spec__ is None whenever the file was run in isolation.

Before your PR is "Ready for review"

Pre checks:

• Make sure you read and followed Contributor guidelines
• Did you write any new necessary tests? No new tests — the fix is to test fixtures themselves; correctness is observable via the existing tests, which now pass when run as a subset.
• Did you add or update any necessary documentation? Not applicable.

Additional Information

To reproduce the fixed errors, run (against main before this PR merges):

uv run --no-sync pytest \
    tests/unit_tests/training/test_timers.py \
    tests/unit_tests/training/test_train_ft_mlflow_logging.py \
    -v --tb=short

Pre-fix: both tests in test_train_ft_mlflow_logging.py fail with ModuleNotFoundError: No module named 'torch.distributed._tensor'; 'torch.distributed' is not a package. After this PR: 17 / 17 pass.

Why CI did not catch this

Full-suite runs collect tests/unit_tests/ alphabetically, so directories like _cli/ (e.g. test_app.py, test_pretrain_cli.py) run long before training/ and import nemo_automodel.recipes.llm.train_ft, which transitively caches torchao.float8 in sys.modules. Once cached, the leaked torch.distributed stub from test_timers.py no longer triggers a re-import of torchao, so the failure is masked. Verified locally:

Scenario (pre-fix)	Result
pytest training/test_timers.py training/test_train_ft_mlflow_logging.py	2 failed, 15 passed
pytest _cli/test_app.py training/test_timers.py training/test_train_ft_mlflow_logging.py	33 passed

Full tests/unit_tests/training/ regression sweep after fix: 55 / 55. ruff format and ruff check clean on the two changed files.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[rob-luke]

@akoumpa & @ko3n1g thank you for the great package. I opened this PR during a GitHub PR outage, and it does not appear in GitHubs Pull Request tab for me, please let me know if I should reopen the PR so it is visible for review. Please disregard this message, the outage appears to be over and the PR is now visible. Sorry

[akoumpa]

/ok to test 11ccaf1