mirror of
https://github.com/elder-plinius/OBLITERATUS.git
synced 2026-08-17 16:37:30 +02:00
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).
This commit is contained in:
Brian Bell
@@ -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.
|
||||
|
||||
@@ -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")
|
||||
Reference in New Issue
Block a user