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

[gemini] gemini support extra-dp (#5043)

Browse Source 
* support ddp

* fix

* fix

* fix

fix

* support ddp

* fix

* fix

* fix

fix

* simplify tests

* fix

* fix

* fix

fix

fix

* fix
This commit is contained in:
flybird11111
2023-11-16 21:03:04 +08:00
committed by GitHub
parent b2ad0d9e8f
commit 3e02154710
10 changed files with 96 additions and 137 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
tests/test_zero/test_gemini/test_chunkv2.py
View File

@@ -39,7 +39,7 @@ def exam_chunk_basic(init_device, keep_gathered, pin_memory):
pg = _get_default_group()
my_chunk = Chunk(
chunk_size=1024,
process_group=pg,
zero_group=pg,
dtype=torch.float32,
init_device=init_device,
cpu_shard_init=True,