github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-11-30 07:14:40 +00:00
Code Issues Packages Projects Releases Wiki Activity

[NFC] polish comments for Chunk class (#2116)

Browse Source
This commit is contained in:
Jiarui Fang
2022-12-12 15:39:31 +08:00
committed by GitHub
parent 09d69e1c25
commit e99edfcb51
6 changed files with 95 additions and 83 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
colossalai/nn/parallel/data_parallel.py
View File

@@ -299,7 +299,7 @@ class ZeroDDP(ColoDDP):
reduced = self.chunk_manager.reduce_chunk(chunk)
if reduced:
if chunk.is_gathered:
chunk.chunk_total.div_(chunk.pg_size)
chunk.cuda_global_chunk.div_(chunk.pg_size)
else:
chunk.cuda_shard.div_(chunk.pg_size)
# check overflow elements
@@ -529,7 +529,7 @@ class ZeroDDP(ColoDDP):
load(parameter_name, tensor, partial(load_fp32_parameter, parameter_slice))
if chunk.is_gathered:
chunk.chunk_total.copy_(temp_chunk)
chunk.cuda_global_chunk.copy_(temp_chunk)
elif chunk.cuda_shard is not None:
chunk.cuda_shard.copy_(temp_chunk[chunk.shard_begin:chunk.shard_end])
else: