File size: 1,387 Bytes
fa2d87b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 | """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
|