ENH: Improve MetaMath training script runtime #2894
Merged
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
The training script of the MetaMathQA PEFT method comparison was calling cuda.empty_cache() and gc.collect() after each step. However, this is not really needed and it also slows down training considerably. It turns out that gc.collect() is not needed at all and thus it has been removed. This results in a big improvement in runtime. As for empty_cache(), not calling it at all leads to an increase in memory usage, but it's not necessary to call it every step. It is instead called every 10th step. Improvement (tested locally, 250 steps): - Removing gc.collect() - 108 sec => 65 sec - memory reserved max stays the same (19.3 GB) - memory reserved 99th percentile stays the same (18.0 GB) - memory reserved avg stays the same (12.0 GB) - Also calling empty_cache() only every 10 steps - 65 sec => 50 sec - memory reserved max stays the same (19.3 GB) - memory reserved avg: 18.0 GB => 19.3 GB - memory reserved avg: 12.0 GB => 14.5 GB Thus gc.collect() can be safely removed. And while calling empty_cache() only every 10th step does increase average memory usage, the peak is unaffected, which is what's most important in this benchmark, so it is a worthwhile tradeoff for the 23% speed improvement we get. Note to maintainers: If this is merged, all MetaMathQA benchmarks should be re-run.
|
The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update. |
githubnemo
reviewed
Nov 4, 2025
method_comparison/MetaMathQA/run.py
Outdated
| } | ||
| print_verbose(json.dumps(log_dict)) | ||
|
|
||
| # # TODO is this needed? |
Collaborator
There was a problem hiding this comment.
I think the comment can be removed if we don't get any more insights :)
githubnemo
approved these changes
Nov 4, 2025
Collaborator
githubnemo
left a comment
githubnemo
left a comment
There was a problem hiding this comment.
LGTM. Thanks for investigating, this is a really nice find.
If there are no new discoveries I think this can be merged.
Also remove gc import
githubnemo
approved these changes
Nov 5, 2025
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
The training script of the MetaMathQA PEFT method comparison was calling
cuda.empty_cache()andgc.collect()after each step. However, this is not really needed and it also slows down training considerably.It turns out that
gc.collect()is not needed at all when it comes to memory and thus it has been removed. This results in a big improvement in runtime. As forempty_cache(), not calling it at all leads to an increase in memory usage, but it's not necessary to call it every step. It is instead called every 10th step.Improvement (tested locally, 250 steps):
gc.collect()empty_cache()only every 10 stepsempty_cache()at all:Thus
gc.collect()can be safely removed, but removingempty_cache()completely is not advisable. Callingempty_cache()only every 10th step does increase average memory usage, but the peak is unaffected, which is what's most important in this benchmark, so it is a worthwhile tradeoff for the 23% speed improvement we get.I also tested how manually deleting certain torch variables (batch, output, loss) would affect training but could not see any difference.
While working on this PR, I also removed an obsolete comment and an unused variable.
Note to maintainers: If this is merged, all MetaMathQA benchmarks should be re-run.