CalvinBackup/OBLITERATUS
1
0
Fork 0
mirror of https://github.com/elder-plinius/OBLITERATUS.git synced 2026-08-17 16:37:30 +02:00
Code Issues Packages Projects Releases Wiki Activity

feat: FP8 and NVFP4 checkpoint support (dequantize, surgery in float, BF16 output)

Browse Source 
Point obliteratus at an FP8 or NVFP4 checkpoint and it just works:
the loader detects the format from config.json + safetensors metadata
(no weight loads), dequantizes shard-by-shard to a temporary BF16 copy,
then runs the normal float pipeline and saves BF16.

Supported layouts:
- FP8 DeepSeek-style block-wise (weight_scale_inv + weight_block_size)
- FP8 per-channel / per-tensor (compressed-tensors, ModelOpt)
- NVFP4 ModelOpt (uint8 nibbles + FP8 group scales + FP32 global),
  including MIXED_PRECISION checkpoints (FP8 mixer + NVFP4 experts)
- NVFP4 compressed-tensors (reciprocal scales)

Design:
- New pure-torch obliteratus/models/quant_dequant.py; no new deps.
  NVFP4 unpack uses torch.float4_e2m1fn_x2 when a runtime probe proves
  it works, else a chunked nibble LUT (bounds transient int64 index
  memory; a naive implementation OOMed at 96GB on a 30B model).
- Scale keys are dropped only when their base weight exists in the
  same shard, so legitimate params ending in _scale (logit_scale et al.)
  survive.
- Unsupported schemes (fbgemm, quanto, W4A4, ...) fail loudly at load,
  naming the scheme.
- Surgery guards: float8 or packed uint8 reaching _dequantize_weight or
  any fused-MoE path raises RuntimeError instead of silently upcasting
  (bitsandbytes quant_state params are explicitly excluded).
- Save path strips quantization metadata and logs that output is BF16;
  re-quantization for serving is out of scope (llm-compressor/modelopt).
- CLI: new --trust-remote-code flag; help text documents auto-detection.

