mirror of
https://github.com/hpcaitech/ColossalAI.git
synced 2025-06-03 20:59:49 +00:00
232142f402
* add model data profiler * add a subclass of torch.profiler.profile * refactor folder structure * remove redundant codes * polish code * use GeminiMemoryManager * fix import path * fix stm profiler ext * polish comments * remove useless file
21 lines
342 B
Python
21 lines
342 B
Python
from abc import ABC, abstractmethod
|
|
|
|
|
|
class ProfilerExtension(ABC):
|
|
|
|
@abstractmethod
|
|
def prepare_trace(self):
|
|
pass
|
|
|
|
@abstractmethod
|
|
def start_trace(self):
|
|
pass
|
|
|
|
@abstractmethod
|
|
def stop_trace(self):
|
|
pass
|
|
|
|
@abstractmethod
|
|
def extend_chrome_trace(self, trace: dict) -> dict:
|
|
pass
|