NVIDIA / NVIDIA/Megatron-LM

Make training-loop-level unit tests easier to write

Open
#7,223 0 comments 0 reactions 0 assignees View on GitHub
Dominant language
Python
Stars
17.9k
Forks
4.5k
Avg merge
4d 6h
Merged PRs (30d)
271

Description

Training-loop-level unit tests are difficult to write because training behavior depends heavily on global state.

For example:

- The test discussed in [PR #7112](https://github.com/NVIDIA/Megatron-LM/pull/7112#discussion_r3981549207) had to parse and mutate global arguments, reset global variables and the microbatch calculator, and repeatedly initialize and destroy model-parallel state.
- [`test_mcore_adapter.py`](https://github.com/NVIDIA/Megatron-LM/blob/b005bf14c46b62169533e755ca2d4e62fe6b7e0a/tests/unit_tests/distributed/mfsdp_v2/test_mcore_adapter.py#L448) manually resets several CUDA graph wrappers and static buffers to prevent state from leaking between tests.

It would be useful to have shared test utilities that initialize, isolate, and clean up this state, making short training-loop tests easier to write and understand.

The current alternative is a functional test. However, functional tests depend on golden values, which introduce maintenance and comparison challenges. They also expose only predefined metrics, making it difficult to inspect intermediate values or assert directly on internal training behavior.

Tag @NVIDIA/mcore-oncall. I talked about this during the standup. I don't think anyone is ready to work on this, but I'm still filing it to keep the record.

Contributor guide

Open the contributing guide

Research direction

Start with the test discussed in PR #7112 and tests/unit_tests/distributed/mfsdp_v2/test_mcore_adapter.py, especially the setup and cleanup around global state, CUDA graph wrappers, and static buffers. Map the state that training-loop tests must initialize and isolate. Done means shared utilities let short tests avoid manual resets while asserting intermediate training behavior directly.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
machine-learning, testing
Issue type
Feature
Difficulty
5/5
Estimated time
Over a week
Activity status
Active
Clarity
Mostly clear
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.