mirror of
https://github.com/elder-plinius/OBLITERATUS.git
synced 2026-08-17 16:37:30 +02:00
fix: harden quantized checkpoint integration
This commit is contained in:
@@ -486,7 +486,7 @@ Quantization roughly halves the GPU count at each step down. A 70B model that ne
|
||||
| FP8 | DeepSeek-style block-wise (`weight_scale_inv` + `weight_block_size`), compressed-tensors per-channel, ModelOpt FP8 |
|
||||
| NVFP4 | ModelOpt (`weight` + `weight_scale` + `weight_scale_2`), compressed-tensors NVFP4 |
|
||||
|
||||
Two things to know: peak VRAM is the **BF16 size** of the model (not the quantized size), and the output is saved as BF16 — re-quantize afterward with llm-compressor or modelopt if you want a quantized serving artifact. Other quantization schemes (fbgemm, quanto, W4A4) fail loudly with a message naming the scheme.
|
||||
Two things to know: peak VRAM is the **BF16 size** of the model (not the quantized size), and the output is saved as BF16 — re-quantize afterward with llm-compressor or modelopt if you want a quantized serving artifact. Unsupported or ambiguous quantization layouts fail loudly with a message naming the scheme.
|
||||
|
||||
### GPU calculator
|
||||
|
||||
|
||||
@@ -128,7 +128,11 @@
|
||||
"runner": "self-hosted, linux, x64, cuda",
|
||||
"prerequisites": "ENABLE_CUDA_GATE=true and a dedicated CUDA runner",
|
||||
"expected_cost": "under 15 self-hosted runner-minutes",
|
||||
"coverage_paths": ["obliteratus/device.py", "obliteratus/models/loader.py"]
|
||||
"coverage_paths": [
|
||||
"obliteratus/device.py",
|
||||
"obliteratus/models/loader.py",
|
||||
"obliteratus/models/quant_dequant.py"
|
||||
]
|
||||
},
|
||||
{
|
||||
"id": "bitsandbytes-runtime",
|
||||
|
||||
@@ -75,6 +75,7 @@
|
||||
"obliteratus/mlx_backend.py",
|
||||
"obliteratus/model_profile.py",
|
||||
"obliteratus/models/loader.py",
|
||||
"obliteratus/models/quant_dequant.py",
|
||||
"obliteratus/presets.py",
|
||||
"obliteratus/runtime_contracts.py",
|
||||
"obliteratus/study_presets.py"
|
||||
@@ -86,6 +87,7 @@
|
||||
"tests/test_config_properties.py",
|
||||
"tests/test_device_boundaries.py",
|
||||
"tests/test_loader_boundaries.py",
|
||||
"tests/test_quant_dequant.py",
|
||||
"tests/test_mlx_backend_boundaries.py",
|
||||
"tests/test_model_profile.py",
|
||||
"tests/test_model_profile_contracts.py",
|
||||
@@ -350,6 +352,13 @@
|
||||
"required_tests": ["tests/test_loader_boundaries.py", "tests/test_offline_integration.py"],
|
||||
"conditional_gates": ["cuda-runtime", "bitsandbytes-runtime"]
|
||||
},
|
||||
{
|
||||
"path": "obliteratus/models/quant_dequant.py",
|
||||
"risk_class": "mixed-runtime",
|
||||
"risk": "FP8 and NVFP4 format detection, numerical dequantization, Hub revision integrity, and checkpoint materialization",
|
||||
"required_tests": ["tests/test_quant_dequant.py"],
|
||||
"conditional_gates": ["cuda-runtime"]
|
||||
},
|
||||
{
|
||||
"path": "obliteratus/mlx_backend.py",
|
||||
"risk_class": "mixed-runtime",
|
||||
|
||||
+23
-12
@@ -1821,10 +1821,13 @@ protocol; anything here that surprised us is marked **(finding)**.
|
||||
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.
|
||||
4. **compressed-tensors inverts only the global NVFP4 scale.** Its documented
|
||||
converter maps ModelOpt `weight` to `weight_packed`, leaves `weight_scale`
|
||||
unchanged, and maps `weight_scale_2` to the reciprocal
|
||||
`weight_global_scale`. Dequantization therefore multiplies by the FP8 block
|
||||
scale and divides by the global scale. Treating both scales as reciprocal
|
||||
produced a self-consistent synthetic test but did not match the published
|
||||
checkpoint contract **(maintainer audit finding)**.
|
||||
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)
|
||||
@@ -1833,9 +1836,11 @@ protocol; anything here that surprised us is marked **(finding)**.
|
||||
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.
|
||||
7. NVFP4 round-trip error on the deterministic random-Gaussian fixtures is
|
||||
below 0.15 relative Frobenius error with cosine above 0.995 for both
|
||||
ModelOpt and compressed-tensors global-scale storage. Seeds, construction,
|
||||
and thresholds are recorded in
|
||||
`tests/fixtures/quant_dequant_provenance.json`.
|
||||
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
|
||||
@@ -1857,7 +1862,12 @@ Smoke-test results on real checkpoints: see Appendix E continuation below
|
||||
|
||||
---
|
||||
|
||||
## Appendix E2: Nemotron-3-Nano-Omni smoke runs (2026-07-24, 1x A100-80GB)
|
||||
## Appendix E2: Contributor-reported Nemotron-3-Nano-Omni smoke runs (2026-07-24)
|
||||
|
||||
The following A100 results were supplied with the contribution. They do not
|
||||
include immutable checkpoint revisions or retained machine-verifiable logs and
|
||||
were not independently reproduced during maintainer integration. They are
|
||||
historical engineering notes, not release evidence or a support claim.
|
||||
|
||||
Real-checkpoint validation of the FP8/NVFP4 dequantization path. Both
|
||||
targets loaded with **zero quantized tensors left, zero NaNs, and
|
||||
@@ -1944,10 +1954,11 @@ Notes:
|
||||
- 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.
|
||||
exercised by the contributor's Nemotron FP8 run, and Llama-3.1 is a gated
|
||||
repo. CT-NVFP4 is covered by deterministic synthetic contract tests only;
|
||||
those tests follow the published `weight_packed`/`weight_scale`/
|
||||
`weight_global_scale` mapping. A pinned public-checkpoint run remains
|
||||
necessary before making a runtime support claim.
|
||||
|
||||
---
|
||||
|
||||
|
||||
@@ -46,6 +46,20 @@ _QUANT_SCALE_ATTRS = (
|
||||
"weight_global_scale",
|
||||
)
|
||||
|
||||
_FUSED_QUANT_SCALE_SUFFIXES = (
|
||||
"_scale",
|
||||
"_scale_2",
|
||||
"_scale_inv",
|
||||
"_global_scale",
|
||||
)
|
||||
|
||||
|
||||
def _has_fused_quant_scale(container: nn.Module, name: str) -> bool:
|
||||
return any(
|
||||
hasattr(container, f"{name}{suffix}")
|
||||
for suffix in _FUSED_QUANT_SCALE_SUFFIXES
|
||||
)
|
||||
|
||||
# Reduce CUDA memory fragmentation for large models. Must be set before any
|
||||
# CUDA allocations, so we do it at import time. This is the PyTorch-recommended
|
||||
# fix for "reserved but unallocated" memory issues.
|
||||
@@ -5021,7 +5035,9 @@ class AbliterationPipeline:
|
||||
data = param.data
|
||||
# FP8/packed fused params are quantized storage — the loader
|
||||
# dequantizes them; raw upcasting here would corrupt scales.
|
||||
if data.dtype in _FP8_DTYPES or data.dtype == torch.uint8:
|
||||
if data.dtype in _FP8_DTYPES or (
|
||||
data.dtype == torch.uint8 and _has_fused_quant_scale(container, name)
|
||||
):
|
||||
raise RuntimeError(
|
||||
f"FP8/NVFP4 fused-expert weight '{name}' reached surgery "
|
||||
f"without dequantization — the loader should have "
|
||||
@@ -5611,8 +5627,13 @@ class AbliterationPipeline:
|
||||
if data.shape[-1] != hidden_dim and data.shape[-2] != hidden_dim:
|
||||
continue
|
||||
|
||||
if (data.dtype in _FP8_DTYPES or data.dtype == torch.uint8) \
|
||||
and not hasattr(param, "quant_state"):
|
||||
if (
|
||||
data.dtype in _FP8_DTYPES
|
||||
or (
|
||||
data.dtype == torch.uint8
|
||||
and _has_fused_quant_scale(container, pname)
|
||||
)
|
||||
) and not hasattr(param, "quant_state"):
|
||||
raise RuntimeError(
|
||||
f"FP8/NVFP4 fused-expert weight '{pname}' reached surgery "
|
||||
f"without dequantization — the loader should have "
|
||||
@@ -5732,8 +5753,13 @@ class AbliterationPipeline:
|
||||
if data.shape[-1] != hidden_dim and data.shape[-2] != hidden_dim:
|
||||
continue
|
||||
|
||||
if (data.dtype in _FP8_DTYPES or data.dtype == torch.uint8) \
|
||||
and not hasattr(param, "quant_state"):
|
||||
if (
|
||||
data.dtype in _FP8_DTYPES
|
||||
or (
|
||||
data.dtype == torch.uint8
|
||||
and _has_fused_quant_scale(container, pname)
|
||||
)
|
||||
) and not hasattr(param, "quant_state"):
|
||||
raise RuntimeError(
|
||||
f"FP8/NVFP4 fused-expert weight '{pname}' reached surgery "
|
||||
f"without dequantization — the loader should have "
|
||||
|
||||
@@ -259,45 +259,88 @@ except Exception:
|
||||
# legacy ``_tied_weights_keys`` (list of names or dict) across modules into
|
||||
# the {param_name: target_name} mapping accelerate expects.
|
||||
# Affected: Nemotron Omni (NemotronH_Nano_Omni_Reasoning_V3)
|
||||
def _install_all_tied_weights_keys_compat(model_cls) -> bool:
|
||||
"""Install the legacy tied-weight adapter only when the class lacks it."""
|
||||
if hasattr(model_cls, "all_tied_weights_keys"):
|
||||
return False
|
||||
|
||||
def _get_all_tied_weights_keys(self):
|
||||
cached = self.__dict__.get("_all_tied_weights_keys_store")
|
||||
if cached is not None:
|
||||
return cached
|
||||
expander = getattr(self, "get_expanded_tied_weights_keys", None)
|
||||
if callable(expander):
|
||||
try:
|
||||
expanded = expander(all_submodels=True)
|
||||
except Exception:
|
||||
# Legacy remote implementations may still expose the modern
|
||||
# method name while keeping an incompatible list contract.
|
||||
pass
|
||||
else:
|
||||
self.__dict__["_all_tied_weights_keys_store"] = expanded
|
||||
return expanded
|
||||
tied: dict = {}
|
||||
for module in self.modules():
|
||||
keys = getattr(module, "_tied_weights_keys", None)
|
||||
if not keys:
|
||||
continue
|
||||
if isinstance(keys, dict):
|
||||
candidates = keys
|
||||
else: # legacy list of patterns/names — one tied group
|
||||
names = list(keys)
|
||||
candidates = {name: names[0] for name in names}
|
||||
for name, target in candidates.items():
|
||||
# Only expose keys that actually resolve on this model —
|
||||
# multimodal wrappers inherit "lm_head.weight" from their
|
||||
# LM submodule's config but have no top-level lm_head.
|
||||
try:
|
||||
self.get_submodule(name.rsplit(".", 1)[0])
|
||||
except AttributeError:
|
||||
continue
|
||||
tied.setdefault(name, target)
|
||||
return tied
|
||||
|
||||
def _set_all_tied_weights_keys(self, value):
|
||||
self.__dict__["_all_tied_weights_keys_store"] = value
|
||||
|
||||
model_cls.all_tied_weights_keys = property(
|
||||
_get_all_tied_weights_keys, _set_all_tied_weights_keys,
|
||||
)
|
||||
return True
|
||||
|
||||
|
||||
try:
|
||||
from transformers import PreTrainedModel as _PTM
|
||||
if not hasattr(_PTM, "all_tied_weights_keys"):
|
||||
def _get_all_tied_weights_keys(self):
|
||||
cached = self.__dict__.get("_all_tied_weights_keys_store")
|
||||
if cached is not None:
|
||||
return cached
|
||||
tied: dict = {}
|
||||
for module in self.modules():
|
||||
keys = getattr(module, "_tied_weights_keys", None)
|
||||
if not keys:
|
||||
continue
|
||||
if isinstance(keys, dict):
|
||||
candidates = keys
|
||||
else: # legacy list of patterns/names — one tied group
|
||||
names = list(keys)
|
||||
candidates = {name: names[0] for name in names}
|
||||
for name, target in candidates.items():
|
||||
# Only expose keys that actually resolve on this model —
|
||||
# multimodal wrappers inherit "lm_head.weight" from their
|
||||
# LM submodule's config but have no top-level lm_head.
|
||||
try:
|
||||
self.get_submodule(name.rsplit(".", 1)[0])
|
||||
except AttributeError:
|
||||
continue
|
||||
tied.setdefault(name, target)
|
||||
return tied
|
||||
|
||||
def _set_all_tied_weights_keys(self, value):
|
||||
# transformers 5.x post_init() assigns this attribute.
|
||||
self.__dict__["_all_tied_weights_keys_store"] = value
|
||||
|
||||
_PTM.all_tied_weights_keys = property(
|
||||
_get_all_tied_weights_keys, _set_all_tied_weights_keys,
|
||||
)
|
||||
_install_all_tied_weights_keys_compat(_PTM)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
|
||||
def _force_eager_config(config) -> None:
|
||||
"""Recursively replace unavailable Flash Attention 2 selections."""
|
||||
from transformers import PretrainedConfig
|
||||
|
||||
if getattr(config, "_attn_implementation", None) == "flash_attention_2":
|
||||
config._attn_implementation = "eager"
|
||||
for subconfig in vars(config).values():
|
||||
if isinstance(subconfig, PretrainedConfig):
|
||||
_force_eager_config(subconfig)
|
||||
|
||||
|
||||
def _config_uses_flash_attention_2(config) -> bool:
|
||||
"""Return whether this config or a nested Transformers config selects FA2."""
|
||||
from transformers import PretrainedConfig
|
||||
|
||||
if getattr(config, "_attn_implementation", None) == "flash_attention_2":
|
||||
return True
|
||||
return any(
|
||||
_config_uses_flash_attention_2(subconfig)
|
||||
for subconfig in vars(config).values()
|
||||
if isinstance(subconfig, PretrainedConfig)
|
||||
)
|
||||
|
||||
|
||||
# ── Deferred shims ──────────────────────────────────────────────────
|
||||
# Some patches must wait until the _LazyModule has fully initialized
|
||||
# (it replaces its __dict__ during bootstrap). We apply these once,
|
||||
@@ -661,7 +704,12 @@ def load_model(
|
||||
# FP8 / NVFP4 checkpoints: dequantize to a plain float copy on disk,
|
||||
# then load through the normal float path. Anything quantized that we
|
||||
# don't explicitly support fails loudly here, before any weight loads.
|
||||
quant_detection = qd.detect_quant_scheme(model_name, token=token)
|
||||
quant_detection = qd.detect_quant_scheme(
|
||||
model_name,
|
||||
token=token,
|
||||
revision=revision,
|
||||
local_files_only=local_files_only,
|
||||
)
|
||||
_dequant_tmp = None
|
||||
_dequant_source = None
|
||||
if quant_detection.scheme is qd.QuantScheme.UNSUPPORTED:
|
||||
@@ -677,7 +725,12 @@ def load_model(
|
||||
)
|
||||
_dequant_source = model_name
|
||||
_dequant_tmp, _ = qd.materialize_dequantized_checkpoint(
|
||||
model_name, quant_detection, out_dtype=torch_dtype, token=token,
|
||||
model_name,
|
||||
quant_detection,
|
||||
out_dtype=torch_dtype,
|
||||
token=token,
|
||||
revision=revision,
|
||||
local_files_only=local_files_only,
|
||||
)
|
||||
model_name = _dequant_tmp
|
||||
config = AutoConfig.from_pretrained(
|
||||
@@ -709,25 +762,7 @@ def load_model(
|
||||
# If flash-attn isn't installed, loading would fail — fall back to eager.
|
||||
# Multimodal wrappers pin it on nested sub-configs too (e.g. Nemotron
|
||||
# Omni's llm_config), so walk those as well.
|
||||
def _force_eager_cfg(cfg) -> None:
|
||||
from transformers import PretrainedConfig as _PTC
|
||||
if getattr(cfg, "_attn_implementation", None) == "flash_attention_2":
|
||||
cfg._attn_implementation = "eager"
|
||||
for sub in vars(cfg).values():
|
||||
if isinstance(sub, _PTC):
|
||||
_force_eager_cfg(sub)
|
||||
|
||||
def _has_flash_attention(cfg) -> bool:
|
||||
from transformers import PretrainedConfig as _PTC
|
||||
if getattr(cfg, "_attn_implementation", None) == "flash_attention_2":
|
||||
return True
|
||||
return any(
|
||||
_has_flash_attention(sub)
|
||||
for sub in vars(cfg).values()
|
||||
if isinstance(sub, _PTC)
|
||||
)
|
||||
|
||||
if _has_flash_attention(config):
|
||||
if _config_uses_flash_attention_2(config):
|
||||
try:
|
||||
from transformers.utils import is_flash_attn_2_available
|
||||
_fa2 = is_flash_attn_2_available()
|
||||
@@ -738,7 +773,7 @@ def load_model(
|
||||
"config.json pins flash_attention_2 but flash-attn is not "
|
||||
"installed — falling back to eager attention."
|
||||
)
|
||||
_force_eager_cfg(config)
|
||||
_force_eager_config(config)
|
||||
|
||||
model_cls = _select_model_class(task, config)
|
||||
load_kwargs: dict = {
|
||||
@@ -861,12 +896,14 @@ def load_model(
|
||||
f"pip install git+https://github.com/huggingface/transformers.git"
|
||||
) from e
|
||||
raise
|
||||
except Exception:
|
||||
# Any other load failure — don't leak the temp dequantized checkpoint.
|
||||
finally:
|
||||
# from_pretrained fully materializes weights before returning. Clean up
|
||||
# on success and across every typed/retry failure path above.
|
||||
if _dequant_tmp is not None:
|
||||
import shutil as _shutil
|
||||
|
||||
_shutil.rmtree(_dequant_tmp, ignore_errors=True)
|
||||
raise
|
||||
logger.info("Removed temporary dequantized checkpoint %s", _dequant_tmp)
|
||||
|
||||
if load_policy.move_to_resolved_device:
|
||||
# Explicit devices and auto-selected MPS/CPU load on CPU before moving.
|
||||
@@ -937,12 +974,7 @@ def load_model(
|
||||
dev.empty_cache()
|
||||
|
||||
if _dequant_tmp is not None:
|
||||
# Weights are fully materialized by from_pretrained at this point —
|
||||
# the temp dequantized copy is no longer needed.
|
||||
import shutil as _shutil
|
||||
model._obliteratus_dequantized_scheme = quant_detection.scheme.value
|
||||
_shutil.rmtree(_dequant_tmp, ignore_errors=True)
|
||||
logger.info("Removed temporary dequantized checkpoint %s", _dequant_tmp)
|
||||
|
||||
try:
|
||||
tokenizer = AutoTokenizer.from_pretrained(
|
||||
|
||||
@@ -29,8 +29,11 @@ NVFP4 (ModelOpt, ``quant_algo: "NVFP4"``):
|
||||
``<name>.weight_scale_2`` float32 scalar — global scale (amax/2688)
|
||||
dequant: e2m1_values * weight_scale * weight_scale_2
|
||||
|
||||
NVFP4 (compressed-tensors): same layout but scales are stored as
|
||||
reciprocals and the global scale may be named ``weight_global_scale``.
|
||||
NVFP4 (compressed-tensors, ``format: "nvfp4-pack-quantized"``):
|
||||
``<name>.weight_packed`` uint8 packed E2M1 values
|
||||
``<name>.weight_scale`` FP8 block scales (multiplicative)
|
||||
``<name>.weight_global_scale`` FP32 reciprocal global scale
|
||||
dequant: e2m1_values * weight_scale / weight_global_scale.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
@@ -97,12 +100,17 @@ class QuantDetection:
|
||||
reason: str = ""
|
||||
block_size: Tuple[int, int] = FP8_DEFAULT_BLOCK
|
||||
group_size: int = NVFP4_GROUP_SIZE
|
||||
scale_is_inverse: bool = False # compressed-tensors stores reciprocals
|
||||
scale_is_inverse: bool = False
|
||||
global_scale_is_inverse: bool = False
|
||||
raw_quant_config: dict = field(default_factory=dict)
|
||||
|
||||
|
||||
def _load_json_from_checkpoint(
|
||||
model_name_or_path: str, filename: str, token: Optional[str] = None,
|
||||
model_name_or_path: str,
|
||||
filename: str,
|
||||
token: Optional[str] = None,
|
||||
revision: Optional[str] = None,
|
||||
local_files_only: bool = False,
|
||||
) -> Optional[dict]:
|
||||
"""Read a JSON metadata file from a local dir or the HF hub."""
|
||||
if os.path.isdir(model_name_or_path):
|
||||
@@ -117,7 +125,13 @@ def _load_json_from_checkpoint(
|
||||
except Exception: # pragma: no cover - huggingface_hub always ships w/ transformers
|
||||
return None
|
||||
try:
|
||||
path = hf_hub_download(model_name_or_path, filename, token=token)
|
||||
path = hf_hub_download(
|
||||
model_name_or_path,
|
||||
filename,
|
||||
token=token,
|
||||
revision=revision,
|
||||
local_files_only=local_files_only,
|
||||
)
|
||||
except EntryNotFoundError:
|
||||
return None
|
||||
except Exception as exc:
|
||||
@@ -128,12 +142,20 @@ def _load_json_from_checkpoint(
|
||||
|
||||
|
||||
def _safetensors_key_names(
|
||||
model_name_or_path: str, config_json: Optional[dict], token: Optional[str] = None,
|
||||
model_name_or_path: str,
|
||||
config_json: Optional[dict],
|
||||
token: Optional[str] = None,
|
||||
revision: Optional[str] = None,
|
||||
local_files_only: bool = False,
|
||||
) -> set:
|
||||
"""Collect tensor key names from the safetensors index (no weight loads)."""
|
||||
keys: set = set()
|
||||
index = _load_json_from_checkpoint(
|
||||
model_name_or_path, "model.safetensors.index.json", token=token,
|
||||
model_name_or_path,
|
||||
"model.safetensors.index.json",
|
||||
token=token,
|
||||
revision=revision,
|
||||
local_files_only=local_files_only,
|
||||
)
|
||||
if index and "weight_map" in index:
|
||||
keys.update(index["weight_map"].keys())
|
||||
@@ -144,7 +166,13 @@ def _safetensors_key_names(
|
||||
else:
|
||||
try:
|
||||
from huggingface_hub import hf_hub_download
|
||||
st_path = hf_hub_download(model_name_or_path, "model.safetensors", token=token)
|
||||
st_path = hf_hub_download(
|
||||
model_name_or_path,
|
||||
"model.safetensors",
|
||||
token=token,
|
||||
revision=revision,
|
||||
local_files_only=local_files_only,
|
||||
)
|
||||
except Exception:
|
||||
return keys
|
||||
if not os.path.exists(st_path):
|
||||
@@ -158,7 +186,12 @@ def _safetensors_key_names(
|
||||
return keys
|
||||
|
||||
|
||||
def detect_quant_scheme(model_name_or_path: str, token: Optional[str] = None) -> QuantDetection:
|
||||
def detect_quant_scheme(
|
||||
model_name_or_path: str,
|
||||
token: Optional[str] = None,
|
||||
revision: Optional[str] = None,
|
||||
local_files_only: bool = False,
|
||||
) -> QuantDetection:
|
||||
"""Classify a checkpoint's quantization without loading any weights.
|
||||
|
||||
Peeks at ``config.json``'s ``quantization_config`` plus safetensors
|
||||
@@ -166,7 +199,13 @@ def detect_quant_scheme(model_name_or_path: str, token: Optional[str] = None) ->
|
||||
reported as UNSUPPORTED with a human-readable reason — the loader
|
||||
turns that into a loud error rather than silent corruption.
|
||||
"""
|
||||
config_json = _load_json_from_checkpoint(model_name_or_path, "config.json", token=token) or {}
|
||||
config_json = _load_json_from_checkpoint(
|
||||
model_name_or_path,
|
||||
"config.json",
|
||||
token=token,
|
||||
revision=revision,
|
||||
local_files_only=local_files_only,
|
||||
) or {}
|
||||
qcfg = config_json.get("quantization_config")
|
||||
if not qcfg:
|
||||
return QuantDetection(QuantScheme.NONE)
|
||||
@@ -182,7 +221,13 @@ def detect_quant_scheme(model_name_or_path: str, token: Optional[str] = None) ->
|
||||
block_size=(int(block[0]), int(block[1])),
|
||||
raw_quant_config=raw,
|
||||
)
|
||||
keys = _safetensors_key_names(model_name_or_path, config_json, token=token)
|
||||
keys = _safetensors_key_names(
|
||||
model_name_or_path,
|
||||
config_json,
|
||||
token=token,
|
||||
revision=revision,
|
||||
local_files_only=local_files_only,
|
||||
)
|
||||
if any(k.endswith("weight_scale_inv") for k in keys):
|
||||
return QuantDetection(QuantScheme.FP8_BLOCKWISE, raw_quant_config=raw)
|
||||
return QuantDetection(
|
||||
@@ -235,25 +280,30 @@ def detect_quant_scheme(model_name_or_path: str, token: Optional[str] = None) ->
|
||||
wtype = str(wcfg.get("type", "")).lower()
|
||||
if num_bits == 8 and wtype == "float":
|
||||
return QuantDetection(QuantScheme.FP8_PER_CHANNEL_CT, raw_quant_config=raw)
|
||||
if num_bits == 4 and wtype == "float":
|
||||
if (
|
||||
num_bits == 4
|
||||
and wtype == "float"
|
||||
and qcfg.get("format") == "nvfp4-pack-quantized"
|
||||
):
|
||||
gs = int(wcfg.get("group_size") or NVFP4_GROUP_SIZE)
|
||||
return QuantDetection(
|
||||
QuantScheme.NVFP4_CT,
|
||||
group_size=gs,
|
||||
scale_is_inverse=True,
|
||||
global_scale_is_inverse=True,
|
||||
raw_quant_config=raw,
|
||||
)
|
||||
return QuantDetection(
|
||||
QuantScheme.UNSUPPORTED,
|
||||
reason=(
|
||||
f"compressed-tensors weights num_bits={num_bits} type={wtype!r} "
|
||||
f"not supported (need float 8-bit or NVFP4 4-bit)"
|
||||
f"format={qcfg.get('format')!r} not supported (need float 8-bit "
|
||||
f"or NVFP4 4-bit with format 'nvfp4-pack-quantized')"
|
||||
),
|
||||
raw_quant_config=raw,
|
||||
)
|
||||
|
||||
if quant_method in ("gptq", "awq", "bitsandbytes", "bitsandbytes_4bit",
|
||||
"bitsandbytes_8bit", "mxfp4", ""):
|
||||
"bitsandbytes_8bit", "mxfp4"):
|
||||
# Handled elsewhere in the loader / surgery layer.
|
||||
return QuantDetection(QuantScheme.NONE, raw_quant_config=raw)
|
||||
|
||||
@@ -416,6 +466,7 @@ def dequantize_nvfp4(
|
||||
global_scale: Optional[torch.Tensor],
|
||||
out_shape: Optional[Tuple[int, ...]] = None,
|
||||
scale_is_inverse: bool = False,
|
||||
global_scale_is_inverse: bool = False,
|
||||
group_size: int = NVFP4_GROUP_SIZE,
|
||||
force_manual: bool = False,
|
||||
) -> torch.Tensor:
|
||||
@@ -424,9 +475,11 @@ def dequantize_nvfp4(
|
||||
``packed_uint8``: (M, N/2) uint8, two E2M1 values per byte, low nibble
|
||||
first along the input dim. ``block_scale``: one FP8-E4M3 scale per
|
||||
``group_size`` elements. ``global_scale``: FP32 scalar (ModelOpt
|
||||
``weight_scale_2`` = amax/2688); None means 1.0. With
|
||||
``scale_is_inverse`` (compressed-tensors) the scales are reciprocals
|
||||
and are divided out instead of multiplied in.
|
||||
``weight_scale_2`` = amax/2688); None means 1.0. ModelOpt stores both
|
||||
scales multiplicatively. Compressed-tensors keeps the block scale but
|
||||
stores the global scale as its reciprocal, selected with
|
||||
``global_scale_is_inverse``. ``scale_is_inverse`` is retained for
|
||||
explicitly reciprocal block-scale layouts.
|
||||
"""
|
||||
vals = unpack_e2m1(packed_uint8, force_manual=force_manual)
|
||||
*lead, N = vals.shape
|
||||
@@ -444,7 +497,7 @@ def dequantize_nvfp4(
|
||||
vals = vals.reshape(*lead, N)
|
||||
if global_scale is not None:
|
||||
gs = global_scale.to(torch.float32)
|
||||
if scale_is_inverse:
|
||||
if global_scale_is_inverse:
|
||||
vals.div_(gs)
|
||||
else:
|
||||
vals.mul_(gs)
|
||||
@@ -482,35 +535,51 @@ def _strip_known_suffix(key: str) -> Optional[str]:
|
||||
return None
|
||||
|
||||
|
||||
def find_scale_siblings(state_dict: Dict[str, torch.Tensor], weight_key: str) -> dict:
|
||||
"""Locate the scale tensors belonging to ``weight_key`` (``foo.weight``)."""
|
||||
base = weight_key[: -len(".weight")] if weight_key.endswith(".weight") else weight_key
|
||||
dotted = weight_key.endswith(".weight")
|
||||
out = {}
|
||||
for name, attr in [
|
||||
(".weight_scale_inv", "scale_inv"),
|
||||
(".weight_scale_2", "global_scale"),
|
||||
(".weight_global_scale", "global_scale"),
|
||||
(".weight_scale", "block_scale"),
|
||||
]:
|
||||
k = base + name
|
||||
if k in state_dict:
|
||||
out[attr] = state_dict[k]
|
||||
if not dotted:
|
||||
# Fused-MoE naming: experts.gate_up_proj{,_scale,_scale_2,_scale_inv}
|
||||
for name, attr in [
|
||||
def _scale_sibling_names(weight_key: str) -> tuple[tuple[str, str], ...]:
|
||||
"""Return candidate ``(key, role)`` pairs for one quantized weight."""
|
||||
if weight_key.endswith(".weight_packed"):
|
||||
base = weight_key[: -len(".weight_packed")]
|
||||
dotted = True
|
||||
elif weight_key.endswith(".weight"):
|
||||
base = weight_key[: -len(".weight")]
|
||||
dotted = True
|
||||
else:
|
||||
base = weight_key
|
||||
dotted = False
|
||||
suffixes = (
|
||||
(
|
||||
(".weight_scale_inv", "scale_inv"),
|
||||
(".weight_scale_2", "global_scale"),
|
||||
(".weight_global_scale", "global_scale"),
|
||||
(".weight_scale", "block_scale"),
|
||||
)
|
||||
if dotted
|
||||
else (
|
||||
# Fused-MoE naming:
|
||||
# experts.gate_up_proj{,_scale,_scale_2,_scale_inv}
|
||||
("_scale_inv", "scale_inv"),
|
||||
("_scale_2", "global_scale"),
|
||||
("_global_scale", "global_scale"),
|
||||
("_scale", "block_scale"),
|
||||
]:
|
||||
k = base + name
|
||||
if k in state_dict and attr not in out:
|
||||
out[attr] = state_dict[k]
|
||||
)
|
||||
)
|
||||
return tuple((base + suffix, role) for suffix, role in suffixes)
|
||||
|
||||
|
||||
def find_scale_siblings(state_dict: Dict[str, torch.Tensor], weight_key: str) -> dict:
|
||||
"""Locate the scale tensors belonging to ``weight_key`` (``foo.weight``)."""
|
||||
out = {}
|
||||
for key, role in _scale_sibling_names(weight_key):
|
||||
if key in state_dict and role not in out:
|
||||
out[role] = state_dict[key]
|
||||
return out
|
||||
|
||||
|
||||
def is_scale_key(key: str, state_dict: Optional[Dict[str, torch.Tensor]] = None) -> bool:
|
||||
def is_scale_key(
|
||||
key: str,
|
||||
state_dict: Optional[Dict[str, torch.Tensor]] = None,
|
||||
known_keys: Optional[set[str]] = None,
|
||||
) -> bool:
|
||||
"""True for keys holding quantization scales, not real parameters.
|
||||
|
||||
With ``state_dict`` given, the key only counts as a scale if its base
|
||||
@@ -521,9 +590,15 @@ def is_scale_key(key: str, state_dict: Optional[Dict[str, torch.Tensor]] = None)
|
||||
base = _strip_known_suffix(key)
|
||||
if base is None:
|
||||
return False
|
||||
if state_dict is not None:
|
||||
# Dotted scales pair with ``base.weight``; fused-MoE scales with ``base``.
|
||||
if base not in state_dict and (base + ".weight") not in state_dict:
|
||||
keys = known_keys if known_keys is not None else state_dict
|
||||
if keys is not None:
|
||||
# Dotted scales pair with dense or packed weights; fused-MoE scales
|
||||
# pair directly with ``base``.
|
||||
if (
|
||||
base not in keys
|
||||
and (base + ".weight") not in keys
|
||||
and (base + ".weight_packed") not in keys
|
||||
):
|
||||
return False
|
||||
return True
|
||||
|
||||
@@ -532,6 +607,8 @@ def dequantize_state_dict(
|
||||
state_dict: Dict[str, torch.Tensor],
|
||||
detection: QuantDetection,
|
||||
out_dtype: torch.dtype = torch.bfloat16,
|
||||
sibling_tensors: Optional[Dict[str, torch.Tensor]] = None,
|
||||
known_keys: Optional[set[str]] = None,
|
||||
) -> Dict[str, torch.Tensor]:
|
||||
"""Dequantize every quantized weight group in one checkpoint shard.
|
||||
|
||||
@@ -540,13 +617,16 @@ def dequantize_state_dict(
|
||||
over naming conventions (dense linears and fused MoE experts alike)
|
||||
because grouping is by key suffix, not by module type.
|
||||
"""
|
||||
scale_lookup = state_dict
|
||||
if sibling_tensors:
|
||||
scale_lookup = {**state_dict, **sibling_tensors}
|
||||
out: Dict[str, torch.Tensor] = {}
|
||||
for key, tensor in state_dict.items():
|
||||
if is_scale_key(key, state_dict):
|
||||
if is_scale_key(key, state_dict, known_keys):
|
||||
# Scales are consumed together with their weight (or dropped if
|
||||
# the weight is absent — e.g. quantization status tensors).
|
||||
continue
|
||||
sib = find_scale_siblings(state_dict, key)
|
||||
sib = find_scale_siblings(scale_lookup, key)
|
||||
if is_fp8_dtype(tensor.dtype):
|
||||
if "scale_inv" in sib:
|
||||
w = dequantize_fp8_blockwise(tensor, sib["scale_inv"], detection.block_size)
|
||||
@@ -567,6 +647,7 @@ def dequantize_state_dict(
|
||||
w = dequantize_nvfp4(
|
||||
tensor, scale, gs,
|
||||
scale_is_inverse=detection.scale_is_inverse,
|
||||
global_scale_is_inverse=detection.global_scale_is_inverse,
|
||||
group_size=detection.group_size,
|
||||
)
|
||||
else:
|
||||
@@ -574,8 +655,9 @@ def dequantize_state_dict(
|
||||
f"packed uint8 tensor {key!r} in a {detection.scheme.value} "
|
||||
f"checkpoint — unsupported layout; report a bug"
|
||||
)
|
||||
out[key] = w.to(out_dtype)
|
||||
elif tensor.dtype == torch.uint8 and key.endswith(".weight"):
|
||||
output_key = key[: -len("_packed")] if key.endswith(".weight_packed") else key
|
||||
out[output_key] = w.to(out_dtype)
|
||||
elif tensor.dtype == torch.uint8 and key.endswith((".weight", ".weight_packed")):
|
||||
raise RuntimeError(
|
||||
f"uint8 weight {key!r} has no recognizable scale siblings "
|
||||
f"(keys present: {[k for k in state_dict if k.startswith(key[:-7])]}) "
|
||||
@@ -596,6 +678,8 @@ def materialize_dequantized_checkpoint(
|
||||
detection: QuantDetection,
|
||||
out_dtype: torch.dtype = torch.bfloat16,
|
||||
token: Optional[str] = None,
|
||||
revision: Optional[str] = None,
|
||||
local_files_only: bool = False,
|
||||
) -> Tuple[str, str]:
|
||||
"""Write a dequantized float copy of a checkpoint to a temp dir.
|
||||
|
||||
@@ -608,6 +692,7 @@ def materialize_dequantized_checkpoint(
|
||||
import shutil
|
||||
import tempfile
|
||||
|
||||
from safetensors import safe_open
|
||||
from safetensors.torch import load_file, save_file
|
||||
|
||||
if os.path.isdir(model_name_or_path):
|
||||
@@ -618,6 +703,8 @@ def materialize_dequantized_checkpoint(
|
||||
src = snapshot_download(
|
||||
model_name_or_path,
|
||||
token=token,
|
||||
revision=revision,
|
||||
local_files_only=local_files_only,
|
||||
allow_patterns=[
|
||||
"*.json", "*.safetensors", "*.py", "tokenizer*", "*.model",
|
||||
"*.txt", "chat_template*", "special_tokens_map.json",
|
||||
@@ -631,52 +718,86 @@ def materialize_dequantized_checkpoint(
|
||||
detection.scheme.value, model_name_or_path, out_dtype, tmp,
|
||||
)
|
||||
|
||||
# Copy non-weight files; strip quantization_config from config.json.
|
||||
for name in os.listdir(src):
|
||||
s = os.path.join(src, name)
|
||||
if not os.path.isfile(s) or name.endswith(".safetensors"):
|
||||
continue
|
||||
d = os.path.join(tmp, name)
|
||||
if name == "config.json":
|
||||
with open(s, "r", encoding="utf-8") as fh:
|
||||
cfg = json.load(fh)
|
||||
cfg.pop("quantization_config", None)
|
||||
with open(d, "w", encoding="utf-8") as fh:
|
||||
json.dump(cfg, fh, indent=2)
|
||||
try:
|
||||
# Copy non-weight files; strip quantization_config from config.json.
|
||||
for name in os.listdir(src):
|
||||
s = os.path.join(src, name)
|
||||
if not os.path.isfile(s) or name.endswith(".safetensors"):
|
||||
continue
|
||||
d = os.path.join(tmp, name)
|
||||
if name == "config.json":
|
||||
with open(s, "r", encoding="utf-8") as fh:
|
||||
cfg = json.load(fh)
|
||||
cfg.pop("quantization_config", None)
|
||||
with open(d, "w", encoding="utf-8") as fh:
|
||||
json.dump(cfg, fh, indent=2)
|
||||
else:
|
||||
shutil.copy2(s, d)
|
||||
|
||||
# Locate shards.
|
||||
index_path = os.path.join(src, "model.safetensors.index.json")
|
||||
old_index = None
|
||||
if os.path.exists(index_path):
|
||||
with open(index_path, "r", encoding="utf-8") as fh:
|
||||
old_index = json.load(fh)
|
||||
shards = sorted(set(old_index["weight_map"].values()))
|
||||
elif os.path.exists(os.path.join(src, "model.safetensors")):
|
||||
shards = ["model.safetensors"]
|
||||
else:
|
||||
shutil.copy2(s, d)
|
||||
raise RuntimeError(
|
||||
f"Quantized checkpoint '{model_name_or_path}' has no safetensors "
|
||||
f"weights (pytorch_model.bin quantized checkpoints are not "
|
||||
f"supported) — use a BF16 checkpoint"
|
||||
)
|
||||
|
||||
# Locate shards.
|
||||
index_path = os.path.join(src, "model.safetensors.index.json")
|
||||
old_index = None
|
||||
if os.path.exists(index_path):
|
||||
with open(index_path, "r", encoding="utf-8") as fh:
|
||||
old_index = json.load(fh)
|
||||
shards = sorted(set(old_index["weight_map"].values()))
|
||||
elif os.path.exists(os.path.join(src, "model.safetensors")):
|
||||
shards = ["model.safetensors"]
|
||||
else:
|
||||
raise RuntimeError(
|
||||
f"Quantized checkpoint '{model_name_or_path}' has no safetensors "
|
||||
f"weights (pytorch_model.bin quantized checkpoints are not "
|
||||
f"supported) — use a BF16 checkpoint"
|
||||
)
|
||||
weight_map = dict(old_index.get("weight_map") or {}) if old_index else {}
|
||||
known_keys = set(weight_map) if weight_map else None
|
||||
new_weight_map: Dict[str, str] = {}
|
||||
new_total_size = 0
|
||||
for i, shard in enumerate(shards, 1):
|
||||
sd = load_file(os.path.join(src, shard), device="cpu")
|
||||
external_scales: Dict[str, torch.Tensor] = {}
|
||||
external_by_shard: Dict[str, list[str]] = {}
|
||||
for weight_key in sd:
|
||||
for scale_key, _role in _scale_sibling_names(weight_key):
|
||||
scale_shard = weight_map.get(scale_key)
|
||||
if scale_shard and scale_shard != shard:
|
||||
external_by_shard.setdefault(scale_shard, []).append(scale_key)
|
||||
for scale_shard, scale_keys in external_by_shard.items():
|
||||
with safe_open(
|
||||
os.path.join(src, scale_shard), framework="pt", device="cpu",
|
||||
) as handle:
|
||||
for scale_key in scale_keys:
|
||||
external_scales[scale_key] = handle.get_tensor(scale_key)
|
||||
|
||||
new_weight_map: Dict[str, str] = {}
|
||||
for i, shard in enumerate(shards, 1):
|
||||
sd = load_file(os.path.join(src, shard), device="cpu")
|
||||
out_sd = dequantize_state_dict(sd, detection, out_dtype=out_dtype)
|
||||
save_file(out_sd, os.path.join(tmp, shard), metadata={"format": "pt"})
|
||||
for k in out_sd:
|
||||
new_weight_map[k] = shard
|
||||
logger.info(
|
||||
"dequantized shard %d/%d (%s): %d tensors -> %d float tensors",
|
||||
i, len(shards), shard, len(sd), len(out_sd),
|
||||
)
|
||||
out_sd = dequantize_state_dict(
|
||||
sd,
|
||||
detection,
|
||||
out_dtype=out_dtype,
|
||||
sibling_tensors=external_scales,
|
||||
known_keys=known_keys,
|
||||
)
|
||||
if out_sd:
|
||||
save_file(out_sd, os.path.join(tmp, shard), metadata={"format": "pt"})
|
||||
for k, tensor in out_sd.items():
|
||||
new_weight_map[k] = shard
|
||||
new_total_size += tensor.numel() * tensor.element_size()
|
||||
logger.info(
|
||||
"dequantized shard %d/%d (%s): %d tensors -> %d float tensors",
|
||||
i, len(shards), shard, len(sd), len(out_sd),
|
||||
)
|
||||
|
||||
if old_index is not None:
|
||||
metadata = dict(old_index.get("metadata") or {})
|
||||
with open(os.path.join(tmp, "model.safetensors.index.json"), "w", encoding="utf-8") as fh:
|
||||
json.dump({"metadata": metadata, "weight_map": new_weight_map}, fh, indent=2)
|
||||
if old_index is not None:
|
||||
metadata = dict(old_index.get("metadata") or {})
|
||||
metadata["total_size"] = new_total_size
|
||||
with open(
|
||||
os.path.join(tmp, "model.safetensors.index.json"),
|
||||
"w",
|
||||
encoding="utf-8",
|
||||
) as fh:
|
||||
json.dump({"metadata": metadata, "weight_map": new_weight_map}, fh, indent=2)
|
||||
except BaseException:
|
||||
shutil.rmtree(tmp, ignore_errors=True)
|
||||
raise
|
||||
|
||||
return tmp, src
|
||||
|
||||
+87
-42
@@ -1,55 +1,100 @@
|
||||
"""Smoke test: load a quantized checkpoint through the OBLITERATUS loader.
|
||||
"""Smoke-load a quantized checkpoint through the OBLITERATUS loader."""
|
||||
|
||||
Verifies: detection, dequantization to BF16, no quantized tensors left in
|
||||
the model, config stripped of quantization_config, and a short generation.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import sys
|
||||
import argparse
|
||||
|
||||
import torch
|
||||
|
||||
from obliteratus.models import quant_dequant as qd
|
||||
from obliteratus.models.loader import load_model
|
||||
|
||||
repo = sys.argv[1]
|
||||
|
||||
det = qd.detect_quant_scheme(repo)
|
||||
print(f"[detection] {det.scheme.value} (inverse={det.scale_is_inverse}, group={det.group_size})")
|
||||
assert det.scheme not in (qd.QuantScheme.NONE, qd.QuantScheme.UNSUPPORTED), det
|
||||
def build_parser() -> argparse.ArgumentParser:
|
||||
parser = argparse.ArgumentParser(description=__doc__)
|
||||
parser.add_argument("repo", help="Local checkpoint path or Hugging Face repository")
|
||||
parser.add_argument(
|
||||
"--revision",
|
||||
help="Immutable Hub commit, tag, or branch to load",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--local-files-only",
|
||||
action="store_true",
|
||||
help="Refuse network access and use only locally cached files",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--trust-remote-code",
|
||||
action="store_true",
|
||||
help=(
|
||||
"Explicitly allow checkpoint-provided Python code to execute. "
|
||||
"Off by default; review and pin the repository revision first."
|
||||
),
|
||||
)
|
||||
return parser
|
||||
|
||||
handle = load_model(
|
||||
repo,
|
||||
task="causal_lm",
|
||||
device="auto",
|
||||
dtype="bfloat16",
|
||||
trust_remote_code=True,
|
||||
)
|
||||
model = handle.model
|
||||
print(f"[load] ok — scheme tag: {getattr(model, '_obliteratus_dequantized_scheme', None)}")
|
||||
|
||||
bad = []
|
||||
n_params = 0
|
||||
for name, p in model.named_parameters():
|
||||
n_params += 1
|
||||
if p.dtype in qd.FP8_DTYPES or p.dtype == torch.uint8:
|
||||
bad.append((name, str(p.dtype)))
|
||||
if torch.isnan(p.data).any().item() if p.data.is_floating_point() else False:
|
||||
bad.append((name, "NaN"))
|
||||
print(f"[check] {n_params} params, quantized/NaN leftovers: {bad[:10] or 'NONE'}")
|
||||
assert not bad, bad
|
||||
def main(argv: list[str] | None = None) -> int:
|
||||
args = build_parser().parse_args(argv)
|
||||
det = qd.detect_quant_scheme(
|
||||
args.repo,
|
||||
revision=args.revision,
|
||||
local_files_only=args.local_files_only,
|
||||
)
|
||||
print(
|
||||
f"[detection] {det.scheme.value} "
|
||||
f"(block_inverse={det.scale_is_inverse}, "
|
||||
f"global_inverse={det.global_scale_is_inverse}, group={det.group_size})"
|
||||
)
|
||||
if det.scheme in (qd.QuantScheme.NONE, qd.QuantScheme.UNSUPPORTED):
|
||||
raise RuntimeError(det.reason or f"checkpoint is not FP8/NVFP4: {args.repo}")
|
||||
|
||||
assert getattr(model.config, "quantization_config", None) is None, "quantization_config survived"
|
||||
handle = load_model(
|
||||
args.repo,
|
||||
task="causal_lm",
|
||||
device="auto",
|
||||
dtype="bfloat16",
|
||||
trust_remote_code=args.trust_remote_code,
|
||||
revision=args.revision,
|
||||
local_files_only=args.local_files_only,
|
||||
)
|
||||
model = handle.model
|
||||
print(
|
||||
"[load] ok — scheme tag: "
|
||||
f"{getattr(model, '_obliteratus_dequantized_scheme', None)}"
|
||||
)
|
||||
|
||||
tok = handle.tokenizer
|
||||
prompt = "The capital of France is"
|
||||
try:
|
||||
inputs = tok(prompt, return_tensors="pt").to(model.device)
|
||||
except Exception:
|
||||
# Omni processors: fall back to bare tokenizer encode
|
||||
ids = tok.encode(prompt, return_tensors="pt").to(model.device)
|
||||
inputs = {"input_ids": ids}
|
||||
with torch.no_grad():
|
||||
out = model.generate(**inputs, max_new_tokens=16, do_sample=False)
|
||||
text = tok.decode(out[0][-16:] if out.dim() > 1 else out[-16:], skip_special_tokens=True)
|
||||
print(f"[generate] {text!r}")
|
||||
print("SMOKE_LOAD_OK")
|
||||
bad = []
|
||||
n_params = 0
|
||||
for name, param in model.named_parameters():
|
||||
n_params += 1
|
||||
if param.dtype in qd.FP8_DTYPES or param.dtype == torch.uint8:
|
||||
bad.append((name, str(param.dtype)))
|
||||
if param.data.is_floating_point() and torch.isnan(param.data).any().item():
|
||||
bad.append((name, "NaN"))
|
||||
print(f"[check] {n_params} params, quantized/NaN leftovers: {bad[:10] or 'NONE'}")
|
||||
if bad:
|
||||
raise RuntimeError(f"quantized or non-finite tensors remain: {bad[:10]}")
|
||||
|
||||
if getattr(model.config, "quantization_config", None) is not None:
|
||||
raise RuntimeError("quantization_config survived dequantization")
|
||||
|
||||
tokenizer = handle.tokenizer
|
||||
prompt = "The capital of France is"
|
||||
try:
|
||||
inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
|
||||
except Exception:
|
||||
ids = tokenizer.encode(prompt, return_tensors="pt").to(model.device)
|
||||
inputs = {"input_ids": ids}
|
||||
with torch.no_grad():
|
||||
output = model.generate(**inputs, max_new_tokens=16, do_sample=False)
|
||||
text = tokenizer.decode(
|
||||
output[0][-16:] if output.dim() > 1 else output[-16:],
|
||||
skip_special_tokens=True,
|
||||
)
|
||||
print(f"[generate] {text!r}")
|
||||
print("SMOKE_LOAD_OK")
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(main())
|
||||
|
||||
+21
@@ -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."
|
||||
}
|
||||
@@ -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)
|
||||
|
||||
|
||||
@@ -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
@@ -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()
|
||||
|
||||
Reference in New Issue
Block a user