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

[tensor] refine linear and add gather for laynorm (#893)

Browse Source 
* refine linear and add function to ColoTensor

* add gather for layernorm

* polish

* polish
This commit is contained in:
Ziyue Jiang
2022-04-28 10:55:40 +08:00
committed by GitHub
parent 26c49639d8
commit cb182da7c5
7 changed files with 225 additions and 123 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
colossalai/tensor/__init__.py
View File

@@ -1,4 +1,4 @@
from .spec import ComputePattern, ParallelAction, TensorSpec
from .spec import ComputePattern, ParallelAction, TensorSpec, ShardPattern
from .op_wrapper import (
colo_op_impl,)
from .colo_tensor import ColoTensor
@@ -7,5 +7,5 @@ from ._ops import *
__all__ = [
'ColoTensor', 'convert_parameter', 'colo_op_impl', 'ComputePattern', 'TensorSpec', 'ParallelAction',
'named_params_with_colotensor'
'named_params_with_colotensor', 'ShardPattern'
]