From 10685a53c2e0b5d6ad104e7ca2f1a36731b0536a Mon Sep 17 00:00:00 2001 From: Joseph Magly <1159087+jmagly@users.noreply.github.com> Date: Sat, 15 Aug 2026 20:52:47 -0400 Subject: [PATCH] test: enforce loader and weight decision contracts --- obliteratus/abliterate.py | 107 ++++++------- obliteratus/models/loader.py | 85 +++++----- obliteratus/runtime_contracts.py | 115 ++++++++++++++ pyproject.toml | 1 + tests/test_loader_boundaries.py | 11 ++ tests/test_projection_math_contracts.py | 83 ++++++++++ tests/test_runtime_contracts.py | 203 ++++++++++++++++++++++++ 7 files changed, 502 insertions(+), 103 deletions(-) diff --git a/obliteratus/abliterate.py b/obliteratus/abliterate.py index 1bf833d..e677d43 100644 --- a/obliteratus/abliterate.py +++ b/obliteratus/abliterate.py @@ -54,6 +54,11 @@ from obliteratus.persistence_contracts import ( # noqa: E402 serialize_checkpoint_metadata, state_dict_size_bytes, ) +from obliteratus.runtime_contracts import ( # noqa: E402 + attention_projection_names, + classify_weight_storage, + is_quantized_parameter, +) from obliteratus.strategies.utils import ( # noqa: E402 get_attention_module, get_ffn_module, @@ -3664,48 +3669,22 @@ class AbliterationPipeline: _text_cfg or config, "num_kv_shared_layers", getattr(config, "num_kv_shared_layers", 0) ) or 0 - _n_layers = len(layers) - _kv_share_start = _n_layers - _kv_shared_layers - _is_shared_kv_layer = _kv_shared_layers > 0 and idx >= _kv_share_start - _is_shared_kv_owner = _is_shared_kv_layer and idx == _kv_share_start - - if self.projection_target == "output": - count += self._project_out_advanced( - attn, d, _ATTN_OUT_NAMES, - norm_preserve=dir_norm_preserve, - regularization=attn_reg, - projection_row_fraction=self.projection_row_fraction, - ) - if self.project_biases: - count += self._project_bias(attn, d, _ATTN_OUT_NAMES) - elif _is_shared_kv_layer and not _is_shared_kv_owner: - # Borrowing layer — skip k/v/k_norm (already projected - # when the owner layer was processed). - _safe_attn_names = [ - n for n in _ATTN_OUT_NAMES + _ATTN_IN_NAMES - if n not in ("k_proj", "v_proj", "k_norm") - ] - count += self._project_out_advanced( - attn, d, _safe_attn_names, - norm_preserve=dir_norm_preserve, - regularization=attn_reg, - projection_row_fraction=self.projection_row_fraction, - ) - if self.project_biases: - count += self._project_bias(attn, d, _safe_attn_names) - else: - # Owner layer (or non-shared layer) — project ALL - # attention weights including k/v. For the shared KV - # owner, this single projection propagates to all - # borrowing layers automatically. - count += self._project_out_advanced( - attn, d, _ATTN_OUT_NAMES + _ATTN_IN_NAMES, - norm_preserve=dir_norm_preserve, - regularization=attn_reg, - projection_row_fraction=self.projection_row_fraction, - ) - if self.project_biases: - count += self._project_bias(attn, d, _ATTN_OUT_NAMES + _ATTN_IN_NAMES) + _attn_names = attention_projection_names( + projection_target=self.projection_target, + layer_index=idx, + num_layers=len(layers), + num_kv_shared_layers=_kv_shared_layers, + output_names=_ATTN_OUT_NAMES, + input_names=_ATTN_IN_NAMES, + ) + count += self._project_out_advanced( + attn, d, _attn_names, + norm_preserve=dir_norm_preserve, + regularization=attn_reg, + projection_row_fraction=self.projection_row_fraction, + ) + if self.project_biases: + count += self._project_bias(attn, d, _attn_names) # Additional head surgery: second-pass precision targeting # on the top safety heads to remove residual refusal signal. @@ -4598,18 +4577,10 @@ class AbliterationPipeline: @staticmethod def _is_quantized_param(param) -> bool: """Check if a parameter is quantized (bitsandbytes, GPTQ, or AWQ).""" - # bitsandbytes NF4/Int8 - if hasattr(param, "quant_state"): - return True - if hasattr(param, "__class__"): - name = param.__class__.__name__ - # bitsandbytes: Params4bit, Int8Params - # GPTQ (auto-gptq / exllamav2): QuantLinear packs weights into qweight - # AWQ (autoawq): WQLinear variants pack weights similarly - if name in ("Params4bit", "Int8Params", "QuantLinear", - "WQLinear", "WQLinear_GEMM", "WQLinear_GEMV"): - return True - return False + return is_quantized_parameter( + class_name=param.__class__.__name__, + has_quant_state=hasattr(param, "quant_state"), + ) @staticmethod def _dequantize_weight(proj_module) -> tuple[torch.Tensor, bool]: @@ -4631,7 +4602,16 @@ class AbliterationPipeline: # These formats pack weights into qweight (not weight), so we # detect at the module level rather than parameter level. module_cls = proj_module.__class__.__name__ - if module_cls in ("QuantLinear", "WQLinear", "WQLinear_GEMM", "WQLinear_GEMV"): + weight = getattr(proj_module, "weight", None) + storage_kind = classify_weight_storage( + module_class_name=module_cls, + parameter_class_name=weight.__class__.__name__ if weight is not None else "", + has_quant_state=hasattr(weight, "quant_state"), + data_is_floating_point=( + weight.data.is_floating_point() if weight is not None else False + ), + ) + if storage_kind == "packed_module": # Both GPTQ and AWQ store packed int weights in qweight with # separate scales/zeros. Use their built-in dequantization. if hasattr(proj_module, "dequantize"): @@ -4649,7 +4629,7 @@ class AbliterationPipeline: # ── bitsandbytes parameter-level detection ───────────────── weight = proj_module.weight - if AbliterationPipeline._is_quantized_param(weight): + if storage_kind == "quantized_parameter": try: import bitsandbytes as bnb W_float = bnb.functional.dequantize_4bit( @@ -4670,7 +4650,7 @@ class AbliterationPipeline: # Some architectures store weights as non-float types (e.g. uint8 from # custom quantization schemes). Projections require float math, so # convert and treat as "quantized" so the caller writes back properly. - if not weight.data.is_floating_point(): + if storage_kind == "integer": return weight.data.to(torch.float32), True return weight.data, False @@ -4683,9 +4663,18 @@ class AbliterationPipeline: inference. """ module_cls = proj_module.__class__.__name__ + weight = getattr(proj_module, "weight", None) + storage_kind = classify_weight_storage( + module_class_name=module_cls, + parameter_class_name=weight.__class__.__name__ if weight is not None else "", + has_quant_state=hasattr(weight, "quant_state"), + data_is_floating_point=( + weight.data.is_floating_point() if weight is not None else False + ), + ) # ── GPTQ/AWQ re-quantization ────────────────────────────── - if module_cls in ("QuantLinear", "WQLinear", "WQLinear_GEMM", "WQLinear_GEMV"): + if storage_kind == "packed_module": if hasattr(proj_module, "pack") and callable(proj_module.pack): # auto-gptq QuantLinear.pack() re-packs float weights try: @@ -4715,7 +4704,7 @@ class AbliterationPipeline: # If the original weight isn't a bitsandbytes/GPTQ/AWQ param, just # replace with the float version so projections are preserved. weight = proj_module.weight - if not AbliterationPipeline._is_quantized_param(weight): + if storage_kind in ("integer", "float"): proj_module.weight = nn.Parameter( W_modified.to(device=weight.device), requires_grad=weight.requires_grad, diff --git a/obliteratus/models/loader.py b/obliteratus/models/loader.py index bc65d80..ab2e1b2 100644 --- a/obliteratus/models/loader.py +++ b/obliteratus/models/loader.py @@ -15,6 +15,7 @@ from obliteratus import device as dev from obliteratus.runtime_contracts import ( effective_model_memory_gb, quantized_model_fits_gpu, + resolve_model_load_policy, should_snapshot_model, validate_model_load_request, ) @@ -613,6 +614,15 @@ def load_model( # Memory estimation and warnings (skip for natively quantized models — estimate is wrong) native_quant = getattr(config, "quantization_config", None) + load_policy = resolve_model_load_policy( + device=device, + resolved_device=resolved_device, + dtype=dtype, + quantization=quantization, + has_native_quantization=native_quant is not None, + device_map_auto_supported=dev.supports_device_map_auto(resolved_device), + bitsandbytes_supported=dev.supports_bitsandbytes(resolved_device), + ) est_gb = _estimate_model_memory_gb(config, torch_dtype) if native_quant is None else 0.0 gpu_gb = _available_gpu_memory_gb() if est_gb > 0 and gpu_gb > 0: @@ -627,65 +637,54 @@ def load_model( load_kwargs: dict = { "pretrained_model_name_or_path": model_name, "config": config, - "torch_dtype": torch_dtype, **hf_kwargs, } + if load_policy.include_torch_dtype: + load_kwargs["torch_dtype"] = torch_dtype if task == "classification": config.num_labels = num_labels load_kwargs["config"] = config # Quantization support (requires bitsandbytes) - if native_quant is not None: + if load_policy.quantization_backend == "native": # Model ships with native quantization (e.g. Mxfp4Config) — don't layer BitsAndBytes # on top, and don't override its compute dtype with our torch_dtype logger.info( f"Model has native quantization ({type(native_quant).__name__}), " f"skipping BitsAndBytes and using model's native dtype" ) - load_kwargs.pop("torch_dtype", None) load_kwargs["device_map"] = "auto" - elif quantization in ("4bit", "8bit"): - # BitsAndBytes only works on NVIDIA CUDA GPUs. - if not dev.supports_bitsandbytes(resolved_device): + elif load_policy.quantization_backend == "bitsandbytes": + try: + import bitsandbytes # noqa: F401 + except ImportError: raise RuntimeError( - f"Quantization '{quantization}' requires an available NVIDIA CUDA device; " - f"resolved device was '{resolved_device}'. Remove --quantization to load in {dtype}.", + f"Quantization '{quantization}' requires bitsandbytes: " + f"pip install -U bitsandbytes>=0.46.1" + ) + from transformers import BitsAndBytesConfig + + # Enable fp32 CPU offload so that models too large to fit entirely on + # GPU (even quantized) can spill to CPU without crashing bitsandbytes. + # This is critical for frontier MoE models (GLM-5 744B, DeepSeek-V3 685B, + # Mistral Large 3 675B, etc.) on single-GPU setups. + if quantization == "4bit": + load_kwargs["quantization_config"] = BitsAndBytesConfig( + load_in_4bit=True, + bnb_4bit_compute_dtype=torch_dtype, + bnb_4bit_quant_type="nf4", + llm_int8_enable_fp32_cpu_offload=True, ) else: - try: - import bitsandbytes # noqa: F401 - except ImportError: - raise RuntimeError( - f"Quantization '{quantization}' requires bitsandbytes: " - f"pip install -U bitsandbytes>=0.46.1" - ) - from transformers import BitsAndBytesConfig - - # Enable fp32 CPU offload so that models too large to fit entirely on - # GPU (even quantized) can spill to CPU without crashing bitsandbytes. - # This is critical for frontier MoE models (GLM-5 744B, DeepSeek-V3 685B, - # Mistral Large 3 675B, etc.) on single-GPU setups. - if quantization == "4bit": - load_kwargs["quantization_config"] = BitsAndBytesConfig( - load_in_4bit=True, - bnb_4bit_compute_dtype=torch_dtype, - bnb_4bit_quant_type="nf4", - llm_int8_enable_fp32_cpu_offload=True, - ) - else: - load_kwargs["quantization_config"] = BitsAndBytesConfig( - load_in_8bit=True, - llm_int8_enable_fp32_cpu_offload=True, - ) - load_kwargs["device_map"] = "auto" + load_kwargs["quantization_config"] = BitsAndBytesConfig( + load_in_8bit=True, + llm_int8_enable_fp32_cpu_offload=True, + ) + load_kwargs["device_map"] = "auto" # device_map="auto" is only reliable on CUDA (accelerate doesn't support MPS). - if "device_map" not in load_kwargs and device == "auto": - if dev.supports_device_map_auto(resolved_device): - load_kwargs["device_map"] = "auto" - else: - # MPS / CPU: load to CPU first, then .to(device) after loading. - pass + if "device_map" not in load_kwargs and load_policy.use_device_map_auto: + load_kwargs["device_map"] = "auto" # Offload support: provide a folder for disk offloading when GPU memory is insufficient _offload_dir = None @@ -756,10 +755,8 @@ def load_model( ) from e raise - if device not in ("auto",) and quantization is None and native_quant is None: - model = model.to(device) - elif device == "auto" and not dev.supports_device_map_auto(): - # MPS / CPU: device_map wasn't used, move model to best device. + if load_policy.move_to_resolved_device: + # Explicit devices and auto-selected MPS/CPU load on CPU before moving. model = model.to(resolved_device) model.eval() diff --git a/obliteratus/runtime_contracts.py b/obliteratus/runtime_contracts.py index 0b86088..b97384d 100644 --- a/obliteratus/runtime_contracts.py +++ b/obliteratus/runtime_contracts.py @@ -3,12 +3,30 @@ from __future__ import annotations from collections.abc import Collection +from dataclasses import dataclass from typing import Literal VALID_MODEL_TASKS = ("causal_lm", "classification") VALID_QUANTIZATIONS = (None, "4bit", "8bit") VALID_DTYPES = ("float32", "float16", "bfloat16") +PACKED_QUANTIZED_MODULE_TYPES = frozenset( + {"QuantLinear", "WQLinear", "WQLinear_GEMM", "WQLinear_GEMV"}, +) +QUANTIZED_PARAMETER_TYPES = frozenset( + {"Params4bit", "Int8Params", *PACKED_QUANTIZED_MODULE_TYPES}, +) +SHARED_KV_PROJECTION_NAMES = frozenset({"k_proj", "v_proj", "k_norm"}) + + +@dataclass(frozen=True) +class ModelLoadPolicy: + """Provider-independent decisions that control model placement and loading.""" + + quantization_backend: Literal["none", "native", "bitsandbytes"] + include_torch_dtype: bool + use_device_map_auto: bool + move_to_resolved_device: bool def validate_model_load_request( @@ -32,6 +50,103 @@ def validate_model_load_request( raise ValueError(f"Unknown dtype {dtype!r}. Choose from {list(VALID_DTYPES)}") +def resolve_model_load_policy( + *, + device: str, + resolved_device: str, + dtype: str, + quantization: str | None, + has_native_quantization: bool, + device_map_auto_supported: bool, + bitsandbytes_supported: bool, +) -> ModelLoadPolicy: + """Resolve dtype, quantization, device-map, and post-load move precedence.""" + if quantization not in VALID_QUANTIZATIONS: + raise ValueError( + "Unknown quantization {!r}. Choose None, '4bit', or '8bit'".format(quantization), + ) + if has_native_quantization: + return ModelLoadPolicy( + quantization_backend="native", + include_torch_dtype=False, + use_device_map_auto=True, + move_to_resolved_device=False, + ) + if quantization in ("4bit", "8bit"): + if not bitsandbytes_supported: + raise RuntimeError( + f"Quantization '{quantization}' requires an available NVIDIA CUDA device; " + f"resolved device was '{resolved_device}'. Remove --quantization to load in {dtype}.", + ) + return ModelLoadPolicy( + quantization_backend="bitsandbytes", + include_torch_dtype=True, + use_device_map_auto=True, + move_to_resolved_device=False, + ) + + use_device_map_auto = device == "auto" and device_map_auto_supported + return ModelLoadPolicy( + quantization_backend="none", + include_torch_dtype=True, + use_device_map_auto=use_device_map_auto, + move_to_resolved_device=not use_device_map_auto, + ) + + +def attention_projection_names( + *, + projection_target: str, + layer_index: int, + num_layers: int, + num_kv_shared_layers: int, + output_names: Collection[str], + input_names: Collection[str], +) -> tuple[str, ...]: + """Choose attention weights while projecting shared KV storage exactly once.""" + outputs = tuple(output_names) + if projection_target == "output": + return outputs + + all_names = outputs + tuple(input_names) + if num_kv_shared_layers <= 0: + return all_names + if num_layers <= 0 or not 0 <= layer_index < num_layers: + raise ValueError("shared-KV projection requires a valid layer index and layer count") + if num_kv_shared_layers > num_layers: + raise ValueError("num_kv_shared_layers cannot exceed the model layer count") + + owner_index = num_layers - num_kv_shared_layers + if layer_index <= owner_index: + return all_names + return tuple(name for name in all_names if name not in SHARED_KV_PROJECTION_NAMES) + + +def is_quantized_parameter(*, class_name: str, has_quant_state: bool) -> bool: + """Return whether a parameter carries a supported packed-quantization marker.""" + return has_quant_state or class_name in QUANTIZED_PARAMETER_TYPES + + +def classify_weight_storage( + *, + module_class_name: str, + parameter_class_name: str, + has_quant_state: bool, + data_is_floating_point: bool, +) -> Literal["packed_module", "quantized_parameter", "integer", "float"]: + """Classify a projection weight so callers select a safe read/write path.""" + if module_class_name in PACKED_QUANTIZED_MODULE_TYPES: + return "packed_module" + if is_quantized_parameter( + class_name=parameter_class_name, + has_quant_state=has_quant_state, + ): + return "quantized_parameter" + if not data_is_floating_point: + return "integer" + return "float" + + def effective_model_memory_gb(estimate_gb: float, quantization: str | None) -> float: """Adjust a full-precision weight estimate for runtime quantization.""" factor = {"4bit": 4, "8bit": 2}.get(quantization, 1) diff --git a/pyproject.toml b/pyproject.toml index 7391e5e..9744bc8 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -128,6 +128,7 @@ only_mutate = [ required_mutation_targets = [ "obliteratus/analysis/numerical_contracts.py", "obliteratus/analysis/whitened_svd.py", + "obliteratus/runtime_contracts.py", ] pytest_add_cli_args = ["--no-cov", "-q"] pytest_add_cli_args_test_selection = [ diff --git a/tests/test_loader_boundaries.py b/tests/test_loader_boundaries.py index 79f27e1..c2b0507 100644 --- a/tests/test_loader_boundaries.py +++ b/tests/test_loader_boundaries.py @@ -3,6 +3,7 @@ from __future__ import annotations import builtins +from pathlib import Path from types import SimpleNamespace from unittest.mock import MagicMock, Mock @@ -552,6 +553,16 @@ def test_cached_path_compatibility_resolves_local_hub_and_fallback(monkeypatch, download.assert_not_called() +def test_working_or_temp_dir_compatibility_preserves_and_cleans_paths(tmp_path): + with loader._working_or_temp_dir(tmp_path) as selected: + assert selected == tmp_path + + with loader._working_or_temp_dir() as generated: + generated_path = Path(generated) + assert generated_path.is_dir() + assert not generated_path.exists() + + def test_deferred_transformers_shims_are_additive_and_idempotent(monkeypatch): fake_transformers = SimpleNamespace(_objects={}) monkeypatch.setitem(loader._sys.modules, "transformers", fake_transformers) diff --git a/tests/test_projection_math_contracts.py b/tests/test_projection_math_contracts.py index c54169a..e8433ad 100644 --- a/tests/test_projection_math_contracts.py +++ b/tests/test_projection_math_contracts.py @@ -709,3 +709,86 @@ def test_project_out_advanced_replaces_quantized_weight_after_successful_project assert len(replacement_calls) == 1 assert replacement_calls[0][0] is linear assert torch.allclose(replacement_calls[0][1][:, 0], torch.zeros(2)) + + +def test_quantized_parameter_wrapper_uses_declared_markers_only(): + params4bit = type("Params4bit", (), {})() + ordinary = SimpleNamespace() + marked = SimpleNamespace(quant_state=object()) + + assert AbliterationPipeline._is_quantized_param(params4bit) is True + assert AbliterationPipeline._is_quantized_param(marked) is True + assert AbliterationPipeline._is_quantized_param(ordinary) is False + + +def test_dequantize_weight_returns_float_storage_without_copying(): + linear = torch.nn.Linear(2, 2, bias=False) + + weight, requires_replacement = AbliterationPipeline._dequantize_weight(linear) + + assert weight.data_ptr() == linear.weight.data.data_ptr() + assert requires_replacement is False + + +def test_dequantize_weight_promotes_integer_storage_for_safe_projection(): + module = SimpleNamespace( + weight=torch.nn.Parameter( + torch.tensor([[1, 2], [3, 4]], dtype=torch.uint8), + requires_grad=False, + ), + ) + + weight, requires_replacement = AbliterationPipeline._dequantize_weight(module) + + assert weight.dtype is torch.float32 + assert torch.equal(weight, torch.tensor([[1.0, 2.0], [3.0, 4.0]])) + assert requires_replacement is True + + +def test_packed_module_dequantization_is_cloned_before_projection(): + packed_type = type("QuantLinear", (), {}) + packed = packed_type() + source = torch.tensor([[1.0, 2.0]]) + packed.dequantize = lambda: source + + weight, requires_replacement = AbliterationPipeline._dequantize_weight(packed) + + assert torch.equal(weight, source) + assert weight.data_ptr() != source.data_ptr() + assert requires_replacement is True + + +def test_packed_module_without_safe_dequantizer_fails_closed(): + packed_type = type("WQLinear_GEMM", (), {}) + packed = packed_type() + packed.qweight = torch.ones(1, dtype=torch.int32) + packed.scales = torch.ones(1) + + with pytest.raises(RuntimeError, match=r"no dequantize\(\) method available"): + AbliterationPipeline._dequantize_weight(packed) + + +def test_weight_replacement_preserves_float_parameter_contract(): + linear = torch.nn.Linear(2, 2, bias=False) + modified = torch.tensor([[1.0, 2.0], [3.0, 4.0]]) + + AbliterationPipeline._replace_quantized_weight(linear, modified) + + assert torch.equal(linear.weight, modified) + assert linear.weight.requires_grad is True + + +def test_packed_weight_replacement_uses_the_module_repacker(): + packed_type = type("QuantLinear", (), {}) + packed = packed_type() + packed.qweight = torch.ones(1) + packed.scales = torch.tensor([0.5]) + calls = [] + packed.pack = lambda weight, scales: calls.append((weight.clone(), scales.clone())) + modified = torch.tensor([[1.0, 2.0]]) + + AbliterationPipeline._replace_quantized_weight(packed, modified) + + assert len(calls) == 1 + assert torch.equal(calls[0][0], modified) + assert torch.equal(calls[0][1], packed.scales) diff --git a/tests/test_runtime_contracts.py b/tests/test_runtime_contracts.py index 9e9b2d4..85dd390 100644 --- a/tests/test_runtime_contracts.py +++ b/tests/test_runtime_contracts.py @@ -6,9 +6,14 @@ import pytest from hypothesis import given, strategies as st from obliteratus.runtime_contracts import ( + ModelLoadPolicy, + attention_projection_names, + classify_weight_storage, classify_architecture_size, effective_model_memory_gb, + is_quantized_parameter, quantized_model_fits_gpu, + resolve_model_load_policy, should_snapshot_model, supports_bfloat16_target, validate_model_load_request, @@ -60,6 +65,204 @@ def test_loader_task_validation_can_follow_the_provider_registry(): ) +@pytest.mark.parametrize( + ("overrides", "expected"), + [ + ( + {}, + ModelLoadPolicy("none", True, False, True), + ), + ( + {"device": "auto", "device_map_auto_supported": True}, + ModelLoadPolicy("none", True, True, False), + ), + ( + {"device": "auto", "resolved_device": "mps"}, + ModelLoadPolicy("none", True, False, True), + ), + ( + {"quantization": "4bit", "bitsandbytes_supported": True}, + ModelLoadPolicy("bitsandbytes", True, True, False), + ), + ( + {"quantization": "8bit", "bitsandbytes_supported": True}, + ModelLoadPolicy("bitsandbytes", True, True, False), + ), + ( + {"quantization": "4bit", "has_native_quantization": True}, + ModelLoadPolicy("native", False, True, False), + ), + ], +) +def test_model_load_policy_precedence_is_explicit(overrides, expected): + inputs = { + "device": "cpu", + "resolved_device": "cpu", + "dtype": "float16", + "quantization": None, + "has_native_quantization": False, + "device_map_auto_supported": False, + "bitsandbytes_supported": False, + } + inputs.update(overrides) + assert resolve_model_load_policy(**inputs) == expected + + +def test_model_load_policy_rejects_unsupported_quantization_before_branching(): + with pytest.raises(ValueError, match="Unknown quantization '3bit'"): + resolve_model_load_policy( + device="auto", + resolved_device="cuda", + dtype="float16", + quantization="3bit", + has_native_quantization=False, + device_map_auto_supported=True, + bitsandbytes_supported=True, + ) + + +@pytest.mark.parametrize("quantization", ["4bit", "8bit"]) +def test_requested_quantization_requires_a_compatible_resolved_device(quantization): + with pytest.raises(RuntimeError) as error: + resolve_model_load_policy( + device="auto", + resolved_device="mps", + dtype="bfloat16", + quantization=quantization, + has_native_quantization=False, + device_map_auto_supported=False, + bitsandbytes_supported=False, + ) + assert str(error.value) == ( + f"Quantization '{quantization}' requires an available NVIDIA CUDA device; " + "resolved device was 'mps'. Remove --quantization to load in bfloat16." + ) + + +_OUTPUT_NAMES = ("o_proj",) +_INPUT_NAMES = ("q_proj", "k_proj", "v_proj", "k_norm") + + +@pytest.mark.parametrize("layer_index", range(6)) +def test_nonshared_attention_projects_every_configured_weight(layer_index): + assert attention_projection_names( + projection_target="all", + layer_index=layer_index, + num_layers=6, + num_kv_shared_layers=0, + output_names=_OUTPUT_NAMES, + input_names=_INPUT_NAMES, + ) == _OUTPUT_NAMES + _INPUT_NAMES + + +def test_output_only_projection_is_independent_of_shared_kv_layout(): + assert attention_projection_names( + projection_target="output", + layer_index=5, + num_layers=6, + num_kv_shared_layers=3, + output_names=_OUTPUT_NAMES, + input_names=_INPUT_NAMES, + ) == _OUTPUT_NAMES + + +@pytest.mark.parametrize( + ("layer_index", "expected"), + [ + (2, _OUTPUT_NAMES + _INPUT_NAMES), + (3, _OUTPUT_NAMES + _INPUT_NAMES), + (4, ("o_proj", "q_proj")), + (5, ("o_proj", "q_proj")), + ], +) +def test_shared_kv_owner_projects_storage_once_and_borrowers_skip_it(layer_index, expected): + assert attention_projection_names( + projection_target="all", + layer_index=layer_index, + num_layers=6, + num_kv_shared_layers=3, + output_names=_OUTPUT_NAMES, + input_names=_INPUT_NAMES, + ) == expected + + +@pytest.mark.parametrize( + ("layer_index", "num_layers", "shared_layers", "message"), + [ + (-1, 6, 3, "valid layer index"), + (6, 6, 3, "valid layer index"), + (0, 0, 1, "valid layer index"), + (0, 6, 7, "cannot exceed"), + ], +) +def test_shared_kv_layout_rejects_impossible_ownership( + layer_index, + num_layers, + shared_layers, + message, +): + with pytest.raises(ValueError, match=message): + attention_projection_names( + projection_target="all", + layer_index=layer_index, + num_layers=num_layers, + num_kv_shared_layers=shared_layers, + output_names=_OUTPUT_NAMES, + input_names=_INPUT_NAMES, + ) + + +@pytest.mark.parametrize( + ("class_name", "has_quant_state", "expected"), + [ + ("Parameter", False, False), + ("Parameter", True, True), + ("Params4bit", False, True), + ("Int8Params", False, True), + ("QuantLinear", False, True), + ("WQLinear_GEMM", False, True), + ], +) +def test_quantized_parameter_markers_are_closed_and_explicit( + class_name, + has_quant_state, + expected, +): + assert is_quantized_parameter( + class_name=class_name, + has_quant_state=has_quant_state, + ) is expected + + +@pytest.mark.parametrize( + ("overrides", "expected"), + [ + ({}, "float"), + ({"data_is_floating_point": False}, "integer"), + ({"has_quant_state": True}, "quantized_parameter"), + ({"parameter_class_name": "Int8Params"}, "quantized_parameter"), + ({"module_class_name": "QuantLinear"}, "packed_module"), + ( + { + "module_class_name": "WQLinear_GEMV", + "has_quant_state": True, + "data_is_floating_point": False, + }, + "packed_module", + ), + ], +) +def test_weight_storage_classification_uses_safe_precedence(overrides, expected): + inputs = { + "module_class_name": "Linear", + "parameter_class_name": "Parameter", + "has_quant_state": False, + "data_is_floating_point": True, + } + inputs.update(overrides) + assert classify_weight_storage(**inputs) == expected + + @given(st.floats(min_value=0, max_value=10_000, allow_nan=False, allow_infinity=False)) def test_effective_memory_respects_quantization_ratios(estimate_gb): assert effective_model_memory_gb(estimate_gb, None) == estimate_gb