scripts/orbitquant_h3_compat.py

"""Narrow MiniMax-H3 dtype-sentinel compatibility for packed linear modules."""

from __future__ import annotations

from dataclasses import dataclass

import torch

from orbitquant.adaln import RTNInt4Linear
from orbitquant.layers import OrbitQuantLinear


@dataclass(frozen=True)
class _WeightDTypeView:
    dtype: torch.dtype


def _attach_dtype_view(module: torch.nn.Module) -> bool:
    if hasattr(module, "weight"):
        return False
    for name in ("bias", "row_norms", "scales"):
        tensor = getattr(module, name, None)
        if isinstance(tensor, torch.Tensor):
            # MiniMax-H3 only reads `.weight.dtype` on these boundary/AdaLN
            # projections. Keep that read cheap and never dequantize a weight.
            object.__setattr__(module, "weight", _WeightDTypeView(tensor.dtype))
            return True
    raise RuntimeError(f"cannot infer source dtype for {type(module).__name__}")


def enable_h3_orbitquant_compat(transformer: torch.nn.Module) -> int:
    """Attach dtype-only views to the packed modules H3 inspects directly."""
    count = 0
    for name, module in transformer.named_modules():
        needs_view = name in {"proj_in", "audio_proj_in"} or name.endswith("adaln_proj.linear")
        if needs_view and isinstance(module, (OrbitQuantLinear, RTNInt4Linear)):
            count += int(_attach_dtype_view(module))
    return count