wehub-resource-sync
94 lines
2.7 KiB
Python
94 lines
2.7 KiB
Python
"""Statistics of the loading process of parameter loaders"""
|
|
|
|
import dataclasses
|
|
import time
|
|
from contextlib import contextmanager
|
|
|
|
from mlc_llm.support import logging
|
|
from mlc_llm.support.style import green
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
@dataclasses.dataclass
|
|
class Stats:
|
|
"""Statistics of the loading process of parameter loaders.
|
|
|
|
Attributes
|
|
----------
|
|
load_time_sec : float
|
|
Time used in loading the parameters.
|
|
|
|
map_time_sec : float
|
|
Time used in applying the mapping function, i.e. `ExternMapping.map_func`.
|
|
|
|
quant_time_sec : float
|
|
Time used in quantizing the parameters, i.e. `QuantizeMapping.quant_func`.
|
|
|
|
current_memory_gb : float
|
|
The current RAM usage in GB.
|
|
|
|
total_memory_gb : float
|
|
The total size data loaded from disk in GB.
|
|
|
|
max_memory_gb : float
|
|
The maximum RAM usage in GB.
|
|
|
|
total_param_num: int
|
|
Total number of parameters (original non-MLC model weights), excluding unused params.
|
|
"""
|
|
|
|
load_time_sec: float = 0.0
|
|
map_time_sec: float = 0.0
|
|
quant_time_sec: float = 0.0
|
|
|
|
current_memory_gb: float = 0.0
|
|
total_memory_gb: float = 0.0
|
|
max_memory_gb: float = 0.0
|
|
|
|
total_param_num: int = 0
|
|
|
|
def timer(self, attr):
|
|
"""A context manager to time the scope and add the time to the attribute."""
|
|
|
|
@contextmanager
|
|
def timed_scope():
|
|
start_time = time.time()
|
|
yield
|
|
elapsed_time = time.time() - start_time
|
|
setattr(self, attr, getattr(self, attr) + elapsed_time)
|
|
|
|
return timed_scope()
|
|
|
|
def mem_add(self, nbytes: int):
|
|
"""Add the memory usage by the given number of bytes."""
|
|
mem_gb = float(nbytes) / float(1024**3)
|
|
self.current_memory_gb += mem_gb
|
|
self.total_memory_gb += mem_gb
|
|
self.max_memory_gb = max(self.max_memory_gb, self.current_memory_gb)
|
|
|
|
def mem_rm(self, nbytes: int):
|
|
"""Remove the memory usage by the given number of bytes."""
|
|
mem_gb = float(nbytes) / float(1024**3)
|
|
self.current_memory_gb -= mem_gb
|
|
|
|
def log_time_info(self, weight_format: str):
|
|
"""Log the time used in loading, pre-quantization and quantization."""
|
|
logger.info(
|
|
"%s: %s loading: %.3f sec; Pre-quantization mapping: %.3f sec; Quantization: %.3f sec",
|
|
green("Time usage"),
|
|
weight_format,
|
|
self.load_time_sec,
|
|
self.map_time_sec,
|
|
self.quant_time_sec,
|
|
)
|
|
|
|
def log_mem_usage(self):
|
|
"""Log the Memory usage information."""
|
|
logger.info(
|
|
"%s: Peak RAM: %.3f GB. Total bytes loaded from disk: %.3f GB",
|
|
green("RAM usage"),
|
|
self.max_memory_gb,
|
|
self.total_memory_gb,
|
|
)
|