github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-06-20 04:32:47 +00:00
Code Issues Packages Projects Releases Wiki Activity

[pre-commit.ci] auto fixes from pre-commit.com hooks

Browse Source 
for more information, see https://pre-commit.ci
This commit is contained in:
pre-commit-ci[bot] 2024-05-20 03:00:06 +00:00
parent a55a9e298b
commit f1918e18a5
1 changed files with 5 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
colossalai/zero/gemini/placement_policy.py
View File

@ -237,12 +237,14 @@ class AutoPlacementPolicy(PlacementPolicy):
if self.gemini_manager.is_warmup(): # no prefetch during warmup since we need compute_list if self.gemini_manager.is_warmup(): # no prefetch during warmup since we need compute_list
return [] return []
# modified from self.evict_tensors # modified from self.evict_tensors
cuda_capacity = self._steady_cuda_cap_ratio * colo_device_memory_capacity(get_accelerator().get_current_device()) 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") 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"] used_cuda_model_data = self.chunk_manager.total_mem["cuda"]
total_cuda_model_data = cuda_capacity - max_cuda_non_model_data_per_period 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 avail_cuda_model_data = total_cuda_model_data - used_cuda_model_data
prefetch_chunk_memory = 0 prefetch_chunk_memory = 0
can_prefetch = self.max_prefetch - len(self.gemini_manager._async_works) can_prefetch = self.max_prefetch - len(self.gemini_manager._async_works)
prefetch = [] prefetch = []
@ -259,6 +261,7 @@ class AutoPlacementPolicy(PlacementPolicy):
break break
return prefetch return prefetch
class PlacementPolicyFactory: class PlacementPolicyFactory:
policies: Dict[str, Type[PlacementPolicy]] = { policies: Dict[str, Type[PlacementPolicy]] = {
"auto": AutoPlacementPolicy, "auto": AutoPlacementPolicy,