CalvinBackup/OBLITERATUS
1
0
Fork 0
mirror of https://github.com/elder-plinius/OBLITERATUS.git synced 2026-08-17 16:37:30 +02:00
Code Issues Packages Projects Releases Wiki Activity

test: enforce loader and weight decision contracts

Browse Source
This commit is contained in:
Joseph Magly
2026-08-15 22:03:22 -04:00
parent 2c2b38b501
commit 10685a53c2
7 changed files with 502 additions and 103 deletions
Download Patch File Download Diff File Expand all files Collapse all files
obliteratus/abliterate.py
+48 -59
View File
@@ -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,
obliteratus/models/loader.py
+41 -44
View File
@@ -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()
obliteratus/runtime_contracts.py
+115
View File
@@ -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)
pyproject.toml
+1
View File
@@ -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 = [
tests/test_loader_boundaries.py
+11
View File
@@ -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)
tests/test_projection_math_contracts.py
+83
View File
@@ -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)
tests/test_runtime_contracts.py
+203
View File
@@ -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