mmgp 3.6.1__tar.gz → 3.6.3__tar.gz

{mmgp-3.6.1/src/mmgp.egg-info → mmgp-3.6.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: mmgp
-Version: 3.6.1
+Version: 3.6.3
 Summary: Memory Management for the GPU Poor
 Author-email: deepbeepmeep <deepbeepmeep@yahoo.com>
 Requires-Python: >=3.10
@@ -15,7 +15,7 @@ Dynamic: license-file
 
 
 <p align="center">
-  <H2>Memory Management 3.6.1 for the GPU Poor by DeepBeepMeep</H2>	
+  <H2>Memory Management 3.6.3 for the GPU Poor by DeepBeepMeep</H2>	
 </p>
 
 

{mmgp-3.6.1 → mmgp-3.6.3}/README.md

@@ -1,6 +1,6 @@
 
 <p align="center">
-  <H2>Memory Management 3.6.1 for the GPU Poor by DeepBeepMeep</H2>	
+  <H2>Memory Management 3.6.3 for the GPU Poor by DeepBeepMeep</H2>	
 </p>
 
 

{mmgp-3.6.1 → mmgp-3.6.3}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "mmgp"
-version = "3.6.1"
+version = "3.6.3"
 authors = [
   { name = "deepbeepmeep", email = "deepbeepmeep@yahoo.com" },
 ]

mmgp-3.6.3/src/mmgp/fp8_quanto_bridge.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())

{mmgp-3.6.1 → mmgp-3.6.3}/src/mmgp/offload.py

@@ -1,4 +1,4 @@
-# ------------------ Memory Management 3.6.1 for the GPU Poor by DeepBeepMeep (mmgp)------------------
+# ------------------ Memory Management 3.6.3 for the GPU Poor by DeepBeepMeep (mmgp)------------------
 #
 # This module contains multiples optimisations so that models such as Flux (and derived), Mochi, CogView, HunyuanVideo, ...  can run smoothly on a 24 GB GPU limited card. 
 # This a replacement for the accelerate library that should in theory manage offloading, but doesn't work properly with models that are loaded / unloaded several
@@ -71,7 +71,7 @@ import torch
 
 from mmgp import safetensors2
 from mmgp import profile_type
-
+from .fp8_quanto_bridge import convert_scaled_fp8_to_quanto, detect_safetensors_format , enable_fp8_fp32_scale_support
 from optimum.quanto import freeze,  qfloat8, qint4 , qint8, quantize, QModuleMixin, QLinear, QTensor,  quantize_module, register_qmodule
 
 # support for Embedding module quantization that is not supported by default by quanto
@@ -688,7 +688,7 @@ def _welcome():
     if welcome_displayed:
          return 
     welcome_displayed = True
-    print(f"{BOLD}{HEADER}************ Memory Management for the GPU Poor (mmgp 3.6.1) by DeepBeepMeep ************{ENDC}{UNBOLD}")
+    print(f"{BOLD}{HEADER}************ Memory Management for the GPU Poor (mmgp 3.6.3) by DeepBeepMeep ************{ENDC}{UNBOLD}")
 
 def change_dtype(model, new_dtype, exclude_buffers = False):
     for submodule_name, submodule in model.named_modules():  
@@ -1319,7 +1319,7 @@ def move_loras_to_device(model, device="cpu" ):
         if ".lora_" in k:
             m.to(device)
 
