Make training-loop-level unit tests easier to write
- 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
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