github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-09-03 01:55:12 +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

5
colossalai/auto_parallel/meta_profiler/constants.py
View File

@@ -5,8 +5,11 @@ import torch.nn as nn
from ..tensor_shard.constants import *
# list of inplace operations
# list of inplace module
INPLACE_MODULE = [nn.ReLU]
# list of inplace operations
INPLACE_OPS = [torch.flatten]
# list of operations that do not save forward activations
NO_SAVE_ACTIVATION = [torch.add, torch.sub, operator.add, operator.sub]