diff --git a/README.md b/README.md index 3667a60..03cf5a3 100644 --- a/README.md +++ b/README.md @@ -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 diff --git a/ci/conditional-test-policy.json b/ci/conditional-test-policy.json index 2312d65..043bde4 100644 --- a/ci/conditional-test-policy.json +++ b/ci/conditional-test-policy.json @@ -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", diff --git a/ci/test-risk-map.json b/ci/test-risk-map.json index 31d0bb6..9de29aa 100644 --- a/ci/test-risk-map.json +++ b/ci/test-risk-map.json @@ -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", diff --git a/docs/theory_journal.md b/docs/theory_journal.md index 6f2c180..0c47cc6 100644 --- a/docs/theory_journal.md +++ b/docs/theory_journal.md @@ -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. --- diff --git a/obliteratus/abliterate.py b/obliteratus/abliterate.py index a547f95..a6a8b19 100644 --- a/obliteratus/abliterate.py +++ b/obliteratus/abliterate.py @@ -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 " diff --git a/obliteratus/models/loader.py b/obliteratus/models/loader.py index 9341f4e..b0eb208 100644 --- a/obliteratus/models/loader.py +++ b/obliteratus/models/loader.py @@ -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( diff --git a/obliteratus/models/quant_dequant.py b/obliteratus/models/quant_dequant.py index a82bb0e..9e6a95f 100644 --- a/obliteratus/models/quant_dequant.py +++ b/obliteratus/models/quant_dequant.py @@ -29,8 +29,11 @@ NVFP4 (ModelOpt, ``quant_algo: "NVFP4"``): ``.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"``): + ``.weight_packed`` uint8 packed E2M1 values + ``.weight_scale`` FP8 block scales (multiplicative) + ``.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 diff --git a/scripts/smoke_load_quant.py b/scripts/smoke_load_quant.py index 69a0c2b..f17c80e 100644 --- a/scripts/smoke_load_quant.py +++ b/scripts/smoke_load_quant.py @@ -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()) diff --git a/tests/fixtures/quant_dequant_provenance.json b/tests/fixtures/quant_dequant_provenance.json new file mode 100644 index 0000000..7e41a45 --- /dev/null +++ b/tests/fixtures/quant_dequant_provenance.json @@ -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." +} diff --git a/tests/test_cli_boundaries.py b/tests/test_cli_boundaries.py index 310f513..cc6285a 100644 --- a/tests/test_cli_boundaries.py +++ b/tests/test_cli_boundaries.py @@ -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) diff --git a/tests/test_loader_boundaries.py b/tests/test_loader_boundaries.py index d60acdf..c1f2871 100644 --- a/tests/test_loader_boundaries.py +++ b/tests/test_loader_boundaries.py @@ -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) diff --git a/tests/test_quant_dequant.py b/tests/test_quant_dequant.py index 401559a..a7792b1 100644 --- a/tests/test_quant_dequant.py +++ b/tests/test_quant_dequant.py @@ -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 # --------------------------------------------------------------------------- diff --git a/tests/test_strategy_navigation_contracts.py b/tests/test_strategy_navigation_contracts.py index 40eb2ca..ee1b37e 100644 --- a/tests/test_strategy_navigation_contracts.py +++ b/tests/test_strategy_navigation_contracts.py @@ -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()