hiera-optim 0.2.0

Drop-in throughput and memory optimisations for FAIR Hiera. 0.2: graph-safe MAE for torch.compile(reduce-overhead) — 2.37x on Hiera-Base GH200 over eager.

hiera-optim

Drop-in throughput optimisations for FAIR's Hiera and its MAE variant. Two lines for the layout fix, one extra flag for graph-safe torch.compile(reduce-overhead). Numerically equivalent within bf16 noise; weights preserved.

from hiera_optim import optimize
optimize(model)                                  # 1.3–1.9x: layout + gather
optimize(model, graph_safe=True)                 # 2.2–2.4x e2e once you also torch.compile

Results

GH200 (Hopper), bf16, full forward + backward, B=128 in-chans=8.

variant	eager	optimize(model)	optimize(graph_safe=True) + compile(mode="reduce-overhead")
Hiera-Tiny	1.00x	1.75x	2.20x
Hiera-Small	1.00x	1.46x	2.34x
Hiera-Base	1.00x	1.32x	2.37x

Hiera-Base step time: 112 ms → 85 ms → 47 ms. Same loss within 5e-3 rel diff in bf16; same gradient flow (worst grad RMS in tests: 4.4e-5).

RTX 4090, Hiera-Base, in-chans=8, B=8: eager 54.3 ms → 13.1 ms (4.13x) under graph-safe + reduce-overhead.

Multi-GPU (DDP) — GH200 4-GPU, B=128 / rank

variant	ms / step (4-GPU)	total samp / s	DDP overhead vs 1-GPU
Hiera-Base	52.3	9,787	+10% (all-reduce on 4× Hopper)
Hiera-Tiny	40.2	12,753	similar

Validated with static_graph=True on DDPStrategy, fused AdamW, and mem_efficient SDPA. Gradients agree across ranks (cross-rank rel diff 0.0).

Variant × q_pool sweep — GH200 B=128 in-chans=8

	q_pool=1	q_pool=2	q_pool=3
Hiera-Tiny	2.35x	2.16x	3.17x
Hiera-Small	2.31x	2.27x	3.33x
Hiera-Base	2.25x	2.31x	3.12x

No regressions across the architecture sweep. q_pool=3 configs see >3× speedup.

Full Wave-1 layout-fix matrix: MATRIX_RESULTS.md. Changelog: CHANGELOG.md.

Install

pip install hiera-optim

From source:

git clone https://github.com/avocardio/hiera-optim.git
cd hiera-optim
pip install -e .

PyTorch >= 2.5, Triton >= 2.3. Recognises FAIR Hiera in-tree (models.hiera) or PyPI (hiera-transformer).

Usage — minimal (1.3-1.9x)

import torch
from hiera_optim import optimize
from hiera import mae_hiera_base_224

model = mae_hiera_base_224(pretrained=False, in_chans=8, input_size=(224, 224))
optimize(model)
model = torch.compile(model, mode="default", dynamic=False)

x = torch.randn(128, 8, 224, 224, device="cuda", dtype=torch.bfloat16)
loss, *_ = model(x, mask_ratio=0.6)
loss.backward()

optimize(model) does two things, in place:

1. Swap every MaskUnitAttention for a 4-D Q/K/V variant so PyTorch SDPA dispatches to FlashAttention / cuDNN-attn / mem-efficient instead of math. FAIR's original feeds SDPA a 5-D tensor that the fused kernels reject (~13x per call on Ada, ~6x on Hopper).
2. Swap x[mask.tile(...)] and x_dec[mask] = ... for explicit torch.gather / scatter_. Removes the indexing_backward_kernel and the aten::nonzero graph break.

Usage — graph-safe + reduce-overhead (2.2–2.4x on GH200)

