Skip to content

Delete extra tensor objects after restoring float8 tensors#1500

Merged
ptrendx merged 13 commits into
NVIDIA:mainfrom
sudhakarsingh27:fix_memory_leak_te_2.0
Feb 28, 2025
Merged

Delete extra tensor objects after restoring float8 tensors#1500
ptrendx merged 13 commits into
NVIDIA:mainfrom
sudhakarsingh27:fix_memory_leak_te_2.0

Conversation

@sudhakarsingh27
Copy link
Copy Markdown
Member

@sudhakarsingh27 sudhakarsingh27 commented Feb 21, 2025

Description

After restoring the float8 tensors in the backward passes of LayernormMLP, LayernormLinear and Linear, the tensor objects are not needed but ctx.tensor_objects still holds the reference and hence it results in extra memory usage. This fixes it.

Fixes # (issue)

Type of change

  • Documentation change (change only to the documentation, either a fix or a new content)
  • Bug fix (non-breaking change which fixes an issue)
  • New feature (non-breaking change which adds functionality)
  • Breaking change (fix or feature that would cause existing functionality to not work as expected)
  • Infra/Build change
  • Code refactoring

Changes

Please list the changes introduced in this PR:

  • Delete extra reference to tensor_objects once they're used in the backwards of LayernormMLP, LayernormLinear and Linear modules.

Checklist:

  • I have read and followed the contributing guidelines
  • The functionality is complete
  • I have commented my code, particularly in hard-to-understand areas
  • I have made corresponding changes to the documentation
  • My changes generate no new warnings
  • I have added tests that prove my fix is effective or that my feature works
  • New and existing unit tests pass locally with my changes

@sudhakarsingh27 sudhakarsingh27 self-assigned this Feb 21, 2025
ksivaman
ksivaman reviewed Feb 21, 2025
View reviewed changes
Comment thread transformer_engine/pytorch/tensor/_internal/float8_tensor_base.py
@ptrendx ptrendx added the 2.1.0 label Feb 21, 2025
ptrendx
ptrendx reviewed Feb 21, 2025
View reviewed changes
Comment thread transformer_engine/pytorch/module/layernorm_linear.py Outdated
@sudhakarsingh27
Copy link
Copy Markdown
Member Author

sudhakarsingh27 commented Feb 24, 2025

/te-ci

@sudhakarsingh27
Copy link
Copy Markdown
Member Author

sudhakarsingh27 commented Feb 27, 2025

/te-ci pytorch

@ptrendx ptrendx mentioned this pull request Feb 27, 2025
Merged
34 tasks
sudhakarsingh27 and others added 5 commits February 27, 2025 14:24
@sudhakarsingh27
uncomment the fix
b3643cd 
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
@sudhakarsingh27
uncomment the fix
e793400 
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
@pre-commit-ci
[pre-commit.ci] auto fixes from pre-commit.com hooks
0c10c5f 
for more information, see https://pre-commit.ci
@sudhakarsingh27
fix lint
0c3aa43 
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
@sudhakarsingh27
Merge branch 'fix_memory_leak_te_2.0' of https://github.com/sudhakars…
32e2a3d 
…ingh27/TransformerEngine into fix_memory_leak_te_2.0
@sudhakarsingh27
Copy link
Copy Markdown
Member Author

sudhakarsingh27 commented Feb 27, 2025

/te-ci pytorch

@ptrendx ptrendx mentioned this pull request Feb 27, 2025
Merged
@ptrendx ptrendx merged commit d3efaeb into NVIDIA:main Feb 28, 2025
ptrendx pushed a commit that referenced this pull request Feb 28, 2025
@sudhakarsingh27 @pre-commit-ci @ptrendx
Delete extra tensor objects after restoring float8 tensors (#1500)
4f9cd42 
* delete extra tensor objects after restoring float8 tensors

Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* nit fix

Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>

* fix the leak in float8tensor and mxfloat8tensor classes

Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* uncomment the fix

Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* fix lint

Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>

---------

Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
@timmoon10 timmoon10 mentioned this pull request Mar 7, 2025
Merged
13 tasks
@timmoon10 timmoon10 mentioned this pull request Apr 15, 2025
Merged
13 tasks
@hungryGeek16 hungryGeek16 mentioned this pull request May 31, 2026
Open
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@ksivaman ksivaman ksivaman left review comments

@pggPL pggPL pggPL left review comments

@ptrendx ptrendx ptrendx approved these changes

Assignees

@sudhakarsingh27 sudhakarsingh27

Labels

2.1.0

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

4 participants

@sudhakarsingh27 @ptrendx @ksivaman @pggPL