wehub-resource-sync
76 lines
2.7 KiB
Python
76 lines
2.7 KiB
Python
"""Common utilities for loading parameters"""
|
|
|
|
import functools
|
|
import operator
|
|
from collections.abc import Iterator
|
|
from pathlib import Path
|
|
from typing import TYPE_CHECKING, Set, Tuple # noqa: UP035
|
|
|
|
import numpy as np
|
|
|
|
from mlc_llm.support import logging
|
|
|
|
if TYPE_CHECKING:
|
|
from .mapping import ExternMapping
|
|
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
def check_parameter_usage(param_map: "ExternMapping", extern_weights: Set[str]): # noqa: UP006
|
|
"""Check that all external parameters have been used and are stored in the weights file."""
|
|
used_extern_names = set(functools.reduce(operator.iadd, param_map.param_map.values(), []))
|
|
# Check 1. All extern parameters in the weight files are used unless explicitly specified
|
|
unused_extern_names = extern_weights - used_extern_names - param_map.unused_params
|
|
if unused_extern_names:
|
|
logger.warning(
|
|
"Unused extern parameters: %s",
|
|
", ".join(sorted(unused_extern_names)),
|
|
)
|
|
# Check 2. All extern parameters required are stored in the weight files
|
|
nonexistent_extern_names = used_extern_names - extern_weights
|
|
if nonexistent_extern_names:
|
|
raise ValueError(
|
|
"The following extern parameters do not exist in the weight files:\n "
|
|
+ "\n ".join(sorted(nonexistent_extern_names)),
|
|
)
|
|
|
|
|
|
def load_torch_shard(path: Path) -> Iterator[Tuple[str, np.ndarray]]: # noqa: UP006
|
|
"""Load and yield PyTorch format parameters."""
|
|
import torch
|
|
|
|
for name, param in torch.load(path, map_location=torch.device("cpu")).items():
|
|
if param is None:
|
|
logger.warning("Encountered None param, skipping it: %s", name)
|
|
continue
|
|
param = param.detach().cpu()
|
|
dtype = str(param.dtype)
|
|
if dtype == "torch.bfloat16":
|
|
param = param.float()
|
|
param = param.numpy()
|
|
yield name, param
|
|
|
|
|
|
def load_safetensor_shard(path: Path) -> Iterator[Tuple[str, np.ndarray]]: # noqa: UP006
|
|
"""Load and yield SafeTensor format parameters."""
|
|
import safetensors
|
|
import torch
|
|
|
|
with safetensors.safe_open(path, framework="pt", device="cpu") as in_file:
|
|
for name in in_file.keys():
|
|
param = in_file.get_tensor(name)
|
|
param = param.detach().cpu()
|
|
dtype = str(param.dtype)
|
|
if dtype == "torch.bfloat16":
|
|
import ml_dtypes
|
|
|
|
param = param.view(torch.float16).cpu().numpy().view(ml_dtypes.bfloat16)
|
|
elif dtype == "torch.float8_e4m3fn":
|
|
import ml_dtypes
|
|
|
|
param = param.view(torch.uint8).cpu().numpy().view(ml_dtypes.float8_e4m3fn)
|
|
else:
|
|
param = param.numpy()
|
|
yield name, param
|