hiera-optim 0.1.0__py3-none-any.whl

hiera_optim/__init__.py

@@ -0,0 +1,44 @@
+"""hiera_optim — training-throughput optimisations for FAIR's Hiera.
+
+Quick start::
+
+    from models.hiera import mae_hiera_base_224         # FAIR Hiera
+    from hiera_optim import optimize
+
+    model = mae_hiera_base_224(pretrained=False, in_chans=8)
+    optimize(model)                                      # in-place
+    # optionally: model = torch.compile(model, mode="default", dynamic=False)
+
+What `optimize()` does:
+  1. Swaps every `MaskUnitAttention` for a FlashAttention/cuDNN-friendly
+     4-D variant (`MaskUnitAttentionFast`). Restores math-fallback SDPA to
+     fused kernel paths — 5-12× per-call attention speedup.
+  2. Replaces the boolean `x[mask.tile(...)]` and `x_dec[mask] = ...`
+     indexing patterns with explicit `torch.gather` / `scatter_`. Removes
+     the `aten::nonzero` graph break (compile-friendly).
+
+Optional add-ons (opt-in, not invoked by default):
+  - `optimize(model, sdpa_backend="auto" | "cudnn" | ...)`: pin the SDPA
+    backend per-block. Sometimes helps, sometimes hurts — see RESULTS.md.
+  - `enable_stage_checkpointing(model, stages=(2,))`: trade compute for
+    activation memory at chosen stages. OOM lever, not a throughput tool.
+"""
+from .patch import optimize, swap_mask_unit_attention, recommended_backend
+from .attention import MaskUnitAttentionFast, BACKEND_NAMES
+from .checkpoint import enable_stage_checkpointing, disable_stage_checkpointing
+from .adapters import HieraAdapter, get_hiera_adapter, auto_detect
+
+__version__ = "0.1.0"
+
+__all__ = [
+    "optimize",
+    "swap_mask_unit_attention",
+    "recommended_backend",
+    "MaskUnitAttentionFast",
+    "BACKEND_NAMES",
+    "enable_stage_checkpointing",
+    "disable_stage_checkpointing",
+    "HieraAdapter",
+    "get_hiera_adapter",
+    "auto_detect",
+]

hiera_optim/adapters/__init__.py

@@ -0,0 +1,28 @@
+"""Model adapters.
+
+The rest of the package is intentionally model-name-free: it operates on
+PyTorch `nn.Module` graphs and looks up attention/block classes through
+adapters. Each adapter teaches `optimize()` how to find the right submodules
+on a specific model family.
+
+Currently bundled:
+  - hiera: FAIR's Hiera / MaskedAutoencoderHiera (https://github.com/facebookresearch/hiera)
+
+Other adapters (Swin, MViTv2, JEPA-Hiera, custom architectures) can plug in by
+implementing the same `ModelAdapter` protocol.
+"""
+from __future__ import annotations
+from typing import Optional
+from .hiera import HieraAdapter, get_hiera_adapter
+
+__all__ = ["HieraAdapter", "get_hiera_adapter", "auto_detect"]
+
+
+def auto_detect(model) -> Optional["ModelAdapter"]:  # type: ignore[name-defined]
+    """Best-effort: return the right adapter for a given model. Returns None
+    if no bundled adapter matches.
+    """
+    a = get_hiera_adapter(model)
+    if a is not None:
+        return a
+    return None

hiera_optim/adapters/hiera.py

