From 714c1adb1a9ed5b6e46312bc6e67c52675cc0be9 Mon Sep 17 00:00:00 2001 From: Baizhou Zhang Date: Thu, 10 Aug 2023 14:52:55 +0800 Subject: [PATCH] [gemini] fix tensor storage cleaning in state dict collection --- colossalai/zero/gemini/gemini_optimizer.py | 6 ------ 1 file changed, 6 deletions(-) diff --git a/colossalai/zero/gemini/gemini_optimizer.py b/colossalai/zero/gemini/gemini_optimizer.py index 7d0db6b1fa23..a2085323f83e 100644 --- a/colossalai/zero/gemini/gemini_optimizer.py +++ b/colossalai/zero/gemini/gemini_optimizer.py @@ -1,6 +1,5 @@ # this code is inspired by the DeepSpeed library and implemented with our own design from scratch import copy -import gc import math import warnings from typing import Any, Dict, Iterator, OrderedDict, Set, Tuple @@ -468,11 +467,6 @@ def collect_states(self, param_id: int, only_rank_0: bool = True) -> dict: self.load_from_compacted_states(compacted_states, collected_states, state_names, shard_offset, shard_size) - # Clean gathered states - for state_shard in gathered_state_shards: - del state_shard[0] - gc.collect() - # Reshape tensors if is_collector: for state_name, state_tensor in collected_states.items():