Validated end-to-end on 1x A100-80GB (see PR description):
Nemotron-3-Nano-Omni-30B NVFP4 (mixed) and FP8, Qwen3-8B-FP8
(block-wise) vs Qwen3-8B BF16 baseline (perplexity 4.23 vs 4.33).
This commit is contained in:
Brian Bell
2026-08-16 12:11:07 -04:00
committed by Joseph Magly
parent 84b97f6620
commit 62b006f6c4
6 changed files with 879 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files
README.md
+9
View File
@@ -479,6 +479,15 @@ obliteratus obliterate meta-llama/Llama-3.1-405B-Instruct \
Quantization roughly halves the GPU count at each step down. A 70B model that needs 3x A100-80GB in bf16 fits on 2 in int8 or 1 in int4.
**FP8 and NVFP4 checkpoints are supported automatically.** No flag needed — the loader detects the format from the checkpoint's `quantization_config`, dequantizes the weights to float (BF16 by default) shard-by-shard, runs the normal pipeline, and saves the output as plain BF16:
| Format | Schemes detected |
|--------|------------------|
| FP8 | DeepSeek-style block-wise (`weight_scale_inv` + `weight_block_size`), compressed-tensors per-channel, ModelOpt FP8 |
| NVFP4 | ModelOpt (`weight` + `weight_scale` + `weight_scale_2`), compressed-tensors NVFP4 |
Two things to know: peak VRAM is the **BF16 size** of the model (not the quantized size), and the output is saved as BF16 — re-quantize afterward with llm-compressor or modelopt if you want a quantized serving artifact. Other quantization schemes (fbgemm, quanto, W4A4) fail loudly with a message naming the scheme.
### GPU calculator
Not sure how many GPUs you need? The `gpu-calc` command estimates the minimum GPU count for any model, accounting for weight memory, activation overhead, and CUDA context:
obliteratus/abliterate.py
+74 -3
View File
@@ -34,6 +34,17 @@ import torch
import torch.nn as nn
from obliteratus import device as dev # noqa: E402 — must import before CUDA setup
from obliteratus.models.quant_dequant import FP8_DTYPES as _FP8_DTYPES
# Module attributes that hold quantization scale tensors; their presence on a
# module means a uint8/float8 ``weight`` is packed quantized data, not a plain
# non-float weight.
_QUANT_SCALE_ATTRS = (
"weight_scale",
"weight_scale_2",
"weight_scale_inv",
"weight_global_scale",
)
# Reduce CUDA memory fragmentation for large models. Must be set before any
# CUDA allocations, so we do it at import time. This is the PyTorch-recommended
@@ -4582,7 +4593,7 @@ class AbliterationPipeline:
return is_quantized_parameter(
class_name=param.__class__.__name__,
has_quant_state=hasattr(param, "quant_state"),
)
) or getattr(param, "dtype", None) in _FP8_DTYPES
@staticmethod
def _dequantize_weight(proj_module) -> tuple[torch.Tensor, bool]:
@@ -4629,8 +4640,22 @@ class AbliterationPipeline:
f"the model in float16/bfloat16 for abliteration."
)
# ── bitsandbytes parameter-level detection ─────────────────
# ── FP8/NVFP4: must have been dequantized by the loader ─────
# Checked before bitsandbytes because _is_quantized_param also
# flags float8 tensors — upcasting here without the scale tensors
# would silently corrupt the model, so fail loudly instead.
weight = proj_module.weight
if weight.data.dtype in _FP8_DTYPES or (
weight.data.dtype == torch.uint8
and any(hasattr(proj_module, a) for a in _QUANT_SCALE_ATTRS)
):
raise RuntimeError(
"FP8/NVFP4 weight reached surgery without dequantization — "
"the loader should have dequantized this checkpoint to float. "
"This is a bug; please report it."
)
# ── bitsandbytes parameter-level detection ─────────────────
if storage_kind == "quantized_parameter":
try:
import bitsandbytes as bnb
@@ -4967,7 +4992,15 @@ class AbliterationPipeline:
if param is None or not isinstance(param, (nn.Parameter, torch.Tensor)):
continue
# Dequantize fused param if necessary
# Dequantize fused param if necessary. FP8 is quantized storage
# the loader should have dequantized — fail loudly, never skip.
if param.data.dtype in _FP8_DTYPES:
raise RuntimeError(
f"FP8 fused-expert weight '{name}' reached surgery "
f"without dequantization — the loader should have "
f"dequantized this checkpoint to float. This is a bug; "
f"please report it."
)
is_quantized = AbliterationPipeline._is_quantized_param(param)
if is_quantized:
try:
@@ -4986,6 +5019,15 @@ class AbliterationPipeline:
continue
else:
data = param.data
# FP8/packed fused params are quantized storage — the loader
# dequantizes them; raw upcasting here would corrupt scales.
if data.dtype in _FP8_DTYPES or data.dtype == torch.uint8:
raise RuntimeError(
f"FP8/NVFP4 fused-expert weight '{name}' reached surgery "
f"without dequantization — the loader should have "
f"dequantized this checkpoint to float. This is a bug; "
f"please report it."
)
# Non-float (e.g. uint8) fused params need float conversion
if not data.is_floating_point():
data = data.float()
@@ -5569,6 +5611,14 @@ class AbliterationPipeline:
if data.shape[-1] != hidden_dim and data.shape[-2] != hidden_dim:
continue
if (data.dtype in _FP8_DTYPES or data.dtype == torch.uint8) \
and not hasattr(param, "quant_state"):
raise RuntimeError(
f"FP8/NVFP4 fused-expert weight '{pname}' reached surgery "
f"without dequantization — the loader should have "
f"dequantized this checkpoint to float. This is a bug; "
f"please report it."
)
is_quantized = AbliterationPipeline._is_quantized_param(param)
if is_quantized:
try:
@@ -5682,6 +5732,14 @@ class AbliterationPipeline:
if data.shape[-1] != hidden_dim and data.shape[-2] != hidden_dim:
continue
if (data.dtype in _FP8_DTYPES or data.dtype == torch.uint8) \
and not hasattr(param, "quant_state"):
raise RuntimeError(
f"FP8/NVFP4 fused-expert weight '{pname}' reached surgery "
f"without dequantization — the loader should have "
f"dequantized this checkpoint to float. This is a bug; "
f"please report it."
)
is_quantized = AbliterationPipeline._is_quantized_param(param)
if is_quantized:
try:
@@ -6630,6 +6688,19 @@ class AbliterationPipeline:
self.log("Stripping native quantization config (weights are now float16)")
model.hf_quantizer.remove_quantization_config(model)
# Runs that started from an FP8/NVFP4 checkpoint were dequantized by
# the loader; make sure no quantization metadata survives into the
# saved config (the temp config already strips it — belt and braces).
_deq_scheme = getattr(model, "_obliteratus_dequantized_scheme", None)
if _deq_scheme is not None:
if getattr(model.config, "quantization_config", None) is not None:
del model.config.quantization_config
self.log(
f"Input checkpoint was {_deq_scheme}; output is saved as plain "
f"float weights. To re-quantize for serving, use llm-compressor "
f"or modelopt on the output directory."
)
# Avoid unsupported reverse conversions when saving a new HF-native artifact.
if hasattr(model, "_weight_conversions"):
del model._weight_conversions
obliteratus/cli.py
+10 -2
View File
@@ -288,7 +288,13 @@ def main(argv: list[str] | None = None):
)
p.add_argument(
"--quantization", type=str, default=None, choices=["4bit", "8bit"],
help="Load model with quantization (4bit or 8bit). Requires bitsandbytes.",
help="Load model with quantization (4bit or 8bit). Requires bitsandbytes. "
"FP8/NVFP4 checkpoints are detected and dequantized automatically.",
)
p.add_argument(
"--trust-remote-code", action="store_true", default=False,
help="Trust custom modeling code from the model repo (required for "
"checkpoints with custom architectures, e.g. Nemotron Omni).",
)
p.add_argument(
"--gpu-memory-utilization",
@@ -439,7 +445,8 @@ def main(argv: list[str] | None = None):
tourney_parser.add_argument("--dataset", type=str, default="builtin", help="Dataset source (default: builtin)")
tourney_parser.add_argument(
"--quantization", type=str, default=None, choices=["4bit", "8bit"],
help="Load model with quantization",
help="Load model with quantization (FP8/NVFP4 checkpoints are "
"detected and dequantized automatically)",
)
tourney_parser.add_argument(
"--output-dir",
@@ -1172,6 +1179,7 @@ def _cmd_abliterate(args):
projection_row_fraction=getattr(args, "projection_row_fraction", None),
quantization=args.quantization,
gpu_memory_utilization=getattr(args, "gpu_memory_utilization", None),
trust_remote_code=getattr(args, "trust_remote_code", False),
large_model_mode=getattr(args, "large_model", False),
verify_sample_size=getattr(args, "verify_sample_size", None),
refusal_max_tokens=getattr(args, "refusal_max_tokens", None),
obliteratus/models/loader.py
+49 -3
View File
@@ -12,6 +12,7 @@ import sys as _sys
import torch
from obliteratus import device as dev
from obliteratus.models import quant_dequant as qd
from obliteratus.runtime_contracts import (
effective_model_memory_gb,
quantized_model_fits_gpu,
@@ -657,6 +658,32 @@ def load_model(
f"If this model requires custom code, pass trust_remote_code=True explicitly."
) from e
# FP8 / NVFP4 checkpoints: dequantize to a plain float copy on disk,
# then load through the normal float path. Anything quantized that we
# don't explicitly support fails loudly here, before any weight loads.
quant_detection = qd.detect_quant_scheme(model_name, token=token)
_dequant_tmp = None
_dequant_source = None
if quant_detection.scheme is qd.QuantScheme.UNSUPPORTED:
raise RuntimeError(
f"Unsupported quantization in '{model_name}': {quant_detection.reason}"
)
if quant_detection.scheme is not qd.QuantScheme.NONE:
logger.warning(
"Quantized checkpoint detected (%s) — dequantizing to %s for "
"surgery. Peak memory is the full %s model size; output is "
"saved as %s.",
quant_detection.scheme.value, dtype, dtype, dtype,
)
_dequant_source = model_name
_dequant_tmp, _ = qd.materialize_dequantized_checkpoint(
model_name, quant_detection, out_dtype=torch_dtype, token=token,
)
model_name = _dequant_tmp
config = AutoConfig.from_pretrained(
model_name, trust_remote_code=trust_remote_code, token=token,
)
# Memory estimation and warnings (skip for natively quantized models — estimate is wrong)
native_quant = getattr(config, "quantization_config", None)
load_policy = resolve_model_load_policy(
@@ -834,6 +861,12 @@ def load_model(
f"pip install git+https://github.com/huggingface/transformers.git"
) from e
raise
except Exception:
# Any other load failure — don't leak the temp dequantized checkpoint.
if _dequant_tmp is not None:
import shutil as _shutil
_shutil.rmtree(_dequant_tmp, ignore_errors=True)
raise
if load_policy.move_to_resolved_device:
# Explicit devices and auto-selected MPS/CPU load on CPU before moving.
@@ -903,12 +936,25 @@ def load_model(
# Free accelerator cache after loading
dev.empty_cache()
if _dequant_tmp is not None:
# Weights are fully materialized by from_pretrained at this point —
# the temp dequantized copy is no longer needed.
import shutil as _shutil
model._obliteratus_dequantized_scheme = quant_detection.scheme.value
_shutil.rmtree(_dequant_tmp, ignore_errors=True)
logger.info("Removed temporary dequantized checkpoint %s", _dequant_tmp)
try:
tokenizer = AutoTokenizer.from_pretrained(model_name, **hf_kwargs)
tokenizer = AutoTokenizer.from_pretrained(
_dequant_source or model_name,
**hf_kwargs,
)
except PermissionError:
fallback_cache = os.path.join(tempfile.gettempdir(), "hf_home", "hub")
tokenizer = AutoTokenizer.from_pretrained(
model_name, cache_dir=fallback_cache, **hf_kwargs,
_dequant_source or model_name,
cache_dir=fallback_cache,
**hf_kwargs,
)
if tokenizer.pad_token is None:
tokenizer.pad_token = tokenizer.eos_token
@@ -917,7 +963,7 @@ def load_model(
model=model,
tokenizer=tokenizer,
config=config,
model_name=model_name,
model_name=_dequant_source or model_name,
task=task,
_offload_dir=_offload_dir,
_owns_offload_dir=_owns_offload_dir,
obliteratus/models/quant_dequant.py
+682
View File
@@ -0,0 +1,682 @@
"""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): same layout but scales are stored as
reciprocals and the global scale may be named ``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 # compressed-tensors stores reciprocals
raw_quant_config: dict = field(default_factory=dict)
def _load_json_from_checkpoint(
model_name_or_path: str, filename: str, token: Optional[str] = None,
) -> 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)
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,
) -> 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,
)
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)
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) -> 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) 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)
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":
gs = int(wcfg.get("group_size") or NVFP4_GROUP_SIZE)
return QuantDetection(
QuantScheme.NVFP4_CT,
group_size=gs,
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"not supported (need float 8-bit or NVFP4 4-bit)"
),
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,
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. With
``scale_is_inverse`` (compressed-tensors) the scales are reciprocals
and are divided out instead of multiplied in.
"""
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 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 find_scale_siblings(state_dict: Dict[str, torch.Tensor], weight_key: str) -> dict:
"""Locate the scale tensors belonging to ``weight_key`` (``foo.weight``)."""
base = weight_key[: -len(".weight")] if weight_key.endswith(".weight") else weight_key
dotted = weight_key.endswith(".weight")
out = {}
for name, attr in [
(".weight_scale_inv", "scale_inv"),
(".weight_scale_2", "global_scale"),
(".weight_global_scale", "global_scale"),
(".weight_scale", "block_scale"),
]:
k = base + name
if k in state_dict:
out[attr] = state_dict[k]
if not dotted:
# Fused-MoE naming: experts.gate_up_proj{,_scale,_scale_2,_scale_inv}
for name, attr in [
("_scale_inv", "scale_inv"),
("_scale_2", "global_scale"),
("_global_scale", "global_scale"),
("_scale", "block_scale"),
]:
k = base + name
if k in state_dict and attr not in out:
out[attr] = state_dict[k]
return out
def is_scale_key(key: str, state_dict: Optional[Dict[str, torch.Tensor]] = 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
if state_dict is not None:
# Dotted scales pair with ``base.weight``; fused-MoE scales with ``base``.
if base not in state_dict and (base + ".weight") not in state_dict:
return False
return True
def dequantize_state_dict(
state_dict: Dict[str, torch.Tensor],
detection: QuantDetection,
out_dtype: torch.dtype = torch.bfloat16,
) -> 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.
"""
out: Dict[str, torch.Tensor] = {}
for key, tensor in state_dict.items():
if is_scale_key(key, state_dict):
# Scales are consumed together with their weight (or dropped if
# the weight is absent — e.g. quantization status tensors).
continue
sib = find_scale_siblings(state_dict, 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,
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"
)
out[key] = w.to(out_dtype)
elif tensor.dtype == torch.uint8 and key.endswith(".weight"):
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,
) -> 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.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,
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,
)
# 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"
)
new_weight_map: Dict[str, str] = {}
for i, shard in enumerate(shards, 1):
sd = load_file(os.path.join(src, shard), device="cpu")
out_sd = dequantize_state_dict(sd, detection, out_dtype=out_dtype)
save_file(out_sd, os.path.join(tmp, shard), metadata={"format": "pt"})
for k in out_sd:
new_weight_map[k] = shard
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 {})
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)
return tmp, src
scripts/smoke_load_quant.py
+55
View File
@@ -0,0 +1,55 @@
"""Smoke test: load a quantized checkpoint through the OBLITERATUS loader.
Verifies: detection, dequantization to BF16, no quantized tensors left in
the model, config stripped of quantization_config, and a short generation.
"""
import sys
import torch
from obliteratus.models import quant_dequant as qd
from obliteratus.models.loader import load_model
repo = sys.argv[1]
det = qd.detect_quant_scheme(repo)
print(f"[detection] {det.scheme.value} (inverse={det.scale_is_inverse}, group={det.group_size})")
assert det.scheme not in (qd.QuantScheme.NONE, qd.QuantScheme.UNSUPPORTED), det
handle = load_model(
repo,
task="causal_lm",
device="auto",
dtype="bfloat16",
trust_remote_code=True,
)
model = handle.model
print(f"[load] ok — scheme tag: {getattr(model, '_obliteratus_dequantized_scheme', None)}")
bad = []
n_params = 0
for name, p in model.named_parameters():
n_params += 1
if p.dtype in qd.FP8_DTYPES or p.dtype == torch.uint8:
bad.append((name, str(p.dtype)))
if torch.isnan(p.data).any().item() if p.data.is_floating_point() else False:
bad.append((name, "NaN"))
print(f"[check] {n_params} params, quantized/NaN leftovers: {bad[:10] or 'NONE'}")
assert not bad, bad
assert getattr(model.config, "quantization_config", None) is None, "quantization_config survived"
tok = handle.tokenizer
prompt = "The capital of France is"
try:
inputs = tok(prompt, return_tensors="pt").to(model.device)
except Exception:
# Omni processors: fall back to bare tokenizer encode
ids = tok.encode(prompt, return_tensors="pt").to(model.device)
inputs = {"input_ids": ids}
with torch.no_grad():
out = model.generate(**inputs, max_new_tokens=16, do_sample=False)
text = tok.decode(out[0][-16:] if out.dim() > 1 else out[-16:], skip_special_tokens=True)
print(f"[generate] {text!r}")
print("SMOKE_LOAD_OK")