"""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"``): ``.weight`` float8_e4m3fn, shape (M, N) ``.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"``): ``.weight`` float8_e4m3fn, shape (M, N) ``.weight_scale`` float32, shape (M, 1) or scalar dequant: w * weight_scale. NVFP4 (ModelOpt, ``quant_algo: "NVFP4"``): ``.weight`` uint8, shape (M, N/2) — two E2M1 nibbles per byte along the input dim, low nibble first ``.weight_scale`` float8_e4m3fn, shape (M, N/16) — one scale per 16-element group ``.weight_scale_2`` float32 scalar — global scale (amax/2688) dequant: e2m1_values * weight_scale * weight_scale_2 NVFP4 (compressed-tensors, ``format: "nvfp4-pack-quantized"``): ``.weight_packed`` uint8 packed E2M1 values ``.weight_scale`` FP8 block scales (multiplicative) ``.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