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

[autoparallel] add binary elementwise metainfo for auto parallel (#2058)

Browse Source 
* [fx] metainfo class for auto parallel

* [fx] add unit test for linear metainfo

* [fx] fix bwd param for linear

* [fx] modify unit test

* [fx] modify unit test

* [fx] modify import

* [fx] modify import

* [fx] modify import

* [fx] move meta profiler to auto parallel

* [fx] add conv metainfo class

* [fx] restore profiler

* [fx] restore meta profiler

* [autoparallel] modify unit test

* [fx] modify unit test

* [autoparallel] add batchnorm metainfo class

* [autoparallel] fix batchnorm unit test function declaration

* [fx] restore profiler

* [fx] add relu metainfo class

* [fx] restore profiler

* [autoparallel] modify metainfo input

* [autoparallel] add pooling metainfo

* [autoparallel] add F.linear metainfo generator

* [autoparallel] add binary elementwise metainfo

* [fx] recover profiler

* [autoparallel] fix forward memory calculation

* [autoparallel] modify constants.py

* [autoparallel] remove redundant print
This commit is contained in:
Boyuan Yao
2022-12-04 15:18:51 +08:00
committed by GitHub
parent 4b40fbd743
commit 616da17fab
9 changed files with 164 additions and 11 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

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