ColossalAI/colossalai/legacy/zero/gemini/gemini_context.py

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",
            )