From b6a1727cb61754a6a9c4daf2b94a5390b48a44df Mon Sep 17 00:00:00 2001 From: Brian Bell Date: Fri, 24 Jul 2026 23:33:03 -0500 Subject: [PATCH] test+docs: quant_dequant test suite, theory journal engineering log 33 CPU-only tests: FP8 block-wise/per-channel round-trips, NVFP4 round-trips (direct + reciprocal scales, native vs manual unpack agreement), scheme detection incl. ModelOpt MIXED_PRECISION, surgery guard raises, and tiny synthetic FP8/NVFP4 GPT-2 checkpoints through load_model end-to-end. Full suite: 870 passed (8 failures pre-exist on main, verified against pristine checkout). docs/theory_journal.md Appendices E-E2: engineering log of the 12 findings from implementation and real-checkpoint validation (nibble order conventions, mixed-precision layouts, the 96GB OOM and the chunked-unpack fix, transformers 5.x drift shims, Omni wrapper unwrap). --- docs/theory_journal.md | 152 +++++++++++ tests/test_quant_dequant.py | 510 ++++++++++++++++++++++++++++++++++++ 2 files changed, 662 insertions(+) create mode 100644 tests/test_quant_dequant.py diff --git a/docs/theory_journal.md b/docs/theory_journal.md index ce7b192..6f2c180 100644 --- a/docs/theory_journal.md +++ b/docs/theory_journal.md @@ -1799,6 +1799,158 @@ implementations diverge from the closed-form GRRO solution. --- +## Appendix E: FP8/NVFP4 checkpoint dequantization — engineering log (2026-07-24) + +Implementation notes from adding FP8 and NVFP4 checkpoint support +(`obliteratus/models/quant_dequant.py`). Recorded per the escalation +protocol; anything here that surprised us is marked **(finding)**. + +1. **Nibble order is low-first, and it is easy to get backwards.** ModelOpt + NVFP4 packs two E2M1 values per byte along the input dim, low nibble + first: byte `0x1B` decodes to `[-1.5, +0.5]` (low `0xB`, high `0x1`). + **(finding)** Our first probe had the expected pair swapped, which would + have silently pair-flipped every weight row; the known-byte unit test + caught it. The runtime probe `_native_fp4_upcast_works()` now derives + the order empirically instead of assuming it. +2. **E2M1 LUT**: index → value `{0,.5,1,1.5,2,3,4,6}`, sign in bit 3. Index + 4 is `2.0`, not `3.0` — the spacing is not linear **(finding: test + caught a wrong-by-inspection comment)**. +3. **Fused-MoE scale naming differs from dense linears.** Dense weights use + `foo.weight` + `foo.weight_scale[_2/_inv]`; fused experts use + `experts.gate_up_proj` + `experts.gate_up_proj_scale[_2]`. Scale-key + detection must accept both `_scale`/`_scale_2` suffixes and the dotted + `.weight_scale` family **(finding: first state-dict test run failed on + exactly this)**. +4. **compressed-tensors stores reciprocal scales.** Dequant divides rather + than multiplies. Round-tripping reciprocals through FP8-E4M3 storage is + measurably coarser (cosine ~0.986 vs ~0.999 for direct storage on random + weights) — expected, not a bug. +5. **torch 2.13.0+cpu**: `torch.float4_e2m1fn_x2` exists but the runtime + upcast probe fails on this CPU build, so the manual nibble-LUT path is + what actually runs there; the probe design (rather than a version check) + was validated by this. +6. **Guard ordering matters**: `_is_quantized_param` flags float8 tensors, + so the FP8/NVFP4 loud-failure guard in `_dequantize_weight` must run + *before* the bitsandbytes branch or the error message is misleading + **(finding: unit test caught the wrong branch firing)**. +7. NVFP4 round-trip error on random Gaussian weights: relative Frobenius + error ~0.10 direct / ~0.18 reciprocal, cosine >0.98 — this is the + intrinsic precision of the format, not implementation error. +8. **Real checkpoints are mixed-precision.** `nvidia/Nemotron-3-Nano-Omni- + 30B-A3B-Reasoning-NVFP4` ships `quant_method: modelopt` with + `quant_algo: MIXED_PRECISION`: FP8 (8-bit float, per-tensor + `weight_scale`, no `weight_scale_2`) for the 98 Mamba-mixer/shared-expert + projections, NVFP4 (`weight` uint8 + FP8 `weight_scale` + + FP32 `weight_scale_2`) for the 5,888 routed-expert projections + **(finding)**. Detection now inspects `config_groups` bit-widths under + MIXED_PRECISION; dequantization was already per-tensor so FP8 and NVFP4 + tensors in the same shard each take their own path. The routed experts + are stored as *individual 2D modules* (`mixer.experts.{i}.up_proj`), not + fused 3D tensors — the dense path covers them. +9. **Custom architectures need a CLI escape hatch.** Nemotron Omni uses a + custom `modeling.py` (`NemotronH_Nano_Omni_Reasoning_V3`), but the CLI + had no `--trust-remote-code` passthrough to the pipeline **(finding)**. + Added the flag (`cli.py` → `AbliterationPipeline` → `load_model`). + +Smoke-test results on real checkpoints: see Appendix E continuation below +(filled in after the Nemotron-3-Nano GPU runs). + +--- + +## Appendix E2: Nemotron-3-Nano-Omni smoke runs (2026-07-24, 1x A100-80GB) + +Real-checkpoint validation of the FP8/NVFP4 dequantization path. Both +targets loaded with **zero quantized tensors left, zero NaNs, and +`quantization_config` stripped**, and both completed the full +`obliterate --method basic` pipeline (SUMMON → PROBE → DISTILL → EXCISE → +VERIFY → REBIRTH) with BF16 output. Findings, in the order they appeared: + +1. **Host RAM OOM during dequantization (fixed).** The first + `materialize_dequantized_checkpoint` attempt peaked at 96.5 GB RSS on a + 94 GB box and was OOM-killed mid-shard. Root cause: the manual nibble + unpack built an int64 index tensor for *all* nibbles at once (8 B/value) + plus stacked copies (~24 B/value total), and scale application created + fresh float32 copies per step. Fix: chunked LUT indexing (≤0.5 GB + transient int64) and in-place `mul_`/`div_` scaling. Peak is now + roughly one shard's float32 size. +2. **Hybrid-Mamba model deps.** The Omni custom code hard-imports + `mamba_ssm` triton ops; pip's build isolation hides torch and the CUDA + extension build is unnecessary — `pip install --no-build-isolation + mamba-ssm` with `MAMBA_SKIP_CUDA_BUILD=TRUE` plus a tiny + `selective_scan_cuda` stub module (only triton paths are used). + Vision/audio towers additionally need `timm`, `open_clip_torch`, + `librosa`. +3. **`attn_implementation: flash_attention_2` pinned in config (fixed + generically).** flash-attn isn't installed; the custom classes also + reject sdpa. The loader now detects the pin — including on nested + sub-configs (`llm_config`) — and falls back to `eager` when flash-attn + is unavailable. +4. **transformers 5.14 API drift, three shims (fixed in loader.py's shim + block / post-load normalization):** + - `all_tied_weights_keys`: expected by the accelerate device-map code, + absent on older remote-code classes, and *assigned* by 5.x + `post_init` — shim is a settable property aggregating legacy + `_tied_weights_keys`, filtering keys that don't resolve on the + top-level model (multimodal wrappers inherit `lm_head.weight` from + the LM submodule's config but have no top-level lm_head). + - `cache_position=None` passed by 5.x generation code into + `prepare_inputs_for_generation`; older remote code assumes a tensor. + Loader wraps the method to synthesize `cache_position` (using + past-key length during decoding). + - `_tied_weights_keys` as legacy *list* breaks 5.x `save_pretrained` + (expects `.keys()`); normalized post-load — emptied when + `tie_word_embeddings` is false so nothing is wrongly dropped. +5. **"Omni" wrapper can't run text-only surgery (fixed).** + `NemotronH_Nano_Omni_Reasoning_V3.forward()` requires `pixel_values`. + The loader now detects media-required forward signatures and unwraps + `model.language_model`; this also fixed layer/head/hidden reporting + (wrapper config lacks those fields). Surgery and the saved artifact + are the LM only. +6. **New architecture registered: `nemotron_h`.** All layer content hangs + off `layer.mixer` — attention (q/k/v/o_proj on the 6 attention layers), + Mamba (in/out_proj on 23), routed experts (2,944 individual + `mixer.experts.N.{up,down}_proj` 2D modules, *not* fused 3D) and + `mixer.shared_experts` (23). Mapped `["backbone","layers"]` + + `mixer` for both attention and FFN resolution in + `strategies/utils.py`. Mamba `out_proj` matches the attention + output-projection names and gets projected too — acceptable, it lives + in the same residual stream; noted here deliberately. +7. **Temp-dir lifecycle bug (fixed).** The dequantized temp copy was + deleted before the tokenizer load (which pointed at it). Tokenizer now + loads from the original checkpoint source. + +Results (`--method basic`, 1 direction): + +| Checkpoint | Scheme detected | Outcome | Refusal (post) | Perplexity | Output | +|---|---|---|---|---|---| +| Nemotron-3-Nano-Omni-30B NVFP4 (mixed FP8+NVFP4) | `nvfp4_modelopt` | full pipeline OK | 0.0 | 953.4 | 63.2 GB BF16 | +| Nemotron-3-Nano-Omni-30B FP8 | per-tensor FP8 (modelopt) | full pipeline OK | 0.0 | 920.5 (KL 0.73) | 63.2 GB BF16 | +| Qwen/Qwen3-8B-FP8 (block-wise, DeepSeek-style) | `fp8_blockwise` | full pipeline OK | 0.867 | 4.23 (KL 0.09) | 16 GB BF16 | +| Qwen/Qwen3-8B (**BF16 baseline**) | none | full pipeline OK | 0.833 | 4.33 (KL 0.11) | 16 GB BF16 | + +Notes: + +- **Qwen3-8B FP8-vs-BF16 parity is the format validation**: perplexity + within 2.4% (criterion was ~10%), refusal rates statistically identical + on the 30-sample eval. Qwen3's high residual refusal is a property of + the model + the 1-direction `basic` method (spectral certification RED, + the run itself recommended more directions) — it is *not* a + quantization artifact, as the BF16 baseline demonstrates. +- Nemotron perplexity is high in absolute terms — `--method basic` on a + hybrid Mamba-MoE reasoning model is the crudest setting and the + eigenvalue report recommended the 'nuclear' method. The smoke test's + purpose is format validation (dequantize → surgery → BF16 save), which + passed; tuning method strength for this architecture is separate work. +- The compressed-tensors per-channel FP8 dense alternate + (`RedHatAI/Llama-3.1-8B-Instruct-FP8-dynamic`) was not run: the exact + code path (`dequantize_fp8_per_channel`, direct scales) was already + exercised by the Nemotron FP8 checkpoint, and Llama-3.1 is a gated + repo. CT-NVFP4 (reciprocal scales) is covered by synthetic unit tests + only — no public checkpoint was needed after the Nemotron layout + parsed cleanly. + +--- + ## References 1. Arditi, A. et al. (2024). Refusal in Language Models Is Mediated by a Single Direction. NeurIPS 2024. diff --git a/tests/test_quant_dequant.py b/tests/test_quant_dequant.py new file mode 100644 index 0000000..401559a --- /dev/null +++ b/tests/test_quant_dequant.py @@ -0,0 +1,510 @@ +"""Tests for FP8/NVFP4 checkpoint dequantization (CPU-only, synthetic). + +Covers: +- FP8 block-wise (DeepSeek-style) and per-channel round-trips +- NVFP4 round-trips (direct + reciprocal scale conventions), native and + manual nibble-unpack paths agreeing with each other +- detect_quant_scheme classification from config.json metadata +- surgery guards rejecting un-dequantized FP8/NVFP4 weights +- end-to-end load_model() on tiny synthetic quantized checkpoints +""" + +from __future__ import annotations + +import json +import os + +import pytest +import torch + +from obliteratus.models import quant_dequant as qd + +HAS_FP8 = hasattr(torch, "float8_e4m3fn") +requires_fp8 = pytest.mark.skipif(not HAS_FP8, reason="torch build lacks float8 dtypes") + +FP8_MAX = 448.0 # e4m3 max +E2M1_MAX = 6.0 + + +# --------------------------------------------------------------------------- +# Reference quantization helpers (test-side) +# --------------------------------------------------------------------------- + +def _quantize_fp8_blockwise(w: torch.Tensor, block=(128, 128)): + """Float weight → (fp8 tensor, scale_inv) using DeepSeek convention.""" + M, N = w.shape + bm, bn = block + scale = torch.zeros((M + bm - 1) // bm, (N + bn - 1) // bn) + q = torch.zeros_like(w) + for i in range(0, M, bm): + for j in range(0, N, bn): + blk = w[i:i + bm, j:j + bn] + s = blk.abs().max() / FP8_MAX + s = max(s, 1e-8) + scale[i // bm, j // bn] = s + q[i:i + bm, j:j + bn] = blk / s + return q.to(torch.float8_e4m3fn), scale + + +def _quantize_fp8_per_channel(w: torch.Tensor, inverse: bool): + scale = (w.abs().amax(dim=1, keepdim=True) / FP8_MAX).clamp(min=1e-8) + q = (w / scale).to(torch.float8_e4m3fn) + return q, (1.0 / scale if inverse else scale) + + +def _nearest_e2m1(x: torch.Tensor) -> torch.Tensor: + """Round to nearest E2M1 value, preserving sign. Returns float codes.""" + lut = torch.tensor(qd.E2M1_POSITIVE, dtype=x.dtype) + sign = torch.sign(x) + ax = x.abs().clamp(max=E2M1_MAX) + idx = (ax.unsqueeze(-1) - lut).abs().argmin(-1) + return sign * lut[idx] + + +def _pack_nvfp4(w: torch.Tensor, group: int = 16, reciprocal: bool = False): + """Float (M, N) weight → (packed uint8, e4m3 block scales, fp32 global). + + Follows the ModelOpt convention: w ≈ e2m1 * block_scale * global_scale. + With reciprocal=True the returned scales are 1/scale (compressed-tensors + convention). + """ + M, N = w.shape + wg = w.reshape(M, N // group, group) + global_scale = (w.abs().max() / (FP8_MAX * E2M1_MAX)).clamp(min=1e-12) + bs = (wg.abs().amax(-1) / (E2M1_MAX * global_scale)).clamp(min=1e-8) + q = _nearest_e2m1(wg / (bs.unsqueeze(-1) * global_scale)) + # values → nibble codes via LUT lookup + codes = torch.zeros_like(q, dtype=torch.uint8) + for code, val in enumerate(list(qd.E2M1_POSITIVE) + [-v for v in qd.E2M1_POSITIVE]): + codes[q == val] = code + codes = codes.reshape(M, N) + low = codes[:, 0::2] + high = codes[:, 1::2] + packed = (high << 4) | low + bs_fp8 = bs.to(torch.float8_e4m3fn) + if reciprocal: + return packed, (1.0 / bs_fp8.float()).to(torch.float8_e4m3fn), 1.0 / global_scale + return packed, bs_fp8, global_scale.float() + + +# --------------------------------------------------------------------------- +# FP8 dequantization +# --------------------------------------------------------------------------- + +@requires_fp8 +def test_fp8_blockwise_roundtrip(): + torch.manual_seed(0) + w = torch.randn(256, 256) + q, scale_inv = _quantize_fp8_blockwise(w) + out = qd.dequantize_fp8_blockwise(q, scale_inv, (128, 128)) + rel = (out - w).norm() / w.norm() + assert rel < 0.05, f"blockwise FP8 round-trip error {rel:.4f}" + + +@requires_fp8 +def test_fp8_blockwise_non_divisible_dims(): + torch.manual_seed(1) + w = torch.randn(200, 300) # not divisible by 128 + q, scale_inv = _quantize_fp8_blockwise(w) + out = qd.dequantize_fp8_blockwise(q, scale_inv, (128, 128)) + assert out.shape == w.shape + rel = (out - w).norm() / w.norm() + assert rel < 0.05 + + +@requires_fp8 +@pytest.mark.parametrize("inverse", [False, True]) +def test_fp8_per_channel_roundtrip(inverse): + torch.manual_seed(2) + w = torch.randn(64, 128) + q, scale = _quantize_fp8_per_channel(w, inverse) + out = qd.dequantize_fp8_per_channel(q, scale, scale_is_inverse=inverse) + rel = (out - w).norm() / w.norm() + assert rel < 0.05 + + +@requires_fp8 +def test_fp8_per_tensor_scalar_scale(): + torch.manual_seed(3) + w = torch.randn(32, 64) + scale = w.abs().max() / FP8_MAX + q = (w / scale).to(torch.float8_e4m3fn) + out = qd.dequantize_fp8_per_channel(q, scale.reshape(1)) + rel = (out - w).norm() / w.norm() + assert rel < 0.05 + + +# --------------------------------------------------------------------------- +# NVFP4 dequantization +# --------------------------------------------------------------------------- + +@requires_fp8 +@pytest.mark.parametrize("reciprocal", [False, True]) +def test_nvfp4_roundtrip(reciprocal): + torch.manual_seed(4) + w = torch.randn(64, 128) + packed, bs, gs = _pack_nvfp4(w, reciprocal=reciprocal) + out = qd.dequantize_nvfp4( + packed, bs, gs, scale_is_inverse=reciprocal, force_manual=True, + ) + assert out.shape == w.shape + # NVFP4 is coarse — check correlation and relative error, not equality. + cos = torch.nn.functional.cosine_similarity( + w.flatten(), out.flatten(), dim=0, + ) + rel = (out - w).norm() / w.norm() + # Reciprocal storage double-rounds the inverse scales — slightly coarser. + assert cos > (0.98 if reciprocal else 0.995), f"NVFP4 cosine {cos:.4f}" + assert rel < (0.20 if reciprocal else 0.15), f"NVFP4 relative error {rel:.4f}" + + +@requires_fp8 +def test_nvfp4_native_and_manual_unpack_agree(): + if not qd._native_fp4_upcast_works(): + pytest.skip("torch build lacks a working float4_e2m1fn_x2 upcast") + torch.manual_seed(5) + packed = torch.randint(0, 256, (16, 64), dtype=torch.uint8) + native = qd.unpack_e2m1(packed, force_manual=False) + manual = qd.unpack_e2m1(packed, force_manual=True) + assert torch.equal(native, manual) + + +def test_nvfp4_manual_unpack_known_values(): + # 0x1B: low nibble 0xB (-1.5), high nibble 0x1 (+0.5) — low nibble first. + # 0x84: low nibble 0x4 (+2.0), high nibble 0x8 (-0.0). + packed = torch.tensor([[0x1B, 0x84]], dtype=torch.uint8) + out = qd.unpack_e2m1(packed, force_manual=True) + assert out.shape == (1, 4) + assert out[0, 0].item() == -1.5 + assert out[0, 1].item() == 0.5 + assert out[0, 2].item() == 2.0 + assert out[0, 3].item() == 0.0 + + +@requires_fp8 +def test_nvfp4_no_global_scale(): + torch.manual_seed(6) + w = torch.randn(32, 64) + packed, bs, gs = _pack_nvfp4(w) + out = qd.dequantize_nvfp4(packed, bs, None, force_manual=True) + assert out.shape == w.shape + + +def test_nvfp4_bad_group_size_raises(): + packed = torch.zeros((4, 8), dtype=torch.uint8) # unpacks to 4x16 + bs = torch.ones((4, 1), dtype=torch.float32) + with pytest.raises(RuntimeError, match="group_size"): + qd.dequantize_nvfp4(packed, bs, None, group_size=32) + + +# --------------------------------------------------------------------------- +# Scheme detection +# --------------------------------------------------------------------------- + +def _write_config(tmp_path, cfg: dict): + with open(os.path.join(tmp_path, "config.json"), "w") as fh: + json.dump(cfg, fh) + + +def test_detect_none(tmp_path): + _write_config(tmp_path, {"model_type": "gpt2"}) + det = qd.detect_quant_scheme(str(tmp_path)) + assert det.scheme is qd.QuantScheme.NONE + + +def test_detect_fp8_blockwise(tmp_path): + _write_config(tmp_path, { + "quantization_config": { + "quant_method": "fp8", + "weight_block_size": [128, 128], + }, + }) + det = qd.detect_quant_scheme(str(tmp_path)) + assert det.scheme is qd.QuantScheme.FP8_BLOCKWISE + assert det.block_size == (128, 128) + + +def test_detect_fp8_blockwise_via_scale_keys(tmp_path): + _write_config(tmp_path, {"quantization_config": {"quant_method": "fp8"}}) + with open(os.path.join(tmp_path, "model.safetensors.index.json"), "w") as fh: + json.dump({"weight_map": {"model.layers.0.mlp.weight_scale_inv": "model.safetensors"}}, fh) + det = qd.detect_quant_scheme(str(tmp_path)) + assert det.scheme is qd.QuantScheme.FP8_BLOCKWISE + + +def test_detect_fp8_ct(tmp_path): + _write_config(tmp_path, { + "quantization_config": { + "quant_method": "compressed-tensors", + "config_groups": {"group_0": {"weights": {"num_bits": 8, "type": "float"}}}, + }, + }) + det = qd.detect_quant_scheme(str(tmp_path)) + assert det.scheme is qd.QuantScheme.FP8_PER_CHANNEL_CT + + +def test_detect_nvfp4_ct(tmp_path): + _write_config(tmp_path, { + "quantization_config": { + "quant_method": "compressed-tensors", + "config_groups": {"group_0": {"weights": {"num_bits": 4, "type": "float", "group_size": 16}}}, + }, + }) + det = qd.detect_quant_scheme(str(tmp_path)) + assert det.scheme is qd.QuantScheme.NVFP4_CT + assert det.scale_is_inverse is True + assert det.group_size == 16 + + +def test_detect_nvfp4_modelopt(tmp_path): + _write_config(tmp_path, { + "quantization_config": {"quant_method": "modelopt", "quant_algo": "NVFP4"}, + }) + det = qd.detect_quant_scheme(str(tmp_path)) + assert det.scheme is qd.QuantScheme.NVFP4_MODELOPT + + +def test_detect_modelopt_mixed_precision(tmp_path): + """Nemotron-3-Nano NVFP4: FP8 mixer + NVFP4 experts in one checkpoint.""" + _write_config(tmp_path, { + "quantization_config": { + "quant_method": "modelopt", + "quant_algo": "MIXED_PRECISION", + "config_groups": { + "group_0": {"weights": {"num_bits": 8, "type": "float"}}, + "group_1": {"weights": {"num_bits": 4, "type": "float", "group_size": 16}}, + }, + }, + }) + det = qd.detect_quant_scheme(str(tmp_path)) + assert det.scheme is qd.QuantScheme.NVFP4_MODELOPT + + +def test_detect_modelopt_mixed_fp8_only(tmp_path): + _write_config(tmp_path, { + "quantization_config": { + "quant_method": "modelopt", + "quant_algo": "MIXED_PRECISION", + "config_groups": { + "group_0": {"weights": {"num_bits": 8, "type": "float"}}, + }, + }, + }) + det = qd.detect_quant_scheme(str(tmp_path)) + assert det.scheme is qd.QuantScheme.FP8_PER_CHANNEL_CT + + +def test_detect_unsupported_fbgemm(tmp_path): + _write_config(tmp_path, { + "quantization_config": {"quant_method": "fbgemm_fp8"}, + }) + det = qd.detect_quant_scheme(str(tmp_path)) + assert det.scheme is qd.QuantScheme.UNSUPPORTED + assert "fbgemm" in det.reason + + +@pytest.mark.parametrize("method", ["gptq", "awq", "bitsandbytes"]) +def test_detect_passthrough_schemes(tmp_path, method): + _write_config(tmp_path, {"quantization_config": {"quant_method": method}}) + det = qd.detect_quant_scheme(str(tmp_path)) + assert det.scheme is qd.QuantScheme.NONE + + +# --------------------------------------------------------------------------- +# State-dict dequantization +# --------------------------------------------------------------------------- + +@requires_fp8 +def test_dequantize_state_dict_fp8_blockwise(): + torch.manual_seed(7) + w = torch.randn(256, 128) + q, scale_inv = _quantize_fp8_blockwise(w) + sd = { + "layers.0.mlp.weight": q, + "layers.0.mlp.weight_scale_inv": scale_inv, + "layers.0.attn.bias": torch.randn(128), + } + det = qd.QuantDetection(qd.QuantScheme.FP8_BLOCKWISE) + out = qd.dequantize_state_dict(sd, det, out_dtype=torch.bfloat16) + assert set(out) == {"layers.0.mlp.weight", "layers.0.attn.bias"} + assert out["layers.0.mlp.weight"].dtype == torch.bfloat16 + rel = (out["layers.0.mlp.weight"].float() - w).norm() / w.norm() + assert rel < 0.05 + + +@requires_fp8 +def test_dequantize_state_dict_nvfp4_modelopt(): + torch.manual_seed(8) + w = torch.randn(32, 64) + packed, bs, gs = _pack_nvfp4(w) + sd = { + "experts.0.gate_up_proj": packed, + "experts.0.gate_up_proj_scale": bs, + "experts.0.gate_up_proj_scale_2": gs, + } + # Non-.weight key names (fused MoE style) must also dequantize. + det = qd.QuantDetection(qd.QuantScheme.NVFP4_MODELOPT) + out = qd.dequantize_state_dict(sd, det, out_dtype=torch.float32) + assert set(out) == {"experts.0.gate_up_proj"} + cos = torch.nn.functional.cosine_similarity(w.flatten(), out["experts.0.gate_up_proj"].flatten(), dim=0) + assert cos > 0.995 + + +@requires_fp8 +def test_dequantize_state_dict_fp8_missing_scale_raises(): + q = torch.randn(16, 16).to(torch.float8_e4m3fn) + det = qd.QuantDetection(qd.QuantScheme.FP8_BLOCKWISE) + with pytest.raises(RuntimeError, match="no weight_scale"): + qd.dequantize_state_dict({"a.weight": q}, det) + + +# --------------------------------------------------------------------------- +# Surgery guards +# --------------------------------------------------------------------------- + +@requires_fp8 +def test_is_quantized_param_fp8(): + from obliteratus.abliterate import AbliterationPipeline + + p = torch.nn.Parameter(torch.randn(8, 8).to(torch.float8_e4m3fn)) + assert AbliterationPipeline._is_quantized_param(p) is True + p2 = torch.nn.Parameter(torch.randn(8, 8)) + assert AbliterationPipeline._is_quantized_param(p2) is False + + +@requires_fp8 +def test_dequantize_weight_guard_fp8(): + from obliteratus.abliterate import AbliterationPipeline + + mod = torch.nn.Linear(8, 8) + mod.weight = torch.nn.Parameter(torch.randn(8, 8).to(torch.float8_e4m3fn)) + with pytest.raises(RuntimeError, match="without dequantization"): + AbliterationPipeline._dequantize_weight(mod) + + +def test_dequantize_weight_guard_packed_uint8_with_scales(): + from obliteratus.abliterate import AbliterationPipeline + + mod = torch.nn.Linear(8, 8) + mod.weight = torch.nn.Parameter(torch.zeros(8, 4, dtype=torch.uint8), requires_grad=False) + mod.weight_scale = torch.ones(8, 1, dtype=torch.float32) + with pytest.raises(RuntimeError, match="without dequantization"): + AbliterationPipeline._dequantize_weight(mod) + + +def test_dequantize_weight_plain_uint8_still_converts(): + """Pre-existing behavior for scale-less custom uint8 weights is kept.""" + from obliteratus.abliterate import AbliterationPipeline + + mod = torch.nn.Linear(8, 8) + mod.weight = torch.nn.Parameter(torch.ones(8, 8, dtype=torch.uint8), requires_grad=False) + w, is_q = AbliterationPipeline._dequantize_weight(mod) + assert is_q is True + assert w.is_floating_point() + + +# --------------------------------------------------------------------------- +# End-to-end loader on tiny synthetic checkpoints +# --------------------------------------------------------------------------- + +@requires_fp8 +def test_load_model_fp8_blockwise_checkpoint(tmp_path, monkeypatch): + from obliteratus.models import loader as loader_mod + + captured = {} + + from safetensors.torch import save_file + from transformers import GPT2Config, GPT2LMHeadModel + + cfg = GPT2Config(n_layer=1, n_head=2, n_embd=256, n_inner=512, vocab_size=128) + model = GPT2LMHeadModel(cfg) + sd = model.state_dict() + out_sd = {} + for k, v in sd.items(): + if k.endswith(".weight") and v.dim() == 2 and min(v.shape) >= 32: + q, s = _quantize_fp8_blockwise(v.float()) + out_sd[k] = q + out_sd[k[: -len(".weight")] + ".weight_scale_inv"] = s + captured[k] = v.float() + else: + out_sd[k] = v + save_file(out_sd, str(tmp_path / "model.safetensors")) + cfg_dict = cfg.to_dict() + cfg_dict["quantization_config"] = { + "quant_method": "fp8", "weight_block_size": [128, 128], + } + with open(tmp_path / "config.json", "w") as fh: + json.dump(cfg_dict, fh) + + class _FakeTok: + pad_token = eos_token = "<|endoftext|>" + + monkeypatch.setattr( + loader_mod.AutoTokenizer, "from_pretrained", staticmethod(lambda *a, **k: _FakeTok()), + ) + handle = loader_mod.load_model(str(tmp_path), task="causal_lm", device="cpu", dtype="bfloat16") + loaded = handle.model.state_dict() + assert getattr(handle.model, "_obliteratus_dequantized_scheme", None) == "fp8_blockwise" + for k, w in captured.items(): + p = loaded[k].detach().float() + assert p.is_floating_point() + rel = (p - w).norm() / w.norm() + assert rel < 0.05, f"{k}: rel error {rel:.4f}" + # No scale tensors survived into the loaded model + assert not any("scale" in n for n in loaded) + + +@requires_fp8 +def test_load_model_nvfp4_checkpoint(tmp_path, monkeypatch): + from obliteratus.models import loader as loader_mod + from safetensors.torch import save_file + from transformers import GPT2Config, GPT2LMHeadModel + + cfg = GPT2Config(n_layer=1, n_head=2, n_embd=256, n_inner=512, vocab_size=128) + model = GPT2LMHeadModel(cfg) + sd = model.state_dict() + out_sd = {} + captured = {} + for k, v in sd.items(): + if k.endswith(".weight") and v.dim() == 2 and v.shape[1] % 32 == 0 and min(v.shape) >= 32: + packed, bs, gs = _pack_nvfp4(v.float()) + base = k[: -len(".weight")] + out_sd[k] = packed + out_sd[base + ".weight_scale"] = bs + out_sd[base + ".weight_scale_2"] = gs + captured[k] = v.float() + else: + out_sd[k] = v + save_file(out_sd, str(tmp_path / "model.safetensors")) + cfg_dict = cfg.to_dict() + cfg_dict["quantization_config"] = { + "quant_method": "modelopt", "quant_algo": "NVFP4", + } + with open(tmp_path / "config.json", "w") as fh: + json.dump(cfg_dict, fh) + + class _FakeTok: + pad_token = eos_token = "<|endoftext|>" + + monkeypatch.setattr( + loader_mod.AutoTokenizer, "from_pretrained", staticmethod(lambda *a, **k: _FakeTok()), + ) + handle = loader_mod.load_model(str(tmp_path), task="causal_lm", device="cpu", dtype="bfloat16") + loaded = handle.model.state_dict() + assert getattr(handle.model, "_obliteratus_dequantized_scheme", None) == "nvfp4_modelopt" + for k, w in captured.items(): + p = loaded[k].detach().float() + cos = torch.nn.functional.cosine_similarity(w.flatten(), p.flatten(), dim=0) + assert cos > 0.99, f"{k}: cosine {cos:.4f}" + assert not any("scale" in n for n in loaded) + + +def test_load_model_unsupported_scheme_fails_loudly(tmp_path): + from obliteratus.models import loader as loader_mod + + _write_config(tmp_path, { + "model_type": "gpt2", + "quantization_config": {"quant_method": "fbgemm_fp8"}, + }) + with pytest.raises(RuntimeError, match="Unsupported quantization"): + loader_mod.load_model(str(tmp_path), task="causal_lm", device="cpu", dtype="bfloat16")