mirror of https://github.com/hpcaitech/ColossalAI
52 lines
1.5 KiB
Python
52 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",
|
|
)
|