Skip to content

Avoid retaining saved tensors in fused norm custom ops under no_grad #2012

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
LeSingh1 wants to merge 1 commit into
base: master
Choose a base branch
from
+30 −12
Open

Avoid retaining saved tensors in fused norm custom ops under no_grad #2012

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
42 changes: 30 additions & 12 deletions apex/normalization/fused_layer_norm.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -182,16 +182,20 @@ def _fused_layer_norm_affine_backward(ctx, grad_output, grad_mean, grad_invvar):
def _fused_layer_norm_affine_setup_context(ctx, inputs, output):
input, weight, bias, normalized_shape, eps, memory_efficient = inputs
output, mean, invvar = output
ctx.normalized_shape = normalized_shape
ctx.eps = eps
ctx.memory_efficient = memory_efficient
# See _fused_rms_norm_affine_setup_context: skip the saves under no_grad to
# avoid leaking the retained tensors (issue #1999).
if not torch.is_grad_enabled():
return
input_ = input.contiguous()
weight_ = weight.contiguous()
bias_ = bias.contiguous()
if memory_efficient:
ctx.save_for_backward(output, weight_, bias_, None, invvar)
else:
ctx.save_for_backward(input_, weight_, bias_, mean, invvar)
ctx.normalized_shape = normalized_shape
ctx.eps = eps
ctx.memory_efficient = memory_efficient

fused_layer_norm_affine_fwd.register_autograd(
_fused_layer_norm_affine_backward,
Expand Down Expand Up @@ -337,15 +341,21 @@ def _fused_rms_norm_affine_backward(ctx, grad_output, grad_invvar):
def _fused_rms_norm_affine_setup_context(ctx, inputs, output):
input_, weight_, normalized_shape, eps, memory_efficient = inputs
output_, invvar = output
ctx.normalized_shape = normalized_shape
ctx.eps = eps
ctx.memory_efficient = memory_efficient
# Under torch.no_grad() backward will never run, so don't retain the
# activation/invvar tensors. For custom ops these saves are held in autograd
# metadata that isn't released after the call returns, which otherwise leaks
# two CUDA tensors per forward in inference loops (issue #1999).
if not torch.is_grad_enabled():
return
input_ = input_.contiguous()
weight_ = weight_.contiguous()
if memory_efficient:
ctx.save_for_backward(output_, weight_, invvar)
else:
ctx.save_for_backward(input_, weight_, invvar)
ctx.normalized_shape = normalized_shape
ctx.eps = eps
ctx.memory_efficient = memory_efficient

fused_rms_norm_affine_fwd.register_autograd(
_fused_rms_norm_affine_backward,
Expand Down Expand Up @@ -515,14 +525,18 @@ def _fused_layer_norm_backward(ctx, grad_output, grad_mean, grad_invvar):
def _fused_layer_norm_setup_context(ctx, inputs, output):
input, normalized_shape, eps, memory_efficient = inputs
output, mean, invvar = output
ctx.normalized_shape = normalized_shape
ctx.eps = eps
ctx.memory_efficient = memory_efficient
# See _fused_rms_norm_affine_setup_context: skip the saves under no_grad to
# avoid leaking the retained tensors (issue #1999).
if not torch.is_grad_enabled():
return
input_ = input.contiguous()
if memory_efficient:
ctx.save_for_backward(output, None, invvar)
else:
ctx.save_for_backward(input_, mean, invvar)
ctx.normalized_shape = normalized_shape
ctx.eps = eps
ctx.memory_efficient = memory_efficient

fused_layer_norm_fwd.register_autograd(
_fused_layer_norm_backward,
Expand Down Expand Up @@ -653,14 +667,18 @@ def _fused_rms_norm_backward(ctx, grad_output, grad_invvar):
def _fused_rms_norm_setup_context(ctx, inputs, output):
input_, normalized_shape, eps, memory_efficient = inputs
output_, invvar = output
ctx.normalized_shape = normalized_shape
ctx.eps = eps
ctx.memory_efficient = memory_efficient
# See _fused_rms_norm_affine_setup_context: skip the saves under no_grad to
# avoid leaking the retained tensors (issue #1999).
if not torch.is_grad_enabled():
return
input_ = input_.contiguous()
if memory_efficient:
ctx.save_for_backward(output_, invvar)
else:
ctx.save_for_backward(input_, invvar)
ctx.normalized_shape = normalized_shape
ctx.eps = eps
ctx.memory_efficient = memory_efficient

fused_rms_norm_fwd.register_autograd(
_fused_rms_norm_backward, setup_context=_fused_rms_norm_setup_context
Expand Down