We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
[tensor] fix some unittests

[tensor] fix some unittests

[tensor] fix some unittests


--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `colossalai/zero/gemini/placement_policy.py`
Content:
```
1 import functools
2 import warnings
3 from abc import ABC, abstractmethod
4 from time import time
5 from typing import Dict, List, Optional, Tuple, Type
6 
7 import torch
8 
9 from colossalai.accelerator import get_accelerator
10 from colossalai.legacy.utils.memory import colo_device_memory_capacity
11 from colossalai.zero.gemini.chunk import Chunk
12 
13 from .chunk import Chunk, ChunkManager
14 from .memory_tracer import ChunkMemStatsCollector
15 
16 
17 class PlacementPolicy(ABC):
18     need_mem_stats: bool = False
19 
20     def __init__(
21         self,
22         gemini_manager: "GeminiManager",
23         chunk_manager: ChunkManager,
24         mem_stats_collector: Optional[ChunkMemStatsCollector] = None,
25         max_prefetch: int = 0,
26         **kwargs,
27     ) -> None:
28         self.gemini_manager = gemini_manager
29         self.chunk_manager = chunk_manager
30         self.mem_stats_collector: Optional[ChunkMemStatsCollector] = mem_stats_collector
31         self.max_prefetch = max_prefetch
32 
33     @abstractmethod
34     def evict_tensors(self, can_evict_chunks: List[Chunk], **kwargs) -> Tuple[int, float]:
35         raise NotImplementedError
36 
37     @abstractmethod
38     def setup_grads_device(
39         self, params: List[torch.Tensor], grads_device_map: Dict[torch.Tensor, torch.device]
40     ) -> None:
41         raise NotImplementedError
42 
43     @abstractmethod
44     def get_prefetch_chunks(self) -> List[Chunk]:
45         raise NotImplementedError
46 
47 
48 class StaticPlacementPolicy(PlacementPolicy):
49     def __init__(
50         self,
51         gemini_manager: "GeminiManager",
52         chunk_manager: ChunkManager,
53         mem_stats_collector: Optional[ChunkMemStatsCollector] = None,
54         max_prefetch: int = 0,
55         shard_param_frac: float = 1.0,
56         offload_optim_frac: float = 0.0,
57         offload_param_frac: float = 0.0,
58         **kwargs,
59     ) -> None:
60         super().__init__(
61             gemini_manager, chunk_manager, mem_stats_collector=mem_stats_collector, max_prefetch=max_prefetch
62         )
63         if offload_param_frac > 0.0 and (shard_param_frac != 1.0 or offload_optim_frac != 1.0):
64             warnings.warn("offload_param_frac is ignored when shard_param_frac != 1.0 or offload_optim_frac != 1.0")
65             offload_param_frac = 0.0
66         self.shard_param_frac = shard_param_frac
67         self.offload_optim_frac = offload_optim_frac
68         self.offload_param_frac = offload_param_frac
69         # these should be initialized in setup_grads_device
70         self.keep_gathered_chunk_mem = 0.0
71         self.keep_cuda_chunk_mem = 0.0
72 
73     def evict_tensors(self, can_evict_chunks: List[Chunk], **kwargs) -> Tuple[int, float]:
74         can_shard_chunk_mem = sum(chunk.chunk_mem for chunk in can_evict_chunks)
75         can_offload_chunk_mem = can_shard_chunk_mem
76         for chunk in can_evict_chunks:
77             if can_shard_chunk_mem <= self.keep_gathered_chunk_mem:
78                 break
79             self.chunk_manager.release_chunk(chunk)
80             # real saved mem is chunk_mem - shard_mem, for simplicity we use chunk_mem
81             can_shard_chunk_mem -= chunk.chunk_mem
82         for chunk in can_evict_chunks:
83             if can_offload_chunk_mem <= self.keep_cuda_chunk_mem:
84                 break
85             self.chunk_manager.move_chunk(chunk, torch.device("cpu"))
86             # real saved mem is shard_mem, for simplicity we use chunk_mem
87             can_offload_chunk_mem -= chunk.chunk_mem
88         return 0, 0.0
89 
90     def setup_grads_device(
91         self, params: List[torch.Tensor], grads_device_map: Dict[torch.Tensor, torch.device]
92     ) -> None:
93         total_chunk_mem = sum(self.chunk_manager.get_chunk(p).chunk_mem for p in params)
94 
95         offload_optim_chunk_mem = total_chunk_mem * self.offload_optim_frac
96         offloaded_optim_chunk_mem = 0
97         chunks = set(self.chunk_manager.get_chunk(p) for p in params)
98         for chunk in chunks:
99             params = chunk.get_tensors()
100             # init offload optim settings
101             # keep gathered chunks are in CUDA
102             if chunk.keep_gathered or offloaded_optim_chunk_mem >= offload_optim_chunk_mem:
103                 device = get_accelerator().get_current_device()
104             else:
105                 device = torch.device("cpu")
106                 # real offloaded mem is chunk.shard_mem, for simplicity we use chunk mem here
107                 offloaded_optim_chunk_mem += chunk.chunk_mem
108             for p in params:
109                 grads_device_map[p] = device
110         self.keep_gathered_chunk_mem = total_chunk_mem * (1 - self.shard_param_frac)
111         self.keep_cuda_chunk_mem = total_chunk_mem * (1 - self.offload_param_frac)
112 
113     def get_prefetch_chunks(self) -> List[Chunk]:
114         if self.gemini_manager.is_warmup():  # no prefetch during warmup since we need compute_list
115             return []
116         can_prefetch = self.max_prefetch - len(self.gemini_manager._async_works)
117         prefetch = []
118         for i in range(self.gemini_manager.compute_idx + 1, len(self.gemini_manager.compute_list)):
119             for chunk in self.gemini_manager.compute_list[i]:
120                 if len(prefetch) >= can_prefetch:
121                     break
122                 if chunk not in prefetch and chunk not in self.chunk_manager.accessed_chunks:
123                     prefetch.append(chunk)
124             if len(prefetch) >= can_prefetch:
125                 break
126         return prefetch
127 
128 
129 class AutoPlacementPolicy(PlacementPolicy):
130     need_mem_stats: bool = True
131 
132     def __init__(
133         self,
134         gemini_manager: "GeminiManager",
135         chunk_manager: ChunkManager,
136         mem_stats_collector: Optional[ChunkMemStatsCollector] = None,
137         max_prefetch: int = 0,
138         warmup_non_model_data_ratio: float = 0.8,
139         steady_cuda_cap_ratio: float = 0.9,
140         **kwargs,
141     ) -> None:
142         super().__init__(
143             gemini_manager, chunk_manager, mem_stats_collector=mem_stats_collector, max_prefetch=max_prefetch
144         )
145         # model data will use 1-_warmup_non_model_data_ratio CUDA memory in warmup phase
146         # you can set them by AutoPlacementPolicy.set_warmup_non_model_data_ratio()
147         # and AutoPlacementPolicy.set_steady_cuda_cap_ratio()
148         self._warmup_non_model_data_ratio = warmup_non_model_data_ratio
149         self._steady_cuda_cap_ratio = steady_cuda_cap_ratio
150 
151     def evict_tensors(
152         self,
153         can_evict_chunks: List[Chunk],
154         cuda_demand: int = 0,
155         warmup: bool = True,
156         compute_list: Optional[List[Tuple[Chunk, ...]]] = None,
157         compute_idx: int = 0,
158         **kwargs,
159     ) -> Tuple[int, float]:
160         """
161         Evict tensors from CUDA device.
162 
163         Args:
164             can_evict_chunks (List[StatefulTensor]): the list of tensors that can be evicted.
165             cuda_demand (int, optional): the volume of data needed on cuda device. Defaults to 0.
166             warmup (bool, optional): a flag indicates whether in the phase of warmup. Defaults to True.
167             compute_list (List[StatefulTensor], optional): TODO. Defaults to [].
168             compute_idx (int, optional): the idx of computing device. Defaults to 0.
169 
170         Raises:
171             RuntimeError:
172 
173         Returns:
174             int: the volume of memory that is evicted
175         """
176         start = time()
177         cuda_capacity = colo_device_memory_capacity(get_accelerator().get_current_device())
178         used_cuda_model_data = self.chunk_manager.total_mem["cuda"]
179         if warmup:
180             # We designate a part of CUDA memory for model data in warmup iterations.
181             max_cuda_non_model_data_per_period = cuda_capacity * self._warmup_non_model_data_ratio
182         else:
183             # max non-model-data cuda memory consumption of this sampling moment and the next sampling moment.
184             max_cuda_non_model_data_per_period = self.mem_stats_collector.next_period_non_model_data_usage("cuda")
185             cuda_capacity *= self._steady_cuda_cap_ratio
186         total_cuda_model_data = cuda_capacity - max_cuda_non_model_data_per_period
187         avail_cuda_model_data = total_cuda_model_data - used_cuda_model_data
188         freed_cuda_model_data = 0
189 
190         if avail_cuda_model_data < cuda_demand:
191             # Move cuda_demand - avail_cuda_model_data volume of tensors
192             # to_free_cuda_model_data = cuda_demand - avail_cuda_model_data
193             to_free_cuda_model_data = cuda_demand - avail_cuda_model_data
194             to_free_chunks = can_evict_chunks
195             if not warmup:
196                 to_free_chunks = self._sort_can_evict_chunks(tuple(to_free_chunks), compute_idx, tuple(compute_list))
197                 # print(self._sort_can_evict_chunks.cache_info())
198             for chunk in to_free_chunks:
199                 if freed_cuda_model_data >= to_free_cuda_model_data:
200                     break
201 
202                 self.chunk_manager.release_chunk(chunk)
203                 self.chunk_manager.move_chunk(chunk, torch.device("cpu"))
204                 freed_cuda_model_data += chunk.chunk_mem
205             if freed_cuda_model_data < to_free_cuda_model_data:
206                 raise RuntimeError(
207                     f"Adjust layout failed! No enough CUDA memory! "
208                     f"Need {to_free_cuda_model_data}, freed {freed_cuda_model_data}"
209                 )
210         return freed_cuda_model_data, time() - start
211 
212     @staticmethod
213     @functools.lru_cache(maxsize=None)
214     def _sort_can_evict_chunks(can_evict_chunks: tuple, compute_idx: int, compute_list: tuple) -> list:
215         next_compute_idx = {chunk: len(compute_list) for chunk in can_evict_chunks}
216         for i in range(len(compute_list) - 1, compute_idx, -1):
217             for chunk in compute_list[i]:
218                 if chunk in next_compute_idx:
219                     next_compute_idx[chunk] = i
220         next_compute_idx = sorted(next_compute_idx.items(), key=lambda pair: pair[1], reverse=True)
221         return [t for (t, idx) in next_compute_idx]
222 
223     def setup_grads_device(
224         self, params: List[torch.Tensor], grads_device_map: Dict[torch.Tensor, torch.device]
225     ) -> None:
226         for p in params:
227             chunk = self.chunk_manager.get_chunk(p)
228             # init offload optim settings
229             # keep gathered chunks are in CUDA
230             if chunk.keep_gathered:
231                 grads_device_map[p] = get_accelerator().get_current_device()
232             else:
233                 grads_device_map[p] = torch.device("cpu")
234 
235     def get_prefetch_chunks(self, max_prefetch: int) -> List[Chunk]:
236         return []  # TODO @botbw: implement prefetching for auto
237 
238 
239 class PlacementPolicyFactory:
240     policies: Dict[str, Type[PlacementPolicy]] = {
241         "auto": AutoPlacementPolicy,
242         "static": StaticPlacementPolicy,
243     }
244 
245     @staticmethod
246     def create(policy_name: str) -> Type[PlacementPolicy]:
247         if policy_name not in PlacementPolicyFactory.policies:
248             raise TypeError(f"Unknown tensor placement policy {policy_name}")
249         return PlacementPolicyFactory.policies[policy_name]
250 
251     @staticmethod
252     def get_policy_names():
253         return tuple(PlacementPolicyFactory.policies.keys())
254
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/colossalai/zero/gemini/placement_policy.py b/colossalai/zero/gemini/placement_policy.py
--- a/colossalai/zero/gemini/placement_policy.py
+++ b/colossalai/zero/gemini/placement_policy.py
@@ -19,7 +19,7 @@
 
     def __init__(
         self,
-        gemini_manager: "GeminiManager",
+        gemini_manager: "GeminiManager",  # TODO @botbw: solve circular import
         chunk_manager: ChunkManager,
         mem_stats_collector: Optional[ChunkMemStatsCollector] = None,
         max_prefetch: int = 0,
@@ -40,9 +40,8 @@
     ) -> None:
         raise NotImplementedError
 
-    @abstractmethod
     def get_prefetch_chunks(self) -> List[Chunk]:
-        raise NotImplementedError
+        return []  # no prefetch by default
 
 
 class StaticPlacementPolicy(PlacementPolicy):
@@ -116,12 +115,14 @@
         can_prefetch = self.max_prefetch - len(self.gemini_manager._async_works)
         prefetch = []
         for i in range(self.gemini_manager.compute_idx + 1, len(self.gemini_manager.compute_list)):
+            break_flag = False
             for chunk in self.gemini_manager.compute_list[i]:
                 if len(prefetch) >= can_prefetch:
+                    break_flag = True
                     break
                 if chunk not in prefetch and chunk not in self.chunk_manager.accessed_chunks:
                     prefetch.append(chunk)
-            if len(prefetch) >= can_prefetch:
+            if break_flag:
                 break
         return prefetch
 
@@ -232,8 +233,33 @@
             else:
                 grads_device_map[p] = torch.device("cpu")
 
-    def get_prefetch_chunks(self, max_prefetch: int) -> List[Chunk]:
-        return []  # TODO @botbw: implement prefetching for auto
+    def get_prefetch_chunks(self) -> List[Chunk]:
+        if self.gemini_manager.is_warmup():  # no prefetch during warmup since we need compute_list
+            return []
+        # modified from self.evict_tensors
+        cuda_capacity = self._steady_cuda_cap_ratio * colo_device_memory_capacity(
+            get_accelerator().get_current_device()
+        )
+        max_cuda_non_model_data_per_period = self.mem_stats_collector.next_period_non_model_data_usage("cuda")
+        used_cuda_model_data = self.chunk_manager.total_mem["cuda"]
+        total_cuda_model_data = cuda_capacity - max_cuda_non_model_data_per_period
+        avail_cuda_model_data = total_cuda_model_data - used_cuda_model_data
+
+        prefetch_chunk_memory = 0
+        can_prefetch = self.max_prefetch - len(self.gemini_manager._async_works)
+        prefetch = []
+        for i in range(self.gemini_manager.compute_idx + 1, len(self.gemini_manager.compute_list)):
+            break_flag = False
+            for chunk in self.gemini_manager.compute_list[i]:
+                chunk: Chunk
+                if len(prefetch) >= can_prefetch or prefetch_chunk_memory + chunk.chunk_mem > avail_cuda_model_data:
+                    break_flag = True
+                    break
+                if chunk not in prefetch and chunk not in self.chunk_manager.accessed_chunks:
+                    prefetch.append(chunk)
+            if break_flag:
+                break
+        return prefetch
 
 
 class PlacementPolicyFactory:

{"golden_diff": "diff --git a/colossalai/zero/gemini/placement_policy.py b/colossalai/zero/gemini/placement_policy.py\n--- a/colossalai/zero/gemini/placement_policy.py\n+++ b/colossalai/zero/gemini/placement_policy.py\n@@ -19,7 +19,7 @@\n \n     def __init__(\n         self,\n-        gemini_manager: \"GeminiManager\",\n+        gemini_manager: \"GeminiManager\",  # TODO @botbw: solve circular import\n         chunk_manager: ChunkManager,\n         mem_stats_collector: Optional[ChunkMemStatsCollector] = None,\n         max_prefetch: int = 0,\n@@ -40,9 +40,8 @@\n     ) -> None:\n         raise NotImplementedError\n \n-    @abstractmethod\n     def get_prefetch_chunks(self) -> List[Chunk]:\n-        raise NotImplementedError\n+        return []  # no prefetch by default\n \n \n class StaticPlacementPolicy(PlacementPolicy):\n@@ -116,12 +115,14 @@\n         can_prefetch = self.max_prefetch - len(self.gemini_manager._async_works)\n         prefetch = []\n         for i in range(self.gemini_manager.compute_idx + 1, len(self.gemini_manager.compute_list)):\n+            break_flag = False\n             for chunk in self.gemini_manager.compute_list[i]:\n                 if len(prefetch) >= can_prefetch:\n+                    break_flag = True\n                     break\n                 if chunk not in prefetch and chunk not in self.chunk_manager.accessed_chunks:\n                     prefetch.append(chunk)\n-            if len(prefetch) >= can_prefetch:\n+            if break_flag:\n                 break\n         return prefetch\n \n@@ -232,8 +233,33 @@\n             else:\n                 grads_device_map[p] = torch.device(\"cpu\")\n \n-    def get_prefetch_chunks(self, max_prefetch: int) -> List[Chunk]:\n-        return []  # TODO @botbw: implement prefetching for auto\n+    def get_prefetch_chunks(self) -> List[Chunk]:\n+        if self.gemini_manager.is_warmup():  # no prefetch during warmup since we need compute_list\n+            return []\n+        # modified from self.evict_tensors\n+        cuda_capacity = self._steady_cuda_cap_ratio * colo_device_memory_capacity(\n+            get_accelerator().get_current_device()\n+        )\n+        max_cuda_non_model_data_per_period = self.mem_stats_collector.next_period_non_model_data_usage(\"cuda\")\n+        used_cuda_model_data = self.chunk_manager.total_mem[\"cuda\"]\n+        total_cuda_model_data = cuda_capacity - max_cuda_non_model_data_per_period\n+        avail_cuda_model_data = total_cuda_model_data - used_cuda_model_data\n+\n+        prefetch_chunk_memory = 0\n+        can_prefetch = self.max_prefetch - len(self.gemini_manager._async_works)\n+        prefetch = []\n+        for i in range(self.gemini_manager.compute_idx + 1, len(self.gemini_manager.compute_list)):\n+            break_flag = False\n+            for chunk in self.gemini_manager.compute_list[i]:\n+                chunk: Chunk\n+                if len(prefetch) >= can_prefetch or prefetch_chunk_memory + chunk.chunk_mem > avail_cuda_model_data:\n+                    break_flag = True\n+                    break\n+                if chunk not in prefetch and chunk not in self.chunk_manager.accessed_chunks:\n+                    prefetch.append(chunk)\n+            if break_flag:\n+                break\n+        return prefetch\n \n \n class PlacementPolicyFactory:\n", "issue": "[tensor] fix some unittests\n\n[tensor] fix some unittests\n\n[tensor] fix some unittests\n\n", "before_files": [{"content": "import functools\nimport warnings\nfrom abc import ABC, abstractmethod\nfrom time import time\nfrom typing import Dict, List, Optional, Tuple, Type\n\nimport torch\n\nfrom colossalai.accelerator import get_accelerator\nfrom colossalai.legacy.utils.memory import colo_device_memory_capacity\nfrom colossalai.zero.gemini.chunk import Chunk\n\nfrom .chunk import Chunk, ChunkManager\nfrom .memory_tracer import ChunkMemStatsCollector\n\n\nclass PlacementPolicy(ABC):\n    need_mem_stats: bool = False\n\n    def __init__(\n        self,\n        gemini_manager: \"GeminiManager\",\n        chunk_manager: ChunkManager,\n        mem_stats_collector: Optional[ChunkMemStatsCollector] = None,\n        max_prefetch: int = 0,\n        **kwargs,\n    ) -> None:\n        self.gemini_manager = gemini_manager\n        self.chunk_manager = chunk_manager\n        self.mem_stats_collector: Optional[ChunkMemStatsCollector] = mem_stats_collector\n        self.max_prefetch = max_prefetch\n\n    @abstractmethod\n    def evict_tensors(self, can_evict_chunks: List[Chunk], **kwargs) -> Tuple[int, float]:\n        raise NotImplementedError\n\n    @abstractmethod\n    def setup_grads_device(\n        self, params: List[torch.Tensor], grads_device_map: Dict[torch.Tensor, torch.device]\n    ) -> None:\n        raise NotImplementedError\n\n    @abstractmethod\n    def get_prefetch_chunks(self) -> List[Chunk]:\n        raise NotImplementedError\n\n\nclass StaticPlacementPolicy(PlacementPolicy):\n    def __init__(\n        self,\n        gemini_manager: \"GeminiManager\",\n        chunk_manager: ChunkManager,\n        mem_stats_collector: Optional[ChunkMemStatsCollector] = None,\n        max_prefetch: int = 0,\n        shard_param_frac: float = 1.0,\n        offload_optim_frac: float = 0.0,\n        offload_param_frac: float = 0.0,\n        **kwargs,\n    ) -> None:\n        super().__init__(\n            gemini_manager, chunk_manager, mem_stats_collector=mem_stats_collector, max_prefetch=max_prefetch\n        )\n        if offload_param_frac > 0.0 and (shard_param_frac != 1.0 or offload_optim_frac != 1.0):\n            warnings.warn(\"offload_param_frac is ignored when shard_param_frac != 1.0 or offload_optim_frac != 1.0\")\n            offload_param_frac = 0.0\n        self.shard_param_frac = shard_param_frac\n        self.offload_optim_frac = offload_optim_frac\n        self.offload_param_frac = offload_param_frac\n        # these should be initialized in setup_grads_device\n        self.keep_gathered_chunk_mem = 0.0\n        self.keep_cuda_chunk_mem = 0.0\n\n    def evict_tensors(self, can_evict_chunks: List[Chunk], **kwargs) -> Tuple[int, float]:\n        can_shard_chunk_mem = sum(chunk.chunk_mem for chunk in can_evict_chunks)\n        can_offload_chunk_mem = can_shard_chunk_mem\n        for chunk in can_evict_chunks:\n            if can_shard_chunk_mem <= self.keep_gathered_chunk_mem:\n                break\n            self.chunk_manager.release_chunk(chunk)\n            # real saved mem is chunk_mem - shard_mem, for simplicity we use chunk_mem\n            can_shard_chunk_mem -= chunk.chunk_mem\n        for chunk in can_evict_chunks:\n            if can_offload_chunk_mem <= self.keep_cuda_chunk_mem:\n                break\n            self.chunk_manager.move_chunk(chunk, torch.device(\"cpu\"))\n            # real saved mem is shard_mem, for simplicity we use chunk_mem\n            can_offload_chunk_mem -= chunk.chunk_mem\n        return 0, 0.0\n\n    def setup_grads_device(\n        self, params: List[torch.Tensor], grads_device_map: Dict[torch.Tensor, torch.device]\n    ) -> None:\n        total_chunk_mem = sum(self.chunk_manager.get_chunk(p).chunk_mem for p in params)\n\n        offload_optim_chunk_mem = total_chunk_mem * self.offload_optim_frac\n        offloaded_optim_chunk_mem = 0\n        chunks = set(self.chunk_manager.get_chunk(p) for p in params)\n        for chunk in chunks:\n            params = chunk.get_tensors()\n            # init offload optim settings\n            # keep gathered chunks are in CUDA\n            if chunk.keep_gathered or offloaded_optim_chunk_mem >= offload_optim_chunk_mem:\n                device = get_accelerator().get_current_device()\n            else:\n                device = torch.device(\"cpu\")\n                # real offloaded mem is chunk.shard_mem, for simplicity we use chunk mem here\n                offloaded_optim_chunk_mem += chunk.chunk_mem\n            for p in params:\n                grads_device_map[p] = device\n        self.keep_gathered_chunk_mem = total_chunk_mem * (1 - self.shard_param_frac)\n        self.keep_cuda_chunk_mem = total_chunk_mem * (1 - self.offload_param_frac)\n\n    def get_prefetch_chunks(self) -> List[Chunk]:\n        if self.gemini_manager.is_warmup():  # no prefetch during warmup since we need compute_list\n            return []\n        can_prefetch = self.max_prefetch - len(self.gemini_manager._async_works)\n        prefetch = []\n        for i in range(self.gemini_manager.compute_idx + 1, len(self.gemini_manager.compute_list)):\n            for chunk in self.gemini_manager.compute_list[i]:\n                if len(prefetch) >= can_prefetch:\n                    break\n                if chunk not in prefetch and chunk not in self.chunk_manager.accessed_chunks:\n                    prefetch.append(chunk)\n            if len(prefetch) >= can_prefetch:\n                break\n        return prefetch\n\n\nclass AutoPlacementPolicy(PlacementPolicy):\n    need_mem_stats: bool = True\n\n    def __init__(\n        self,\n        gemini_manager: \"GeminiManager\",\n        chunk_manager: ChunkManager,\n        mem_stats_collector: Optional[ChunkMemStatsCollector] = None,\n        max_prefetch: int = 0,\n        warmup_non_model_data_ratio: float = 0.8,\n        steady_cuda_cap_ratio: float = 0.9,\n        **kwargs,\n    ) -> None:\n        super().__init__(\n            gemini_manager, chunk_manager, mem_stats_collector=mem_stats_collector, max_prefetch=max_prefetch\n        )\n        # model data will use 1-_warmup_non_model_data_ratio CUDA memory in warmup phase\n        # you can set them by AutoPlacementPolicy.set_warmup_non_model_data_ratio()\n        # and AutoPlacementPolicy.set_steady_cuda_cap_ratio()\n        self._warmup_non_model_data_ratio = warmup_non_model_data_ratio\n        self._steady_cuda_cap_ratio = steady_cuda_cap_ratio\n\n    def evict_tensors(\n        self,\n        can_evict_chunks: List[Chunk],\n        cuda_demand: int = 0,\n        warmup: bool = True,\n        compute_list: Optional[List[Tuple[Chunk, ...]]] = None,\n        compute_idx: int = 0,\n        **kwargs,\n    ) -> Tuple[int, float]:\n        \"\"\"\n        Evict tensors from CUDA device.\n\n        Args:\n            can_evict_chunks (List[StatefulTensor]): the list of tensors that can be evicted.\n            cuda_demand (int, optional): the volume of data needed on cuda device. Defaults to 0.\n            warmup (bool, optional): a flag indicates whether in the phase of warmup. Defaults to True.\n            compute_list (List[StatefulTensor], optional): TODO. Defaults to [].\n            compute_idx (int, optional): the idx of computing device. Defaults to 0.\n\n        Raises:\n            RuntimeError:\n\n        Returns:\n            int: the volume of memory that is evicted\n        \"\"\"\n        start = time()\n        cuda_capacity = colo_device_memory_capacity(get_accelerator().get_current_device())\n        used_cuda_model_data = self.chunk_manager.total_mem[\"cuda\"]\n        if warmup:\n            # We designate a part of CUDA memory for model data in warmup iterations.\n            max_cuda_non_model_data_per_period = cuda_capacity * self._warmup_non_model_data_ratio\n        else:\n            # max non-model-data cuda memory consumption of this sampling moment and the next sampling moment.\n            max_cuda_non_model_data_per_period = self.mem_stats_collector.next_period_non_model_data_usage(\"cuda\")\n            cuda_capacity *= self._steady_cuda_cap_ratio\n        total_cuda_model_data = cuda_capacity - max_cuda_non_model_data_per_period\n        avail_cuda_model_data = total_cuda_model_data - used_cuda_model_data\n        freed_cuda_model_data = 0\n\n        if avail_cuda_model_data < cuda_demand:\n            # Move cuda_demand - avail_cuda_model_data volume of tensors\n            # to_free_cuda_model_data = cuda_demand - avail_cuda_model_data\n            to_free_cuda_model_data = cuda_demand - avail_cuda_model_data\n            to_free_chunks = can_evict_chunks\n            if not warmup:\n                to_free_chunks = self._sort_can_evict_chunks(tuple(to_free_chunks), compute_idx, tuple(compute_list))\n                # print(self._sort_can_evict_chunks.cache_info())\n            for chunk in to_free_chunks:\n                if freed_cuda_model_data >= to_free_cuda_model_data:\n                    break\n\n                self.chunk_manager.release_chunk(chunk)\n                self.chunk_manager.move_chunk(chunk, torch.device(\"cpu\"))\n                freed_cuda_model_data += chunk.chunk_mem\n            if freed_cuda_model_data < to_free_cuda_model_data:\n                raise RuntimeError(\n                    f\"Adjust layout failed! No enough CUDA memory! \"\n                    f\"Need {to_free_cuda_model_data}, freed {freed_cuda_model_data}\"\n                )\n        return freed_cuda_model_data, time() - start\n\n    @staticmethod\n    @functools.lru_cache(maxsize=None)\n    def _sort_can_evict_chunks(can_evict_chunks: tuple, compute_idx: int, compute_list: tuple) -> list:\n        next_compute_idx = {chunk: len(compute_list) for chunk in can_evict_chunks}\n        for i in range(len(compute_list) - 1, compute_idx, -1):\n            for chunk in compute_list[i]:\n                if chunk in next_compute_idx:\n                    next_compute_idx[chunk] = i\n        next_compute_idx = sorted(next_compute_idx.items(), key=lambda pair: pair[1], reverse=True)\n        return [t for (t, idx) in next_compute_idx]\n\n    def setup_grads_device(\n        self, params: List[torch.Tensor], grads_device_map: Dict[torch.Tensor, torch.device]\n    ) -> None:\n        for p in params:\n            chunk = self.chunk_manager.get_chunk(p)\n            # init offload optim settings\n            # keep gathered chunks are in CUDA\n            if chunk.keep_gathered:\n                grads_device_map[p] = get_accelerator().get_current_device()\n            else:\n                grads_device_map[p] = torch.device(\"cpu\")\n\n    def get_prefetch_chunks(self, max_prefetch: int) -> List[Chunk]:\n        return []  # TODO @botbw: implement prefetching for auto\n\n\nclass PlacementPolicyFactory:\n    policies: Dict[str, Type[PlacementPolicy]] = {\n        \"auto\": AutoPlacementPolicy,\n        \"static\": StaticPlacementPolicy,\n    }\n\n    @staticmethod\n    def create(policy_name: str) -> Type[PlacementPolicy]:\n        if policy_name not in PlacementPolicyFactory.policies:\n            raise TypeError(f\"Unknown tensor placement policy {policy_name}\")\n        return PlacementPolicyFactory.policies[policy_name]\n\n    @staticmethod\n    def get_policy_names():\n        return tuple(PlacementPolicyFactory.policies.keys())\n", "path": "colossalai/zero/gemini/placement_policy.py"}], "after_files": [{"content": "import functools\nimport warnings\nfrom abc import ABC, abstractmethod\nfrom time import time\nfrom typing import Dict, List, Optional, Tuple, Type\n\nimport torch\n\nfrom colossalai.accelerator import get_accelerator\nfrom colossalai.legacy.utils.memory import colo_device_memory_capacity\nfrom colossalai.zero.gemini.chunk import Chunk\n\nfrom .chunk import Chunk, ChunkManager\nfrom .memory_tracer import ChunkMemStatsCollector\n\n\nclass PlacementPolicy(ABC):\n    need_mem_stats: bool = False\n\n    def __init__(\n        self,\n        gemini_manager: \"GeminiManager\",  # TODO @botbw: solve circular import\n        chunk_manager: ChunkManager,\n        mem_stats_collector: Optional[ChunkMemStatsCollector] = None,\n        max_prefetch: int = 0,\n        **kwargs,\n    ) -> None:\n        self.gemini_manager = gemini_manager\n        self.chunk_manager = chunk_manager\n        self.mem_stats_collector: Optional[ChunkMemStatsCollector] = mem_stats_collector\n        self.max_prefetch = max_prefetch\n\n    @abstractmethod\n    def evict_tensors(self, can_evict_chunks: List[Chunk], **kwargs) -> Tuple[int, float]:\n        raise NotImplementedError\n\n    @abstractmethod\n    def setup_grads_device(\n        self, params: List[torch.Tensor], grads_device_map: Dict[torch.Tensor, torch.device]\n    ) -> None:\n        raise NotImplementedError\n\n    def get_prefetch_chunks(self) -> List[Chunk]:\n        return []  # no prefetch by default\n\n\nclass StaticPlacementPolicy(PlacementPolicy):\n    def __init__(\n        self,\n        gemini_manager: \"GeminiManager\",\n        chunk_manager: ChunkManager,\n        mem_stats_collector: Optional[ChunkMemStatsCollector] = None,\n        max_prefetch: int = 0,\n        shard_param_frac: float = 1.0,\n        offload_optim_frac: float = 0.0,\n        offload_param_frac: float = 0.0,\n        **kwargs,\n    ) -> None:\n        super().__init__(\n            gemini_manager, chunk_manager, mem_stats_collector=mem_stats_collector, max_prefetch=max_prefetch\n        )\n        if offload_param_frac > 0.0 and (shard_param_frac != 1.0 or offload_optim_frac != 1.0):\n            warnings.warn(\"offload_param_frac is ignored when shard_param_frac != 1.0 or offload_optim_frac != 1.0\")\n            offload_param_frac = 0.0\n        self.shard_param_frac = shard_param_frac\n        self.offload_optim_frac = offload_optim_frac\n        self.offload_param_frac = offload_param_frac\n        # these should be initialized in setup_grads_device\n        self.keep_gathered_chunk_mem = 0.0\n        self.keep_cuda_chunk_mem = 0.0\n\n    def evict_tensors(self, can_evict_chunks: List[Chunk], **kwargs) -> Tuple[int, float]:\n        can_shard_chunk_mem = sum(chunk.chunk_mem for chunk in can_evict_chunks)\n        can_offload_chunk_mem = can_shard_chunk_mem\n        for chunk in can_evict_chunks:\n            if can_shard_chunk_mem <= self.keep_gathered_chunk_mem:\n                break\n            self.chunk_manager.release_chunk(chunk)\n            # real saved mem is chunk_mem - shard_mem, for simplicity we use chunk_mem\n            can_shard_chunk_mem -= chunk.chunk_mem\n        for chunk in can_evict_chunks:\n            if can_offload_chunk_mem <= self.keep_cuda_chunk_mem:\n                break\n            self.chunk_manager.move_chunk(chunk, torch.device(\"cpu\"))\n            # real saved mem is shard_mem, for simplicity we use chunk_mem\n            can_offload_chunk_mem -= chunk.chunk_mem\n        return 0, 0.0\n\n    def setup_grads_device(\n        self, params: List[torch.Tensor], grads_device_map: Dict[torch.Tensor, torch.device]\n    ) -> None:\n        total_chunk_mem = sum(self.chunk_manager.get_chunk(p).chunk_mem for p in params)\n\n        offload_optim_chunk_mem = total_chunk_mem * self.offload_optim_frac\n        offloaded_optim_chunk_mem = 0\n        chunks = set(self.chunk_manager.get_chunk(p) for p in params)\n        for chunk in chunks:\n            params = chunk.get_tensors()\n            # init offload optim settings\n            # keep gathered chunks are in CUDA\n            if chunk.keep_gathered or offloaded_optim_chunk_mem >= offload_optim_chunk_mem:\n                device = get_accelerator().get_current_device()\n            else:\n                device = torch.device(\"cpu\")\n                # real offloaded mem is chunk.shard_mem, for simplicity we use chunk mem here\n                offloaded_optim_chunk_mem += chunk.chunk_mem\n            for p in params:\n                grads_device_map[p] = device\n        self.keep_gathered_chunk_mem = total_chunk_mem * (1 - self.shard_param_frac)\n        self.keep_cuda_chunk_mem = total_chunk_mem * (1 - self.offload_param_frac)\n\n    def get_prefetch_chunks(self) -> List[Chunk]:\n        if self.gemini_manager.is_warmup():  # no prefetch during warmup since we need compute_list\n            return []\n        can_prefetch = self.max_prefetch - len(self.gemini_manager._async_works)\n        prefetch = []\n        for i in range(self.gemini_manager.compute_idx + 1, len(self.gemini_manager.compute_list)):\n            break_flag = False\n            for chunk in self.gemini_manager.compute_list[i]:\n                if len(prefetch) >= can_prefetch:\n                    break_flag = True\n                    break\n                if chunk not in prefetch and chunk not in self.chunk_manager.accessed_chunks:\n                    prefetch.append(chunk)\n            if break_flag:\n                break\n        return prefetch\n\n\nclass AutoPlacementPolicy(PlacementPolicy):\n    need_mem_stats: bool = True\n\n    def __init__(\n        self,\n        gemini_manager: \"GeminiManager\",\n        chunk_manager: ChunkManager,\n        mem_stats_collector: Optional[ChunkMemStatsCollector] = None,\n        max_prefetch: int = 0,\n        warmup_non_model_data_ratio: float = 0.8,\n        steady_cuda_cap_ratio: float = 0.9,\n        **kwargs,\n    ) -> None:\n        super().__init__(\n            gemini_manager, chunk_manager, mem_stats_collector=mem_stats_collector, max_prefetch=max_prefetch\n        )\n        # model data will use 1-_warmup_non_model_data_ratio CUDA memory in warmup phase\n        # you can set them by AutoPlacementPolicy.set_warmup_non_model_data_ratio()\n        # and AutoPlacementPolicy.set_steady_cuda_cap_ratio()\n        self._warmup_non_model_data_ratio = warmup_non_model_data_ratio\n        self._steady_cuda_cap_ratio = steady_cuda_cap_ratio\n\n    def evict_tensors(\n        self,\n        can_evict_chunks: List[Chunk],\n        cuda_demand: int = 0,\n        warmup: bool = True,\n        compute_list: Optional[List[Tuple[Chunk, ...]]] = None,\n        compute_idx: int = 0,\n        **kwargs,\n    ) -> Tuple[int, float]:\n        \"\"\"\n        Evict tensors from CUDA device.\n\n        Args:\n            can_evict_chunks (List[StatefulTensor]): the list of tensors that can be evicted.\n            cuda_demand (int, optional): the volume of data needed on cuda device. Defaults to 0.\n            warmup (bool, optional): a flag indicates whether in the phase of warmup. Defaults to True.\n            compute_list (List[StatefulTensor], optional): TODO. Defaults to [].\n            compute_idx (int, optional): the idx of computing device. Defaults to 0.\n\n        Raises:\n            RuntimeError:\n\n        Returns:\n            int: the volume of memory that is evicted\n        \"\"\"\n        start = time()\n        cuda_capacity = colo_device_memory_capacity(get_accelerator().get_current_device())\n        used_cuda_model_data = self.chunk_manager.total_mem[\"cuda\"]\n        if warmup:\n            # We designate a part of CUDA memory for model data in warmup iterations.\n            max_cuda_non_model_data_per_period = cuda_capacity * self._warmup_non_model_data_ratio\n        else:\n            # max non-model-data cuda memory consumption of this sampling moment and the next sampling moment.\n            max_cuda_non_model_data_per_period = self.mem_stats_collector.next_period_non_model_data_usage(\"cuda\")\n            cuda_capacity *= self._steady_cuda_cap_ratio\n        total_cuda_model_data = cuda_capacity - max_cuda_non_model_data_per_period\n        avail_cuda_model_data = total_cuda_model_data - used_cuda_model_data\n        freed_cuda_model_data = 0\n\n        if avail_cuda_model_data < cuda_demand:\n            # Move cuda_demand - avail_cuda_model_data volume of tensors\n            # to_free_cuda_model_data = cuda_demand - avail_cuda_model_data\n            to_free_cuda_model_data = cuda_demand - avail_cuda_model_data\n            to_free_chunks = can_evict_chunks\n            if not warmup:\n                to_free_chunks = self._sort_can_evict_chunks(tuple(to_free_chunks), compute_idx, tuple(compute_list))\n                # print(self._sort_can_evict_chunks.cache_info())\n            for chunk in to_free_chunks:\n                if freed_cuda_model_data >= to_free_cuda_model_data:\n                    break\n\n                self.chunk_manager.release_chunk(chunk)\n                self.chunk_manager.move_chunk(chunk, torch.device(\"cpu\"))\n                freed_cuda_model_data += chunk.chunk_mem\n            if freed_cuda_model_data < to_free_cuda_model_data:\n                raise RuntimeError(\n                    f\"Adjust layout failed! No enough CUDA memory! \"\n                    f\"Need {to_free_cuda_model_data}, freed {freed_cuda_model_data}\"\n                )\n        return freed_cuda_model_data, time() - start\n\n    @staticmethod\n    @functools.lru_cache(maxsize=None)\n    def _sort_can_evict_chunks(can_evict_chunks: tuple, compute_idx: int, compute_list: tuple) -> list:\n        next_compute_idx = {chunk: len(compute_list) for chunk in can_evict_chunks}\n        for i in range(len(compute_list) - 1, compute_idx, -1):\n            for chunk in compute_list[i]:\n                if chunk in next_compute_idx:\n                    next_compute_idx[chunk] = i\n        next_compute_idx = sorted(next_compute_idx.items(), key=lambda pair: pair[1], reverse=True)\n        return [t for (t, idx) in next_compute_idx]\n\n    def setup_grads_device(\n        self, params: List[torch.Tensor], grads_device_map: Dict[torch.Tensor, torch.device]\n    ) -> None:\n        for p in params:\n            chunk = self.chunk_manager.get_chunk(p)\n            # init offload optim settings\n            # keep gathered chunks are in CUDA\n            if chunk.keep_gathered:\n                grads_device_map[p] = get_accelerator().get_current_device()\n            else:\n                grads_device_map[p] = torch.device(\"cpu\")\n\n    def get_prefetch_chunks(self) -> List[Chunk]:\n        if self.gemini_manager.is_warmup():  # no prefetch during warmup since we need compute_list\n            return []\n        # modified from self.evict_tensors\n        cuda_capacity = self._steady_cuda_cap_ratio * colo_device_memory_capacity(\n            get_accelerator().get_current_device()\n        )\n        max_cuda_non_model_data_per_period = self.mem_stats_collector.next_period_non_model_data_usage(\"cuda\")\n        used_cuda_model_data = self.chunk_manager.total_mem[\"cuda\"]\n        total_cuda_model_data = cuda_capacity - max_cuda_non_model_data_per_period\n        avail_cuda_model_data = total_cuda_model_data - used_cuda_model_data\n\n        prefetch_chunk_memory = 0\n        can_prefetch = self.max_prefetch - len(self.gemini_manager._async_works)\n        prefetch = []\n        for i in range(self.gemini_manager.compute_idx + 1, len(self.gemini_manager.compute_list)):\n            break_flag = False\n            for chunk in self.gemini_manager.compute_list[i]:\n                chunk: Chunk\n                if len(prefetch) >= can_prefetch or prefetch_chunk_memory + chunk.chunk_mem > avail_cuda_model_data:\n                    break_flag = True\n                    break\n                if chunk not in prefetch and chunk not in self.chunk_manager.accessed_chunks:\n                    prefetch.append(chunk)\n            if break_flag:\n                break\n        return prefetch\n\n\nclass PlacementPolicyFactory:\n    policies: Dict[str, Type[PlacementPolicy]] = {\n        \"auto\": AutoPlacementPolicy,\n        \"static\": StaticPlacementPolicy,\n    }\n\n    @staticmethod\n    def create(policy_name: str) -> Type[PlacementPolicy]:\n        if policy_name not in PlacementPolicyFactory.policies:\n            raise TypeError(f\"Unknown tensor placement policy {policy_name}\")\n        return PlacementPolicyFactory.policies[policy_name]\n\n    @staticmethod\n    def get_policy_names():\n        return tuple(PlacementPolicyFactory.policies.keys())\n", "path": "colossalai/zero/gemini/placement_policy.py"}]}