mirror of https://github.com/hpcaitech/ColossalAI
170 lines
7.2 KiB
Python
170 lines
7.2 KiB
Python
from abc import ABC, abstractmethod
|
|
from time import time
|
|
from typing import List, Optional, Tuple, Dict
|
|
import torch
|
|
from colossalai.utils import get_current_device
|
|
from colossalai.utils.memory import colo_device_memory_capacity
|
|
|
|
from colossalai.gemini.memory_tracer.memstats_collector import MemStatsCollectorV2
|
|
from typing import Type
|
|
import functools
|
|
from colossalai.gemini import Chunk, ChunkManager
|
|
|
|
|
|
class PlacementPolicy(ABC):
|
|
need_mem_stats: bool = False
|
|
|
|
def __init__(self, chunk_manager: ChunkManager, mem_stats_collector: Optional[MemStatsCollectorV2] = None) -> None:
|
|
self.chunk_manager = chunk_manager
|
|
self.mem_stats_collector: Optional[MemStatsCollectorV2] = mem_stats_collector
|
|
|
|
@abstractmethod
|
|
def evict_tensors(self, can_evict_chunks: List[Chunk], **kwargs) -> None:
|
|
raise NotImplementedError
|
|
|
|
@staticmethod
|
|
def get_default_device() -> torch.device:
|
|
return torch.device('cpu')
|
|
|
|
|
|
class CPUPlacementPolicy(PlacementPolicy):
|
|
|
|
def __init__(self, chunk_manager: ChunkManager, mem_stats_collector: Optional[MemStatsCollectorV2] = None) -> None:
|
|
super().__init__(chunk_manager, mem_stats_collector=mem_stats_collector)
|
|
|
|
def evict_tensors(self, can_evict_chunks: List[Chunk], **kwargs) -> int:
|
|
volume = 0
|
|
start = time()
|
|
for chunk in can_evict_chunks:
|
|
self.chunk_manager.move_chunk(chunk, torch.device('cpu'), update_ptr=False)
|
|
volume += chunk.mem
|
|
return volume, time() - start
|
|
|
|
|
|
class CUDAPlacementPolicy(PlacementPolicy):
|
|
|
|
def __init__(self, chunk_manager: ChunkManager, mem_stats_collector: Optional[MemStatsCollectorV2] = None) -> None:
|
|
assert torch.cuda.is_available(), 'Cannot use CUDATensorPlacementPolicy when CUDA is not available'
|
|
super().__init__(chunk_manager, mem_stats_collector=mem_stats_collector)
|
|
|
|
def evict_tensors(self, can_evict_chunks: List[Chunk], **kwargs) -> int:
|
|
return 0, 0
|
|
|
|
@staticmethod
|
|
def get_default_device() -> torch.device:
|
|
return get_current_device()
|
|
|
|
|
|
class AutoPlacementPolicy(PlacementPolicy):
|
|
|
|
need_mem_stats: bool = True
|
|
# model data will use 1-_warmup_non_model_data_ratio CUDA memory in warmup phase
|
|
# you can set them by AutoPlacementPolicy.set_warmup_non_model_data_ratio() and AutoPlacementPolicy.set_steady_cuda_cap_ratio()
|
|
_warmup_non_model_data_ratio: float = 0.8
|
|
_steady_cuda_cap_ratio: float = 0.9
|
|
|
|
def __init__(self, chunk_manager: ChunkManager, mem_stats_collector: Optional[MemStatsCollectorV2] = None) -> None:
|
|
super().__init__(chunk_manager, mem_stats_collector=mem_stats_collector)
|
|
|
|
def evict_tensors(self,
|
|
can_evict_chunks: List[Chunk],
|
|
cuda_demand: int = 0,
|
|
warmup: bool = True,
|
|
compute_list: List[Tuple[Chunk, ...]] = [],
|
|
compute_idx: int = 0,
|
|
**kwargs) -> int:
|
|
"""
|
|
Evict tensors from CUDA device.
|
|
|
|
Args:
|
|
hold_cuda_tensor_list (List[StatefulTensor]): the list of tensor in state of HOLD-like
|
|
cuda_demand (int, optional): the volume of data needed on cuda device. Defaults to 0.
|
|
warmup (bool, optional): a flag indicates whether in the phase of warmup. Defaults to True.
|
|
compute_list (List[StatefulTensor], optional): TODO. Defaults to [].
|
|
compute_idx (int, optional): the idx of computing device. Defaults to 0.
|
|
|
|
Raises:
|
|
RuntimeError:
|
|
|
|
Returns:
|
|
int: the volume of memory that is evicted
|
|
"""
|
|
start = time()
|
|
cuda_capacity = colo_device_memory_capacity(get_current_device())
|
|
used_cuda_model_data = self.chunk_manager.total_mem['cuda']
|
|
if warmup:
|
|
# We designate a part of CUDA memory for model data in warmup iterations.
|
|
max_cuda_non_model_data_per_period = cuda_capacity * AutoPlacementPolicy._warmup_non_model_data_ratio
|
|
else:
|
|
# max non-model-data cuda memory consumption of this sampling moment and the next sampling moment.
|
|
max_cuda_non_model_data_per_period = self.mem_stats_collector.next_period_non_model_data_usage('cuda')
|
|
cuda_capacity *= AutoPlacementPolicy._steady_cuda_cap_ratio
|
|
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
|
|
freed_cuda_model_data = 0
|
|
|
|
if avail_cuda_model_data < cuda_demand:
|
|
# Move cuda_demand - avail_cuda_model_data volume of tensors
|
|
# to_free_cuda_model_data = cuda_demand - avail_cuda_model_data
|
|
to_free_cuda_model_data = cuda_demand - avail_cuda_model_data
|
|
to_free_chunks = can_evict_chunks
|
|
if not warmup:
|
|
to_free_chunks = self._sort_can_evict_chunks(tuple(to_free_chunks), compute_idx, tuple(compute_list))
|
|
# print(self._sort_can_evict_chunks.cache_info())
|
|
for chunk in to_free_chunks:
|
|
if freed_cuda_model_data >= to_free_cuda_model_data:
|
|
break
|
|
freed_cuda_model_data += chunk.mem
|
|
self.chunk_manager.move_chunk(chunk, torch.device('cpu'), update_ptr=False)
|
|
if freed_cuda_model_data < to_free_cuda_model_data:
|
|
raise RuntimeError(
|
|
f"Adjust layout failed! No enough CUDA memory! Need {to_free_cuda_model_data}, freed {freed_cuda_model_data}"
|
|
)
|
|
return freed_cuda_model_data, time() - start
|
|
|
|
@staticmethod
|
|
@functools.lru_cache(maxsize=None)
|
|
def _sort_can_evict_chunks(can_evict_chunks: tuple, compute_idx: int, compute_list: tuple) -> list:
|
|
next_compute_idx = {chunk: len(compute_list) for chunk in can_evict_chunks}
|
|
for i in range(len(compute_list) - 1, compute_idx, -1):
|
|
for chunk in compute_list[i]:
|
|
if chunk in next_compute_idx:
|
|
next_compute_idx[chunk] = i
|
|
next_compute_idx = sorted(next_compute_idx.items(), key=lambda pair: pair[1], reverse=True)
|
|
return [t for (t, idx) in next_compute_idx]
|
|
|
|
@staticmethod
|
|
def set_warmup_non_model_data_ratio(ratio: float) -> None:
|
|
ratio = float(ratio)
|
|
assert 0.0 < ratio < 1.0
|
|
AutoPlacementPolicy._warmup_non_model_data_ratio = ratio
|
|
|
|
@staticmethod
|
|
def set_steady_cuda_cap_ratio(ratio: float) -> None:
|
|
ratio = float(ratio)
|
|
assert 0.0 < ratio < 1.0
|
|
AutoPlacementPolicy._steady_cuda_cap_ratio = ratio
|
|
|
|
|
|
class PlacementPolicyFactory:
|
|
policies: Dict[str, PlacementPolicy] = {
|
|
'cpu': CPUPlacementPolicy,
|
|
'cuda': CUDAPlacementPolicy,
|
|
'auto': AutoPlacementPolicy
|
|
}
|
|
|
|
@staticmethod
|
|
def create(policy_name: str) -> Type[PlacementPolicy]:
|
|
if policy_name not in PlacementPolicyFactory.policies:
|
|
raise TypeError(f"Unknown tensor placement policy {policy_name}")
|
|
return PlacementPolicyFactory.policies[policy_name]
|
|
|
|
@staticmethod
|
|
def get_polocy_names():
|
|
return tuple(PlacementPolicyFactory.policies.keys())
|
|
|
|
@staticmethod
|
|
def get_default_device(policy_name: str) -> torch.device:
|
|
policy_cls = PlacementPolicyFactory.create(policy_name)
|
|
return policy_cls.get_default_device()
|