mirror of https://github.com/hpcaitech/ColossalAI
46 lines
1.4 KiB
Python
46 lines
1.4 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 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')
|