NVIDIA / NVIDIA/Megatron-LM

[ENHANCEMENT] Discard the intermediate recomputed outputs while restoring the original outputs

Open
#1,825 0 comments 0 reactions 0 assignees View on GitHub
enhancement
Dominant language
Python
Stars
17.9k
Forks
4.5k
Avg merge
4d 6h
Merged PRs (30d)
271

Description

**Is your feature request related to a problem? Please describe.**

mcore has provided a [discard_output_checkpoint](https://github.com/NVIDIA/Megatron-LM/blob/core_v0.15.0rc5/megatron/core/tensor_parallel/random.py#L521C7-L521C30). It will do the recomputation via a backward hook. However, it does store the recomputation result and then restore the storage of the original outputs , followed by an inplace copy. See https://github.com/NVIDIA/Megatron-LM/blob/core_v0.15.0rc5/megatron/core/tensor_parallel/random.py#L587-L594. Note that the intermediate result can be released right after the copy to save memory aggressively. [However](https://github.com/NVIDIA/Megatron-LM/blob/core_v0.15.0rc5/megatron/core/tensor_parallel/random.py#L593) , mcore simply sets the intermediate tensors `None`, so that the memory management is returned to Pytorch/Python.

I'm wondering if we can mannualy release those tensors.

**Describe the solution you'd like**

Release the output within [the loop](https://github.com/NVIDIA/Megatron-LM/blob/core_v0.15.0rc5/megatron/core/tensor_parallel/random.py#L588-L590)
```
output_size = recomputation_output.untyped_storage().size()
output.untyped_storage().resize_(output_size)
output.untyped_storage().copy_(recomputation_output.untyped_storage())
recomputation_output.untyped_storage.resize_(0)
```

**Proposed implementation**
If you have a proposed implementation for the feature state it here or link to a PR.

**Additional context**
Add any other context or screenshots about the feature request here.

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.