@@ -0,0 +1,142 @@
+"""Adapter for FAIR Hiera (https://github.com/facebookresearch/hiera).
+
+This is the ONE module in `hiera_optim` allowed to import FAIR classes by name.
+Everything else works through this adapter's protocol so the package is
+trivially portable to derivative architectures.
+
+The adapter resolves at import time and degrades gracefully if FAIR Hiera
+isn't installed — `get_hiera_adapter(model)` returns None instead of crashing.
+"""
+from __future__ import annotations
+import importlib
+from dataclasses import dataclass
+from typing import Any, Optional, Type
+
+import torch
+import torch.nn as nn
+
+
+@dataclass(frozen=True)
+class _HieraSymbols:
+    """Resolved references to the FAIR Hiera classes we touch."""
+    Hiera: Type[nn.Module]
+    MaskedAutoencoderHiera: Type[nn.Module]
+    HieraBlock: Type[nn.Module]
+    MaskUnitAttention: Type[nn.Module]
+    apply_fusion_head: callable
+    undo_windowing: callable
+
+
+def _resolve() -> Optional[_HieraSymbols]:
+    """Try a few import paths. Returns None if Hiera isn't available."""
+    candidates = [
+        # In-tree (brain_atlas, the development environment)
+        ("models.hiera", "utils.hiera_utils"),
+        # PyPI package (`pip install hiera-transformer`)
+        ("hiera.hiera", "hiera.hiera_utils"),
+        ("hiera", "hiera.hiera_utils"),
+    ]
+    for hmod, hutil in candidates:
+        try:
+            hm = importlib.import_module(hmod)
+            hu = importlib.import_module(hutil)
+            return _HieraSymbols(
+                Hiera=hm.Hiera,
+                MaskedAutoencoderHiera=hm.MaskedAutoencoderHiera,
+                HieraBlock=hm.HieraBlock,
+                MaskUnitAttention=hm.MaskUnitAttention,
+                apply_fusion_head=hm.apply_fusion_head,
+                undo_windowing=hu.undo_windowing,
+            )
+        except (ImportError, AttributeError):
+            continue
+    return None
+
+
+_SYMS = _resolve()
+
+
+def is_available() -> bool:
+    return _SYMS is not None
+
+
+def symbols() -> _HieraSymbols:
+    """Return resolved FAIR symbols. Raises if Hiera isn't installed."""
+    if _SYMS is None:
+        raise ImportError(
+            "FAIR Hiera not installed. Add `pip install hiera-transformer`, "
+            "or ensure `models.hiera` is importable in this project."
+        )
+    return _SYMS
+
+
+class HieraAdapter:
+    """Describes how to introspect / patch a FAIR Hiera model.
+
+    The adapter is the bridge between FAIR's class hierarchy and our
+    framework-agnostic patching code. Methods are pure introspection — they
+    do not import FAIR classes unless explicitly needed.
+    """
+
+    def __init__(self):
+        if not is_available():
+            raise ImportError(
+                "HieraAdapter requires FAIR Hiera to be importable."
+            )
+        self._syms = symbols()
+
+    # ---- Introspection -----------------------------------------------------
+
+    def matches(self, model: nn.Module) -> bool:
+        """True if `model` is a Hiera-family model."""
+        return isinstance(model, (self._syms.Hiera, self._syms.MaskedAutoencoderHiera))
+
+    def is_mae(self, model: nn.Module) -> bool:
+        return isinstance(model, self._syms.MaskedAutoencoderHiera)
+
+    def block_class(self) -> Type[nn.Module]:
+        return self._syms.HieraBlock
+
+    def attention_class(self) -> Type[nn.Module]:
+        return self._syms.MaskUnitAttention
+
+    def encoder_blocks(self, model: nn.Module) -> nn.ModuleList:
+        """Returns the encoder block list (`model.blocks`)."""
+        return model.blocks
+
+    def decoder_blocks(self, model: nn.Module) -> Optional[nn.ModuleList]:
+        """Returns the MAE decoder block list, or None for non-MAE models."""
+        return getattr(model, "decoder_blocks", None)
+
+    def stage_ends(self, model: nn.Module) -> list[int]:
+        return list(model.stage_ends)
+
+    # ---- Layout convention ------------------------------------------------
+    #
+    # FAIR's `Unroll` produces a (T outer, nw inner) token layout: token n in
+    # the N axis maps to (t = n // nw, w = n % nw). This is what FAIR's
+    # `do_pool(x, stride)` (which expects stride as the OUTER dim) and the
+    # MaskUnitAttention reshape pattern both rely on.
+    layout = "T_outer_nw_inner"
+
+    # ---- Helpers used by the patched forwards -----------------------------
+
+    def apply_fusion_head(self, head: nn.Module, x: torch.Tensor) -> torch.Tensor:
+        return self._syms.apply_fusion_head(head, x)
+
+    def undo_windowing(self, x, shape, mu_shape):
+        return self._syms.undo_windowing(x, shape, mu_shape)
+
+
+# Module-level singleton resolution (cheap; just an isinstance check)
+_DEFAULT_ADAPTER: Optional[HieraAdapter] = None
+
+
+def get_hiera_adapter(model: nn.Module) -> Optional[HieraAdapter]:
+    """Return a HieraAdapter if the model is a Hiera-family model, else None."""
+    global _DEFAULT_ADAPTER
+    if not is_available():
+        return None
+    if _DEFAULT_ADAPTER is None:
+        _DEFAULT_ADAPTER = HieraAdapter()
+    return _DEFAULT_ADAPTER if _DEFAULT_ADAPTER.matches(model) else None

