orbitquant 0.1.0__py3-none-any.whl

orbitquant/artifacts/comparisons.py

@@ -0,0 +1,140 @@
+from __future__ import annotations
+
+import json
+import re
+from pathlib import Path
+from typing import Any
+
+from orbitquant.artifacts.assets import record_artifact_asset
+from orbitquant.artifacts.validator import validate_orbitquant_artifact
+
+_IMAGE_SUFFIXES = {".jpeg", ".jpg", ".png", ".webp"}
+_SAFE_NAME_RE = re.compile(r"[^A-Za-z0-9_.-]+")
+
+
+def _read_jsonl(path: Path) -> list[dict[str, Any]]:
+    if not path.is_file():
+        return []
+    return [
+        json.loads(line)
+        for line in path.read_text(encoding="utf-8").splitlines()
+        if line.strip()
+    ]
+
+
+def _safe_name(value: Any) -> str:
+    return _SAFE_NAME_RE.sub("-", str(value)).strip("-") or "unknown"
+
+
+def _prompt_id(metadata: dict[str, Any]) -> str:
+    prompt_record = metadata.get("prompt_record")
+    if isinstance(prompt_record, dict) and prompt_record.get("id") is not None:
+        return str(prompt_record["id"])
+    return "prompt"
+
+
+def _output_path(artifact_path: Path, metadata: dict[str, Any]) -> Path | None:
+    value = metadata.get("output_path")
+    if not value:
+        return None
+    path = Path(str(value))
+    return path if path.is_absolute() else artifact_path / path
+
+
+def _artifact_path(artifact_path: Path, value: Any) -> Path:
+    path = Path(str(value))
+    return path if path.is_absolute() else artifact_path / path
+
+
+def _comparison_source_path(artifact_path: Path, metadata: dict[str, Any]) -> Path | None:
+    output_path = _output_path(artifact_path, metadata)
+    if output_path is not None and output_path.suffix.lower() in _IMAGE_SUFFIXES:
+        return output_path
+
+    contact_sheet_path = metadata.get("contact_sheet_path")
+    if contact_sheet_path:
+        path = _artifact_path(artifact_path, contact_sheet_path)
+        if path.suffix.lower() in _IMAGE_SUFFIXES:
+            return path
+
+    for asset_path in metadata.get("asset_paths", []):
+        path = _artifact_path(artifact_path, asset_path)
+        if "contact_sheet" in path.name and path.suffix.lower() in _IMAGE_SUFFIXES:
+            return path
+    return None
+
+
+def _comparison_key(metadata: dict[str, Any]) -> tuple[str, str, str]:
+    return (
+        str(metadata.get("suite", "")),
+        str(metadata.get("seed", "")),
+        _prompt_id(metadata),
+    )
+
+
+def _indexed_records(artifact_path: Path, split: str) -> dict[tuple[str, str, str], dict[str, Any]]:
+    records = {}
+    for record in _read_jsonl(artifact_path / "benchmark" / f"{split}.metrics.jsonl"):
+        metadata = dict(record.get("metadata", {}))
+        path = _comparison_source_path(artifact_path, metadata)
+        if path is None or not path.is_file():
+            continue
+        records[_comparison_key(metadata)] = {"metadata": metadata, "path": path}
+    return records
+
+
+def create_artifact_image_comparisons(
+    artifact_dir: str | Path,
+    *,
+    comparison_keys: set[tuple[Any, Any, Any]] | None = None,
+    validate_checksums_enabled: bool = True,
+    refresh_checksums_enabled: bool = True,
+) -> list[str]:
+    from orbitquant.eval.assets import create_image_comparison_sheet
+
+    artifact_path = Path(artifact_dir)
+    validate_orbitquant_artifact(
+        artifact_path,
+        validate_checksums_enabled=validate_checksums_enabled,
+        validate_tensors=validate_checksums_enabled,
+    )
+    original_records = _indexed_records(artifact_path, "original")
+    orbitquant_records = _indexed_records(artifact_path, "orbitquant")
+    selected_keys = (
+        None
+        if comparison_keys is None
+        else {(str(suite), str(seed), str(prompt_id)) for suite, seed, prompt_id in comparison_keys}
+    )
+    created = []
+    for key, orbitquant in sorted(orbitquant_records.items()):
+        if selected_keys is not None and key not in selected_keys:
+            continue
+        original = original_records.get(key)
+        if original is None:
+            continue
+        suite, seed, prompt_id = key
+        bit_setting = orbitquant["metadata"].get("bit_setting", "orbitquant")
+        output_path = (
+            artifact_path
+            / "assets"
+            / (
+                "original_vs_orbitquant_"
+                f"{_safe_name(suite)}_seed{_safe_name(seed)}_"
+                f"{_safe_name(bit_setting)}_{_safe_name(prompt_id)}.webp"
+            )
+        )
+        create_image_comparison_sheet(
+            original["path"],
+            orbitquant["path"],
+            output_path,
+            labels=("BF16", f"OrbitQuant {bit_setting}"),
+        )
+        created.append(
+            record_artifact_asset(
+                artifact_path,
+                output_path,
+                validate_checksums_enabled=validate_checksums_enabled,
+                refresh_checksums_enabled=refresh_checksums_enabled,
+            )
+        )
+    return created