graph_safe=True rewrites forward_loss and get_pixel_label_* with mask-weighted-mean reductions (no pred[mask] boolean indexing — the data-dependent shape would otherwise crash CUDA Graphs on replay) and auto-pins a CUDA-Graph-safe SDPA backend on Hopper (cuDNN-attention isn't graph-safe in PyTorch 2.9).

import torch
from hiera_optim import optimize, MAEStepInputs
from hiera import mae_hiera_base_224

model = mae_hiera_base_224(pretrained=False, in_chans=8, input_size=(224, 224))
model = model.to("cuda", torch.bfloat16)
optimize(model, graph_safe=True)
model = torch.compile(model, mode="reduce-overhead", dynamic=False)

# CUDA Graphs need stable input tensor addresses across iterations.
inputs = MAEStepInputs(
    model._orig_mod, batch_size=128, in_chans=8,
    input_size=(224, 224), mask_ratio=0.6,
    device="cuda", dtype=torch.bfloat16,
)
opt = torch.optim.AdamW(model.parameters(), lr=3e-4, fused=True)

for batch in loader:
    inputs.x.copy_(batch.cuda(non_blocking=True))
    inputs.refresh_mask()
    out = model(inputs.x, mask_ratio=0.6,
                mask=inputs.mask, keep_idx=inputs.keep_idx)
    out[0].backward()
    opt.step(); opt.zero_grad()

MAEStepInputs handles the static-buffer pattern that CUDA Graphs require: the mask and keep_idx are sampled outside the captured region and copied into persistent tensors each step.

Production checklist for graph_safe + reduce-overhead

	required	why
PYTORCH_ALLOC_CONF=max_split_size_mb:512,expandable_segments:True	yes at large B	Without this, the captured graph pool fragments and reduce-overhead OOMs at B=1024/GPU on Hiera-Base GH200 even though peak alloc is only 67 GiB of 95 GiB.
drop_last=True on dataloader	yes	a partial last batch triggers a ~50-second recompile (whole-graph re-capture for the new shape). drop_last keeps every step at the captured B.
static_graph=True on DDPStrategy	yes	DDP needs a stable comm pattern for the captured allreduce hooks.
static batch shape across steps	yes	feed inputs.x.copy_(batch) — never re-allocate.
in-chans / image size / mask ratio constant for the run	yes	changing any retriggers compile.
Muon / manual_optimization, accumulate_grad_batches=1	ok — validated on GH200 (loss matches FAIR ref within 6.2e-3 bf16).
Muon / manual_optimization, accumulate_grad_batches > 1	broken under reduce-overhead in PyTorch 2.9 — 4 workarounds tried (mark_step_begin / +sync / +clone / +del), all hit accessing tensor output of CUDAGraphs that has been overwritten. Use accum=1 with a larger per-GPU batch instead.
gradient checkpointing under reduce-overhead	don't — enable_stage_checkpointing increases memory under CUDA Graphs (recompute path also captured into pool). Tested: stage-2 ckpt goes 67 → 77 GiB peak.
compile cache	persistent across runs if you set TRITON_CACHE_DIR and TORCHINDUCTOR_CACHE_DIR to a stable path. First step still pays the autotune cost.

Production scenarios on 4×GH200, effective batch 4096

config	per-GPU B	accum	mode	samp/s/GPU	total samp/s
user's current production	512	2	default	1675	6700
recommended	1024	1	reduce-overhead + graph_safe	2717	10869

1.62× production throughput at the same effective batch.

Optional

from hiera_optim import optimize, enable_stage_checkpointing

optimize(model, sdpa_backend="auto")             # per-block SDPA hint
optimize(model, sdpa_backend="mem_efficient")    # pin one backend everywhere
enable_stage_checkpointing(model, stages=(2,))   # OOM lever

Flexibility

Validated across the Hiera matrix (86/86 tests). Equivalence to FAIR baseline holds for:

	tested
variants	Tiny / Small / Base / Large / Huge
q_pool	1 / 2 / 3
mask ratio	0.5 / 0.6 / 0.75
dtypes	fp32 / bf16
in_chans	1 / 3 / 8
input sizes	128 / 224
1-D / 2-D / 3-D MAE	all paths
production raw_reshape_224 (HieraMAEv5)	yes

graph_safe=True adds:

• mask-weighted-mean loss (no bool indexing in label/loss)
• Hopper-safe SDPA backend auto-pin (avoid cuDNN-attention under CUDA Graphs)
• MAEStepInputs helper for the static-buffer pattern

It does not change the 0.1.0 layout / gather fixes — the new flag is opt-in. Default optimize(model) calls work exactly as before.

GPU support

Architecture	SM	Layout fix	graph_safe + reduce-overhead
Ada (RTX 4090, L40)	SM89	Tested	Tested (4.13x on Base B=8)
Hopper (H100, GH200)	SM90	Tested	Tested (2.37x on Base B=128)
Ampere (A100)	SM80	Should work	Likely works (cuDNN-attn not used on Ampere)
Blackwell (B200)	SM100	Should work	Should work

Tests

pip install -e .[test]
pytest

86 tests cover all 5 Hiera variants × q_pool {1, 2, 3} × mask ratios × bf16/fp16/fp32 × 1D/2D/3D inputs × classification + MAE × graph-safe path.

License

MIT.