github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-09-20 00:55:29 +00:00
Code Issues Packages Projects Releases Wiki Activity

[util] fixed activation checkpointing on torch 1.9 (#719)

Browse Source
This commit is contained in:
Frank Lee
2022-04-12 09:35:45 +08:00
committed by GitHub
parent 04ff5ea546
commit 2412429d54
1 changed files with 9 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
colossalai/utils/activation_checkpoint.py
View File

@@ -68,6 +68,9 @@ class CheckpointFunction(torch.autograd.Function):
else: else:
ctx.inputs.append(arg) ctx.inputs.append(arg)
if activation_offload:
ctx.tensor_inputs = tensor_inputs
else:
ctx.save_for_backward(*tensor_inputs) ctx.save_for_backward(*tensor_inputs)
return outputs return outputs
@@ -79,6 +82,10 @@ class CheckpointFunction(torch.autograd.Function):
# Copy the list to avoid modifying original list. # Copy the list to avoid modifying original list.
inputs = list(ctx.inputs) inputs = list(ctx.inputs)
tensor_indices = ctx.tensor_indices tensor_indices = ctx.tensor_indices
if ctx.activation_offload:
tensors = ctx.tensor_inputs
else:
tensors = ctx.saved_tensors tensors = ctx.saved_tensors
# store the current states # store the current states