mirror of
https://github.com/elder-plinius/OBLITERATUS.git
synced 2026-08-17 16:37:30 +02:00
804 lines
30 KiB
Python
804 lines
30 KiB
Python
"""Dequantize FP8 and NVFP4 checkpoints to plain float weights.
|
|
|
|
OBLITERATUS performs weight surgery in float space. Checkpoints stored in
|
|
FP8 (DeepSeek block-wise or compressed-tensors per-channel) or NVFP4
|
|
(ModelOpt or compressed-tensors) are detected at load time and dequantized
|
|
to the requested float dtype before the model is handed to the pipeline.
|
|
Output is always saved as plain float (BF16 by default) — re-quantization
|
|
is deliberately out of scope.
|
|
|
|
Pure torch + safetensors; no new dependencies.
|
|
|
|
Layout conventions handled
|
|
--------------------------
|
|
FP8 block-wise (DeepSeek-style, ``quant_method: "fp8"``):
|
|
``<name>.weight`` float8_e4m3fn, shape (M, N)
|
|
``<name>.weight_scale_inv`` float32, shape (ceil(M/128), ceil(N/128))
|
|
dequant: w * scale_inv broadcast over ``weight_block_size`` blocks.
|
|
|
|
FP8 per-channel (compressed-tensors, ``num_bits: 8, type: "float"``):
|
|
``<name>.weight`` float8_e4m3fn, shape (M, N)
|
|
``<name>.weight_scale`` float32, shape (M, 1) or scalar
|
|
dequant: w * weight_scale.
|
|
|
|
NVFP4 (ModelOpt, ``quant_algo: "NVFP4"``):
|
|
``<name>.weight`` uint8, shape (M, N/2) — two E2M1 nibbles per
|
|
byte along the input dim, low nibble first
|
|
``<name>.weight_scale`` float8_e4m3fn, shape (M, N/16) — one scale
|
|
per 16-element group
|
|
``<name>.weight_scale_2`` float32 scalar — global scale (amax/2688)
|
|
dequant: e2m1_values * weight_scale * weight_scale_2
|
|
|
|
NVFP4 (compressed-tensors, ``format: "nvfp4-pack-quantized"``):
|
|
``<name>.weight_packed`` uint8 packed E2M1 values
|
|
``<name>.weight_scale`` FP8 block scales (multiplicative)
|
|
``<name>.weight_global_scale`` FP32 reciprocal global scale
|
|
dequant: e2m1_values * weight_scale / weight_global_scale.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import logging
|
|
import os
|
|
from dataclasses import dataclass, field
|
|
from enum import Enum
|
|
from typing import Dict, Optional, Tuple
|
|
|
|
import torch
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Dtypes / constants
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def _fp8_dtypes() -> frozenset:
|
|
"""torch float8 dtypes available in this build (torch >= 2.1)."""
|
|
out = set()
|
|
for name in ("float8_e4m3fn", "float8_e5m2", "float8_e4m3fnuz", "float8_e5m2fnuz"):
|
|
dt = getattr(torch, name, None)
|
|
if dt is not None:
|
|
out.add(dt)
|
|
return frozenset(out)
|
|
|
|
|
|
FP8_DTYPES = _fp8_dtypes()
|
|
|
|
# E2M1 magnitude values indexed by the low 3 bits; bit 3 is the sign.
|
|
E2M1_POSITIVE = (0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0)
|
|
E2M1_LUT = torch.tensor(
|
|
[v for v in E2M1_POSITIVE] + [-v for v in E2M1_POSITIVE],
|
|
dtype=torch.float32,
|
|
)
|
|
|
|
NVFP4_GROUP_SIZE = 16
|
|
FP8_DEFAULT_BLOCK = (128, 128)
|
|
|
|
|
|
def is_fp8_dtype(dtype: torch.dtype) -> bool:
|
|
return dtype in FP8_DTYPES
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Scheme detection
|
|
# ---------------------------------------------------------------------------
|
|
|
|
class QuantScheme(Enum):
|
|
NONE = "none"
|
|
FP8_BLOCKWISE = "fp8_blockwise"
|
|
FP8_PER_CHANNEL_CT = "fp8_per_channel_compressed_tensors"
|
|
NVFP4_MODELOPT = "nvfp4_modelopt"
|
|
NVFP4_CT = "nvfp4_compressed_tensors"
|
|
UNSUPPORTED = "unsupported"
|
|
|
|
|
|
@dataclass
|
|
class QuantDetection:
|
|
scheme: QuantScheme
|
|
reason: str = ""
|
|
block_size: Tuple[int, int] = FP8_DEFAULT_BLOCK
|
|
group_size: int = NVFP4_GROUP_SIZE
|
|
scale_is_inverse: bool = False
|
|
global_scale_is_inverse: bool = False
|
|
raw_quant_config: dict = field(default_factory=dict)
|
|
|
|
|
|
def _load_json_from_checkpoint(
|
|
model_name_or_path: str,
|
|
filename: str,
|
|
token: Optional[str] = None,
|
|
revision: Optional[str] = None,
|
|
local_files_only: bool = False,
|
|
) -> Optional[dict]:
|
|
"""Read a JSON metadata file from a local dir or the HF hub."""
|
|
if os.path.isdir(model_name_or_path):
|
|
path = os.path.join(model_name_or_path, filename)
|
|
if not os.path.exists(path):
|
|
return None
|
|
with open(path, "r", encoding="utf-8") as fh:
|
|
return json.load(fh)
|
|
try:
|
|
from huggingface_hub import hf_hub_download
|
|
from huggingface_hub.utils import EntryNotFoundError
|
|
except Exception: # pragma: no cover - huggingface_hub always ships w/ transformers
|
|
return None
|
|
try:
|
|
path = hf_hub_download(
|
|
model_name_or_path,
|
|
filename,
|
|
token=token,
|
|
revision=revision,
|
|
local_files_only=local_files_only,
|
|
)
|
|
except EntryNotFoundError:
|
|
return None
|
|
except Exception as exc:
|
|
logger.debug("could not fetch %s for %s: %s", filename, model_name_or_path, exc)
|
|
return None
|
|
with open(path, "r", encoding="utf-8") as fh:
|
|
return json.load(fh)
|
|
|
|
|
|
def _safetensors_key_names(
|
|
model_name_or_path: str,
|
|
config_json: Optional[dict],
|
|
token: Optional[str] = None,
|
|
revision: Optional[str] = None,
|
|
local_files_only: bool = False,
|
|
) -> set:
|
|
"""Collect tensor key names from the safetensors index (no weight loads)."""
|
|
keys: set = set()
|
|
index = _load_json_from_checkpoint(
|
|
model_name_or_path,
|
|
"model.safetensors.index.json",
|
|
token=token,
|
|
revision=revision,
|
|
local_files_only=local_files_only,
|
|
)
|
|
if index and "weight_map" in index:
|
|
keys.update(index["weight_map"].keys())
|
|
return keys
|
|
# Single-file checkpoint: open headers only.
|
|
if os.path.isdir(model_name_or_path):
|
|
st_path = os.path.join(model_name_or_path, "model.safetensors")
|
|
else:
|
|
try:
|
|
from huggingface_hub import hf_hub_download
|
|
st_path = hf_hub_download(
|
|
model_name_or_path,
|
|
"model.safetensors",
|
|
token=token,
|
|
revision=revision,
|
|
local_files_only=local_files_only,
|
|
)
|
|
except Exception:
|
|
return keys
|
|
if not os.path.exists(st_path):
|
|
return keys
|
|
try:
|
|
from safetensors import safe_open
|
|
with safe_open(st_path, framework="pt", device="cpu") as fh:
|
|
keys.update(fh.keys())
|
|
except Exception as exc:
|
|
logger.debug("could not read safetensors header %s: %s", st_path, exc)
|
|
return keys
|
|
|
|
|
|
def detect_quant_scheme(
|
|
model_name_or_path: str,
|
|
token: Optional[str] = None,
|
|
revision: Optional[str] = None,
|
|
local_files_only: bool = False,
|
|
) -> QuantDetection:
|
|
"""Classify a checkpoint's quantization without loading any weights.
|
|
|
|
Peeks at ``config.json``'s ``quantization_config`` plus safetensors
|
|
key names. Anything quantized that we do not explicitly support is
|
|
reported as UNSUPPORTED with a human-readable reason — the loader
|
|
turns that into a loud error rather than silent corruption.
|
|
"""
|
|
config_json = _load_json_from_checkpoint(
|
|
model_name_or_path,
|
|
"config.json",
|
|
token=token,
|
|
revision=revision,
|
|
local_files_only=local_files_only,
|
|
) or {}
|
|
qcfg = config_json.get("quantization_config")
|
|
if not qcfg:
|
|
return QuantDetection(QuantScheme.NONE)
|
|
|
|
quant_method = str(qcfg.get("quant_method", "")).lower()
|
|
raw = dict(qcfg)
|
|
|
|
if quant_method == "fp8":
|
|
block = qcfg.get("weight_block_size")
|
|
if block:
|
|
return QuantDetection(
|
|
QuantScheme.FP8_BLOCKWISE,
|
|
block_size=(int(block[0]), int(block[1])),
|
|
raw_quant_config=raw,
|
|
)
|
|
keys = _safetensors_key_names(
|
|
model_name_or_path,
|
|
config_json,
|
|
token=token,
|
|
revision=revision,
|
|
local_files_only=local_files_only,
|
|
)
|
|
if any(k.endswith("weight_scale_inv") for k in keys):
|
|
return QuantDetection(QuantScheme.FP8_BLOCKWISE, raw_quant_config=raw)
|
|
return QuantDetection(
|
|
QuantScheme.UNSUPPORTED,
|
|
reason=(
|
|
f"quant_method 'fp8' without weight_block_size/weight_scale_inv "
|
|
f"(activation scheme {qcfg.get('activation_scheme')!r}); only "
|
|
f"DeepSeek-style block-wise FP8 and compressed-tensors FP8 are "
|
|
f"supported — use a BF16 checkpoint"
|
|
),
|
|
raw_quant_config=raw,
|
|
)
|
|
|
|
if quant_method == "modelopt":
|
|
algo = str(qcfg.get("quant_algo", "")).upper()
|
|
kv = str(qcfg.get("kv_cache_quant_algo", "") or "").upper()
|
|
if "MIXED" in algo:
|
|
# Mixed-precision ModelOpt (e.g. FP8 mixer + NVFP4 experts):
|
|
# inspect config_groups — any 4-bit group means NVFP4 tensors are
|
|
# present; dequantization is per-tensor so FP8 tensors still take
|
|
# the FP8 path automatically.
|
|
bits = {
|
|
(grp or {}).get("weights", {}).get("num_bits")
|
|
for grp in (qcfg.get("config_groups") or {}).values()
|
|
}
|
|
if 4 in bits:
|
|
return QuantDetection(QuantScheme.NVFP4_MODELOPT, raw_quant_config=raw)
|
|
return QuantDetection(QuantScheme.FP8_PER_CHANNEL_CT, raw_quant_config=raw)
|
|
if "NVFP4" in algo or "FP4" in algo:
|
|
return QuantDetection(QuantScheme.NVFP4_MODELOPT, raw_quant_config=raw)
|
|
if "FP8" in algo:
|
|
# ModelOpt FP8 checkpoints are per-tensor; treat like the
|
|
# per-channel path with scalar scales.
|
|
return QuantDetection(QuantScheme.FP8_PER_CHANNEL_CT, raw_quant_config=raw)
|
|
return QuantDetection(
|
|
QuantScheme.UNSUPPORTED,
|
|
reason=f"modelopt quant_algo {algo or kv or 'unknown'!r} not supported",
|
|
raw_quant_config=raw,
|
|
)
|
|
|
|
if quant_method == "compressed-tensors":
|
|
groups = qcfg.get("config_groups") or {}
|
|
wcfg = {}
|
|
for grp in groups.values():
|
|
w = (grp or {}).get("weights") or {}
|
|
if w.get("num_bits") is not None:
|
|
wcfg = w
|
|
break
|
|
num_bits = wcfg.get("num_bits")
|
|
wtype = str(wcfg.get("type", "")).lower()
|
|
if num_bits == 8 and wtype == "float":
|
|
return QuantDetection(QuantScheme.FP8_PER_CHANNEL_CT, raw_quant_config=raw)
|
|
if (
|
|
num_bits == 4
|
|
and wtype == "float"
|
|
and qcfg.get("format") == "nvfp4-pack-quantized"
|
|
):
|
|
gs = int(wcfg.get("group_size") or NVFP4_GROUP_SIZE)
|
|
return QuantDetection(
|
|
QuantScheme.NVFP4_CT,
|
|
group_size=gs,
|
|
global_scale_is_inverse=True,
|
|
raw_quant_config=raw,
|
|
)
|
|
return QuantDetection(
|
|
QuantScheme.UNSUPPORTED,
|
|
reason=(
|
|
f"compressed-tensors weights num_bits={num_bits} type={wtype!r} "
|
|
f"format={qcfg.get('format')!r} not supported (need float 8-bit "
|
|
f"or NVFP4 4-bit with format 'nvfp4-pack-quantized')"
|
|
),
|
|
raw_quant_config=raw,
|
|
)
|
|
|
|
if quant_method in ("gptq", "awq", "bitsandbytes", "bitsandbytes_4bit",
|
|
"bitsandbytes_8bit", "mxfp4"):
|
|
# Handled elsewhere in the loader / surgery layer.
|
|
return QuantDetection(QuantScheme.NONE, raw_quant_config=raw)
|
|
|
|
return QuantDetection(
|
|
QuantScheme.UNSUPPORTED,
|
|
reason=f"quant_method {quant_method!r} not supported — use a BF16 checkpoint",
|
|
raw_quant_config=raw,
|
|
)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# FP8 dequantization
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def _upcast_fp8(t: torch.Tensor) -> torch.Tensor:
|
|
"""float8 (or uint8-viewed-as-float8) tensor → float32."""
|
|
if t.dtype in FP8_DTYPES:
|
|
return t.to(torch.float32)
|
|
if t.dtype == torch.uint8 and getattr(torch, "float8_e4m3fn", None) is not None:
|
|
return t.view(torch.float8_e4m3fn).to(torch.float32)
|
|
raise RuntimeError(f"cannot upcast dtype {t.dtype} as FP8")
|
|
|
|
|
|
def dequantize_fp8_blockwise(
|
|
w_fp8: torch.Tensor,
|
|
scale_inv: torch.Tensor,
|
|
block_size: Tuple[int, int] = FP8_DEFAULT_BLOCK,
|
|
) -> torch.Tensor:
|
|
"""DeepSeek-style block-wise FP8 → float32.
|
|
|
|
``w_fp8``: (M, N) float8. ``scale_inv``: (ceil(M/bm), ceil(N/bn)) — one
|
|
scale per ``block_size`` tile, applied multiplicatively.
|
|
"""
|
|
w = _upcast_fp8(w_fp8)
|
|
bm, bn = int(block_size[0]), int(block_size[1])
|
|
M, N = w.shape[-2], w.shape[-1]
|
|
s = scale_inv.to(torch.float32)
|
|
if s.shape[-2] * bm < M or s.shape[-1] * bn < N:
|
|
raise RuntimeError(
|
|
f"weight_scale_inv shape {tuple(s.shape)} incompatible with "
|
|
f"weight shape {(M, N)} and block {block_size}"
|
|
)
|
|
s = s.repeat_interleave(bm, dim=-2).repeat_interleave(bn, dim=-1)
|
|
s = s[..., :M, :N]
|
|
w.mul_(s) # in-place: float32 weights are already 2x the BF16 size
|
|
return w
|
|
|
|
|
|
def dequantize_fp8_per_channel(
|
|
w_fp8: torch.Tensor,
|
|
scale: torch.Tensor,
|
|
scale_is_inverse: bool = False,
|
|
) -> torch.Tensor:
|
|
"""Per-channel (or per-tensor) FP8 → float32."""
|
|
w = _upcast_fp8(w_fp8)
|
|
s = scale.to(torch.float32)
|
|
while s.ndim < w.ndim:
|
|
s = s.unsqueeze(-1)
|
|
if scale_is_inverse:
|
|
w.div_(s)
|
|
else:
|
|
w.mul_(s)
|
|
return w
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# NVFP4 dequantization
|
|
# ---------------------------------------------------------------------------
|
|
|
|
_NATIVE_FP4_OK: Optional[bool] = None
|
|
_NATIVE_FP4_LOW_FIRST: bool = True
|
|
|
|
|
|
def _native_fp4_upcast_works() -> bool:
|
|
"""Probe: does this torch build upcast float4_e2m1fn_x2 correctly?
|
|
|
|
Op support for the packed-FP4 dtype is spotty across torch versions,
|
|
so we verify with a known byte pattern instead of trusting a version
|
|
check. Byte 0x1B = low nibble 0xB (-1.5), high nibble 0x1 (+0.5).
|
|
Also detects the nibble order the native path uses.
|
|
"""
|
|
global _NATIVE_FP4_OK, _NATIVE_FP4_LOW_FIRST
|
|
if _NATIVE_FP4_OK is not None:
|
|
return _NATIVE_FP4_OK
|
|
_NATIVE_FP4_OK = False
|
|
dt = getattr(torch, "float4_e2m1fn_x2", None)
|
|
if dt is None:
|
|
return False
|
|
try:
|
|
packed = torch.tensor([0x1B], dtype=torch.uint8)
|
|
vals = packed.view(dt).to(torch.float32).flatten()
|
|
if vals.numel() != 2:
|
|
return False
|
|
got = vals.tolist()
|
|
# ModelOpt convention is low-nibble-first: byte 0x1B → [-1.5, +0.5].
|
|
if got == [-1.5, 0.5]:
|
|
_NATIVE_FP4_OK, _NATIVE_FP4_LOW_FIRST = True, True
|
|
elif got == [0.5, -1.5]:
|
|
_NATIVE_FP4_OK, _NATIVE_FP4_LOW_FIRST = True, False
|
|
except Exception:
|
|
_NATIVE_FP4_OK = False
|
|
return _NATIVE_FP4_OK
|
|
|
|
|
|
def _unpack_e2m1_native(packed: torch.Tensor) -> torch.Tensor:
|
|
"""uint8 (…, K) → float32 (…, 2K) via the native FP4 dtype."""
|
|
dt = torch.float4_e2m1fn_x2
|
|
vals = packed.view(dt).to(torch.float32)
|
|
vals = vals.reshape(*packed.shape[:-1], packed.shape[-1] * 2)
|
|
if not _NATIVE_FP4_LOW_FIRST:
|
|
# Native path emitted high nibble first: swap adjacent pairs.
|
|
vals = vals.reshape(*vals.shape[:-1], -1, 2).flip(-1).reshape(vals.shape)
|
|
return vals
|
|
|
|
|
|
def _unpack_e2m1_manual(packed: torch.Tensor) -> torch.Tensor:
|
|
"""uint8 (…, K) → float32 (…, 2K), low nibble first, via LUT.
|
|
|
|
Memory-conscious: LUT indexing requires int64 indices (8 B/value), so
|
|
the leading dim is processed in chunks to bound the transient index
|
|
tensors — indexing a whole 30B-model shard naively peaks at >100 GB.
|
|
"""
|
|
lut = E2M1_LUT.to(device=packed.device)
|
|
out = torch.empty(
|
|
*packed.shape[:-1], packed.shape[-1] * 2,
|
|
dtype=torch.float32, device=packed.device,
|
|
)
|
|
flat = packed.reshape(-1, packed.shape[-1])
|
|
out_flat = out.reshape(-1, packed.shape[-1] * 2)
|
|
# ~64M nibbles per chunk → ≤0.5 GB transient int64
|
|
chunk = max(1, (64 * 1024 * 1024) // (packed.shape[-1] * 2))
|
|
for i in range(0, flat.shape[0], chunk):
|
|
pk = flat[i:i + chunk]
|
|
dst = out_flat[i:i + chunk]
|
|
dst[:, 0::2] = lut[(pk & 0x0F).long()]
|
|
dst[:, 1::2] = lut[(pk >> 4).long()]
|
|
return out
|
|
|
|
|
|
def unpack_e2m1(packed: torch.Tensor, force_manual: bool = False) -> torch.Tensor:
|
|
"""Unpack NVFP4 nibbles to float32, preferring the native dtype."""
|
|
if not force_manual and _native_fp4_upcast_works():
|
|
try:
|
|
return _unpack_e2m1_native(packed)
|
|
except Exception:
|
|
logger.debug("native FP4 upcast failed at runtime; using manual LUT")
|
|
return _unpack_e2m1_manual(packed)
|
|
|
|
|
|
def _upcast_fp8_scales(t: torch.Tensor) -> torch.Tensor:
|
|
"""Block scales may be stored as float8 or as raw uint8 bytes."""
|
|
if t.dtype in FP8_DTYPES or t.dtype == torch.uint8:
|
|
return _upcast_fp8(t)
|
|
return t.to(torch.float32)
|
|
|
|
|
|
def dequantize_nvfp4(
|
|
packed_uint8: torch.Tensor,
|
|
block_scale: torch.Tensor,
|
|
global_scale: Optional[torch.Tensor],
|
|
out_shape: Optional[Tuple[int, ...]] = None,
|
|
scale_is_inverse: bool = False,
|
|
global_scale_is_inverse: bool = False,
|
|
group_size: int = NVFP4_GROUP_SIZE,
|
|
force_manual: bool = False,
|
|
) -> torch.Tensor:
|
|
"""NVFP4 → float32.
|
|
|
|
``packed_uint8``: (M, N/2) uint8, two E2M1 values per byte, low nibble
|
|
first along the input dim. ``block_scale``: one FP8-E4M3 scale per
|
|
``group_size`` elements. ``global_scale``: FP32 scalar (ModelOpt
|
|
``weight_scale_2`` = amax/2688); None means 1.0. ModelOpt stores both
|
|
scales multiplicatively. Compressed-tensors keeps the block scale but
|
|
stores the global scale as its reciprocal, selected with
|
|
``global_scale_is_inverse``. ``scale_is_inverse`` is retained for
|
|
explicitly reciprocal block-scale layouts.
|
|
"""
|
|
vals = unpack_e2m1(packed_uint8, force_manual=force_manual)
|
|
*lead, N = vals.shape
|
|
if N % group_size != 0:
|
|
raise RuntimeError(
|
|
f"unpacked NVFP4 dim {N} not divisible by group_size {group_size}"
|
|
)
|
|
vals = vals.reshape(*lead, N // group_size, group_size)
|
|
bs = _upcast_fp8_scales(block_scale).reshape(*lead, N // group_size, 1)
|
|
# In-place scaling — full-model float32 copies are 2x the BF16 size each.
|
|
if scale_is_inverse:
|
|
vals.div_(bs)
|
|
else:
|
|
vals.mul_(bs)
|
|
vals = vals.reshape(*lead, N)
|
|
if global_scale is not None:
|
|
gs = global_scale.to(torch.float32)
|
|
if global_scale_is_inverse:
|
|
vals.div_(gs)
|
|
else:
|
|
vals.mul_(gs)
|
|
if out_shape is not None and tuple(vals.shape) != tuple(out_shape):
|
|
vals = vals.reshape(out_shape)
|
|
return vals
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# State-dict level dequantization (used by the loader)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
_SCALE_SUFFIXES = (
|
|
# Dense-linears (``foo.weight`` + ``foo.weight_scale``)
|
|
".weight_scale_inv",
|
|
".weight_scale_2",
|
|
".weight_global_scale",
|
|
".weight_scale",
|
|
".input_scale",
|
|
".input_global_scale",
|
|
".activation_scale",
|
|
# Fused MoE experts (``experts.gate_up_proj`` + ``experts.gate_up_proj_scale``)
|
|
"_scale_inv",
|
|
"_scale_2",
|
|
"_global_scale",
|
|
"_input_scale",
|
|
"_scale",
|
|
)
|
|
|
|
|
|
def _strip_known_suffix(key: str) -> Optional[str]:
|
|
for suf in _SCALE_SUFFIXES:
|
|
if key.endswith(suf):
|
|
return key[: -len(suf)]
|
|
return None
|
|
|
|
|
|
def _scale_sibling_names(weight_key: str) -> tuple[tuple[str, str], ...]:
|
|
"""Return candidate ``(key, role)`` pairs for one quantized weight."""
|
|
if weight_key.endswith(".weight_packed"):
|
|
base = weight_key[: -len(".weight_packed")]
|
|
dotted = True
|
|
elif weight_key.endswith(".weight"):
|
|
base = weight_key[: -len(".weight")]
|
|
dotted = True
|
|
else:
|
|
base = weight_key
|
|
dotted = False
|
|
suffixes = (
|
|
(
|
|
(".weight_scale_inv", "scale_inv"),
|
|
(".weight_scale_2", "global_scale"),
|
|
(".weight_global_scale", "global_scale"),
|
|
(".weight_scale", "block_scale"),
|
|
)
|
|
if dotted
|
|
else (
|
|
# Fused-MoE naming:
|
|
# experts.gate_up_proj{,_scale,_scale_2,_scale_inv}
|
|
("_scale_inv", "scale_inv"),
|
|
("_scale_2", "global_scale"),
|
|
("_global_scale", "global_scale"),
|
|
("_scale", "block_scale"),
|
|
)
|
|
)
|
|
return tuple((base + suffix, role) for suffix, role in suffixes)
|
|
|
|
|
|
def find_scale_siblings(state_dict: Dict[str, torch.Tensor], weight_key: str) -> dict:
|
|
"""Locate the scale tensors belonging to ``weight_key`` (``foo.weight``)."""
|
|
out = {}
|
|
for key, role in _scale_sibling_names(weight_key):
|
|
if key in state_dict and role not in out:
|
|
out[role] = state_dict[key]
|
|
return out
|
|
|
|
|
|
def is_scale_key(
|
|
key: str,
|
|
state_dict: Optional[Dict[str, torch.Tensor]] = None,
|
|
known_keys: Optional[set[str]] = None,
|
|
) -> bool:
|
|
"""True for keys holding quantization scales, not real parameters.
|
|
|
|
With ``state_dict`` given, the key only counts as a scale if its base
|
|
tensor exists in the shard — otherwise a legitimate float parameter
|
|
that merely *ends* in ``_scale`` (e.g. ``logit_scale``) would be
|
|
silently dropped from the dequantized checkpoint.
|
|
"""
|
|
base = _strip_known_suffix(key)
|
|
if base is None:
|
|
return False
|
|
keys = known_keys if known_keys is not None else state_dict
|
|
if keys is not None:
|
|
# Dotted scales pair with dense or packed weights; fused-MoE scales
|
|
# pair directly with ``base``.
|
|
if (
|
|
base not in keys
|
|
and (base + ".weight") not in keys
|
|
and (base + ".weight_packed") not in keys
|
|
):
|
|
return False
|
|
return True
|
|
|
|
|
|
def dequantize_state_dict(
|
|
state_dict: Dict[str, torch.Tensor],
|
|
detection: QuantDetection,
|
|
out_dtype: torch.dtype = torch.bfloat16,
|
|
sibling_tensors: Optional[Dict[str, torch.Tensor]] = None,
|
|
known_keys: Optional[set[str]] = None,
|
|
) -> Dict[str, torch.Tensor]:
|
|
"""Dequantize every quantized weight group in one checkpoint shard.
|
|
|
|
Returns a new dict with float tensors only: scale/aux tensors are
|
|
dropped, quantized weights become ``out_dtype``. Works generically
|
|
over naming conventions (dense linears and fused MoE experts alike)
|
|
because grouping is by key suffix, not by module type.
|
|
"""
|
|
scale_lookup = state_dict
|
|
if sibling_tensors:
|
|
scale_lookup = {**state_dict, **sibling_tensors}
|
|
out: Dict[str, torch.Tensor] = {}
|
|
for key, tensor in state_dict.items():
|
|
if is_scale_key(key, state_dict, known_keys):
|
|
# Scales are consumed together with their weight (or dropped if
|
|
# the weight is absent — e.g. quantization status tensors).
|
|
continue
|
|
sib = find_scale_siblings(scale_lookup, key)
|
|
if is_fp8_dtype(tensor.dtype):
|
|
if "scale_inv" in sib:
|
|
w = dequantize_fp8_blockwise(tensor, sib["scale_inv"], detection.block_size)
|
|
elif "block_scale" in sib:
|
|
w = dequantize_fp8_per_channel(
|
|
tensor, sib["block_scale"], detection.scale_is_inverse
|
|
)
|
|
else:
|
|
raise RuntimeError(
|
|
f"FP8 tensor {key!r} has no weight_scale/weight_scale_inv "
|
|
f"sibling — unsupported layout; report a bug"
|
|
)
|
|
out[key] = w.to(out_dtype)
|
|
elif tensor.dtype == torch.uint8 and ("block_scale" in sib or "scale_inv" in sib):
|
|
gs = sib.get("global_scale")
|
|
scale = sib.get("block_scale", sib.get("scale_inv"))
|
|
if detection.scheme in (QuantScheme.NVFP4_MODELOPT, QuantScheme.NVFP4_CT):
|
|
w = dequantize_nvfp4(
|
|
tensor, scale, gs,
|
|
scale_is_inverse=detection.scale_is_inverse,
|
|
global_scale_is_inverse=detection.global_scale_is_inverse,
|
|
group_size=detection.group_size,
|
|
)
|
|
else:
|
|
raise RuntimeError(
|
|
f"packed uint8 tensor {key!r} in a {detection.scheme.value} "
|
|
f"checkpoint — unsupported layout; report a bug"
|
|
)
|
|
output_key = key[: -len("_packed")] if key.endswith(".weight_packed") else key
|
|
out[output_key] = w.to(out_dtype)
|
|
elif tensor.dtype == torch.uint8 and key.endswith((".weight", ".weight_packed")):
|
|
raise RuntimeError(
|
|
f"uint8 weight {key!r} has no recognizable scale siblings "
|
|
f"(keys present: {[k for k in state_dict if k.startswith(key[:-7])]}) "
|
|
f"— unsupported packed layout; use a BF16 checkpoint"
|
|
)
|
|
else:
|
|
out[key] = tensor
|
|
return out
|
|
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Checkpoint materialization (dequantize an on-disk checkpoint to float)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def materialize_dequantized_checkpoint(
|
|
model_name_or_path: str,
|
|
detection: QuantDetection,
|
|
out_dtype: torch.dtype = torch.bfloat16,
|
|
token: Optional[str] = None,
|
|
revision: Optional[str] = None,
|
|
local_files_only: bool = False,
|
|
) -> Tuple[str, str]:
|
|
"""Write a dequantized float copy of a checkpoint to a temp dir.
|
|
|
|
Returns ``(tmp_dir, source_dir)``. The tmp dir is a complete
|
|
standalone checkpoint (config sans ``quantization_config``, tokenizer,
|
|
safetensors shards + index) that ``from_pretrained`` can load through
|
|
the normal float path — no custom module placement needed. Shards are
|
|
processed one at a time so peak RAM stays near one shard.
|
|
"""
|
|
import shutil
|
|
import tempfile
|
|
|
|
from safetensors import safe_open
|
|
from safetensors.torch import load_file, save_file
|
|
|
|
if os.path.isdir(model_name_or_path):
|
|
src = model_name_or_path
|
|
else:
|
|
from huggingface_hub import snapshot_download
|
|
|
|
src = snapshot_download(
|
|
model_name_or_path,
|
|
token=token,
|
|
revision=revision,
|
|
local_files_only=local_files_only,
|
|
allow_patterns=[
|
|
"*.json", "*.safetensors", "*.py", "tokenizer*", "*.model",
|
|
"*.txt", "chat_template*", "special_tokens_map.json",
|
|
"generation_config.json",
|
|
],
|
|
)
|
|
|
|
tmp = tempfile.mkdtemp(prefix="obliteratus_dequant_")
|
|
logger.warning(
|
|
"Dequantizing %s checkpoint %s -> %s (temporary dir %s)",
|
|
detection.scheme.value, model_name_or_path, out_dtype, tmp,
|
|
)
|
|
|
|
try:
|
|
# Copy non-weight files; strip quantization_config from config.json.
|
|
for name in os.listdir(src):
|
|
s = os.path.join(src, name)
|
|
if not os.path.isfile(s) or name.endswith(".safetensors"):
|
|
continue
|
|
d = os.path.join(tmp, name)
|
|
if name == "config.json":
|
|
with open(s, "r", encoding="utf-8") as fh:
|
|
cfg = json.load(fh)
|
|
cfg.pop("quantization_config", None)
|
|
with open(d, "w", encoding="utf-8") as fh:
|
|
json.dump(cfg, fh, indent=2)
|
|
else:
|
|
shutil.copy2(s, d)
|
|
|
|
# Locate shards.
|
|
index_path = os.path.join(src, "model.safetensors.index.json")
|
|
old_index = None
|
|
if os.path.exists(index_path):
|
|
with open(index_path, "r", encoding="utf-8") as fh:
|
|
old_index = json.load(fh)
|
|
shards = sorted(set(old_index["weight_map"].values()))
|
|
elif os.path.exists(os.path.join(src, "model.safetensors")):
|
|
shards = ["model.safetensors"]
|
|
else:
|
|
raise RuntimeError(
|
|
f"Quantized checkpoint '{model_name_or_path}' has no safetensors "
|
|
f"weights (pytorch_model.bin quantized checkpoints are not "
|
|
f"supported) — use a BF16 checkpoint"
|
|
)
|
|
|
|
weight_map = dict(old_index.get("weight_map") or {}) if old_index else {}
|
|
known_keys = set(weight_map) if weight_map else None
|
|
new_weight_map: Dict[str, str] = {}
|
|
new_total_size = 0
|
|
for i, shard in enumerate(shards, 1):
|
|
sd = load_file(os.path.join(src, shard), device="cpu")
|
|
external_scales: Dict[str, torch.Tensor] = {}
|
|
external_by_shard: Dict[str, list[str]] = {}
|
|
for weight_key in sd:
|
|
for scale_key, _role in _scale_sibling_names(weight_key):
|
|
scale_shard = weight_map.get(scale_key)
|
|
if scale_shard and scale_shard != shard:
|
|
external_by_shard.setdefault(scale_shard, []).append(scale_key)
|
|
for scale_shard, scale_keys in external_by_shard.items():
|
|
with safe_open(
|
|
os.path.join(src, scale_shard), framework="pt", device="cpu",
|
|
) as handle:
|
|
for scale_key in scale_keys:
|
|
external_scales[scale_key] = handle.get_tensor(scale_key)
|
|
|
|
out_sd = dequantize_state_dict(
|
|
sd,
|
|
detection,
|
|
out_dtype=out_dtype,
|
|
sibling_tensors=external_scales,
|
|
known_keys=known_keys,
|
|
)
|
|
if out_sd:
|
|
save_file(out_sd, os.path.join(tmp, shard), metadata={"format": "pt"})
|
|
for k, tensor in out_sd.items():
|
|
new_weight_map[k] = shard
|
|
new_total_size += tensor.numel() * tensor.element_size()
|
|
logger.info(
|
|
"dequantized shard %d/%d (%s): %d tensors -> %d float tensors",
|
|
i, len(shards), shard, len(sd), len(out_sd),
|
|
)
|
|
|
|
if old_index is not None:
|
|
metadata = dict(old_index.get("metadata") or {})
|
|
metadata["total_size"] = new_total_size
|
|
with open(
|
|
os.path.join(tmp, "model.safetensors.index.json"),
|
|
"w",
|
|
encoding="utf-8",
|
|
) as fh:
|
|
json.dump({"metadata": metadata, "weight_map": new_weight_map}, fh, indent=2)
|
|
except BaseException:
|
|
shutil.rmtree(tmp, ignore_errors=True)
|
|
raise
|
|
|
|
return tmp, src
|