hiera_optim/attention/__init__.py

@@ -0,0 +1,6 @@
+"""Attention modules. Currently exports MaskUnitAttentionFast — the
+FlashAttention/cuDNN-friendly 4-D variant of FAIR's MaskUnitAttention.
+"""
+from .mask_unit import MaskUnitAttentionFast, copy_weights_from_orig, BACKEND_NAMES
+
+__all__ = ["MaskUnitAttentionFast", "copy_weights_from_orig", "BACKEND_NAMES"]

hiera_optim/attention/mask_unit.py

@@ -0,0 +1,116 @@
+"""Optimized MaskUnitAttention — drop-in replacement.
+
+Key fixes vs FAIR original (models/hiera.py):
+  - Reshape Q/K/V to 4D (B*num_windows, heads, T, D) so SDPA dispatches to
+    FlashAttention (cuDNN/Flash). Original feeds 5D tensors which fall back
+    to the math backend (12-13x slower on stage-0 shapes per microbench).
+  - Match FAIR's N-axis layout exactly: token n in input has n = t*nw + w,
+    so within-window positions are SLOW-varying and num_windows is FAST.
+    That's because the upstream Unroll module produces this interleaved
+    layout (and do_pool relies on the stride axis being outer).
+  - Skip per-tensor .contiguous() — the permute lands flash-friendly.
+  - Optional per-stage SDPA backend hint via `sdpa_backend` attribute. Set to
+    one of {"cudnn", "flash", "mem_efficient", "math", None}; None lets the
+    PyTorch dispatcher pick. Per-stage tuning is useful because Hiera's small
+    Tq stages (stage-1 post q_pool: T=16) often favor mem-efficient while the
+    long-seq global-attention stages favor cuDNN-attn or flash on Hopper.
+
+Numerically identical to FAIR baseline up to bf16 noise (~1e-2 max abs diff,
+~1e-3 rel RMS) — the noise is the difference between SDPA math vs flash.
+"""
+from __future__ import annotations
+
+from typing import Optional
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.nn.attention import SDPBackend, sdpa_kernel
+
+
+_BACKEND_MAP = {
+    "cudnn": SDPBackend.CUDNN_ATTENTION,
+    "flash": SDPBackend.FLASH_ATTENTION,
+    "mem_efficient": SDPBackend.EFFICIENT_ATTENTION,
+    "math": SDPBackend.MATH,
+}
+
+BACKEND_NAMES = tuple(_BACKEND_MAP.keys())
+
+
+class MaskUnitAttentionFast(nn.Module):
+    """Drop-in replacement for models.hiera.MaskUnitAttention with 4D SDPA.
+
+    Args:
+        sdpa_backend: optional SDPA backend hint. One of {"cudnn", "flash",
+            "mem_efficient", "math", None}. None = default dispatcher.
+    """
+
+    def __init__(
+        self,
+        dim: int,
+        dim_out: int,
+        heads: int,
+        q_stride: int = 1,
+        window_size: int = 0,
+        use_mask_unit_attn: bool = False,
+        sdpa_backend: Optional[str] = None,
+    ):
+        super().__init__()
+        self.dim = dim
+        self.dim_out = dim_out
+        self.heads = heads
+        self.q_stride = q_stride
+        self.head_dim = dim_out // heads
+        self.scale = self.head_dim ** -0.5
+        self.qkv = nn.Linear(dim, 3 * dim_out)
+        self.proj = nn.Linear(dim_out, dim_out)
+        self.window_size = window_size
+        self.use_mask_unit_attn = use_mask_unit_attn
+        if sdpa_backend is not None and sdpa_backend not in _BACKEND_MAP:
+            raise ValueError(f"sdpa_backend must be one of {list(_BACKEND_MAP)} or None; got {sdpa_backend!r}")
+        self.sdpa_backend = sdpa_backend
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        B, N, _ = x.shape
+        H, D = self.heads, self.head_dim
+        nw = (N // (self.q_stride * self.window_size)) if self.use_mask_unit_attn else 1
+        T = N // nw  # tokens per window before q-pool
+
+        # qkv: (B, N, 3*dim_out)
+        # FAIR layout: N is interpreted as (T, nw) with nw fast-varying.
+        # We want (3, B*nw, H, T, D) so SDPA gets a 4D Q/K/V.
+        qkv = self.qkv(x).view(B, T, nw, 3, H, D)
+        # permute -> (3, B, nw, H, T, D), then flatten (B, nw) into batch
+        qkv = qkv.permute(3, 0, 2, 4, 1, 5).reshape(3, B * nw, H, T, D)
+        q, k, v = qkv[0], qkv[1], qkv[2]
+
+        if self.q_stride > 1:
+            # Max-pool Q over the q_stride flat axis. T = q_stride * Tq.
+            # FAIR's do_pool: view(B, stride, -1, C).max(dim=1) — stride is OUTER.
+            # So inside T, q_stride is slow-varying. Mirror: view (.., q_stride, Tq, ..)
+            Tq = T // self.q_stride
+            q = q.view(B * nw, H, self.q_stride, Tq, D).amax(dim=2)
+
+        # 4D SDPA → FlashAttention/cuDNN. Optionally pin a backend.
+        if self.sdpa_backend is None:
+            out = F.scaled_dot_product_attention(q, k, v)
+        else:
+            with sdpa_kernel([_BACKEND_MAP[self.sdpa_backend]]):
+                out = F.scaled_dot_product_attention(q, k, v)
+        # out: (B*nw, H, Tq, D)
+        Tq_out = out.shape[2]
+
+        # Back to (B, N_out, dim_out) with N_out indexed as (tq, w) tq slow / w fast
+        # out: (B*nw, H, Tq, D) -> (B, nw, H, Tq, D) -> (B, Tq, nw, H, D) -> reshape
+        out = out.view(B, nw, H, Tq_out, D).permute(0, 3, 1, 2, 4).reshape(B, Tq_out * nw, self.dim_out)
+        return self.proj(out)
+
+
+@torch.no_grad()
+def copy_weights_from_orig(fast: MaskUnitAttentionFast, orig) -> None:
+    """Copy parameters from FAIR's MaskUnitAttention into a fast one."""
+    fast.qkv.weight.copy_(orig.qkv.weight)
+    fast.qkv.bias.copy_(orig.qkv.bias)
+    fast.proj.weight.copy_(orig.proj.weight)
+    fast.proj.bias.copy_(orig.proj.bias)

hiera_optim/checkpoint.py

@@ -0,0 +1,113 @@
+"""Selective gradient checkpointing for Hiera.
+
+Wraps specific blocks (default: stage-2) with `torch.utils.checkpoint` so that
+their activations are *not* stored during forward but recomputed on backward.
+Costs ~33% extra compute on the checkpointed blocks; saves ~50% of activation
+memory total in Hiera-Base (stage-2 holds 16 of 24 enc blocks).
+
+Use case: fit larger per-GPU batch → climb out of launch-overhead regime →
+higher MFU. Especially useful when going from B=128 → B=256 per GPU on H100.
+
+Compatible with `torch.compile(mode="default")`. Use `use_reentrant=False`
+(the modern checkpoint variant) to avoid known autograd issues with reentrant
+checkpointing + compile.
+"""
+from __future__ import annotations
+from typing import Iterable, Optional, Union
+
+import torch
+import torch.nn as nn
+from torch.utils.checkpoint import checkpoint
+
+
+class _CheckpointedBlock(nn.Module):
+    """Wraps a block so its forward goes through `torch.utils.checkpoint`.
+
+    Keeps a reference to the original block (for weight loading / state_dict
+    compatibility) and forwards all attribute access to it. Works for any
+    single-input single-output `nn.Module` whose forward is a function of
+    its parameters and the input — typical for transformer blocks.
+    """
+    def __init__(self, block: HieraBlock):
+        super().__init__()
+        self.block = block
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return checkpoint(self.block, x, use_reentrant=False)
+
+    # Pass-through for `.attn`, `.dim`, etc. that downstream code may inspect
+    def __getattr__(self, name):
+        try:
+            return super().__getattr__(name)
+        except AttributeError:
+            return getattr(self.block, name)
+
+
+def _block_to_stage(model: nn.Module) -> list[int]:
+    """Returns a list mapping block_idx -> stage_idx based on model.stage_ends."""
+    stage_ends = model.stage_ends
+    mapping = []
+    for i in range(stage_ends[-1] + 1):
+        for s, end in enumerate(stage_ends):
+            if i <= end:
+                mapping.append(s)
+                break
+    return mapping
+
+
+def enable_stage_checkpointing(
+    model: nn.Module,
+    stages: Union[Iterable[int], str] = (2,),
+    decoder: bool = False,
+) -> int:
+    """In-place: wrap blocks in `stages` with checkpoint().
+
+    Args:
+        model: a built Hiera / MaskedAutoencoderHiera.
+        stages: iterable of stage indices to checkpoint (default: just stage 2),
+                or the string "all" to checkpoint every stage.
+        decoder: if True (and model is MAE) also checkpoint decoder blocks.
+                 Decoder blocks have a much smaller token count so this is rarely
+                 worth it; default is False.
+
+    Returns the count of wrapped blocks.
+    """
+    if isinstance(stages, str):
+        if stages != "all":
+            raise ValueError(f"stages must be an iterable or 'all'; got {stages!r}")
+        stages = set(range(len(model.stage_ends)))
+    else:
+        stages = set(stages)
+
+    mapping = _block_to_stage(model)
+    n_wrapped = 0
+    for i, b in enumerate(model.blocks):
+        if isinstance(b, _CheckpointedBlock):
+            continue  # idempotent
+        if mapping[i] in stages:
+            model.blocks[i] = _CheckpointedBlock(b)
+            n_wrapped += 1
+
+    if decoder and hasattr(model, "decoder_blocks"):
+        for i, b in enumerate(model.decoder_blocks):
+            if isinstance(b, _CheckpointedBlock):
+                continue
+            model.decoder_blocks[i] = _CheckpointedBlock(b)
+            n_wrapped += 1
+
+    return n_wrapped
+
+
+def disable_stage_checkpointing(model: nn.Module) -> int:
+    """Unwrap any _CheckpointedBlock wrappers. Returns count unwrapped."""
+    n = 0
+    for i, b in enumerate(model.blocks):
+        if isinstance(b, _CheckpointedBlock):
+            model.blocks[i] = b.block
+            n += 1
+    if hasattr(model, "decoder_blocks"):
+        for i, b in enumerate(model.decoder_blocks):
+            if isinstance(b, _CheckpointedBlock):
+                model.decoder_blocks[i] = b.block
+                n += 1
+    return n

hiera_optim/kernels/__init__.py

@@ -0,0 +1,22 @@
+"""Triton kernels.
+
+These are pure shape-parametric kernels with no model dependencies. Each
+provides a PyTorch reference implementation alongside the Triton kernel for
+correctness testing.
+
+Bundled:
+  - mask_gather: Triton mask-unit gather. Correct, but `torch.gather` is
+    faster at the shapes we tested — kept for fusion experiments.
+  - flash_qpool: Triton fused Q-pool + FlashAttention. Correct, but loses
+    to PyTorch's flash on RTX 4090 short-seq workloads. May be useful on
+    Hopper with hand-tuned TMA paths.
+"""
+from .mask_gather import gather_mu_groups_triton, keep_idx_from_bool_mask
+from .flash_qpool import flash_qpool_attention, qpool_attention_ref
+
+__all__ = [
+    "gather_mu_groups_triton",
+    "keep_idx_from_bool_mask",
+    "flash_qpool_attention",
+    "qpool_attention_ref",
+]