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

[polish] add update directory in gemini; rename AgChunk to ChunkV2 (#1432)

Browse Source
This commit is contained in:
HELSON
2022-08-10 16:40:29 +08:00
committed by GitHub
parent f20cb4e893
commit 039b7ed3bc
3 changed files with 4 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
tests/test_gemini/chunk/test_agchunk.py → tests/test_gemini/update/test_chunkv2.py
View File

@@ -9,7 +9,7 @@ from colossalai.utils import free_port, get_current_device
from colossalai.tensor import ProcessGroup as ColoProcessGroup
from colossalai.tensor import ColoParameter
from colossalai.gemini import TensorState
from colossalai.gemini.ag_chunk import AgChunk
from colossalai.gemini.update import ChunkV2
def dist_sum(x):
@@ -38,7 +38,7 @@ def check_euqal(param, param_cp):
def exam_chunk_basic(init_device, keep_gathered, pin_memory):
world_size = torch.distributed.get_world_size()
pg = ColoProcessGroup()
my_chunk = AgChunk(
my_chunk = ChunkV2(
chunk_size=1024,
process_group=pg,
dtype=torch.float32,