github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2026-01-29 21:49:54 +00:00
Code Issues Packages Projects Releases Wiki Activity

[zero] fixed api consistency (#1098)

Browse Source
This commit is contained in:
Frank Lee
2022-06-10 16:59:59 +08:00
committed by GitHub
parent cb18922c47
commit 14e5b11d7f
2 changed files with 3 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
colossalai/zero/zero_optimizer.py
View File

@@ -136,7 +136,7 @@ class ZeroOptimizer(ColossalaiOptimizer):
fp32_params_used_cuda_margin_mem = 0
for fp16_param_chunk, fp32_param_chunk in zip(self.chunk_manager.chunk_groups['fp16_param'],
self.chunk_manager.chunk_groups['fp32_param']):
if fp32_param_chunk.is_free:
if fp32_param_chunk.is_empty:
continue
if fp32_params_used_cuda_margin_mem + fp32_param_chunk.mem < fp32_params_available_cuda_margin_mem:
self.chunk_manager.move_chunk(fp32_param_chunk, get_current_device())