github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-09-05 19:13:01 +00:00
Code Issues Packages Projects Releases Wiki Activity

[autoparallel] align the data_ptr with the old version of auto activation checkpoint pipeline (#2261)

Browse Source
This commit is contained in:
Boyuan Yao
2023-01-03 10:30:15 +08:00
committed by GitHub
parent 3ccf58aa76
commit 1ea99b869e
5 changed files with 51 additions and 43 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
colossalai/auto_parallel/passes/constants.py Normal file
View File

@@ -0,0 +1,8 @@
import torch
OUTPUT_SAVED_OPS = [torch.nn.functional.relu, torch.nn.functional.softmax, torch.flatten]
OUTPUT_SAVED_MOD = [
torch.nn.ReLU,
torch.nn.Softmax,
]