Skip to content

avoid calling gc.collect and cuda.empty_cache#34514

Merged
ydshieh merged 5 commits intomainfrom
speed_no_empty
Oct 31, 2024
Merged

avoid calling gc.collect and cuda.empty_cache#34514
ydshieh merged 5 commits intomainfrom
speed_no_empty

Conversation

@ydshieh
Copy link
Collaborator

@ydshieh ydshieh commented Oct 30, 2024

What does this PR do?

Let's avoid calling gc.collect and cuda.empty_cache while the tests are running on CPU:

  • those operations are slow
  • (actually, in most cases, they are only relevant for integration tests which use large models)

Running on GPT2 tests,

60 seconds on main, 20 seconds on this PR

@ydshieh ydshieh requested a review from Rocketknight1 October 30, 2024 16:44
Copy link
Member

@Rocketknight1 Rocketknight1 left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yes, this seems like a good speed fix! cc @LysandreJik @ArthurZucker for core maintainer review

@ydshieh ydshieh changed the title Speed no empty avoid calling gc.collect and cuda.empty_cache Oct 30, 2024
@ydshieh ydshieh changed the title avoid calling gc.collect and cuda.empty_cache avoid calling gc.collect and cuda.empty_cache Oct 30, 2024
Copy link
Member

@LysandreJik LysandreJik left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Smart! Should a helper method be made that only runs on CPU?

Both the gc.collect and the torch device checks could be moved into the backend_empty_cache method (or an other method that wraps both)

@ydshieh
Copy link
Collaborator Author

ydshieh commented Oct 31, 2024

Yes, a helper method is nice. Will update

@ydshieh
Copy link
Collaborator Author

ydshieh commented Oct 31, 2024

updated.

So far it doesn't call gc.collect() at all (default value False).
I would like to see if this would cause issue.
In general, we don't need to call it after each test method (so in tearDown) as it is slow.

@ydshieh ydshieh merged commit ab98f0b into main Oct 31, 2024
@ydshieh ydshieh deleted the speed_no_empty branch October 31, 2024 15:36
@HuggingFaceDocBuilderDev

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.

BernardZach pushed a commit to BernardZach/transformers that referenced this pull request Dec 5, 2024
* update

* update

* update

* update

* update

---------

Co-authored-by: ydshieh <ydshieh@users.noreply.github.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants