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).
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
+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")