github
/
ColossalAI
mirror of https://github.com/hpcaitech/ColossalAI
1
0
Fork 0
Code Issues Projects Releases Wiki Activity

[test] update shardformer tests

pull/4445/head
ver217 2023-07-05 14:30:17 +08:00 committed by Hongxin Liu
parent b0b8ad2823
commit 2d6cc07feb
2 changed files with 3 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
tests/test_shardformer/test_model/_utils.py
View File

@ -12,8 +12,8 @@ def build_model(model_fn, enable_fused_normalization=True, enable_tensor_paralle
enable_tensor_parallelism=enable_tensor_parallelism) enable_tensor_parallelism=enable_tensor_parallelism)
model_copy = copy.deepcopy(org_model) model_copy = copy.deepcopy(org_model)
shard_former = ShardFormer(shard_config=shard_config) shard_former = ShardFormer(shard_config=shard_config)
sharded_model = shard_former.optimize(model_copy).cuda() sharded_model, shared_params = shard_former.optimize(model_copy)
return org_model, sharded_model return org_model, sharded_model.cuda()
def run_forward(original_model, sharded_model, data_gen_fn, output_transform_fn, loss_fn): def run_forward(original_model, sharded_model, data_gen_fn, output_transform_fn, loss_fn):

2
tests/test_shardformer/test_with_torch_ddp.py
View File

@ -44,7 +44,7 @@ def check_shardformer_with_ddp(rank, world_size, port):
for name, (model_fn, data_gen_fn, output_transform_fn, loss_fn, _) in sub_model_zoo.items(): for name, (model_fn, data_gen_fn, output_transform_fn, loss_fn, _) in sub_model_zoo.items():
# create and shard model # create and shard model
model = model_fn().cuda() model = model_fn().cuda()
sharded_model = shardformer.optimize(model) sharded_model, _ = shardformer.optimize(model)
# add ddp # add ddp
sharded_ddp_model = DDP(sharded_model, process_group=dp_process_group) sharded_ddp_model = DDP(sharded_model, process_group=dp_process_group)