orbitquant/artifacts/loader.py

@@ -0,0 +1,109 @@
+from __future__ import annotations
+
+import json
+from pathlib import Path
+
+import torch
+from safetensors.torch import load_file
+
+from orbitquant.adaln import RTNInt4Linear
+from orbitquant.artifacts.checksums import (
+    validate_checksums as validate_artifact_checksums,
+)
+from orbitquant.artifacts.checksums import (
+    validate_sha256sums,
+)
+from orbitquant.artifacts.manifest import OrbitQuantManifest
+from orbitquant.artifacts.validator import (
+    _validate_config_manifest,
+    validate_required_artifact_files,
+)
+from orbitquant.config import OrbitQuantConfig
+from orbitquant.layers import OrbitQuantLinear
+from orbitquant.modeling import _parent_and_child, _set_child
+
+
+def _get_module(model: torch.nn.Module, module_name: str) -> torch.nn.Module:
+    module = model
+    for part in module_name.split("."):
+        if part.isdigit() and isinstance(module, torch.nn.ModuleList | torch.nn.Sequential):
+            module = module[int(part)]
+        elif isinstance(module, torch.nn.ModuleDict):
+            module = module[part]
+        else:
+            module = getattr(module, part)
+    return module
+
+
+def _first_tensor_device(model: torch.nn.Module) -> torch.device:
+    for parameter in model.parameters():
+        return parameter.device
+    for buffer in model.buffers():
+        return buffer.device
+    return torch.device("cpu")
+
+
+def _safetensors_load_device(device: str | torch.device | None, model: torch.nn.Module) -> str:
+    target_device = _first_tensor_device(model) if device is None else torch.device(device)
+    if target_device.type == "cuda":
+        return str(target_device)
+    return "cpu"
+
+
+def load_orbitquant_artifact(
+    model: torch.nn.Module,
+    artifact_dir: str | Path,
+    *,
+    strict: bool = True,
+    validate_checksums: bool = True,
+    device: str | torch.device | None = None,
+    runtime_mode: str | None = None,
+    activation_kernel_backend: str | None = None,
+) -> OrbitQuantManifest:
+    artifact_path = Path(artifact_dir)
+    validate_required_artifact_files(artifact_path)
+    config = OrbitQuantConfig.from_dict(
+        json.loads((artifact_path / "quantization_config.json").read_text(encoding="utf-8"))
+    )
+    manifest = OrbitQuantManifest.from_dict(
+        json.loads((artifact_path / "orbitquant_manifest.json").read_text(encoding="utf-8"))
+    )
+    _validate_config_manifest(config, manifest)
+    if runtime_mode is not None or activation_kernel_backend is not None:
+        config_payload = config.to_dict()
+        if runtime_mode is not None:
+            config_payload["runtime_mode"] = runtime_mode
+        if activation_kernel_backend is not None:
+            config_payload["activation_kernel_backend"] = activation_kernel_backend
+        config = OrbitQuantConfig.from_dict(config_payload)
+    if validate_checksums:
+        validate_artifact_checksums(artifact_path, manifest.checksums)
+        validate_sha256sums(
+            artifact_path,
+            required_paths=tuple(manifest.checksums) + ("README.md", "orbitquant_manifest.json"),
+        )
+
+    for name in manifest.quantized_modules:
+        module = _get_module(model, name)
+        if not isinstance(module, torch.nn.Linear):
+            raise TypeError(f"expected Linear at {name}, got {type(module).__name__}")
+        replacement = OrbitQuantLinear.empty_from_linear(module, config=config, module_name=name)
+        parent, child_name = _parent_and_child(model, name)
+        _set_child(parent, child_name, replacement)
+
+    for name in manifest.adaln_modules:
+        module = _get_module(model, name)
+        if not isinstance(module, torch.nn.Linear):
+            raise TypeError(f"expected Linear at {name}, got {type(module).__name__}")
+        replacement = RTNInt4Linear.empty_from_linear(module, config=config, module_name=name)
+        parent, child_name = _parent_and_child(model, name)
+        _set_child(parent, child_name, replacement)
+
+    state_dict = load_file(
+        artifact_path / "model.safetensors",
+        device=_safetensors_load_device(device, model),
+    )
+    missing, unexpected = model.load_state_dict(state_dict, strict=strict)
+    if strict and (missing or unexpected):
+        raise RuntimeError(f"artifact state mismatch: missing={missing}, unexpected={unexpected}")
+    return manifest

