Spaces:
Running on Zero
Running on Zero
| from __future__ import annotations | |
| import os | |
| from typing import Any, Dict, List, Tuple | |
| import torch | |
| import torch.nn as nn | |
| from .config import QuantConfig | |
| from .qmodules import QuantConv1d, QuantConv2d, QuantLinear | |
| def _unwrap_state_dict(obj: Any) -> Dict[str, torch.Tensor]: | |
| if isinstance(obj, dict): | |
| for key in ("state_dict", "model", "net", "params"): | |
| value = obj.get(key) | |
| if isinstance(value, dict): | |
| return value | |
| return obj | |
| raise ValueError("Unsupported checkpoint format") | |
| def load_state_dict_flexible( | |
| model: nn.Module, | |
| ckpt_path: str, | |
| device: torch.device, | |
| ) -> Tuple[List[str], List[str]]: | |
| ckpt = torch.load(ckpt_path, map_location="cpu") | |
| sd = _unwrap_state_dict(ckpt) | |
| model_sd = model.state_dict() | |
| if any(k.startswith("module.") for k in sd) and not any(k.startswith("module.") for k in model_sd): | |
| sd = {k[len("module."):]: v for k, v in sd.items()} | |
| missing, unexpected = model.load_state_dict(sd, strict=False) | |
| model.to(device) | |
| return list(missing), list(unexpected) | |
| def save_quant_checkpoint( | |
| path: str, | |
| q_model: nn.Module, | |
| qcfg: QuantConfig, | |
| extra: Dict[str, Any] | None = None, | |
| ) -> None: | |
| os.makedirs(os.path.dirname(os.path.abspath(path)), exist_ok=True) | |
| torch.save( | |
| { | |
| "state_dict": q_model.state_dict(), | |
| "quant_config": qcfg.to_dict(), | |
| "extra": extra or {}, | |
| }, | |
| path, | |
| ) | |
| def _get_parent(root: nn.Module, full_name: str) -> tuple[nn.Module, str]: | |
| parts = full_name.split(".") | |
| if len(parts) == 1: | |
| return root, parts[0] | |
| parent = root | |
| for part in parts[:-1]: | |
| parent = getattr(parent, part) | |
| return parent, parts[-1] | |
| def _has_quant_weight(state_dict: Dict[str, torch.Tensor], module_name: str) -> bool: | |
| prefix = module_name + "." | |
| return any(k.startswith(prefix + "w_") for k in state_dict) | |
| def _infer_weight_bits(state_dict: Dict[str, torch.Tensor], module_name: str, fallback: int) -> int: | |
| prefix = module_name + "." | |
| if (prefix + "w_q_packed") in state_dict: | |
| return 4 | |
| if (prefix + "w_q_int8") in state_dict: | |
| return int(fallback) | |
| return int(fallback) | |
| def rebuild_quant_model_from_checkpoint( | |
| fp_model: nn.Module, | |
| quant_ckpt_path: str, | |
| device: torch.device, | |
| ) -> nn.Module: | |
| payload = torch.load(quant_ckpt_path, map_location="cpu") | |
| if not (isinstance(payload, dict) and "quant_config" in payload and "state_dict" in payload): | |
| raise ValueError("Invalid quant checkpoint format; expected quant_config and state_dict") | |
| qcfg = QuantConfig.from_dict(payload["quant_config"]) | |
| sd: Dict[str, torch.Tensor] = payload["state_dict"] | |
| for full_name, module in list(fp_model.named_modules()): | |
| if not isinstance(module, (nn.Linear, nn.Conv2d, nn.Conv1d)): | |
| continue | |
| if not _has_quant_weight(sd, full_name): | |
| continue | |
| parent, child_name = _get_parent(fp_model, full_name) | |
| w_bits = _infer_weight_bits(sd, full_name, fallback=qcfg.w_bits) | |
| if isinstance(module, nn.Linear): | |
| replacement = QuantLinear( | |
| module, | |
| w_bits=w_bits, | |
| a_bits=qcfg.a_bits, | |
| w_per_channel=qcfg.w_per_channel, | |
| w_group_size=qcfg.w_group_size, | |
| w_clip_percentile=qcfg.w_clip_percentile, | |
| act_mode=qcfg.act_mode, | |
| act_group_size=qcfg.act_group_size, | |
| dstg_group_size=qcfg.dstg_group_size, | |
| act_per_channel=qcfg.a_per_channel, | |
| a_clip_percentile=qcfg.a_clip_percentile, | |
| dstg_clip_percentile=qcfg.dstg_clip_percentile, | |
| token_clip_percentile=qcfg.token_clip_percentile, | |
| token_outlier_bypass=qcfg.token_outlier_bypass, | |
| token_outlier_topk=qcfg.token_outlier_topk, | |
| dcrp_tau=qcfg.dcrp_tau, | |
| dcrp_zero_ratio=qcfg.dcrp_zero_ratio, | |
| token_step_over_std=qcfg.token_step_over_std, | |
| token_zero_ratio=qcfg.token_zero_ratio, | |
| compute_dtype=qcfg.compute_dtype, | |
| full_name=full_name, | |
| ) | |
| elif isinstance(module, nn.Conv2d): | |
| replacement = QuantConv2d( | |
| module, | |
| w_bits=w_bits, | |
| a_bits=qcfg.a_bits, | |
| act_per_channel=qcfg.a_per_channel, | |
| w_per_channel=qcfg.w_per_channel, | |
| w_clip_percentile=qcfg.w_clip_percentile, | |
| a_clip_percentile=qcfg.a_clip_percentile, | |
| compute_dtype=qcfg.compute_dtype, | |
| full_name=full_name, | |
| ) | |
| else: | |
| replacement = QuantConv1d( | |
| module, | |
| w_bits=w_bits, | |
| a_bits=qcfg.a_bits, | |
| act_per_channel=qcfg.a_per_channel, | |
| w_per_channel=qcfg.w_per_channel, | |
| w_clip_percentile=qcfg.w_clip_percentile, | |
| a_clip_percentile=qcfg.a_clip_percentile, | |
| compute_dtype=qcfg.compute_dtype, | |
| full_name=full_name, | |
| ) | |
| setattr(parent, child_name, replacement) | |
| msg = fp_model.load_state_dict(sd, strict=False) | |
| fp_model.to(device) | |
| if str(qcfg.compute_dtype).lower() in ("fp16", "float16", "half") and device.type == "cuda": | |
| fp_model.half() | |
| print( | |
| "[Builder] Rebuilt quant model. " | |
| f"missing={len(getattr(msg, 'missing_keys', []))}, " | |
| f"unexpected={len(getattr(msg, 'unexpected_keys', []))}" | |
| ) | |
| return fp_model | |