-def fast_load_transformers_model(model_path: str,  do_quantize = False, quantizationType =  qint8, pinToMemory = False, partialPinning = False, forcedConfigPath = None, defaultConfigPath = None, modelClass=None, modelPrefix = None, writable_tensors = True, verboseLevel = -1, preprocess_sd  = None, modules = None,  return_shared_modules = None,  configKwargs ={}):
+def fast_load_transformers_model(model_path: str,  do_quantize = False, quantizationType =  qint8, pinToMemory = False, partialPinning = False, forcedConfigPath = None, defaultConfigPath = None, modelClass=None, modelPrefix = None, writable_tensors = True, verboseLevel = -1, preprocess_sd  = None, modules = None,  return_shared_modules = None, default_dtype = torch.bfloat16, ignore_unused_weights = False, configKwargs ={}):
     """
     quick version of .LoadfromPretrained of  the transformers library
     used to build a model and load the corresponding weights (quantized or not)
@@ -1407,13 +1407,13 @@ def fast_load_transformers_model(model_path: str,  do_quantize = False, quantiza
 
     model._config = transformer_config
             
-    load_model_data(model,model_path, do_quantize = do_quantize, quantizationType = quantizationType, pinToMemory= pinToMemory, partialPinning= partialPinning, modelPrefix = modelPrefix, writable_tensors =writable_tensors, preprocess_sd = preprocess_sd , modules = modules, return_shared_modules =  return_shared_modules, verboseLevel=verboseLevel )
+    load_model_data(model,model_path, do_quantize = do_quantize, quantizationType = quantizationType, pinToMemory= pinToMemory, partialPinning= partialPinning, modelPrefix = modelPrefix, writable_tensors =writable_tensors, preprocess_sd = preprocess_sd , modules = modules, return_shared_modules =  return_shared_modules, default_dtype = default_dtype, ignore_unused_weights = ignore_unused_weights, verboseLevel=verboseLevel )
 
     return model
 
 
 
-def load_model_data(model, file_path, do_quantize = False, quantizationType = qint8, pinToMemory = False, partialPinning = False, modelPrefix = None, writable_tensors = True,  preprocess_sd = None, modules = None, return_shared_modules = None, verboseLevel = -1):
+def load_model_data(model, file_path, do_quantize = False, quantizationType = qint8, pinToMemory = False, partialPinning = False, modelPrefix = None, writable_tensors = True,  preprocess_sd = None, postprocess_sd = None, modules = None, return_shared_modules = None, default_dtype = torch.bfloat16, ignore_unused_weights = False, verboseLevel = -1):
     """
     Load a model, detect if it has been previously quantized using quanto and do the extra setup if necessary
     """
@@ -1489,29 +1489,41 @@ def load_model_data(model, file_path, do_quantize = False, quantizationType = qi
                     state_dict.update(sd)
             else:
                 state_dict, metadata = _safetensors_load_file(file, writable_tensors =writable_tensors)
-                
-            if metadata !=  None:
-                quantization_map = metadata.get("quantization_map", None)
-                config = metadata.get("config", None)
-                if config is not None:
-                    model._config = config
-
-                tied_weights_map = metadata.get("tied_weights_map", None)
-                if tied_weights_map != None:
-                    for name, tied_weights_list in tied_weights_map.items():
-                        mapped_weight = state_dict[name]
-                        for tied_weights in tied_weights_list:
-                            state_dict[tied_weights] = mapped_weight
-
-            if quantization_map is None:
-                pos = str.rfind(file, ".")
-                if pos > 0:
-                    quantization_map_path = file[:pos]
-                quantization_map_path += "_map.json"
-
-                if os.path.isfile(quantization_map_path):
-                    with open(quantization_map_path, 'r') as f:
-                        quantization_map = json.load(f)
+
+        if preprocess_sd != None:
+            state_dict = preprocess_sd(state_dict)
+
+        if metadata !=  None:
+            quantization_map = metadata.get("quantization_map", None)
+            config = metadata.get("config", None)
+            if config is not None:
+                model._config = config
+
+            tied_weights_map = metadata.get("tied_weights_map", None)
+            if tied_weights_map != None:
+                for name, tied_weights_list in tied_weights_map.items():
+                    mapped_weight = state_dict[name]
+                    for tied_weights in tied_weights_list:
+                        state_dict[tied_weights] = mapped_weight
+
+        if quantization_map is None:
+            detection_type = detect_safetensors_format(state_dict)
+            if detection_type["kind"] in ['scaled_fp8','fp8']:
+                conv_result = convert_scaled_fp8_to_quanto(state_dict, dtype = default_dtype, in_place= True)
+                state_dict = conv_result["state_dict"]
+                quantization_map = conv_result["quant_map"]
+                conv_result = None
+                # enable_fp8_fp32_scale_support()
+
+        if quantization_map is None:
+            pos = str.rfind(file, ".")
+            if pos > 0:
+                quantization_map_path = file[:pos]
+            quantization_map_path += "_map.json"
+
+            if os.path.isfile(quantization_map_path):
+                with open(quantization_map_path, 'r') as f:
+                    quantization_map = json.load(f)
         
         full_state_dict.update(state_dict)
         if quantization_map != None:
@@ -1530,8 +1542,8 @@ def load_model_data(model, file_path, do_quantize = False, quantizationType = qi
     full_state_dict, full_quantization_map, full_tied_weights_map = None, None, None
 
     # deal if we are trying to load just a sub part of a larger model
-    if preprocess_sd != None:
-        state_dict, quantization_map = preprocess_sd(state_dict, quantization_map)
+    if postprocess_sd != None:
+        state_dict, quantization_map = postprocess_sd(state_dict, quantization_map)
         
     if modelPrefix != None:
         base_model_prefix = modelPrefix + "."
@@ -1562,7 +1574,7 @@ def load_model_data(model, file_path, do_quantize = False, quantizationType = qi
         
     del state_dict
 
-    if len(unexpected_keys) > 0 and verboseLevel >=2:
+    if len(unexpected_keys) > 0 and verboseLevel >=2 and not ignore_unused_weights:
         print(f"Unexpected keys while loading '{file_path}': {unexpected_keys}")
 
     for k,p in model.named_parameters():

{mmgp-3.6.1 → mmgp-3.6.3/src/mmgp.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: mmgp
-Version: 3.6.1
+Version: 3.6.3
 Summary: Memory Management for the GPU Poor
 Author-email: deepbeepmeep <deepbeepmeep@yahoo.com>
 Requires-Python: >=3.10
@@ -15,7 +15,7 @@ Dynamic: license-file
 
 
 <p align="center">
-  <H2>Memory Management 3.6.1 for the GPU Poor by DeepBeepMeep</H2>	
+  <H2>Memory Management 3.6.3 for the GPU Poor by DeepBeepMeep</H2>	
 </p>
 
 

{mmgp-3.6.1 → mmgp-3.6.3}/src/mmgp.egg-info/SOURCES.txt

@@ -3,6 +3,7 @@ README.md
 pyproject.toml
 src/__init__.py
 src/mmgp/__init__.py
+src/mmgp/fp8_quanto_bridge.py
 src/mmgp/offload.py
 src/mmgp/safetensors2.py
 src/mmgp.egg-info/PKG-INFO