mmgp 3.3.1__py3-none-any.whl → 3.6.11__py3-none-any.whl

mmgp/fp8_quanto_bridge_old.py

@@ -0,0 +1,498 @@
+from __future__ import annotations
+import json, re, inspect
+from types import SimpleNamespace
+from typing import Dict, Optional, Tuple, Union, Iterable, Callable
+
+import torch
+from safetensors.torch import safe_open, save_file
+
+# ---------- Constants ----------
+DATA_SUFFIX  = "._data"
+SCALE_SUFFIX = "._scale"          # per-channel, shape [out, 1, ...]
+IN_SCALE    = ".input_scale"      # 1-D placeholder tensor [1]
+OUT_SCALE   = ".output_scale"     # 1-D placeholder tensor [1]
+
+_QTYPE_NAME = {
+    "e4m3fn": "qfloat8_e4m3fn",
+    "e5m2":   "qfloat8_e5m2",
+    "auto":   "qfloat8",
+}
+
+_SCALE_META_KEYS = (
+    "fp8_scale_map", "fp8.scale_map", "scale_map",
+    "quant_scale_map", "weights_scales", "scales",
+)
+
+_DTYPE_ALIASES = {
+    "float32": torch.float32, "fp32": torch.float32,
+    "bfloat16": torch.bfloat16, "bf16": torch.bfloat16,
+    "float16": torch.float16, "fp16": torch.float16, "half": torch.float16,
+}
+
+def _is_weight_key(k: str) -> bool:
+    return k.endswith(".weight")
+
+# ---------- Accessors (unify file vs dict) ----------
+class Accessor:
+    def keys(self) -> Iterable[str]: ...
+    def get_tensor(self, key: str) -> torch.Tensor: ...
+    def metadata(self) -> Dict[str, str]: ...
+    def has(self, key: str) -> bool: ...          # NEW
+    def can_delete(self) -> bool: return False
+    def delete(self, key: str) -> None: raise NotImplementedError
+
+class FileAccessor(Accessor):
+    def __init__(self, path: str):
+        self._fh = safe_open(path, framework="pt")
+        self._keys = list(self._fh.keys())
+        self._keys_set = set(self._keys)          # O(1) membership
+        self._meta = self._fh.metadata() or {}
+    def keys(self) -> Iterable[str]: return self._keys
+    def has(self, key: str) -> bool: return key in self._keys_set
+    def get_tensor(self, key: str) -> torch.Tensor: return self._fh.get_tensor(key)
+    def metadata(self) -> Dict[str, str]: return self._meta
+    def close(self) -> None: self._fh.close()
+
+class DictAccessor(Accessor):
+    def __init__(self, sd: Dict[str, torch.Tensor], meta: Optional[Dict[str, str]] = None,
+                 in_place: bool = False, free_cuda_cache: bool = False, cuda_cache_interval: int = 32):
+        self.sd = sd
+        self._meta = meta or {}
+        self._in_place = in_place
+        self._free = free_cuda_cache
+        self._interval = int(cuda_cache_interval)
+        self._deletions = 0
+    def keys(self) -> Iterable[str]: return list(self.sd.keys())
+    def has(self, key: str) -> bool: return key in self.sd          # dict membership = O(1)
+    def get_tensor(self, key: str) -> torch.Tensor: return self.sd[key]
+    def metadata(self) -> Dict[str, str]: return self._meta
+    def can_delete(self) -> bool: return self._in_place
+    def delete(self, key: str) -> None:
+        if key in self.sd:
+            self.sd.pop(key, None)
+            self._deletions += 1
+            if self._free and (self._deletions % self._interval == 0) and torch.cuda.is_available():
+                torch.cuda.empty_cache()
+def _as_accessor(src: Union[str, Dict[str, torch.Tensor]], **dict_opts) -> Tuple[Accessor, Callable[[], None]]:
+    if isinstance(src, str):
+        acc = FileAccessor(src)
+        return acc, acc.close
+    acc = DictAccessor(src, **dict_opts)
+    return acc, (lambda: None)
+
+# ---------- Shared helpers ----------
+def _normalize_scale_dtype(scale_dtype: Union[str, torch.dtype]) -> torch.dtype:
+    if isinstance(scale_dtype, torch.dtype):
+        return scale_dtype
+    key = str(scale_dtype).lower()
+    if key not in _DTYPE_ALIASES:
+        raise ValueError(f"scale_dtype must be one of {list(_DTYPE_ALIASES.keys())} or a torch.dtype")
+    return _DTYPE_ALIASES[key]
+
+def _json_to_dict(s: str) -> Optional[Dict]:
+    # Strictly catch JSON decoding only
+    try:
+        return json.loads(s)
+    except json.JSONDecodeError:
+        return None
+
+def _maybe_parse_scale_map(meta: Dict[str, str]) -> Optional[Dict[str, float]]:
+    def try_parse(obj) -> Optional[Dict[str, float]]:
+        if not isinstance(obj, dict):
+            return None
+        out: Dict[str, float] = {}
+        for wk, v in obj.items():
+            if isinstance(v, (int, float)):
+                out[wk] = float(v)
+            elif isinstance(v, dict) and "scale" in v:
+                sc = v["scale"]
+                if isinstance(sc, (int, float)):
+                    out[wk] = float(sc)
+                elif isinstance(sc, (list, tuple)) and len(sc) == 1 and isinstance(sc[0], (int, float)):
+                    out[wk] = float(sc[0])
+        if out:
+            return out
+        for sub in ("weights", "tensors", "params", "map"):
+            subobj = obj.get(sub)
+            if isinstance(subobj, dict):
+                got = try_parse(subobj)
+                if got:
+                    return got
+        return None
+
+    # exact keys first
+    for k in _SCALE_META_KEYS:
+        raw = meta.get(k)
+        if isinstance(raw, str) and raw.startswith("{") and raw.endswith("}"):
+            parsed = _json_to_dict(raw)
+            if parsed:
+                got = try_parse(parsed)
+                if got:
+                    return got
+
+    # loose scan of any JSON-looking value
+    for v in meta.values():
+        if isinstance(v, str) and v.startswith("{") and v.endswith("}"):
+            parsed = _json_to_dict(v)
+            if parsed:
+                got = try_parse(parsed)
+                if got:
+                    return got
+    return None
+
+def _quick_fp8_variant_from_sentinel(acc: Accessor) -> Optional[str]:
+    if "scaled_fp8" in set(acc.keys()):
+        dt = acc.get_tensor("scaled_fp8").dtype
+        if dt == torch.float8_e4m3fn: return "e4m3fn"
+        if dt == torch.float8_e5m2:   return "e5m2"
+    return None
+
+def _per_channel_reshape(vec: torch.Tensor, weight: torch.Tensor) -> torch.Tensor:
+    return vec.view(weight.shape[0], *([1] * (weight.ndim - 1)))
+
+# ---------- Unified converter ----------
+class ConvertResult(Dict[str, object]):
+    @property
+    def state_dict(self) -> Dict[str, torch.Tensor]: return self["state_dict"]  # type: ignore
+    @property
+    def quant_map(self) -> Dict[str, Dict]: return self["quant_map"]            # type: ignore
+    @property
+    def fp8_format(self) -> str: return self["fp8_format"]                      # type: ignore
+    @property
+    def patch_needed(self) -> bool: return self["patch_needed"]                 # type: ignore
+
+def convert_scaled_fp8_to_quanto(
+    src: Union[str, Dict[str, torch.Tensor]],
+    *,
+    fp8_format: Optional[str] = None,            # 'e4m3fn' | 'e5m2' | None (auto)
+    require_scale: bool = False,
+    allow_default_scale: bool = True,
+    default_missing_scale: float = 1.0,
+    dtype: Union[str, torch.dtype] = "float32",
+    add_activation_placeholders: bool = True,
+    # dict mode options
+    sd_metadata: Optional[Dict[str, str]] = None,
+    in_place: bool = False,
+    free_cuda_cache: bool = False,
+    cuda_cache_interval: int = 32,
+) -> ConvertResult:
+    sd_scale_dtype = _normalize_scale_dtype(dtype)
+    patch_needed = (sd_scale_dtype == torch.float32)
+
+    acc, closer = _as_accessor(
+        src,
+        meta=sd_metadata,
+        in_place=in_place,
+        free_cuda_cache=free_cuda_cache,
+        cuda_cache_interval=cuda_cache_interval,
+    )
+    if not acc.can_delete(): in_place = False
+    try:
+        meta = acc.metadata() or {}
+        meta_scale_map = _maybe_parse_scale_map(meta) or {}
+
+        keys = list(acc.keys())
+
+        # FP8 variant: sentinel -> first FP8 weight -> 'auto'
+        fmt = fp8_format or _quick_fp8_variant_from_sentinel(acc)
+        if fmt is None:
+            for wk in keys:
+                if not _is_weight_key(wk): continue
+                dt = acc.get_tensor(wk).dtype
+                if dt == torch.float8_e4m3fn: fmt = "e4m3fn"; break
+                if dt == torch.float8_e5m2:   fmt = "e5m2";   break
+        if fmt is None: fmt = "auto"
+
+        # Map '<base>.scale_weight' -> '<base>.weight'
+        scale_weight_map: Dict[str, str] = {}
+        for sk in keys:
+            if sk.endswith(".scale_weight"):
+                base = sk[: -len(".scale_weight")]
+                wk = base + ".weight"
+                if wk in keys:
+                    scale_weight_map[wk] = sk
+
+        def get_scale_vec_for_weight(wk: str, out_ch: int) -> Optional[torch.Tensor]:
+            # 1) explicit tensor
+            sk = scale_weight_map.get(wk)
+            if sk is not None:
+                s_t = acc.get_tensor(sk).to(torch.float32)
+                if in_place: acc.delete(s_t)
+                if s_t.numel() == 1:
+                    return torch.full((out_ch,), float(s_t.item()), dtype=torch.float32)
+                if s_t.numel() == out_ch:
+                    return s_t.reshape(out_ch)
+                if torch.numel(s_t.unique()) == 1:
+                    return torch.full((out_ch,), float(s_t.view(-1)[0].item()), dtype=torch.float32)
+                raise ValueError(f"Unexpected scale length for '{wk}': {s_t.numel()} (out_ch={out_ch})")
+            # 2) metadata exact / normalized
+            if wk in meta_scale_map:
+                return torch.full((out_ch,), float(meta_scale_map[wk]), dtype=torch.float32)
+            for alt in (wk.replace("model.", ""), re.sub(r"(^|\.)weight$", "", wk)):
+                if alt in meta_scale_map:
+                    return torch.full((out_ch,), float(meta_scale_map[alt]), dtype=torch.float32)
+            return None
+
+        out_sd: Dict[str, torch.Tensor] = {}
+        qmap: Dict[str, Dict] = {}
+
+        # Single pass: rewrite FP8 weights, copy-through others
+        for k in keys:
+            # Drop source-only artifacts
+            if k == "scaled_fp8" or k.endswith(".scale_weight") :
+                continue
+
+            t = acc.get_tensor(k)
+            if in_place: acc.delete(k)
+            if _is_weight_key(k) and t.dtype in (torch.float8_e4m3fn, torch.float8_e5m2):
+                # Quantized: keep original FP8 tensor as _data
+                out_sd[k + DATA_SUFFIX] = t
+
+                out_ch = int(t.shape[0])
+                s_vec = get_scale_vec_for_weight(k, out_ch)
+                if s_vec is None:
+                    if require_scale and not allow_default_scale:
+                        raise KeyError(f"No scale found for '{k}' (looked for '.scale_weight' and metadata).")
+                    s_vec = torch.full((out_ch,), float(default_missing_scale), dtype=torch.float32)
+
+                s_grid = _per_channel_reshape(s_vec, t).to(sd_scale_dtype)
+                out_sd[k + SCALE_SUFFIX] = s_grid
+
+                if add_activation_placeholders:
+                    base = k[:-len(".weight")]
+                    out_sd[base + IN_SCALE]  = torch.tensor([1], dtype=sd_scale_dtype)
+                    out_sd[base + OUT_SCALE] = torch.tensor([1], dtype=sd_scale_dtype)
+
+                base = k[:-len(".weight")]
+                qmap[base] = {"weights": _QTYPE_NAME[fmt], "activations": "none"}
+            else:
+                out_sd[k] =  t if t.dtype == dtype or t.dtype == torch.float32 else t.to(dtype)
+            t = None
+        return ConvertResult(state_dict=out_sd, quant_map=qmap, fp8_format=fmt, patch_needed=patch_needed)
+    finally:
+        closer()
+
+def detect_safetensors_format(
+    src: Union[str, Dict[str, torch.Tensor]],
+    *,
+    sd_metadata: Optional[Dict[str, str]] = None,
+    probe_weights: bool = False,   # if True, we may read up to 2 weights total
+    with_hints: bool = False,
+) -> Dict[str, str]:
+    """
+    Returns:
+      {
+        'kind': 'quanto' | 'scaled_fp8' | 'fp8' | 'none',
+        'quant_format': 'qfloat8_e4m3fn' | 'qfloat8_e5m2' | 'qfloat8' | 'qint8' | 'qint4' | 'unknown' | '',
+        'fp8_format': 'e4m3fn' | 'e5m2' | 'unknown' | '',
+        'hint': '...'  # only when with_hints=True
+      }
+    """
+    acc, closer = _as_accessor(src, meta=sd_metadata, in_place=False)
+    try:
+        # --- O(1) sentinel test up-front (no key scan) ---
+        if acc.has("scaled_fp8"):
+            dt = acc.get_tensor("scaled_fp8").dtype
+            fp8_fmt = "e4m3fn" if dt == torch.float8_e4m3fn else ("e5m2" if dt == torch.float8_e5m2 else "unknown")
+            out = {"kind": "scaled_fp8", "quant_format": "", "fp8_format": fp8_fmt}
+            if with_hints: out["hint"] = "sentinel"
+            return out
+
+        # --- Single pass over keys (no re-scans) ---
+        ks = list(acc.keys())
+        has_scale_weight = False
+        saw_quanto_data = False
+        fp8_variant = None
+        fp8_probe_budget = 2 if probe_weights else 1
+
+        for k in ks:
+            # Quanto pack short-circuit
+            if not saw_quanto_data and k.endswith(DATA_SUFFIX):
+                saw_quanto_data = True
+                # we can break here, but keep minimal state setting uniformity
+                break
+
+        if saw_quanto_data:
+            out = {"kind": "quanto", "quant_format": "qfloat8", "fp8_format": ""}
+            if with_hints: out["hint"] = "keys:*._data"
+            return out
+
+        # continue single pass for the rest (scale keys + bounded dtype probe)
+        for k in ks:
+            if not has_scale_weight and k.endswith(".scale_weight"):
+                has_scale_weight = True
+                # don't return yet; we may still probe a dtype to grab variant
+
+            if fp8_probe_budget > 0 and _is_weight_key(k):
+                dt = acc.get_tensor(k).dtype
+                if dt == torch.float8_e4m3fn:
+                    fp8_variant = "e4m3fn"; fp8_probe_budget -= 1
+                elif dt == torch.float8_e5m2:
+                    fp8_variant = "e5m2";   fp8_probe_budget -= 1
+
+        if has_scale_weight:
+            out = {"kind": "scaled_fp8", "quant_format": "", "fp8_format": fp8_variant or "unknown"}
+            if with_hints: out["hint"] = "scale_weight keys"
+            return out
+
+        if fp8_variant is not None:
+            out = {"kind": "fp8", "quant_format": "", "fp8_format": fp8_variant}
+            if with_hints: out["hint"] = "weight dtype (plain fp8)"
+            return out
+
+        # --- Cheap metadata peek only if keys didn't decide it (no JSON parsing) ---
+        meta = acc.metadata() or {}
+        blob = " ".join(v for v in meta.values() if isinstance(v, str)).lower()
+
+        # scaled-fp8 hinted by metadata only
+        has_scale_map = (
+            any(k in meta for k in _SCALE_META_KEYS) or
+            (("scale" in blob) and (("fp8" in blob) or ("float8" in blob)))
+        )
+        if has_scale_map:
+            fmt = "e4m3fn" if "e4m3" in blob else ("e5m2" if "e5m2" in blob else "unknown")
+            out = {"kind": "scaled_fp8", "quant_format": "", "fp8_format": fmt}
+            if with_hints: out["hint"] = "metadata"
+            return out
+
+        # quanto hinted by metadata only (not decisive without keys)
+        qtype_hint = ""
+        for tok in ("qfloat8_e4m3fn", "qfloat8_e5m2", "qfloat8", "qint8", "qint4"):
+            if tok in blob:
+                qtype_hint = tok
+                break
+
+        out = {"kind": "none", "quant_format": qtype_hint, "fp8_format": ""}
+        if with_hints: out["hint"] = "no decisive keys"
+        return out
+
+    finally:
+        closer()
+
+# ---------- Optional Quanto runtime patch (FP32-scale support), enable/disable ----------
+_patch_state = SimpleNamespace(enabled=False, orig=None, scale_index=None)
+
+def enable_fp8_fp32_scale_support():
+    """
+    Version-robust wrapper for WeightQBytesTensor.create:
+      - matches both positional/keyword call styles via *args/**kwargs,
+      - for FP8 + FP32 scales, expands scalar/uniform scales with a VIEW to the needed length,
+      - leaves bf16/fp16 (classic Quanto) untouched.
+    Enable only if you emitted float32 scales.
+    """
+    if _patch_state.enabled:
+        return True
+
+    from optimum.quanto.tensor.weights import qbytes as _qbytes  # late import
+    orig = _qbytes.WeightQBytesTensor.create
+    sig = inspect.signature(orig)
+    params = list(sig.parameters.keys())
+    scale_index = params.index("scale") if "scale" in params else 5  # fallback
+
+    def wrapper(*args, **kwargs):
+        # Extract fields irrespective of signature
+        qtype = kwargs.get("qtype", args[0] if len(args) > 0 else None)
+        axis  = kwargs.get("axis",  args[1] if len(args) > 1 else None)
+        size  = kwargs.get("size",  args[2] if len(args) > 2 else None)
+
+        if "scale" in kwargs:
+            scale = kwargs["scale"]
+            def set_scale(new): kwargs.__setitem__("scale", new)
+        else:
+            scale = args[scale_index] if len(args) > scale_index else None
+            def set_scale(new):
+                nonlocal args
+                args = list(args)
+                if len(args) > scale_index:
+                    args[scale_index] = new
+                else:
+                    kwargs["scale"] = new
+                args = tuple(args)
+
+        is_fp8 = isinstance(qtype, str) and ("float8" in qtype.lower() or "qfloat8" in qtype.lower()) or \
+                 (not isinstance(qtype, str) and "float8" in str(qtype).lower())
+
+        if is_fp8 and isinstance(scale, torch.Tensor) and scale.dtype == torch.float32:
+            need = int(size[axis]) if (isinstance(size, (tuple, list)) and axis is not None and axis >= 0) else None
+            if need is not None:
+                if scale.numel() == 1:
+                    scale = scale.view(1).expand(need, *scale.shape[1:])
+                elif scale.shape[0] != need:
+                    # Expand if uniform; otherwise raise
+                    uniform = (scale.numel() == 1) or (torch.numel(scale.unique()) == 1)
+                    if uniform:
+                        scale = scale.reshape(1, *scale.shape[1:]).expand(need, *scale.shape[1:])
+                    else:
+                        raise ValueError(f"Scale leading dim {scale.shape[0]} != required {need}")
+            set_scale(scale)
+
+        return orig(*args, **kwargs)
+
+    _qbytes.WeightQBytesTensor.create = wrapper
+    _patch_state.enabled = True
+    _patch_state.orig = orig
+    _patch_state.scale_index = scale_index
+    return True
+
+def disable_fp8_fp32_scale_support():
+    """Restore Quanto's original factory."""
+    if not _patch_state.enabled:
+        return False
+    from optimum.quanto.tensor.weights import qbytes as _qbytes
+    _qbytes.WeightQBytesTensor.create = _patch_state.orig
+    _patch_state.enabled = False
+    _patch_state.orig = None
+    _patch_state.scale_index = None
+    return True
+
+# ---------- Tiny CLI (optional) ----------
+def _cli():
+    import argparse, json as _json
+    p = argparse.ArgumentParser("fp8_quanto_bridge")
+    sub = p.add_subparsers(dest="cmd", required=True)
+
+    p_conv = sub.add_parser("convert", help="Convert scaled-FP8 (file) to Quanto artifacts.")
+    p_conv.add_argument("in_path")
+    p_conv.add_argument("out_weights")
+    p_conv.add_argument("out_qmap")
+    p_conv.add_argument("--fp8-format", choices=("e4m3fn", "e5m2"), default=None)
+    p_conv.add_argument("--scale-dtype", default="float32",
+                        choices=("float32","bfloat16","float16","fp32","bf16","fp16","half"))
+    p_conv.add_argument("--no-activation-placeholders", action="store_true")
+    p_conv.add_argument("--default-missing-scale", type=float, default=1.0)
+
+    p_det = sub.add_parser("detect", help="Detect format quickly (path).")
+    p_det.add_argument("path")
+    p_det.add_argument("--probe", action="store_true")
+    p_det.add_argument("--hints", action="store_true")
+
+    p_patch = sub.add_parser("patch", help="Enable/disable FP32-scale runtime patch.")
+    p_patch.add_argument("mode", choices=("enable","disable"))
+
+    args = p.parse_args()
+
+    if args.cmd == "convert":
+        res = convert_scaled_fp8_to_quanto(
+            args.in_path,
+            fp8_format=args.fp8_format,
+            dtype=args.scale_dtype,
+            add_activation_placeholders=not args.no_activation_placeholders,
+            default_missing_scale=args.default_missing_scale,
+        )
+        save_file(res.state_dict, args.out_weights)
+        with open(args.out_qmap, "w") as f:
+            _json.dump(res.quant_map, f)
+        print(f"Wrote: {args.out_weights} and {args.out_qmap}. Patch needed: {res.patch_needed}")
+        return 0
+
+    if args.cmd == "detect":
+        info = detect_safetensors_format(args.path, probe_weights=args.probe, with_hints=args.hints)
+        print(info); return 0
+
+    if args.cmd == "patch":
+        ok = enable_fp8_fp32_scale_support() if args.mode == "enable" else disable_fp8_fp32_scale_support()
+        print(f"patch {args.mode}: {'ok' if ok else 'already in that state'}")
+        return 0
+
+if __name__ == "__main__":
+    raise SystemExit(_cli())