orbitquant/artifacts/manifest.py

@@ -0,0 +1,138 @@
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Any
+
+from orbitquant.config import OrbitQuantConfig
+
+
+@dataclass(frozen=True)
+class OrbitQuantManifest:
+    source_model_id: str
+    source_revision: str
+    source_license: str
+    weight_bits: int
+    activation_bits: int
+    rotation_seed: int
+    block_size: int | str
+    block_size_policy: str
+    codebook_version: int
+    target_policy: str
+    runtime_mode: str
+    activation_kernel_backend: str
+    activation_eps: float = 1e-10
+    adaln_group_size: int = 64
+    quantization_device: str = "unknown"
+    weight_quantization_backend: str = "unknown"
+    quantization_staging_mode: str = "unknown"
+    quantized_modules: list[str] = field(default_factory=list)
+    adaln_modules: list[str] = field(default_factory=list)
+    skipped_modules: list[str] = field(default_factory=list)
+    module_shapes: dict[str, list[int]] = field(default_factory=dict)
+    checksums: dict[str, str] = field(default_factory=dict)
+
+    @classmethod
+    def from_config(
+        cls,
+        config: OrbitQuantConfig,
+        *,
+        source_model_id: str,
+        source_revision: str,
+        source_license: str,
+        quantized_modules: list[str],
+        skipped_modules: list[str],
+        adaln_modules: list[str] | None = None,
+        module_shapes: dict[str, list[int]] | None = None,
+        checksums: dict[str, str] | None = None,
+        quantization_device: str = "unknown",
+        weight_quantization_backend: str = "unknown",
+        quantization_staging_mode: str = "unknown",
+    ) -> OrbitQuantManifest:
+        return cls(
+            source_model_id=source_model_id,
+            source_revision=source_revision,
+            source_license=source_license,
+            weight_bits=config.weight_bits,
+            activation_bits=config.activation_bits,
+            rotation_seed=config.rotation_seed,
+            block_size=config.block_size,
+            block_size_policy="largest_power_of_two_dividing_dim"
+            if config.block_size == "paper"
+            else "explicit",
+            codebook_version=1,
+            target_policy=config.target_policy,
+            runtime_mode=config.runtime_mode,
+            activation_kernel_backend=config.activation_kernel_backend,
+            activation_eps=config.activation_eps,
+            adaln_group_size=config.adaln_group_size,
+            quantization_device=quantization_device,
+            weight_quantization_backend=weight_quantization_backend,
+            quantization_staging_mode=quantization_staging_mode,
+            quantized_modules=list(quantized_modules),
+            adaln_modules=[] if adaln_modules is None else list(adaln_modules),
+            skipped_modules=list(skipped_modules),
+            module_shapes={} if module_shapes is None else dict(module_shapes),
+            checksums={} if checksums is None else dict(checksums),
+        )
+
+    def to_dict(self) -> dict[str, Any]:
+        return {
+            "artifact_format": "orbitquant-v1",
+            "source_model_id": self.source_model_id,
+            "source_revision": self.source_revision,
+            "source_license": self.source_license,
+            "quant_method": "orbitquant",
+            "paper": "https://arxiv.org/abs/2607.02461",
+            "weight_bits": self.weight_bits,
+            "activation_bits": self.activation_bits,
+            "rotation": "rpbh",
+            "rotation_seed": self.rotation_seed,
+            "block_size": self.block_size,
+            "block_size_policy": self.block_size_policy,
+            "codebook": "lloyd_max",
+            "codebook_version": self.codebook_version,
+            "row_norm_dtype": "bfloat16",
+            "runtime_mode": self.runtime_mode,
+            "activation_kernel_backend": self.activation_kernel_backend,
+            "activation_eps": self.activation_eps,
+            "adaln_group_size": self.adaln_group_size,
+            "quantization_device": self.quantization_device,
+            "weight_quantization_backend": self.weight_quantization_backend,
+            "quantization_staging_mode": self.quantization_staging_mode,
+            "target_policy": self.target_policy,
+            "adaln_policy": f"int4_rtn_group{self.adaln_group_size}_bf16_activation",
+            "quantized_modules": self.quantized_modules,
+            "adaln_modules": self.adaln_modules,
+            "skipped_modules": self.skipped_modules,
+            "module_shapes": self.module_shapes,
+            "checksums": self.checksums,
+        }
+
+    @classmethod
+    def from_dict(cls, data: dict[str, Any]) -> OrbitQuantManifest:
+        return cls(
+            source_model_id=data["source_model_id"],
+            source_revision=data["source_revision"],
+            source_license=data["source_license"],
+            weight_bits=int(data["weight_bits"]),
+            activation_bits=int(data["activation_bits"]),
+            rotation_seed=int(data.get("rotation_seed", 0)),
+            block_size=data.get("block_size", "paper"),
+            block_size_policy=data.get(
+                "block_size_policy", "largest_power_of_two_dividing_dim"
+            ),
+            codebook_version=int(data.get("codebook_version", 1)),
+            target_policy=data.get("target_policy", "auto"),
+            runtime_mode=data.get("runtime_mode", "dequant_bf16"),
+            activation_kernel_backend=data.get("activation_kernel_backend", "auto"),
+            activation_eps=float(data.get("activation_eps", 1e-10)),
+            adaln_group_size=int(data.get("adaln_group_size", 64)),
+            quantization_device=data.get("quantization_device", "unknown"),
+            weight_quantization_backend=data.get("weight_quantization_backend", "unknown"),
+            quantization_staging_mode=data.get("quantization_staging_mode", "unknown"),
+            quantized_modules=list(data.get("quantized_modules", [])),
+            adaln_modules=list(data.get("adaln_modules", [])),
+            skipped_modules=list(data.get("skipped_modules", [])),
+            module_shapes=dict(data.get("module_shapes", {})),
+            checksums=dict(data.get("checksums", {})),
+        )