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

[shardformer] add util functions for shardformer tests/fix sync_shared_param (#4366)

Browse Source 
* add util functions for shardformer tests & rewrite gpt2 test

* fix shared_params & embedding/merging

* fix precision
This commit is contained in:
Baizhou Zhang
2023-08-03 17:50:15 +08:00
committed by Hongxin Liu
parent 5c6f183192
commit b1feeced8e
4 changed files with 189 additions and 113 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
colossalai/booster/plugin/hybrid_parallel_plugin.py
View File

@@ -37,7 +37,8 @@ class HybridParallelModule(ModelWrapper):
self.shared_param_process_groups = []
for shared_param in self.shared_params:
if len(shared_param) > 0:
self.stage_manager.init_process_group_by_stages(list(shared_param.keys()))
self.shared_param_process_groups.append(
self.stage_manager.init_process_group_by_stages(list(shared_param.keys())))
if precision == 'fp16':
module = module.half().cuda()
elif precision == 'bf16':