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

[gemini] fixed the gemini checkpoint io (#3934)

Browse Source
This commit is contained in:
Frank Lee
2023-06-09 09:48:49 +08:00
committed by GitHub
parent a98e16ed07
commit bd1ab98158
3 changed files with 19 additions and 11 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
colossalai/zero/gemini/gemini_ddp.py
View File

@@ -716,7 +716,10 @@ class _StateDictSharder:
tensor_size = calculate_tensor_size(tensor)
ret_block = None
ret_block_size = 0
if self.current_block_size + tensor_size > self.max_shard_size:
# before we return the current block and create a new block,
# we need to ensure that the current block is not empty
if self.current_block_size + tensor_size > self.max_shard_size and self.current_block_size > 0:
ret_block = self.current_block
ret_block_size = self.current_block_size
self.current_block = OrderedDict()