KEMBAR78
Reduce the memory leak caused by `torch.jit.trace` by ydshieh · Pull Request #16691 · huggingface/transformers · GitHub
Skip to content

Conversation

@ydshieh
Copy link
Collaborator

@ydshieh ydshieh commented Apr 11, 2022

What does this PR do?

Reduce the memory leak caused by torch.jit.trace.

Without this fix, each call to _create_and_check_torchscript increases RAM usage by ~20MB.
(Even with this call, there are still memory leak by ~0.04MB)

Remark

Since our torchscript tests are slow tests, they are run only in scheduled CI (where the test jobs are organized by models), this memory leak issue is not important (memory is released after each job process is done/exit).

However, for PR like #16679, I need to make sure the modified tests will pass. When I run in a GCP VM, I got OOM issue.
So the change in this PR still have its value, I think.

More information

The method is copied from torch

https://github.com/pytorch/pytorch/blob/bcf6974c207ac0339bfb8bdfdb0b0ec348f7a22f/torch/testing/_internal/jit_utils.py#L59

This is called in torch/test, including

https://github.com/pytorch/pytorch/blob/bcf6974c207ac0339bfb8bdfdb0b0ec348f7a22f/test/test_jit.py#L12962

    # clear the class registry as we will be defining foo multiple times
    jit_utils.clear_class_registry()

@ydshieh ydshieh requested review from LysandreJik and sgugger April 11, 2022 07:42
@HuggingFaceDocBuilderDev
Copy link

HuggingFaceDocBuilderDev commented Apr 11, 2022

The documentation is not available anymore as the PR was closed or merged.

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.

Oh that's a very good catch, and may well have been one of the causes of the out-of-memory errors we were getting when running all tests sequentially.

Copy link
Collaborator

@sgugger sgugger left a comment

Choose a reason for hiding this comment

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

Thanks for fixing this!

@ydshieh ydshieh merged commit 3918d6a into huggingface:main Apr 11, 2022
@ydshieh ydshieh deleted the fix_test_torchscript_mem_leak branch April 11, 2022 16:22
elusenji pushed a commit to elusenji/transformers that referenced this pull request Jun 12, 2022
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