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

[shardformer] refactored layernorm (#4086)

Browse Source
This commit is contained in:
Frank Lee
2023-06-26 18:05:00 +08:00
parent c4b1b65931
commit d33a44e8c3
4 changed files with 51 additions and 77 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
colossalai/shardformer/layer/__init__.py
View File

@@ -1,11 +1,11 @@
from .dropout import Dropout1D
from .embedding import Embedding1D, VocabParallelEmbedding1D
from .layernorm import LayerNorm1D
from .layernorm import FusedLayerNorm
from .linear import Linear1D_Col, Linear1D_Row
from .linear_conv import LinearConv1D_Col, LinearConv1D_Row
from .loss import cross_entropy_1d
__all__ = [
"Embedding1D", "VocabParallelEmbedding1D", "Linear1D_Col", "Linear1D_Row", "LinearConv1D_Col", "LinearConv1D_Row",
"Dropout1D", "cross_entropy_1d", 'LayerNorm1D'
"Dropout1D", "cross_entropy_1d", 'FusedLayerNorm'
]