github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-09-16 06:30:41 +00:00
Code Issues Packages Projects Releases Wiki Activity

[shardformer] rewrite tests for opt/bloom/llama/vit/chatglm (#4395)

Browse Source 
* rewrite opt tests

* rewrite llama tests

* rewrite bloom & vit tests

* rewrite chatglm tests

* fix LinearCol for classfiers

* add judge for other tp layers, fix lazy init in util
This commit is contained in:
Baizhou Zhang
2023-08-11 15:43:23 +08:00
committed by Hongxin Liu
parent 21e0a42fd1
commit 7711bd524a
19 changed files with 1064 additions and 1273 deletions
Download Patch File Download Diff File Expand all files Collapse all files

16
colossalai/shardformer/layer/linear.py
View File

@@ -143,6 +143,14 @@ class Linear1D_Col(ParallelModule):
f'Expected only one process group, got {len(process_group)}.'
process_group = process_group[0]
tp_size = dist.get_world_size(process_group)
if out_features < tp_size:
return module
if out_features % tp_size != 0:
raise ValueError(
f"The size of out_features:{out_features} is not integer multiples of tensor parallel size: {tp_size}!")
linear_1d = Linear1D_Col(in_features=in_features,
out_features=out_features,
bias=bias,
@@ -293,6 +301,14 @@ class Linear1D_Row(ParallelModule):
f'Expected only one process group, got {len(process_group)}.'
process_group = process_group[0]
tp_size = dist.get_world_size(process_group)
if in_features < tp_size:
return module
if in_features % tp_size != 0:
raise ValueError(
f"The size of in_features:{in_features} is not integer multiples of tensor parallel size: {tp_size}!")
linear_1d = Linear1D_Row(in_features=in_features,
out_features=out_features,
bias=bias,