Metadata-Version: 2.4
Name: hiera-optim
Version: 0.2.0
Summary: 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.
Author: Maxi Kalcher
License: MIT
Project-URL: Homepage, https://github.com/avocardio/hiera-optim
Project-URL: Repository, https://github.com/avocardio/hiera-optim
Project-URL: Issues, https://github.com/avocardio/hiera-optim/issues
Keywords: pytorch,transformer,vision,hiera,mae,flash-attention,triton,hopper,h100,gh200
Classifier: Development Status :: 4 - Beta
Classifier: License :: OSI Approved :: MIT License
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.10
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Classifier: Operating System :: POSIX :: Linux
Requires-Python: >=3.10
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: torch>=2.5.0
Requires-Dist: triton>=2.3.0
Provides-Extra: hiera
Requires-Dist: hiera-transformer>=0.1.4; extra == "hiera"
Provides-Extra: test
Requires-Dist: pytest>=7.0; extra == "test"
Provides-Extra: dev
Requires-Dist: pytest>=7.0; extra == "dev"
Requires-Dist: ruff; extra == "dev"
Requires-Dist: build; extra == "dev"
Requires-Dist: twine; extra == "dev"
Dynamic: license-file

# hiera-optim

Drop-in throughput optimisations for [FAIR's Hiera](https://github.com/facebookresearch/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.

```python
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`](MATRIX_RESULTS.md). Changelog: [`CHANGELOG.md`](CHANGELOG.md).

## Install

```bash
pip install hiera-optim
```

From source:

```bash
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)

```python
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).

```python
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

```python
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

```bash
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.
