orbitquant 0.1.0__py3-none-any.whl

orbitquant/__init__.py

@@ -0,0 +1,30 @@
+"""OrbitQuant clean-room implementation for diffusion transformer quantization."""
+
+from orbitquant.config import OrbitQuantConfig
+from orbitquant.layers import OrbitQuantLinear
+from orbitquant.modeling import prewarm_quantized_linear_modules
+from orbitquant.pipeline import (
+    build_diffusers_pipeline_quantization_config,
+    load_quantized_pipeline_component,
+    load_quantized_pipeline_from_artifact,
+    quantize_pipeline,
+    save_quantized_pipeline_component,
+)
+from orbitquant.quantizer import register_hf_quantizers
+
+__version__ = "0.1.0"
+
+register_hf_quantizers()
+
+__all__ = [
+    "OrbitQuantConfig",
+    "OrbitQuantLinear",
+    "__version__",
+    "build_diffusers_pipeline_quantization_config",
+    "load_quantized_pipeline_from_artifact",
+    "load_quantized_pipeline_component",
+    "prewarm_quantized_linear_modules",
+    "quantize_pipeline",
+    "register_hf_quantizers",
+    "save_quantized_pipeline_component",
+]

orbitquant/adaln.py

@@ -0,0 +1,174 @@
+from __future__ import annotations
+
+import torch
+from torch import nn
+from torch.nn import functional as F
+
+from orbitquant.config import OrbitQuantConfig
+from orbitquant.packing import pack_lowbit, unpack_lowbit
+
+
+def _packed_length(value_count: int, bits: int) -> int:
+    return (value_count * bits + 7) // 8
+
+
+def _quantize_adaln_weight_reference(
+    weight: torch.Tensor,
+    *,
+    group_size: int,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    num_groups = (weight.shape[1] + group_size - 1) // group_size
+    padded = torch.zeros(
+        weight.shape[0],
+        num_groups * group_size,
+        dtype=torch.float32,
+        device=weight.device,
+    )
+    padded[:, : weight.shape[1]] = weight.to(torch.float32)
+    grouped = padded.reshape(weight.shape[0], num_groups, group_size)
+    scales = grouped.abs().amax(dim=-1).clamp_min(1e-12) / 7.0
+    quantized_signed = torch.round(grouped / scales[..., None]).clamp(-8, 7).to(torch.int16)
+    quantized_unsigned = (quantized_signed + 8).to(torch.uint8)
+    packed = pack_lowbit(quantized_unsigned.flatten(), bits=4, validate=False)
+    return packed, scales
+
+
+def _quantize_adaln_weight(
+    weight: torch.Tensor,
+    *,
+    group_size: int,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    if weight.is_cuda:
+        try:
+            from orbitquant.kernels.triton_cuda import quantize_adaln_weight_with_triton
+        except Exception:
+            pass
+        else:
+            return quantize_adaln_weight_with_triton(weight, group_size=group_size)
+    return _quantize_adaln_weight_reference(weight, group_size=group_size)
+
+
+class RTNInt4Linear(nn.Module):
+    """Symmetric INT4 round-to-nearest linear used for AdaLN modulation weights."""
+
+    def __init__(
+        self,
+        *,
+        in_features: int,
+        out_features: int,
+        group_size: int,
+        module_name: str,
+        packed_weight: torch.Tensor,
+        scales: torch.Tensor,
+        bias: torch.Tensor | None,
+    ) -> None:
+        super().__init__()
+        self.in_features = in_features
+        self.out_features = out_features
+        self.group_size = group_size
+        self.module_name = module_name
+        self.num_groups = (in_features + group_size - 1) // group_size
+        self.register_buffer("packed_weight", packed_weight)
+        self.register_buffer("scales", scales)
+        if bias is None:
+            self.register_parameter("bias", None)
+        else:
+            self.bias = nn.Parameter(bias.detach().clone(), requires_grad=False)
+        self._dequantized_weight_cache: torch.Tensor | None = None
+        self._dequantized_weight_cache_key: tuple[str, torch.dtype] | None = None
+
+    def clear_dequantized_cache(self) -> None:
+        self._dequantized_weight_cache = None
+        self._dequantized_weight_cache_key = None
+
+    @classmethod
+    def from_linear(
+        cls, layer: nn.Linear, *, config: OrbitQuantConfig, module_name: str
+    ) -> RTNInt4Linear:
+        group_size = config.adaln_group_size
+        weight = layer.weight.detach().to(torch.float32)
+        packed, scales = _quantize_adaln_weight(weight, group_size=group_size)
+        bias = None if layer.bias is None else layer.bias.detach()
+        return cls(
+            in_features=layer.in_features,
+            out_features=layer.out_features,
+            group_size=group_size,
+            module_name=module_name,
+            packed_weight=packed,
+            scales=scales.to(torch.bfloat16),
+            bias=bias,
+        )
+
+    @classmethod
+    def empty_from_linear(
+        cls, layer: nn.Linear, *, config: OrbitQuantConfig, module_name: str
+    ) -> RTNInt4Linear:
+        group_size = config.adaln_group_size
+        num_groups = (layer.in_features + group_size - 1) // group_size
+        packed = torch.empty(
+            _packed_length(layer.out_features * num_groups * group_size, 4),
+            dtype=torch.uint8,
+            device=layer.weight.device,
+        )
+        scales = torch.empty(
+            layer.out_features, num_groups, dtype=torch.bfloat16, device=layer.weight.device
+        )
+        bias = None
+        if layer.bias is not None:
+            bias = torch.zeros(layer.out_features, dtype=layer.bias.dtype, device=layer.bias.device)
+        return cls(
+            in_features=layer.in_features,
+            out_features=layer.out_features,
+            group_size=group_size,
+            module_name=module_name,
+            packed_weight=packed,
+            scales=scales,
+            bias=bias,
+        )
+
+    def _dequantize_weight(self, *, device: torch.device, dtype: torch.dtype) -> torch.Tensor:
+        cache_key = (str(device), dtype)
+        if (
+            self._dequantized_weight_cache is not None
+            and self._dequantized_weight_cache_key == cache_key
+        ):
+            return self._dequantized_weight_cache
+
+        if device.type == "cuda":
+            try:
+                from orbitquant.kernels.triton_cuda import dequantize_adaln_weight_with_triton
+            except Exception:
+                dequantize_adaln_weight_with_triton = None
+            if dequantize_adaln_weight_with_triton is not None:
+                weight = dequantize_adaln_weight_with_triton(
+                    self.packed_weight,
+                    self.scales,
+                    out_features=self.out_features,
+                    in_features=self.in_features,
+                    group_size=self.group_size,
+                    device=device,
+                )
+                dequantized = weight.to(dtype=dtype)
+                self._dequantized_weight_cache = dequantized.detach()
+                self._dequantized_weight_cache_key = cache_key
+                return dequantized
+
+        total = self.out_features * self.num_groups * self.group_size
+        unsigned = unpack_lowbit(self.packed_weight, bits=4, length=total).to(device=device)
+        signed = unsigned.to(torch.int16).sub(8).to(torch.float32)
+        grouped = signed.reshape(self.out_features, self.num_groups, self.group_size)
+        scales = self.scales.to(device=device, dtype=torch.float32)
+        weight = (grouped * scales[..., None]).reshape(
+            self.out_features, self.num_groups * self.group_size
+        )
+        weight = weight[:, : self.in_features]
+        dequantized = weight.to(device=device, dtype=dtype)
+        self._dequantized_weight_cache = dequantized.detach()
+        self._dequantized_weight_cache_key = cache_key
+        return dequantized
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        activation = x.to(torch.bfloat16)
+        weight = self._dequantize_weight(device=x.device, dtype=torch.bfloat16)
+        bias = None if self.bias is None else self.bias.to(device=x.device, dtype=torch.bfloat16)
+        return F.linear(activation, weight, bias)

orbitquant/artifacts/__init__.py

@@ -0,0 +1,32 @@
+from orbitquant.artifacts.assets import record_artifact_asset
+from orbitquant.artifacts.benchmark import record_artifact_metrics
+from orbitquant.artifacts.checksums import sha256_file, write_sha256sums
+from orbitquant.artifacts.comparisons import create_artifact_image_comparisons
+from orbitquant.artifacts.loader import load_orbitquant_artifact
+from orbitquant.artifacts.manifest import OrbitQuantManifest
+from orbitquant.artifacts.model_card import render_model_card
+from orbitquant.artifacts.refresh import refresh_artifact_checksums
+from orbitquant.artifacts.repair import repair_artifact_metadata
+from orbitquant.artifacts.validator import (
+    validate_artifact_policy_inventory,
+    validate_orbitquant_artifact,
+    validate_policy_inventory_payloads,
+)
+from orbitquant.artifacts.writer import save_orbitquant_artifact
+
+__all__ = [
+    "OrbitQuantManifest",
+    "create_artifact_image_comparisons",
+    "load_orbitquant_artifact",
+    "record_artifact_asset",
+    "record_artifact_metrics",
+    "refresh_artifact_checksums",
+    "render_model_card",
+    "repair_artifact_metadata",
+    "save_orbitquant_artifact",
+    "validate_artifact_policy_inventory",
+    "validate_orbitquant_artifact",
+    "validate_policy_inventory_payloads",
+    "sha256_file",
+    "write_sha256sums",
+]

orbitquant/artifacts/assets.py

@@ -0,0 +1,59 @@
+from __future__ import annotations
+
+import json
+from pathlib import Path
+
+from orbitquant.artifacts.checksums import (
+    sha256_file,
+    validate_checksums,
+    write_sha256sums_from_manifest,
+)
+from orbitquant.artifacts.manifest import OrbitQuantManifest
+from orbitquant.artifacts.validator import validate_required_artifact_files
+
+
+def _read_manifest(artifact_path: Path) -> OrbitQuantManifest:
+    return OrbitQuantManifest.from_dict(
+        json.loads((artifact_path / "orbitquant_manifest.json").read_text(encoding="utf-8"))
+    )
+
+
+def _relative_asset_path(artifact_path: Path, asset_path: Path) -> str:
+    try:
+        relative = asset_path.resolve().relative_to(artifact_path.resolve())
+    except ValueError as exc:
+        raise ValueError("artifact asset must be inside the artifact directory") from exc
+    if relative.parts[0] != "assets":
+        raise ValueError("artifact asset must live under assets/")
+    return relative.as_posix()
+
+
+def record_artifact_asset(
+    artifact_dir: str | Path,
+    asset_path: str | Path,
+    *,
+    validate_checksums_enabled: bool = True,
+    refresh_checksums_enabled: bool = True,
+) -> str:
+    artifact_path = Path(artifact_dir)
+    asset = Path(asset_path)
+    if not asset.is_file():
+        raise RuntimeError(f"artifact asset missing: {asset}")
+
+    validate_required_artifact_files(artifact_path)
+    manifest = _read_manifest(artifact_path)
+    if validate_checksums_enabled:
+        validate_checksums(artifact_path, manifest.checksums)
+
+    relative_path = _relative_asset_path(artifact_path, asset)
+    if not refresh_checksums_enabled:
+        return relative_path
+    checksums = dict(manifest.checksums)
+    checksums[relative_path] = sha256_file(asset)
+    payload = manifest.to_dict()
+    payload["checksums"] = checksums
+    (artifact_path / "orbitquant_manifest.json").write_text(
+        json.dumps(payload, indent=2) + "\n", encoding="utf-8"
+    )
+    write_sha256sums_from_manifest(artifact_path, checksums)
+    return relative_path

orbitquant/artifacts/benchmark.py

@@ -0,0 +1,133 @@
+from __future__ import annotations
+
+import csv
+import json
+from pathlib import Path
+from typing import Any
+
+from orbitquant.artifacts.checksums import (
+    sha256_file,
+    validate_checksums,
+    write_sha256sums_from_manifest,
+)
+from orbitquant.artifacts.manifest import OrbitQuantManifest
+from orbitquant.artifacts.validator import validate_required_artifact_files
+
+_VALID_METRIC_SPLITS = {"original", "orbitquant"}
+
+
+def _normalize_json_value(value: Any) -> Any:
+    if isinstance(value, Path):
+        return str(value)
+    if isinstance(value, dict):
+        return {str(key): _normalize_json_value(item) for key, item in value.items()}
+    if isinstance(value, list | tuple):
+        return [_normalize_json_value(item) for item in value]
+    item = getattr(value, "item", None)
+    if callable(item):
+        try:
+            return item()
+        except Exception:
+            pass
+    return value
+
+
+def _csv_value(value: Any) -> Any:
+    normalized = _normalize_json_value(value)
+    if isinstance(normalized, dict | list):
+        return json.dumps(normalized, sort_keys=True, separators=(",", ":"))
+    return normalized
+
+
+def _read_manifest(artifact_path: Path) -> OrbitQuantManifest:
+    return OrbitQuantManifest.from_dict(
+        json.loads((artifact_path / "orbitquant_manifest.json").read_text(encoding="utf-8"))
+    )
+
+
+def _read_summary(summary_path: Path) -> dict[str, Any]:
+    if not summary_path.is_file():
+        return {}
+    return json.loads(summary_path.read_text(encoding="utf-8"))
+
+
+def _record_count(jsonl_path: Path) -> int:
+    return sum(1 for line in jsonl_path.read_text(encoding="utf-8").splitlines() if line.strip())
+
+
+def _write_manifest_with_refreshed_checksums(
+    artifact_path: Path, manifest: OrbitQuantManifest, relative_paths: tuple[str, ...]
+) -> dict[str, str]:
+    checksums = dict(manifest.checksums)
+    for relative_path in relative_paths:
+        checksums[relative_path] = sha256_file(artifact_path / relative_path)
+    payload = manifest.to_dict()
+    payload["checksums"] = checksums
+    (artifact_path / "orbitquant_manifest.json").write_text(
+        json.dumps(payload, indent=2) + "\n", encoding="utf-8"
+    )
+    return checksums
+
+
+def record_artifact_metrics(
+    artifact_dir: str | Path,
+    *,
+    split: str,
+    metrics: dict[str, Any],
+    metadata: dict[str, Any] | None = None,
+    validate_checksums_enabled: bool = True,
+    refresh_checksums_enabled: bool = True,
+) -> dict[str, Any]:
+    normalized_split = split.lower()
+    if normalized_split not in _VALID_METRIC_SPLITS:
+        raise ValueError("split must be one of: original, orbitquant")
+    if not metrics:
+        raise ValueError("metrics must not be empty")
+
+    artifact_path = Path(artifact_dir)
+    validate_required_artifact_files(artifact_path)
+    manifest = _read_manifest(artifact_path)
+    if validate_checksums_enabled:
+        validate_checksums(artifact_path, manifest.checksums)
+
+    record = {
+        "split": normalized_split,
+        "metrics": _normalize_json_value(metrics),
+        "metadata": {} if metadata is None else _normalize_json_value(metadata),
+    }
+    summary_path = artifact_path / "benchmark" / "summary.json"
+    jsonl_path = artifact_path / "benchmark" / f"{normalized_split}.metrics.jsonl"
+    csv_path = artifact_path / "benchmark" / f"{normalized_split}.metrics.csv"
+
+    with jsonl_path.open("a", encoding="utf-8") as handle:
+        handle.write(json.dumps(record, separators=(",", ":")) + "\n")
+
+    write_csv_header = not csv_path.exists()
+    with csv_path.open("a", encoding="utf-8", newline="") as handle:
+        writer = csv.writer(handle)
+        if write_csv_header:
+            writer.writerow(["metric", "value"])
+        for metric, value in record["metrics"].items():
+            writer.writerow([metric, _csv_value(value)])
+
+    summary = _read_summary(summary_path)
+    summary["status"] = "metrics_recorded"
+    split_summaries = summary.setdefault("metrics", {})
+    split_summaries[normalized_split] = {
+        "records": _record_count(jsonl_path),
+        "latest": record,
+    }
+    summary_path.write_text(json.dumps(summary, indent=2) + "\n", encoding="utf-8")
+
+    changed_paths = (
+        "benchmark/summary.json",
+        f"benchmark/{normalized_split}.metrics.jsonl",
+        f"benchmark/{normalized_split}.metrics.csv",
+    )
+    if not refresh_checksums_enabled:
+        return record
+    checksums = _write_manifest_with_refreshed_checksums(
+        artifact_path, manifest, changed_paths
+    )
+    write_sha256sums_from_manifest(artifact_path, checksums)
+    return record

orbitquant/artifacts/checksums.py

@@ -0,0 +1,132 @@
+from __future__ import annotations
+
+import hashlib
+from pathlib import Path
+
+_IGNORED_ARTIFACT_PATH_PARTS = {
+    ".cache",
+    ".gitattributes",
+    ".gitignore",
+    ".pytest_cache",
+    ".ruff_cache",
+    "__pycache__",
+}
+
+
+def is_ignored_artifact_relative_path(relative_path: str | Path) -> bool:
+    return any(part in _IGNORED_ARTIFACT_PATH_PARTS for part in Path(relative_path).parts)
+
+
+def sha256_file(path: str | Path) -> str:
+    digest = hashlib.sha256()
+    with Path(path).open("rb") as handle:
+        for chunk in iter(lambda: handle.read(1024 * 1024), b""):
+            digest.update(chunk)
+    return digest.hexdigest()
+
+
+def write_sha256sums(root: str | Path, output: str | Path | None = None) -> Path:
+    root_path = Path(root)
+    output_path = root_path / "SHA256SUMS" if output is None else Path(output)
+    lines: list[str] = []
+    for path in sorted(root_path.rglob("*")):
+        if not path.is_file() or path == output_path:
+            continue
+        rel = path.relative_to(root_path)
+        if is_ignored_artifact_relative_path(rel):
+            continue
+        lines.append(f"{sha256_file(path)}  {rel.as_posix()}")
+    output_path.write_text("\n".join(lines) + ("\n" if lines else ""), encoding="utf-8")
+    return output_path
+
+
+def read_sha256sums(path: str | Path) -> dict[str, str]:
+    path = Path(path)
+    if not path.is_file():
+        return {}
+    entries: dict[str, str] = {}
+    for line in path.read_text(encoding="utf-8").splitlines():
+        if not line.strip():
+            continue
+        digest, _, relative_path = line.partition("  ")
+        if digest and relative_path:
+            entries[relative_path] = digest
+    return entries
+
+
+def write_sha256sums_from_manifest(
+    root: str | Path,
+    checksums: dict[str, str],
+    output: str | Path | None = None,
+) -> Path:
+    root_path = Path(root)
+    output_path = root_path / "SHA256SUMS" if output is None else Path(output)
+    entries = read_sha256sums(output_path)
+    entries.update(checksums)
+
+    manifest_path = root_path / "orbitquant_manifest.json"
+    if manifest_path.is_file():
+        entries["orbitquant_manifest.json"] = sha256_file(manifest_path)
+    readme_path = root_path / "README.md"
+    if readme_path.is_file():
+        entries["README.md"] = sha256_file(readme_path)
+
+    entries = {
+        relative_path: digest
+        for relative_path, digest in entries.items()
+        if relative_path != "SHA256SUMS"
+        and not is_ignored_artifact_relative_path(relative_path)
+        and (root_path / relative_path).is_file()
+    }
+    output_path.write_text(
+        "\n".join(f"{digest}  {relative_path}" for relative_path, digest in sorted(entries.items()))
+        + ("\n" if entries else ""),
+        encoding="utf-8",
+    )
+    return output_path
+
+
+def validate_checksums(root: str | Path, checksums: dict[str, str]) -> None:
+    root_path = Path(root)
+    for relative_path, expected in checksums.items():
+        if is_ignored_artifact_relative_path(relative_path):
+            continue
+        path = root_path / relative_path
+        if not path.is_file():
+            raise RuntimeError(f"artifact checksum target missing: {relative_path}")
+        actual = sha256_file(path)
+        if actual != expected:
+            raise RuntimeError(
+                f"artifact checksum mismatch for {relative_path}: "
+                f"expected {expected}, got {actual}"
+            )
+
+
+def validate_sha256sums(
+    root: str | Path,
+    *,
+    required_paths: tuple[str, ...] = (),
+    sha256sums_path: str | Path | None = None,
+) -> dict[str, str]:
+    root_path = Path(root)
+    sums_path = root_path / "SHA256SUMS" if sha256sums_path is None else Path(sha256sums_path)
+    entries = read_sha256sums(sums_path)
+    if not entries:
+        raise RuntimeError("SHA256SUMS is empty or missing")
+    missing_entries = sorted(set(required_paths) - set(entries))
+    if missing_entries:
+        raise RuntimeError(f"SHA256SUMS missing entries: {missing_entries}")
+    for relative_path, expected in entries.items():
+        if is_ignored_artifact_relative_path(relative_path):
+            continue
+        if relative_path == "SHA256SUMS":
+            raise RuntimeError("SHA256SUMS must not include itself")
+        path = root_path / relative_path
+        if not path.is_file():
+            raise RuntimeError(f"SHA256SUMS target missing: {relative_path}")
+        actual = sha256_file(path)
+        if actual != expected:
+            raise RuntimeError(
+                f"SHA256SUMS mismatch for {relative_path}: expected {expected}, got {actual}"
+            )
+    return entries