fix: harden quantized checkpoint integration

This commit is contained in:
Joseph Magly
2026-08-16 12:57:17 -04:00
parent b6a1727cb6
commit b25b3f3b66
13 changed files with 1379 additions and 228 deletions
+21
View File
@@ -0,0 +1,21 @@
{
"schema_version": 1,
"origin": "synthetic-runtime-generated",
"external_data": false,
"contract_sources": [
"https://docs.vllm.ai/projects/llm-compressor/en/stable/guides/entrypoints/convert/",
"https://github.com/huggingface/transformers/blob/main/docs/source/en/quantization/compressed_tensors.md"
],
"generator": "tests/test_quant_dequant.py",
"torch_manual_seeds": [0, 1, 2, 3, 4, 5, 6, 7, 8],
"formats": {
"fp8": "IEEE-like E4M3 values generated by torch.float8_e4m3fn",
"nvfp4": "E2M1 values generated from OBLITERATUS's explicit 16-value test oracle"
},
"oracles": {
"fp8_relative_l2_max": 0.05,
"nvfp4_relative_l2_max": 0.15,
"nvfp4_cosine_min": 0.995
},
"reproducibility": "Fixtures are regenerated deterministically during each test run; no model weights, datasets, or network resources are consumed."
}
+13
View File
@@ -80,6 +80,18 @@ def test_refusal_max_tokens_cli_default_and_positive_override(monkeypatch):
assert command.call_args.args[0].refusal_max_tokens == 512
def test_trust_remote_code_requires_explicit_cli_opt_in(monkeypatch):
command = Mock()
monkeypatch.setattr(cli, "_cmd_abliterate", command)
cli.main(["abliterate", "local/model"])
assert command.call_args.args[0].trust_remote_code is False
command.reset_mock()
cli.main(["abliterate", "local/model", "--trust-remote-code"])
assert command.call_args.args[0].trust_remote_code is True
@pytest.mark.parametrize("invalid", ["0", "-1", "not-an-integer"])
def test_refusal_max_tokens_cli_rejects_invalid_values(invalid):
with pytest.raises(SystemExit) as exc:
@@ -510,6 +522,7 @@ def test_abliterate_pipeline_callbacks_residue_and_contribution(monkeypatch, tmp
cli._cmd_abliterate(args)
assert factory.call_args.kwargs["refusal_max_tokens"] == 512
assert factory.call_args.kwargs["gpu_memory_utilization"] == 0.95
assert factory.call_args.kwargs["trust_remote_code"] is False
assert (result_path / "hard_negative_residue.json").is_file()
telemetry.assert_called_once_with(pipeline)
+219
View File
@@ -23,6 +23,7 @@ def test_compatibility_shims_tolerate_independent_missing_or_broken_imports(monk
original_import = builtins.__import__
failure_keys = {
("transformers", ("AutoModelForImageTextToText",)): ImportError,
("transformers", ("PreTrainedModel",)): RuntimeError,
("transformers.utils", ("output_capturing",)): ImportError,
("transformers.utils.import_utils", ()): RuntimeError,
("transformers.utils", ()): RuntimeError,
@@ -109,6 +110,72 @@ def test_compatibility_shims_tolerate_rejected_generic_attribute_patch(monkeypat
assert _execute_loader_source()["TASK_MODEL_MAP"]
def test_all_tied_weights_compat_aggregates_filters_and_caches():
class LegacyModel(torch.nn.Module):
pass
assert loader._install_all_tied_weights_keys_compat(LegacyModel) is True
assert loader._install_all_tied_weights_keys_compat(LegacyModel) is False
model = LegacyModel()
model.lm_head = torch.nn.Linear(2, 2)
model._tied_weights_keys = ["lm_head.weight", "missing.weight"]
child = torch.nn.Module()
child._tied_weights_keys = {"lm_head.bias": "lm_head.weight"}
model.child = child
assert model.all_tied_weights_keys == {
"lm_head.weight": "lm_head.weight",
"lm_head.bias": "lm_head.weight",
}
model.all_tied_weights_keys = {"cached.weight": "source.weight"}
assert model.all_tied_weights_keys == {"cached.weight": "source.weight"}
def test_all_tied_weights_compat_prefers_transformers_expansion_contract():
class ExpandableModel(torch.nn.Module):
def get_expanded_tied_weights_keys(self, *, all_submodels):
assert all_submodels is True
return {"lm_head.weight": "model.embed_tokens.weight"}
assert loader._install_all_tied_weights_keys_compat(ExpandableModel) is True
model = ExpandableModel()
assert model.all_tied_weights_keys == {
"lm_head.weight": "model.embed_tokens.weight",
}
assert model.__dict__["_all_tied_weights_keys_store"] == {
"lm_head.weight": "model.embed_tokens.weight",
}
def test_all_tied_weights_compat_falls_back_when_expansion_is_legacy():
class LegacyExpandableModel(torch.nn.Module):
def get_expanded_tied_weights_keys(self, *, all_submodels):
raise TypeError("legacy list contract")
assert loader._install_all_tied_weights_keys_compat(LegacyExpandableModel) is True
model = LegacyExpandableModel()
model.lm_head = torch.nn.Linear(2, 2)
model._tied_weights_keys = ["lm_head.weight"]
assert model.all_tied_weights_keys == {"lm_head.weight": "lm_head.weight"}
def test_flash_attention_config_helpers_recurse_into_transformers_subconfigs():
from transformers import PretrainedConfig
outer = PretrainedConfig()
inner = PretrainedConfig()
outer.inner = inner
inner._attn_implementation = "flash_attention_2"
assert loader._config_uses_flash_attention_2(outer) is True
loader._force_eager_config(outer)
assert inner._attn_implementation == "eager"
assert loader._config_uses_flash_attention_2(outer) is False
def _config(**overrides):
values = {
"model_type": "gpt2",
@@ -200,6 +267,158 @@ def test_revision_trust_and_offline_flags_reach_every_provider(loader_boundary):
loader.dev.empty_cache.assert_called_once_with()
def test_unavailable_flash_attention_falls_back_to_eager(
loader_boundary, monkeypatch,
):
import transformers.utils
loader_boundary.config._attn_implementation = "flash_attention_2"
monkeypatch.setattr(
transformers.utils,
"is_flash_attn_2_available",
lambda: False,
)
monkeypatch.setattr(
loader.qd,
"detect_quant_scheme",
lambda *args, **kwargs: loader.qd.QuantDetection(loader.qd.QuantScheme.NONE),
)
loader.load_model("local/model", skip_snapshot=True)
assert loader_boundary.config._attn_implementation == "eager"
def test_broken_flash_attention_probe_falls_back_to_eager(
loader_boundary, monkeypatch,
):
loader_boundary.config._attn_implementation = "flash_attention_2"
real_import = builtins.__import__
def rejecting_import(name, globals=None, locals=None, fromlist=(), level=0):
if (
name == "transformers.utils"
and "is_flash_attn_2_available" in tuple(fromlist or ())
):
raise RuntimeError("flash-attention availability probe failed")
return real_import(name, globals, locals, fromlist, level)
monkeypatch.setattr(builtins, "__import__", rejecting_import)
loader.load_model("local/model", skip_snapshot=True)
assert loader_boundary.config._attn_implementation == "eager"
def test_available_flash_attention_preserves_requested_implementation(
loader_boundary, monkeypatch,
):
import transformers.utils
loader_boundary.config._attn_implementation = "flash_attention_2"
monkeypatch.setattr(
transformers.utils,
"is_flash_attn_2_available",
lambda: True,
)
loader.load_model("local/model", skip_snapshot=True)
assert loader_boundary.config._attn_implementation == "flash_attention_2"
def test_multimodal_unwrap_generation_cache_and_legacy_ties(
loader_boundary, monkeypatch,
):
class FakeLanguageModel(torch.nn.Module):
def __init__(self):
super().__init__()
self.config = _config(tie_word_embeddings=False)
self._tied_weights_keys = ["lm_head.weight"]
self.last_cache_position = None
def prepare_inputs_for_generation(
self, input_ids, *args, cache_position=None, **kwargs,
):
self.last_cache_position = cache_position
return {"input_ids": input_ids, "cache_position": cache_position}
class MultimodalWrapper(torch.nn.Module):
def __init__(self):
super().__init__()
self.language_model = FakeLanguageModel()
def forward(self, input_ids, pixel_values):
return input_ids, pixel_values
wrapper = MultimodalWrapper()
loader_boundary.model_class.from_pretrained.return_value = wrapper
monkeypatch.setattr(loader, "PreTrainedModel", FakeLanguageModel)
monkeypatch.setattr(
loader.qd,
"detect_quant_scheme",
lambda *args, **kwargs: loader.qd.QuantDetection(loader.qd.QuantScheme.NONE),
)
handle = loader.load_model("local/model", skip_snapshot=True)
assert handle.model is wrapper.language_model
assert handle.model._tied_weights_keys == {}
input_ids = torch.ones(1, 2, dtype=torch.long)
class Cache:
@staticmethod
def get_seq_length():
return 4
prepared = handle.model.prepare_inputs_for_generation(
input_ids, past_key_values=Cache(),
)
assert prepared["cache_position"].tolist() == [4, 5]
legacy_cache = ((torch.zeros(1, 1, 3, 1),),)
prepared = handle.model.prepare_inputs_for_generation(
input_ids, past_key_values=legacy_cache,
)
assert prepared["cache_position"].tolist() == [3, 4]
prepared = handle.model.prepare_inputs_for_generation(
input_ids, past_key_values=object(),
)
assert prepared["cache_position"].tolist() == [0, 1]
prepared = handle.model.prepare_inputs_for_generation(input_ids)
assert prepared["cache_position"].tolist() == [0, 1]
explicit = torch.tensor([7, 8])
prepared = handle.model.prepare_inputs_for_generation(
input_ids, cache_position=explicit,
)
assert prepared["cache_position"] is explicit
def test_multimodal_signature_probe_failure_preserves_wrapper(
loader_boundary, monkeypatch,
):
class FakeLanguageModel(torch.nn.Module):
def __init__(self):
super().__init__()
self.config = _config()
class WrapperWithUninspectableForward(torch.nn.Module):
def __init__(self):
super().__init__()
self.language_model = FakeLanguageModel()
self.forward = object()
wrapper = WrapperWithUninspectableForward()
loader_boundary.model_class.from_pretrained.return_value = wrapper
monkeypatch.setattr(loader, "PreTrainedModel", FakeLanguageModel)
handle = loader.load_model("local/model", skip_snapshot=True)
assert handle.model is wrapper
def test_hf_token_is_forwarded_without_logging_value(loader_boundary, monkeypatch):
monkeypatch.setenv("HF_TOKEN", "secret-token")
loader.load_model("x", skip_snapshot=True)
+646 -14
View File
@@ -2,7 +2,7 @@
Covers:
- FP8 block-wise (DeepSeek-style) and per-channel round-trips
- NVFP4 round-trips (direct + reciprocal scale conventions), native and
- NVFP4 round-trips (ModelOpt + compressed-tensors global-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
@@ -13,6 +13,8 @@ from __future__ import annotations
import json
import os
from pathlib import Path
from types import SimpleNamespace
import pytest
import torch
@@ -61,12 +63,12 @@ def _nearest_e2m1(x: torch.Tensor) -> torch.Tensor:
return sign * lut[idx]
def _pack_nvfp4(w: torch.Tensor, group: int = 16, reciprocal: bool = False):
def _pack_nvfp4(w: torch.Tensor, group: int = 16, global_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).
Compressed-tensors keeps the block scale but stores the reciprocal of
the global scale.
"""
M, N = w.shape
wg = w.reshape(M, N // group, group)
@@ -82,8 +84,8 @@ def _pack_nvfp4(w: torch.Tensor, group: int = 16, reciprocal: bool = False):
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
if global_reciprocal:
return packed, bs_fp8, 1.0 / global_scale
return packed, bs_fp8, global_scale.float()
@@ -139,13 +141,17 @@ def test_fp8_per_tensor_scalar_scale():
# ---------------------------------------------------------------------------
@requires_fp8
@pytest.mark.parametrize("reciprocal", [False, True])
def test_nvfp4_roundtrip(reciprocal):
@pytest.mark.parametrize("global_reciprocal", [False, True])
def test_nvfp4_roundtrip(global_reciprocal):
torch.manual_seed(4)
w = torch.randn(64, 128)
packed, bs, gs = _pack_nvfp4(w, reciprocal=reciprocal)
packed, bs, gs = _pack_nvfp4(w, global_reciprocal=global_reciprocal)
out = qd.dequantize_nvfp4(
packed, bs, gs, scale_is_inverse=reciprocal, force_manual=True,
packed,
bs,
gs,
global_scale_is_inverse=global_reciprocal,
force_manual=True,
)
assert out.shape == w.shape
# NVFP4 is coarse — check correlation and relative error, not equality.
@@ -153,9 +159,8 @@ def test_nvfp4_roundtrip(reciprocal):
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}"
assert cos > 0.995, f"NVFP4 cosine {cos:.4f}"
assert rel < 0.15, f"NVFP4 relative error {rel:.4f}"
@requires_fp8
@@ -206,6 +211,67 @@ def _write_config(tmp_path, cfg: dict):
json.dump(cfg, fh)
def test_load_json_local_missing_returns_none(tmp_path):
assert qd._load_json_from_checkpoint(str(tmp_path), "missing.json") is None
def test_safetensors_key_names_reads_single_file_and_index(tmp_path):
from safetensors.torch import save_file
save_file({"a.weight": torch.ones(1)}, tmp_path / "model.safetensors")
assert qd._safetensors_key_names(str(tmp_path), None) == {"a.weight"}
(tmp_path / "model.safetensors.index.json").write_text(json.dumps({
"weight_map": {"b.weight": "model.safetensors"},
}))
assert qd._safetensors_key_names(str(tmp_path), None) == {"b.weight"}
def test_safetensors_key_names_handles_missing_and_invalid_files(tmp_path):
assert qd._safetensors_key_names(str(tmp_path), None) == set()
(tmp_path / "model.safetensors").write_text("not safetensors")
assert qd._safetensors_key_names(str(tmp_path), None) == set()
def test_safetensors_key_names_remote_propagates_hub_policy(tmp_path, monkeypatch):
from safetensors.torch import save_file
checkpoint = tmp_path / "remote.safetensors"
save_file({"remote.weight": torch.ones(1)}, checkpoint)
monkeypatch.setattr(qd, "_load_json_from_checkpoint", lambda *args, **kwargs: None)
calls = []
def fake_download(repo, filename, **kwargs):
calls.append((repo, filename, kwargs))
return str(checkpoint)
monkeypatch.setattr("huggingface_hub.hf_hub_download", fake_download)
assert qd._safetensors_key_names(
"org/model",
None,
token="token",
revision="immutable-sha",
local_files_only=True,
) == {"remote.weight"}
assert calls == [(
"org/model",
"model.safetensors",
{
"token": "token",
"revision": "immutable-sha",
"local_files_only": True,
},
)]
def test_safetensors_key_names_remote_download_failure_is_empty(monkeypatch):
monkeypatch.setattr(qd, "_load_json_from_checkpoint", lambda *args, **kwargs: None)
monkeypatch.setattr(
"huggingface_hub.hf_hub_download",
lambda *args, **kwargs: (_ for _ in ()).throw(OSError("offline")),
)
assert qd._safetensors_key_names("org/model", None) == set()
def test_detect_none(tmp_path):
_write_config(tmp_path, {"model_type": "gpt2"})
det = qd.detect_quant_scheme(str(tmp_path))
@@ -232,6 +298,18 @@ def test_detect_fp8_blockwise_via_scale_keys(tmp_path):
assert det.scheme is qd.QuantScheme.FP8_BLOCKWISE
def test_detect_fp8_without_supported_layout_fails_closed(tmp_path):
_write_config(tmp_path, {
"quantization_config": {
"quant_method": "fp8",
"activation_scheme": "static",
},
})
det = qd.detect_quant_scheme(str(tmp_path))
assert det.scheme is qd.QuantScheme.UNSUPPORTED
assert "without weight_block_size" in det.reason
def test_detect_fp8_ct(tmp_path):
_write_config(tmp_path, {
"quantization_config": {
@@ -247,15 +325,33 @@ def test_detect_nvfp4_ct(tmp_path):
_write_config(tmp_path, {
"quantization_config": {
"quant_method": "compressed-tensors",
"format": "nvfp4-pack-quantized",
"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.scale_is_inverse is False
assert det.global_scale_is_inverse is True
assert det.group_size == 16
def test_detect_nvfp4_ct_without_packed_format_fails_closed(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.UNSUPPORTED
assert "nvfp4-pack-quantized" in det.reason
def test_detect_nvfp4_modelopt(tmp_path):
_write_config(tmp_path, {
"quantization_config": {"quant_method": "modelopt", "quant_algo": "NVFP4"},
@@ -294,6 +390,23 @@ def test_detect_modelopt_mixed_fp8_only(tmp_path):
assert det.scheme is qd.QuantScheme.FP8_PER_CHANNEL_CT
def test_detect_modelopt_fp8_and_unknown(tmp_path):
_write_config(tmp_path, {
"quantization_config": {"quant_method": "modelopt", "quant_algo": "FP8"},
})
assert qd.detect_quant_scheme(str(tmp_path)).scheme is qd.QuantScheme.FP8_PER_CHANNEL_CT
_write_config(tmp_path, {
"quantization_config": {
"quant_method": "modelopt",
"kv_cache_quant_algo": "INT8",
},
})
det = qd.detect_quant_scheme(str(tmp_path))
assert det.scheme is qd.QuantScheme.UNSUPPORTED
assert "INT8" in det.reason
def test_detect_unsupported_fbgemm(tmp_path):
_write_config(tmp_path, {
"quantization_config": {"quant_method": "fbgemm_fp8"},
@@ -303,6 +416,15 @@ def test_detect_unsupported_fbgemm(tmp_path):
assert "fbgemm" in det.reason
def test_detect_nonempty_quant_config_without_method_fails_closed(tmp_path):
_write_config(tmp_path, {
"quantization_config": {"bits": 4, "format": "unknown-packed-layout"},
})
det = qd.detect_quant_scheme(str(tmp_path))
assert det.scheme is qd.QuantScheme.UNSUPPORTED
assert "quant_method ''" 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}})
@@ -310,10 +432,335 @@ def test_detect_passthrough_schemes(tmp_path, method):
assert det.scheme is qd.QuantScheme.NONE
def test_detect_remote_checkpoint_pins_revision_and_offline_mode(
tmp_path, monkeypatch,
):
config_path = tmp_path / "config.json"
config_path.write_text(json.dumps({
"quantization_config": {
"quant_method": "fp8",
"weight_block_size": [128, 128],
},
}))
calls = []
def fake_download(repo, filename, **kwargs):
calls.append((repo, filename, kwargs))
return str(config_path)
monkeypatch.setattr("huggingface_hub.hf_hub_download", fake_download)
det = qd.detect_quant_scheme(
"org/model",
token="token",
revision="immutable-sha",
local_files_only=True,
)
assert det.scheme is qd.QuantScheme.FP8_BLOCKWISE
assert calls == [(
"org/model",
"config.json",
{
"token": "token",
"revision": "immutable-sha",
"local_files_only": True,
},
)]
def test_loader_propagates_revision_and_offline_mode_to_detection(monkeypatch):
from obliteratus.models import loader as loader_mod
captured = {}
monkeypatch.setattr(
loader_mod.AutoConfig,
"from_pretrained",
staticmethod(lambda *args, **kwargs: SimpleNamespace(quantization_config=None)),
)
def fake_detect(model_name, **kwargs):
captured["model_name"] = model_name
captured.update(kwargs)
return qd.QuantDetection(qd.QuantScheme.UNSUPPORTED, reason="test sentinel")
monkeypatch.setattr(qd, "detect_quant_scheme", fake_detect)
with pytest.raises(RuntimeError, match="test sentinel"):
loader_mod.load_model(
"org/model",
task="causal_lm",
device="cpu",
dtype="float32",
revision="immutable-sha",
local_files_only=True,
)
assert captured == {
"model_name": "org/model",
"token": None,
"revision": "immutable-sha",
"local_files_only": True,
}
def test_materialize_remote_checkpoint_pins_revision_and_offline_mode(
tmp_path, monkeypatch,
):
from safetensors.torch import save_file
source = tmp_path / "source"
source.mkdir()
_write_config(source, {
"model_type": "gpt2",
"quantization_config": {"quant_method": "fp8"},
})
save_file({"plain.weight": torch.ones(2, 2)}, source / "model.safetensors")
captured = {}
def fake_snapshot(repo, **kwargs):
captured["repo"] = repo
captured.update(kwargs)
return str(source)
monkeypatch.setattr("huggingface_hub.snapshot_download", fake_snapshot)
output, returned_source = qd.materialize_dequantized_checkpoint(
"org/model",
qd.QuantDetection(qd.QuantScheme.FP8_BLOCKWISE),
token="token",
revision="immutable-sha",
local_files_only=True,
)
try:
assert returned_source == str(source)
assert captured["repo"] == "org/model"
assert captured["token"] == "token"
assert captured["revision"] == "immutable-sha"
assert captured["local_files_only"] is True
assert "*.safetensors" in captured["allow_patterns"]
finally:
import shutil
shutil.rmtree(output)
def test_materialize_failure_removes_partial_checkpoint(tmp_path, monkeypatch):
import tempfile
source = tmp_path / "source"
source.mkdir()
_write_config(source, {
"model_type": "gpt2",
"quantization_config": {"quant_method": "fp8"},
})
partial = tmp_path / "partial-output"
def fake_mkdtemp(*, prefix):
assert prefix == "obliteratus_dequant_"
partial.mkdir()
return str(partial)
monkeypatch.setattr(tempfile, "mkdtemp", fake_mkdtemp)
with pytest.raises(RuntimeError, match="no safetensors weights"):
qd.materialize_dequantized_checkpoint(
str(source),
qd.QuantDetection(qd.QuantScheme.FP8_BLOCKWISE),
)
assert not partial.exists()
def test_materialize_sharded_checkpoint_rewrites_index_and_copies_metadata(tmp_path):
import shutil
from safetensors.torch import load_file, save_file
save_file({"a.weight": torch.ones(1)}, tmp_path / "part-1.safetensors")
save_file({"b.weight": torch.ones(1)}, tmp_path / "part-2.safetensors")
(tmp_path / "model.safetensors.index.json").write_text(json.dumps({
"metadata": {"total_size": 8},
"weight_map": {
"a.weight": "part-1.safetensors",
"b.weight": "part-2.safetensors",
},
}))
_write_config(tmp_path, {
"model_type": "gpt2",
"quantization_config": {"quant_method": "fp8"},
})
(tmp_path / "tokenizer_config.json").write_text('{"test": true}')
output, _ = qd.materialize_dequantized_checkpoint(
str(tmp_path),
qd.QuantDetection(qd.QuantScheme.FP8_BLOCKWISE),
out_dtype=torch.float32,
)
try:
index = json.loads(Path(output, "model.safetensors.index.json").read_text())
assert index["metadata"] == {"total_size": 8}
assert index["weight_map"] == {
"a.weight": "part-1.safetensors",
"b.weight": "part-2.safetensors",
}
assert load_file(Path(output, "part-1.safetensors"))["a.weight"].item() == 1
assert json.loads(Path(output, "tokenizer_config.json").read_text()) == {
"test": True,
}
finally:
shutil.rmtree(output)
@requires_fp8
def test_materialize_resolves_scale_tensor_across_shard_boundary(tmp_path):
import shutil
from safetensors.torch import load_file, save_file
weight = torch.tensor([[0.5, 1.0]], dtype=torch.float8_e4m3fn)
save_file({"a.weight": weight}, tmp_path / "part-1.safetensors")
save_file(
{
"a.weight_scale": torch.tensor([2.0]),
"norm.weight": torch.ones(1),
},
tmp_path / "part-2.safetensors",
)
(tmp_path / "model.safetensors.index.json").write_text(json.dumps({
"metadata": {"total_size": 6},
"weight_map": {
"a.weight": "part-1.safetensors",
"a.weight_scale": "part-2.safetensors",
"norm.weight": "part-2.safetensors",
},
}))
_write_config(tmp_path, {
"model_type": "gpt2",
"quantization_config": {"quant_method": "modelopt", "quant_algo": "FP8"},
})
output, _ = qd.materialize_dequantized_checkpoint(
str(tmp_path),
qd.QuantDetection(qd.QuantScheme.FP8_PER_CHANNEL_CT),
out_dtype=torch.float32,
)
try:
index = json.loads(Path(output, "model.safetensors.index.json").read_text())
assert index["metadata"]["total_size"] == 12
assert index["weight_map"] == {
"a.weight": "part-1.safetensors",
"norm.weight": "part-2.safetensors",
}
assert load_file(Path(output, "part-1.safetensors"))["a.weight"].tolist() == [
[1.0, 2.0],
]
assert set(load_file(Path(output, "part-2.safetensors"))) == {"norm.weight"}
finally:
shutil.rmtree(output)
def test_loader_failure_removes_materialized_checkpoint(tmp_path, monkeypatch):
from obliteratus.models import loader as loader_mod
materialized = tmp_path / "materialized"
materialized.mkdir()
config = SimpleNamespace(quantization_config=None)
monkeypatch.setattr(
loader_mod.AutoConfig,
"from_pretrained",
staticmethod(lambda *args, **kwargs: config),
)
monkeypatch.setattr(
qd,
"detect_quant_scheme",
lambda *args, **kwargs: qd.QuantDetection(qd.QuantScheme.FP8_BLOCKWISE),
)
monkeypatch.setattr(
qd,
"materialize_dequantized_checkpoint",
lambda *args, **kwargs: (str(materialized), "source"),
)
monkeypatch.setattr(loader_mod, "_estimate_model_memory_gb", lambda *args: 0.0)
class FailingModelClass:
@staticmethod
def from_pretrained(**kwargs):
raise OSError("load failed")
monkeypatch.setattr(loader_mod, "_select_model_class", lambda *args: FailingModelClass)
with pytest.raises(OSError, match="load failed"):
loader_mod.load_model(
"org/model",
task="causal_lm",
device="cpu",
dtype="float32",
)
assert not materialized.exists()
def test_smoke_loader_requires_explicit_remote_code_trust():
from scripts import smoke_load_quant
parser = smoke_load_quant.build_parser()
assert parser.parse_args(["org/model"]).trust_remote_code is False
assert parser.parse_args([
"org/model", "--trust-remote-code",
]).trust_remote_code is True
def test_numerical_fixture_provenance_is_machine_readable():
provenance = json.loads(
(Path(__file__).parent / "fixtures" / "quant_dequant_provenance.json").read_text()
)
assert provenance["origin"] == "synthetic-runtime-generated"
assert provenance["external_data"] is False
assert provenance["torch_manual_seeds"] == list(range(9))
assert provenance["oracles"]["fp8_relative_l2_max"] == 0.05
assert provenance["oracles"]["nvfp4_relative_l2_max"] == 0.15
# ---------------------------------------------------------------------------
# State-dict dequantization
# ---------------------------------------------------------------------------
@requires_fp8
def test_upcast_raw_fp8_bytes_and_reject_plain_float():
fp8 = torch.tensor([0.5, -1.0], dtype=torch.float8_e4m3fn)
raw = fp8.view(torch.uint8)
assert torch.equal(qd._upcast_fp8(raw), fp8.float())
with pytest.raises(RuntimeError, match="cannot upcast"):
qd._upcast_fp8(torch.ones(2, dtype=torch.float32))
@requires_fp8
def test_fp8_blockwise_rejects_incompatible_scale_shape():
weight = torch.ones(4, 4, dtype=torch.float8_e4m3fn)
with pytest.raises(RuntimeError, match="incompatible"):
qd.dequantize_fp8_blockwise(weight, torch.ones(1, 1), block_size=(2, 2))
def test_unpack_native_failure_falls_back_to_manual(monkeypatch):
packed = torch.tensor([[0x1B]], dtype=torch.uint8)
monkeypatch.setattr(qd, "_native_fp4_upcast_works", lambda: True)
monkeypatch.setattr(
qd,
"_unpack_e2m1_native",
lambda value: (_ for _ in ()).throw(RuntimeError("unsupported")),
)
assert qd.unpack_e2m1(packed).tolist() == [[-1.5, 0.5]]
def test_nvfp4_inverse_block_scale_and_output_shape():
packed = torch.tensor([[0x11] * 8], dtype=torch.uint8)
output = qd.dequantize_nvfp4(
packed,
torch.full((1, 1), 2.0),
None,
out_shape=(2, 8),
scale_is_inverse=True,
force_manual=True,
)
assert output.shape == (2, 8)
assert torch.all(output == 0.25)
@requires_fp8
def test_dequantize_state_dict_fp8_blockwise():
torch.manual_seed(7)
@@ -350,6 +797,75 @@ def test_dequantize_state_dict_nvfp4_modelopt():
assert cos > 0.995
@requires_fp8
def test_dequantize_state_dict_nvfp4_compressed_tensors_contract():
torch.manual_seed(8)
w = torch.randn(32, 64)
packed, bs, gs = _pack_nvfp4(w, global_reciprocal=True)
sd = {
"layers.0.mlp.weight_packed": packed,
"layers.0.mlp.weight_scale": bs,
"layers.0.mlp.weight_global_scale": gs,
"layers.0.mlp.input_global_scale": torch.tensor(1.0),
}
det = qd.QuantDetection(
qd.QuantScheme.NVFP4_CT,
global_scale_is_inverse=True,
)
out = qd.dequantize_state_dict(sd, det, out_dtype=torch.float32)
assert set(out) == {"layers.0.mlp.weight"}
cos = torch.nn.functional.cosine_similarity(
w.flatten(), out["layers.0.mlp.weight"].flatten(), dim=0,
)
rel = (out["layers.0.mlp.weight"] - w).norm() / w.norm()
assert cos > 0.995
assert rel < 0.15
@requires_fp8
def test_materialize_nvfp4_compressed_tensors_checkpoint_contract(tmp_path):
import shutil
from safetensors.torch import load_file, save_file
torch.manual_seed(8)
w = torch.randn(32, 64)
packed, bs, gs = _pack_nvfp4(w, global_reciprocal=True)
save_file({
"layers.0.mlp.weight_packed": packed,
"layers.0.mlp.weight_scale": bs,
"layers.0.mlp.weight_global_scale": gs,
"layers.0.mlp.input_global_scale": torch.tensor(1.0),
"layers.0.norm.weight": torch.ones(32),
}, tmp_path / "model.safetensors")
_write_config(tmp_path, {
"model_type": "gpt2",
"quantization_config": {
"quant_method": "compressed-tensors",
"format": "nvfp4-pack-quantized",
},
})
det = qd.QuantDetection(
qd.QuantScheme.NVFP4_CT,
global_scale_is_inverse=True,
)
output, source = qd.materialize_dequantized_checkpoint(
str(tmp_path), det, out_dtype=torch.float32,
)
try:
assert source == str(tmp_path)
state = load_file(os.path.join(output, "model.safetensors"))
assert set(state) == {"layers.0.mlp.weight", "layers.0.norm.weight"}
rel = (state["layers.0.mlp.weight"] - w).norm() / w.norm()
assert rel < 0.15
materialized_config = json.loads(
Path(output, "config.json").read_text(encoding="utf-8")
)
assert "quantization_config" not in materialized_config
finally:
shutil.rmtree(output)
@requires_fp8
def test_dequantize_state_dict_fp8_missing_scale_raises():
q = torch.randn(16, 16).to(torch.float8_e4m3fn)
@@ -358,6 +874,47 @@ def test_dequantize_state_dict_fp8_missing_scale_raises():
qd.dequantize_state_dict({"a.weight": q}, det)
@requires_fp8
def test_dequantize_state_dict_fp8_per_channel():
state = {
"a.weight": torch.tensor([[0.5, 1.0]], dtype=torch.float8_e4m3fn),
"a.weight_scale": torch.tensor([2.0]),
}
out = qd.dequantize_state_dict(
state,
qd.QuantDetection(qd.QuantScheme.FP8_PER_CHANNEL_CT),
out_dtype=torch.float32,
)
assert set(out) == {"a.weight"}
assert out["a.weight"].tolist() == [[1.0, 2.0]]
def test_dequantize_state_dict_rejects_packed_weight_for_wrong_scheme():
state = {
"a.weight": torch.zeros(1, 8, dtype=torch.uint8),
"a.weight_scale": torch.ones(1, 1),
}
with pytest.raises(RuntimeError, match="unsupported layout"):
qd.dequantize_state_dict(
state,
qd.QuantDetection(qd.QuantScheme.FP8_BLOCKWISE),
)
def test_dequantize_state_dict_rejects_unscaled_packed_weight():
with pytest.raises(RuntimeError, match="no recognizable scale siblings"):
qd.dequantize_state_dict(
{"a.weight_packed": torch.zeros(1, 8, dtype=torch.uint8)},
qd.QuantDetection(qd.QuantScheme.NVFP4_CT),
)
def test_scale_key_requires_matching_weight():
state = {"logit_scale": torch.tensor(1.0)}
assert qd.is_scale_key("logit_scale", state) is False
assert qd.is_scale_key("ordinary.weight", state) is False
# ---------------------------------------------------------------------------
# Surgery guards
# ---------------------------------------------------------------------------
@@ -403,6 +960,81 @@ def test_dequantize_weight_plain_uint8_still_converts():
assert w.is_floating_point()
def test_fused_uint8_guard_requires_quantization_scale_sibling():
from obliteratus.abliterate import AbliterationPipeline
container = torch.nn.Module()
container.w1 = torch.nn.Parameter(
torch.zeros(2, 4, 2, dtype=torch.uint8), requires_grad=False,
)
container.w1_scale = torch.ones(2, 4, 1)
with pytest.raises(RuntimeError, match="without dequantization"):
AbliterationPipeline._project_fused_3d(
container,
torch.ones(2, 1),
["w1"],
norm_preserve=False,
scale=1.0,
)
@requires_fp8
def test_fused_fp8_guard_rejects_raw_storage():
from obliteratus.abliterate import AbliterationPipeline
container = torch.nn.Module()
container.w1 = torch.nn.Parameter(
torch.zeros(2, 4, 2).to(torch.float8_e4m3fn), requires_grad=False,
)
with pytest.raises(RuntimeError, match="without dequantization"):
AbliterationPipeline._project_fused_3d(
container,
torch.ones(2, 1),
["w1"],
norm_preserve=False,
scale=1.0,
)
def test_granular_fused_uint8_guard_rejects_packed_storage():
from obliteratus.abliterate import AbliterationPipeline
container = torch.nn.Module()
container.w1 = torch.nn.Parameter(
torch.zeros(2, 4, 2, dtype=torch.uint8), requires_grad=False,
)
container.w1_scale = torch.ones(2, 4, 1)
with pytest.raises(RuntimeError, match="without dequantization"):
AbliterationPipeline._project_fused_3d_granular(
container,
torch.ones(2),
{},
["w1"],
norm_preserve=False,
scale=1.0,
)
def test_selective_fused_uint8_guard_rejects_packed_storage():
from obliteratus.abliterate import AbliterationPipeline
container = torch.nn.Module()
container.w1 = torch.nn.Parameter(
torch.zeros(2, 4, 2, dtype=torch.uint8), requires_grad=False,
)
container.w1_global_scale = torch.ones(1)
with pytest.raises(RuntimeError, match="without dequantization"):
AbliterationPipeline._project_fused_3d_selective_inversion(
container,
torch.ones(2),
["w1"],
safety_indices={0},
reflect_scale=2.0,
remove_scale=1.0,
norm_preserve=False,
)
# ---------------------------------------------------------------------------
# End-to-end loader on tiny synthetic checkpoints
# ---------------------------------------------------------------------------
@@ -63,6 +63,15 @@ class _NoEmbeddingModel(nn.Module):
self.model.layers = nn.ModuleList([nn.Module()])
class _NemotronHModel(nn.Module):
def __init__(self):
super().__init__()
self.backbone = nn.Module()
layer = nn.Module()
layer.mixer = nn.Module()
self.backbone.layers = nn.ModuleList([layer])
def _handle(model: nn.Module, *, architecture: str, hidden_size: int = 8, num_layers: int = 1, num_heads: int = 2):
return ModelHandle(
model=model,
@@ -93,6 +102,15 @@ def test_strategy_navigation_resolves_fallback_layers_and_missing_attention():
get_attention_module(broken, "qwen3_5_moe")
def test_nemotron_h_navigation_uses_backbone_layers_and_mixer():
handle = _handle(_NemotronHModel(), architecture="nemotron_h")
layers = get_layer_modules(handle)
assert layers is handle.model.backbone.layers
assert get_attention_module(layers[0], handle.architecture) is layers[0].mixer
assert get_ffn_module(layers[0], handle.architecture) is layers[0].mixer
def test_head_pruning_zeros_qkv_and_output_slices_for_standard_attention():
model = nn.Module()
model.model = nn.Module()