HI, during my usage of the tutel fast_dispatch, I noticed that when gradient accumulation is enabled, the GPU memory consumption significantly exceeds that when gradient accumulation is OFF, and the consumed memory increases with the growth of gradient accumulation steps.
Upon reviewing the relevant code, I found that in the Encode/Decode implementation, the input tensor is directly placed in the ctx without invoking save_for_backward.
According to the PyTorch documentation, this practice may lead to a memory leak.
I attempted to fix this issue, and the memory consumption no longer grows with the number of gradient accumulation steps.
ghostplant
commented
2 months ago
HI, during my usage of the tutel fast_dispatch, I noticed that when
gradient accumulationis enabled, the GPU memory consumption significantly exceeds that whengradient accumulationis OFF, and the consumed memory increases with the growth ofgradient accumulationsteps.Upon reviewing the relevant code, I found that in the
Encode/Decodeimplementation, the input tensor is directly placed in thectxwithout invokingsave_for_backward. According to the PyTorch documentation, this practice may lead to a memory leak.I attempted to fix this issue, and the memory consumption no longer grows with the number of gradient accumulation steps.