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

[gemini] update ddp strict mode (#2518)

Browse Source 
* [zero] add strict ddp mode for chunk init

* [gemini] update gpt example
This commit is contained in:
HELSON
2023-01-28 14:35:25 +08:00
committed by GitHub
parent 0af793836c
commit 707b11d4a0
16 changed files with 133 additions and 54 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
examples/language/gpt/gemini/train_gpt_demo.py
View File

@@ -263,7 +263,7 @@ def main():
if args.distplan == "colossalai":
# all param must use the same process group.
world_size = torch.distributed.get_world_size()
shard_pg = ProcessGroup(tp_degree=world_size)
shard_pg = ProcessGroup(tp_degree=world_size) if args.shardinit else None
default_dist_spec = ShardSpec([-1], [world_size]) if args.shardinit else None
# build GPT model