mirror of https://github.com/hpcaitech/ColossalAI
33 lines
713 B
Python
33 lines
713 B
Python
from contextlib import contextmanager
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
|
|
|
|
@contextmanager
|
|
def low_precision_init(target_dtype: torch.dtype = torch.float16):
|
|
dtype = torch.get_default_dtype()
|
|
try:
|
|
torch.set_default_dtype(target_dtype)
|
|
yield
|
|
finally:
|
|
torch.set_default_dtype(dtype)
|
|
|
|
|
|
def get_model_numel(model: nn.Module) -> int:
|
|
return sum(p.numel() for p in model.parameters())
|
|
|
|
|
|
def format_numel_str(numel: int) -> str:
|
|
B = 1024**3
|
|
M = 1024**2
|
|
K = 1024
|
|
if numel >= B:
|
|
return f"{numel / B:.2f} B"
|
|
elif numel >= M:
|
|
return f"{numel / M:.2f} M"
|
|
elif numel >= K:
|
|
return f"{numel / K:.2f} K"
|
|
else:
|
|
return f"{numel}"
|