mirror of https://github.com/hpcaitech/ColossalAI
49 lines
1.5 KiB
Python
49 lines
1.5 KiB
Python
from enum import EnumMeta
|
|
|
|
|
|
class GeminiMemoryManager(object):
|
|
|
|
def __init__(self, states_cls: EnumMeta):
|
|
super().__init__()
|
|
self.states_cls = states_cls
|
|
self._cnter = 0 # the counter of instances
|
|
|
|
self.total_mem = dict()
|
|
self.state_mem = dict()
|
|
self.state_mem['cpu'] = dict()
|
|
self.state_mem['cuda'] = dict()
|
|
|
|
self.reset()
|
|
|
|
@property
|
|
def total_number(self):
|
|
return self._cnter
|
|
|
|
def reset(self):
|
|
self._cnter = 0 # the counter of instances
|
|
|
|
self.total_mem['cpu'] = 0 # memory occupation of instances in cpu
|
|
self.total_mem['cuda'] = 0 # memory of occupation of instances in cuda
|
|
|
|
# memory conditions for all states
|
|
for state in self.states_cls:
|
|
self.state_mem['cpu'][state] = 0
|
|
self.state_mem['cuda'][state] = 0
|
|
|
|
def register_new_instance(self):
|
|
self._cnter += 1
|
|
|
|
def delete_instance(self):
|
|
self._cnter -= 1
|
|
|
|
def print_info(self):
|
|
print(f"Total number: {self.total_number}",
|
|
f"Total CPU memory occupation: {self.total_mem['cpu']}",
|
|
f"Total CUDA memory occupation: {self.total_mem['cuda']}\n",
|
|
sep='\n')
|
|
|
|
for state in self.states_cls:
|
|
print(f"{state}: CPU memory occupation: {self.state_mem['cpu'][state]}",
|
|
f"{state}: CUDA memory occupation: {self.state_mem['cuda'][state]}\n",
|
|
sep='\n')
|