From 62b006f6c4f249a5d5ed86fddb4e237329aaca5a Mon Sep 17 00:00:00 2001 From: Brian Bell Date: Fri, 24 Jul 2026 23:32:48 -0500 Subject: [PATCH] feat: FP8 and NVFP4 checkpoint support (dequantize, surgery in float, BF16 output) 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). --- README.md | 9 + obliteratus/abliterate.py | 77 +++- obliteratus/cli.py | 12 +- obliteratus/models/loader.py | 52 ++- obliteratus/models/quant_dequant.py | 682 ++++++++++++++++++++++++++++ scripts/smoke_load_quant.py | 55 +++ 6 files changed, 879 insertions(+), 8 deletions(-) create mode 100644 obliteratus/models/quant_dequant.py create mode 100644 scripts/smoke_load_quant.py diff --git a/README.md b/README.md index a61214d..3667a60 100644 --- a/README.md +++ b/README.md @@ -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: diff --git a/obliteratus/abliterate.py b/obliteratus/abliterate.py index 3a6a212..a547f95 100644 --- a/obliteratus/abliterate.py +++ b/obliteratus/abliterate.py @@ -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 diff --git a/obliteratus/cli.py b/obliteratus/cli.py index 333dbea..c72fc88 100644 --- a/obliteratus/cli.py +++ b/obliteratus/cli.py @@ -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), diff --git a/obliteratus/models/loader.py b/obliteratus/models/loader.py index c623e5e..9341f4e 100644 --- a/obliteratus/models/loader.py +++ b/obliteratus/models/loader.py @@ -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, diff --git a/obliteratus/models/quant_dequant.py b/obliteratus/models/quant_dequant.py new file mode 100644 index 0000000..a82bb0e --- /dev/null +++ b/obliteratus/models/quant_dequant.py @@ -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"``): + ``.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): 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 diff --git a/scripts/smoke_load_quant.py b/scripts/smoke_load_quant.py new file mode 100644 index 0000000..69a0c2b --- /dev/null +++ b/scripts/smoke_load_quant.py @@ -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")