[Pipelining] Free memory usage earlier in last stage #138504
Conversation
[ghstack-poisoned]
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/138504
Note: Links to docs will display an error until the docs builds have been completed. ❗ 1 Active SEVsThere are 1 currently active SEVs. If your PR is affected, please view them below: ❌ 5 New FailuresAs of commit bed5bb5 with merge base deaf041 ( NEW FAILURES - The following jobs have failed:
This comment was automatically generated by Dr. CI and updates every 15 minutes. |
pytorch-bot
bot
added
the
oncall: distributed
H-Huang
added
release notes: distributed (pipeline)
| # to return to the user in merge_output_chunks, therefore | ||
| # this should be detached to release autograd graph context and free memory earlier | ||
| for t in stage_output: | ||
| t.detach_() |
There was a problem hiding this comment.
i'm confused about the difference in the 2 PRs. The first PR detaches 'stage_output' in stage_backward_input, which iiuc runs inside this function (backward_one_chunk) but above this point, in the case of non-full-backward. If this is the last stage, wouldn't we have already run stage_backward_input and detached stage_outputs by the time we get here? or are there multiple copies of 'stage_output' and we have to detach them both?
There was a problem hiding this comment.
For the last stage backward_one_chunk operates on the loss. The dependencies look like:
rest_of_autograd_graph -> stage_output -> loss
The "stage_outputs" detached in the first PR was actually the loss for the last stage (i guess it should be renamed to be stage_output_or_loss. So we also need to also detach the stage_output in the last stage to allow the memory to be freed
There was a problem hiding this comment.
gotcha. yea, rename for clarity. and maybe even put that simple illustration in to the comment rest_of_autograd_graph -> stage_output -> loss
|
@pytorchbot merge |
Merge startedYour change will be merged once all checks pass (ETA 0-4 Hours). Learn more about merging in the wiki. Questions? Feedback? Please reach out to the PyTorch DevX Team |
Merge failedReason: 1 jobs have failed, first few of them are: linux-binary-libtorch-cxx11-abi / libtorch-cpu-shared-with-deps-cxx11-abi-build / build Details for Dev Infra teamRaised by workflow job |
|
@pytorchbot merge -i |
Merge startedYour change will be merged while ignoring the following 5 checks: linux-binary-libtorch-cxx11-abi / libtorch-cpu-shared-with-deps-cxx11-abi-build / build, linux-binary-libtorch-pre-cxx11 / libtorch-cpu-shared-with-deps-pre-cxx11-build / build, linux-binary-manywheel / manywheel-py3_9-cuda12_4-build / build, linux-binary-manywheel / manywheel-py3_9-cuda11_8-build / build, linux-binary-manywheel / manywheel-py3_9-cuda12_1-build / build Learn more about merging in the wiki. Questions? Feedback? Please reach out to the PyTorch DevX Team |
Addressing the comments in previous PRs to update the variable names and add additional code comments Pull Request resolved: #138735 Approved by: https://github.com/wconstab ghstack dependencies: #138119, #138504
Pull Request resolved: #138720 Approved by: https://github.com/wconstab ghstack dependencies: #138119, #138504, #138735
3 participants
Stack from ghstack (oldest at bottom):
This fix is similar to that done in #138119, except this is an edge case for the last stage. For the last stage we perform backward on the
losswhich we detached in the previous PR. However, we also hold thestage_outputsalive because we return all the output chunks inmerge_output_chunks()after the step is over. This will also still keep the autograd graph alive, so detaching these tensors frees the memory earlier.pre-fix:
post-fix:
cc @awgu @kwen2501 @wanchaol @fegin @fduwjj @wz337 @wconstab @d4l3k @c-p-i-o