Metadata-Version: 2.4
Name: hiera-optim
Version: 0.2.1
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: lightning
Requires-Dist: lightning>=2.0; extra == "lightning"
Provides-Extra: test
Requires-Dist: pytest>=7.0; extra == "test"
Requires-Dist: lightning>=2.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"
Requires-Dist: lightning>=2.0; 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.

### Hiera-Large — GH200 in-chans=8 bf16

| | eager | `optimize(model)` | graph_safe + reduce-overhead |
|---|---|---|---|
| Hiera-Large B=64 | 1.00x (140 ms) | 1.40x | **2.72x (52 ms)** |
| Hiera-Large B=128 | 1.00x (169 ms) | 1.47x | **1.73x (98 ms)** |

Peak memory: 9.2 GiB (B=64), 16.8 GiB (B=128). Large is compute-heavy so the launch-amortization win is smaller at high B than for Hiera-Base, but still 1.7–2.7×.

### 1-D Hiera (launch-bound) — GH200 bf16, base, fwd+bwd

A genuine 1-D MAE (`q_stride=(2,)`, 1024 samples → 64 tokens, 8 mask units) is in a
different regime than 2-D: it's **launch-bound** (tiny per-channel GEMMs, thousands of
small kernels), not attention-bound.

| | eager | `optimize(model)` (layout) | graph_safe + reduce-overhead |
|---|---|---|---|
| B=128 | 1.00x (1931 samp/s, 0.9 GB) | **0.92x** (no-op) | **8.46x (16,333 samp/s, 0.8 GB)** |
| B=512 | 1.00x (7643 samp/s, 3.5 GB) | 1.22x | **3.12x (23,829 samp/s, 1.0 GB)** |

**On 1-D the eager layout swap does nothing** (0.92x at B=128 — it can even regress);
the entire win is `graph_safe=True` + `compile(mode="reduce-overhead")`. It also cuts
memory ~3.5× at B=512 (3.5→1.0 GB), which is what lets you push much larger batches.
Loss is bit-identical to FAIR for a matched mask. **For 1-D training, graph_safe +
reduce-overhead is the recommended default** — see [`CI_SPEEDUP.md`](CI_SPEEDUP.md).

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.

### With PyTorch Lightning (one-line training_step)

For Lightning users, `MAEStaticInputsCallback` owns the buffer lifecycle so the training_step stays clean. Install with `pip install hiera-optim[lightning]` or just `pip install hiera-optim lightning`.

```python
from lightning.pytorch import Trainer
from lightning.pytorch.strategies import DDPStrategy
from hiera_optim import optimize, MAEStaticInputsCallback, feed_batch

# in LightningModule.__init__:
optimize(self.model, graph_safe=True)
self.model = torch.compile(self.model, mode="reduce-overhead", dynamic=False)

# in LightningModule.training_step:
def training_step(self, batch, batch_idx):
    inputs = self._mae_inputs                           # attached by callback
    feed_batch(inputs, batch, self._batch_key)          # 1-line buffer fill + mask refresh
    out = self.model(inputs.x, mask_ratio=self.mask_ratio,
                     mask=inputs.mask, keep_idx=inputs.keep_idx)
    loss = out[0]
    ...

# Trainer:
cb = MAEStaticInputsCallback(
    batch_size=B, in_chans=8, input_size=(224, 224),
    mask_ratio=0.6, dtype=torch.bfloat16, batch_key="eeg_raw",
)
trainer = Trainer(
    callbacks=[cb],
    strategy=DDPStrategy(static_graph=True),
    precision="bf16-mixed",
    devices=4,
)
```

The callback:
- allocates static buffers on `setup(stage="fit")` (knows the device by then)
- walks `_orig_mod` / `.module` to find the inner Hiera model
- refreshes the mask on `on_train_batch_start` (so the captured CUDA Graph sees stable input addresses every step)
- exposes the buffers as `pl_module._mae_inputs` for any callback or hook to read

### 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

The bench numbers above are **single-GPU, bare PyTorch, AdamW**. Real Lightning + Muon + DDP stacks pay extra shell overhead that the lab bench doesn't capture. Realistic per-GPU samp/s for the same effective-batch=4096 recipe:

| stack | optimizer | per-GPU B | mode | samp/s/GPU |
|---|---|---|---|---|
| bare PT, single GPU (lab) | AdamW | 1024 | reduce-overhead + graph_safe | **2839** |
| bare PT, single GPU (lab) | Muon | 1024 | reduce-overhead + graph_safe | 2221 (Muon is 78% of AdamW) |
| 4×GH200 DDP, Lightning, `sync_dist=True` per step (current prod) | Muon | 1024 | reduce-overhead + graph_safe | ~932 |
| 4×GH200 DDP, Lightning, `sync_dist=False` (one-line fix) | Muon | 1024 | reduce-overhead + graph_safe | **~1500-1800** |
| 4×GH200 DDP, Lightning, no per-step log + AdamW | AdamW | 1024 | reduce-overhead + graph_safe | **~2500** |

Lightning's `sync_dist=True` on the per-step loss log is a hidden ~50% throughput tax at this it/s, because it does an all-reduce on the loss scalar *outside* the captured CUDA Graph and serializes every step. The user-side fix is one line: change `sync_dist=True` → `sync_dist=False` (Lightning still aggregates at epoch end if `on_epoch=True`). See [`cluster/REALITY_CHECK.md`](cluster/REALITY_CHECK.md) for the full diagnosis.

If you want raw throughput in a Lightning + DDP stack, take the headline as **1.5×** over the previous production baseline, not 2.31×.

## 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.
