diff --git a/colossalai/gemini/gemini_mgr.py b/colossalai/gemini/gemini_mgr.py index 068af1257..6b5d23252 100644 --- a/colossalai/gemini/gemini_mgr.py +++ b/colossalai/gemini/gemini_mgr.py @@ -73,11 +73,11 @@ class GeminiManager: def _get_layout_info(self, compute_idx: int, warmup: bool, chunks: Tuple[Chunk, ...], group_name: str): cuda_demand = 0 for chunk in chunks: - if chunk.device_type == 'cpu' or chunk.is_free: + if chunk.device_type == 'cpu' or chunk.is_empty: cuda_demand += chunk.mem can_evict_chunks = [] for chunk in self._chunk_manager.chunk_groups[group_name]: - if not chunk.is_free and chunk.device_type == 'cuda' and chunk.can_move_device: + if not chunk.is_empty and chunk.device_type == 'cuda' and chunk.can_move_device: can_evict_chunks.append(chunk) return cuda_demand, can_evict_chunks diff --git a/colossalai/zero/zero_optimizer.py b/colossalai/zero/zero_optimizer.py index b9208bace..d263da59d 100644 --- a/colossalai/zero/zero_optimizer.py +++ b/colossalai/zero/zero_optimizer.py @@ -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())