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

[shardformer] refactored some doc and api (#4137)

Browse Source 
* [shardformer] refactored some doc and api

* polish code
This commit is contained in:
Frank Lee
2023-07-03 15:29:11 +08:00
parent 7f9b30335b
commit 74257cb446
15 changed files with 355 additions and 490 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -11,7 +11,7 @@ def build_model(model_fn):
shard_config = ShardConfig(enable_fused_normalization=True)
model_copy = copy.deepcopy(org_model)
shard_former = ShardFormer(shard_config=shard_config)
sharded_model = shard_former.shard_model(model_copy).cuda()
sharded_model = shard_former.optimize(model_copy).cuda()
return org_model, sharded_model