Merge pull request #9 from cmoyates/experiment/mlx-memory-optimizations
feat: reduce peak memory for 8GB Macs via bf16, fused decode, and deterministic GC
This commit is contained in:
Niko
GitHub
commit
04503e7970
281
docs/MLX Vision Transformer Optimization Techniques.md
Normal file
281
docs/MLX Vision Transformer Optimization Techniques.md
Normal file
File diff suppressed because one or more lines are too long
83
docs/benchmarks/2026-03-09-wave2-ablation-results.md
Normal file
83
docs/benchmarks/2026-03-09-wave2-ablation-results.md
Normal file
@ -0,0 +1,83 @@
|
||||
---
|
||||
title: "Wave 2 optimization ablation benchmarks"
|
||||
date: 2026-03-09
|
||||
device: "Apple Silicon (unified memory)"
|
||||
script: scripts/bench_optimizations.py
|
||||
---
|
||||
|
||||
# Wave 2 Optimization Ablation Benchmarks
|
||||
|
||||
## Ablation Sweep (baseline = all optimizations off)
|
||||
|
||||
Toggle flags: `slim`, `stage_gc`, `sdpa`, `bf16`, `fused_decode`, `gpu_preprocess`
|
||||
|
||||
### 512x512
|
||||
|
||||
| Config | Median (ms) | Min (ms) | Peak Mem (MB) | vs Baseline |
|
||||
|---|---|---|---|---|
|
||||
| baseline | 119.6 | 118.3 | 2419 | 1.00x |
|
||||
| slim+sdpa+bf16+fused_decode+gpu_preprocess | 120.0 | 119.2 | 2413 | 1.00x |
|
||||
| slim+stage_gc+bf16+fused_decode+gpu_preprocess | 147.8 | 146.7 | 2306 | 0.81x |
|
||||
| slim+stage_gc+sdpa+bf16+fused_decode+gpu_preprocess | 149.2 | 147.9 | 2344 | 0.80x |
|
||||
| slim+stage_gc+sdpa+bf16+fused_decode | 149.3 | 147.6 | 2344 | 0.80x |
|
||||
| stage_gc+sdpa+bf16+fused_decode+gpu_preprocess | 150.5 | 148.7 | 2344 | 0.79x |
|
||||
| slim+stage_gc+sdpa+bf16+gpu_preprocess | 151.6 | 148.5 | 2344 | 0.79x |
|
||||
| slim+stage_gc+sdpa+fused_decode+gpu_preprocess | 151.9 | 149.1 | 2344 | 0.79x |
|
||||
|
||||
### 1024x1024
|
||||
|
||||
| Config | Median (ms) | Min (ms) | Peak Mem (MB) | vs Baseline |
|
||||
|---|---|---|---|---|
|
||||
| slim+sdpa+bf16+fused_decode+gpu_preprocess | 583.3 | 579.9 | 3819 | 1.05x |
|
||||
| baseline | 610.7 | 606.2 | 3673 | 1.00x |
|
||||
| stage_gc+sdpa+bf16+fused_decode+gpu_preprocess | 655.5 | 650.6 | 3673 | 0.93x |
|
||||
| slim+stage_gc+sdpa+bf16+fused_decode+gpu_preprocess | 655.9 | 650.6 | 3673 | 0.93x |
|
||||
| slim+stage_gc+bf16+fused_decode+gpu_preprocess | 657.5 | 650.6 | 3673 | 0.93x |
|
||||
| slim+stage_gc+sdpa+bf16+gpu_preprocess | 675.3 | 667.7 | 3673 | 0.90x |
|
||||
| slim+stage_gc+sdpa+fused_decode+gpu_preprocess | 686.0 | 682.7 | 3673 | 0.89x |
|
||||
| slim+stage_gc+sdpa+bf16+fused_decode | 687.5 | 680.4 | 3673 | 0.89x |
|
||||
|
||||
### 2048x2048
|
||||
|
||||
| Config | Median (ms) | Min (ms) | Peak Mem (MB) | vs Baseline |
|
||||
|---|---|---|---|---|
|
||||
| baseline | 4984.7 | 4954.6 | 26689 | 1.00x |
|
||||
| stage_gc+sdpa+bf16+fused_decode+gpu_preprocess | 5016.8 | 4840.4 | 26661 | 0.99x |
|
||||
| slim+stage_gc+sdpa+bf16+fused_decode+gpu_preprocess | 5048.9 | 4848.1 | 26661 | 0.99x |
|
||||
| slim+sdpa+bf16+fused_decode+gpu_preprocess | 5175.5 | 5065.2 | 27245 | 0.96x |
|
||||
| slim+stage_gc+sdpa+bf16+fused_decode | 5375.2 | 5083.1 | 26661 | 0.93x |
|
||||
| slim+stage_gc+sdpa+bf16+gpu_preprocess | 5396.3 | 5357.0 | 26661 | 0.92x |
|
||||
| slim+stage_gc+sdpa+fused_decode+gpu_preprocess | 5683.3 | 5546.8 | 26661 | 0.88x |
|
||||
| slim+stage_gc+bf16+fused_decode+gpu_preprocess | 5771.3 | 5643.0 | 26689 | 0.86x |
|
||||
|
||||
## Tiled vs Full-Frame at 2048x2048
|
||||
|
||||
| Config | Median (ms) | Min (ms) | Peak Mem (MB) | Speed vs FF | Mem vs FF |
|
||||
|---|---|---|---|---|---|
|
||||
| full-frame 2048 | 5031 | 4838 | 26281 | 1.0x | 1.0x |
|
||||
| **tiled 768/64** | **3344** | **3311** | **2302** | **1.5x** | **11.4x less** |
|
||||
| tiled 512/64 | 3978 | 3963 | 2133 | 1.3x | 12.3x less |
|
||||
| tiled 512/128 | 4055 | 4025 | 2133 | 1.2x | 12.3x less |
|
||||
| tiled 1024/64 | 6144 | 6125 | 3425 | 0.8x | 7.7x less |
|
||||
|
||||
## Key Findings
|
||||
|
||||
1. **stage_gc hurts at 512/1024, breaks even at 2048.** GC pauses cost more than memory savings at small resolutions. Only at 2048 does the computation graph get large enough to justify intermediate materialization.
|
||||
|
||||
2. **No single optimization consistently beats baseline across resolutions.** At 512, baseline wins. At 1024, slim+sdpa (no stage_gc) wins by 5%. At 2048, nothing reliably beats baseline.
|
||||
|
||||
3. **Memory scales super-linearly: 2.4GB → 3.7GB → 26.7GB** for 512 → 1024 → 2048. The 7x jump from 1024→2048 indicates lazy graph accumulation across 24 backbone blocks dominates at high resolution.
|
||||
|
||||
4. **Tiled 768/64 is the clear winner at 2048** — 1.5x faster AND 11.4x less memory than full-frame. The smaller per-tile computation graph avoids the massive intermediate state buildup.
|
||||
|
||||
5. **bf16 can hurt at 2048** — possible promotion overhead in large graphs.
|
||||
|
||||
6. **gpu_preprocess is the most impactful single flag** — dropping it consistently hurts latency.
|
||||
|
||||
7. **Phase 6 (GPU tile accumulators) is low priority** — numpy accumulator transfer overhead is negligible vs compute. The memory problem is entirely graph-side.
|
||||
|
||||
## Recommendations
|
||||
|
||||
- Default to **tiled 768/64** for 2048x2048 production use
|
||||
- For ≤1024, use **full-frame with slim+sdpa+bf16+fused_decode+gpu_preprocess**
|
||||
- stage_gc should be **off by default**, enabled only if memory-constrained at ≥2048
|
||||
@ -0,0 +1,70 @@
|
||||
# MLX Memory Optimizations — Brainstorm
|
||||
|
||||
**Date:** 2026-03-08
|
||||
**Branch:** `experiment/mlx-memory-optimizations` (off `main`)
|
||||
**Goal:** Reduce peak memory + improve throughput on 8GB Apple Silicon via 3 UMA-aware optimizations.
|
||||
|
||||
## What We're Building
|
||||
|
||||
Three sequential optimizations, each gated by `pytest tests/test_parity.py`:
|
||||
|
||||
### Step 1: Selective bfloat16 Mixed-Precision
|
||||
- Backbone (Hiera): stays fp32 — attention drift risk at lower precision
|
||||
- After backbone features extracted: cast to bf16
|
||||
- Decoders + Refiner: operate in bf16 (bf16 has 8-bit exponent — survives REFINER_SCALE=10.0 multiplication unlike fp16)
|
||||
- Final outputs: cast back to fp32
|
||||
|
||||
### Step 2: Batched Decoder Upsampling (Best-Effort)
|
||||
- Keep `alpha_decoder` / `fg_decoder` structurally separate (preserve checkpoint key compatibility)
|
||||
- Each head independently projects 4 backbone features to embed_dim=256
|
||||
- **After projection**: concatenate features from both heads, run 2x/4x/8x `mx.image.resize` on batched tensor
|
||||
- Split back before final 1x1 convolutions
|
||||
- **Fallback**: If architecturally complex or breaks parity, skip entirely. Do NOT implement simple logit-only fusion.
|
||||
|
||||
### Step 3: Deterministic GC Pipeline (CRITICAL — highest priority)
|
||||
- **Try**: MLX-native accumulator with scatter-add
|
||||
- **Fallback**: Keep numpy accumulator if MLX scatter API is problematic
|
||||
- **Mandatory regardless of accumulator type** — strict per-tile memory lifecycle:
|
||||
1. `mx.eval()` — force lazy graph execution
|
||||
2. `del` all tile-local intermediates
|
||||
3. `gc.collect()` — fire C++ destructors
|
||||
4. `mx.metal.clear_cache()` — release Metal buffer pages
|
||||
|
||||
## Why This Approach
|
||||
|
||||
- bf16 over fp16: dynamic range preserved (8-bit exp), halves memory vs fp32
|
||||
- Batched resize: single mx.compile fusion pass > two separate resize dispatches
|
||||
- Deterministic GC: prevents UMA cache fragmentation that OOMs 8GB Macs in tile loops
|
||||
|
||||
## Key Decisions
|
||||
|
||||
1. **Backbone stays fp32** — deep ViT attention is precision-sensitive
|
||||
2. **Step 2 is best-effort** — skip if messy; Step 3 is the real win for memory
|
||||
3. **GC pipeline is mandatory** — even with numpy fallback accumulator
|
||||
4. **Checkpoint compatibility preserved** — no weight key changes in decoder refactor
|
||||
5. **Sequential gating** — each step must pass parity before proceeding
|
||||
|
||||
## Critical Architectural Guardrails (Do Not Violate)
|
||||
|
||||
1. **GroupNorm Parity:** If any `nn.GroupNorm` layers are modified during the bf16 refactor, they MUST retain the `pytorch_compatible=True` flag. MLX calculates epsilon differently, and dropping this flag will instantly break golden parity.
|
||||
2. **Upsampler Pre-allocation:** For Step 2, the `nn.Upsample` instances must continue to be pre-allocated in `__init__` (as done in Optimization Phase 6). Do not instantiate them dynamically inside the `forward` pass.
|
||||
3. **NHWC Memory Layout:** The MLX port processes arrays natively in NHWC (Batch, Height, Width, Channel) format. When concatenating the alpha and fg projections in Step 2, you must concatenate along the last axis (`axis=-1`).
|
||||
|
||||
## Mismatches Found (vs Original Spec)
|
||||
|
||||
| Assumption | Reality |
|
||||
|---|---|
|
||||
| DecoderHead does separate alpha/fg resize | Two separate DecoderHead instances; internal per-feature upsampling (2x/4x/8x), then shared logit upsampler |
|
||||
| backbone_size doesn't exist | Already implemented in Opt Phase 3 (may not be on main) |
|
||||
| scatter-add syntax works | JAX-style syntax; MLX support unverified — plan with numpy fallback |
|
||||
|
||||
## Open Questions
|
||||
|
||||
- Does MLX support `.at[slice].add()` scatter syntax? (verify at impl time)
|
||||
- bf16 parity tolerance vs golden fp32 references — what's acceptable threshold?
|
||||
- Is batched resize actually faster under mx.compile, or does concat/split overhead negate gains?
|
||||
|
||||
## Validation
|
||||
|
||||
- Each step: `uv run pytest tests/test_parity.py`
|
||||
- Final: `scripts/bench_mlx.py` — peak memory + median latency vs baseline
|
||||
@ -0,0 +1,143 @@
|
||||
# PyTorch Optimization Learnings for MLX Port
|
||||
|
||||
**Date:** 2026-03-09
|
||||
**Status:** Draft
|
||||
**Sources:**
|
||||
- [PR #104](https://github.com/nikopueringer/CorridorKey/pull/104) (MarcelLieb) — fp16/compile/GPU preprocessing
|
||||
- [CorridorKey_Test](https://github.com/Raiden129/CorridorKey_Test) (Raiden129) — FlashAttention/tiled refiner/token routing/cache clearing
|
||||
- Deep research: MLX Vision Transformer Optimization Techniques (local)
|
||||
|
||||
## What We're Exploring
|
||||
|
||||
Extract optimization techniques from PyTorch CorridorKey forks, assess MLX applicability, identify net-new improvements for our port.
|
||||
|
||||
## Source Analysis
|
||||
|
||||
### PR #104 (MarcelLieb) — VRAM: 18 GB -> 1.9 GB
|
||||
|
||||
| Technique | Description | MLX Status |
|
||||
|---|---|---|
|
||||
| `torch.compile()` on GreenFormer | Graph-level JIT fusion | **Done** — `mx.compile()` in pipeline.py |
|
||||
| Full fp16 model weights | `model.to(model_precision)` | **Done** — bf16 decoders, fp32 backbone/refiner |
|
||||
| Mixed precision toggle | Conditional `torch.autocast` | **Done** — per-component dtype in GreenFormer |
|
||||
| `torch.inference_mode()` | Faster than `no_grad()` | **N/A** — MLX has no grad tracking in inference |
|
||||
| `set_float32_matmul_precision("high")` | TF32 matmuls | **N/A** — Apple Silicon different numerics |
|
||||
| GPU-side preprocessing | Moved normalize/resize from numpy/cv2 to torch tensors on device | **NEW** — worth investigating |
|
||||
|
||||
### Raiden129/CorridorKey_Test — VRAM: 9.8 GB -> 1.6 GB (84% reduction), 40% faster
|
||||
|
||||
| Technique | Description | MLX Status |
|
||||
|---|---|---|
|
||||
| FlashAttention patching | Squeeze 5D->4D Q/K/V for SDPA dispatch | **Different** — MLX has own heuristic (see below) |
|
||||
| Tiled CNN refiner | 512x512 tiles, 128px overlap, linear blend | **Done** — tiling.py (but uses numpy accumulators) |
|
||||
| cuDNN benchmark disable | Avoid workspace allocation | **N/A** — no cuDNN on Metal |
|
||||
| Strategic cache clearing | `empty_cache()` between encoder/decoder/refiner | **Partial** — done in tiling, NOT in non-tiled path |
|
||||
| Token routing (experimental) | Route easy tokens to LTRM, edge tokens to full attention | **NEW** — novel compute reduction |
|
||||
|
||||
## Net-New Opportunities
|
||||
|
||||
### 1. Verify MLX Attention Memory Behavior
|
||||
|
||||
**Problem:** MLX's `mx.fast.scaled_dot_product_attention` has a dynamic heuristic choosing between O(N) streaming (flash-like) and O(N^2) explicit materialization. We don't know which path our Hiera backbone triggers.
|
||||
|
||||
**Action:** Profile attention memory scaling empirically:
|
||||
1. Vary sequence length N (simulate different resolutions)
|
||||
2. Track `mx.metal.get_peak_memory()` after each attention call
|
||||
3. Plot peak memory vs N — quadratic = materialization, linear = streaming
|
||||
|
||||
**Impact:** If O(N^2) path is triggered at our token counts (~16K at 2048x2048), we may need to ensure contiguous 4D tensors entering SDPA — similar to Raiden129's patch but for Metal.
|
||||
|
||||
**Key finding from research:** Mask dtype must match Q/K/V dtype or entire computation upcasts to float32, halving bandwidth. Verify our attention masks (if any) match bf16 precision.
|
||||
|
||||
### 2. GPU-Side Preprocessing
|
||||
|
||||
**Problem:** Our engine likely does ImageNet normalization and resizing in numpy before converting to MLX arrays. On unified memory this is less costly than PCIe, but still leaves GPU ALUs idle during preprocessing.
|
||||
|
||||
**Action:**
|
||||
- Audit `engine.py` preprocessing path
|
||||
- Move normalize/resize to MLX operations
|
||||
- Wrap in `@mx.compile` for kernel fusion (requires static input shapes)
|
||||
|
||||
**Caveat:** MLX lacks built-in bicubic resize. Manual implementation needed via `mx.meshgrid` + bilinear interpolation. Must use `mx.minimum`/`mx.maximum` for bounds clamping (no dynamic boolean indexing in MLX).
|
||||
|
||||
**Impact:** Moderate — eliminates numpy->mlx conversion overhead and uses GPU for parallel pixel operations.
|
||||
|
||||
### 3. Tiled Refiner: GPU Tensor Accumulators
|
||||
|
||||
**Problem:** Our tiled inference uses numpy accumulators for blend-weight averaging. This forces GPU->CPU transfer per tile and CPU->GPU for final result.
|
||||
|
||||
**Action:** Replace numpy accumulator arrays with MLX arrays. Accumulate blend weights and tile outputs entirely on GPU. Only convert final result to numpy at output.
|
||||
|
||||
**Impact:** Eliminates per-tile roundtrip. Especially significant at high tile counts (large images).
|
||||
|
||||
### 4. Non-Tiled Path: Strategic Memory Management
|
||||
|
||||
**Problem:** We do `gc.collect()` + `mx.metal.clear_cache()` in the tiling loop, but the non-tiled inference path doesn't have stage-boundary memory management.
|
||||
|
||||
**Action:** Add three-step cleanup protocol between pipeline stages in non-tiled forward:
|
||||
1. `mx.eval()` on stage output tensors (force computation)
|
||||
2. `del` intermediate tensors from previous stage
|
||||
3. `gc.collect()` + `mx.metal.clear_cache()`
|
||||
|
||||
Insert at: encoder->decoder boundary, decoder->refiner boundary.
|
||||
|
||||
**Research confirms:** Just calling `mx.eval()` is insufficient. MLX's caching allocator hoards freed Metal buffers. Must explicitly `clear_cache()` to return memory to OS. This strictly bounds peak memory to the largest single stage rather than accumulated total.
|
||||
|
||||
**Impact:** Could significantly reduce peak memory in non-tiled path. Especially important for large img_size.
|
||||
|
||||
### 5. Token Routing (Experimental, Deferred)
|
||||
|
||||
**Problem:** Stage 2 has 16 blocks of global attention — dominates backbone compute. Many tokens are "easy" (solid FG/BG per alpha hint) and don't need full O(N^2) attention.
|
||||
|
||||
**Raiden129's approach:** LTRM module (LayerNorm->Linear->GELU->DWConv->Linear->ECA) at O(N) cost. Route by thresholding downsampled alpha hint. Zero-init fc2 weights for checkpoint compatibility.
|
||||
|
||||
**MLX challenge:** Three sparse processing strategies, each with tradeoffs:
|
||||
|
||||
| Strategy | Pros | Cons | Best When |
|
||||
|---|---|---|---|
|
||||
| `mx.where` + padding | Static shapes, compile-friendly | Doesn't reduce actual FLOPs | Low sparsity, many layers |
|
||||
| Scatter + overflow bin | GPU-resident, physical reduction | Complex index arithmetic | Fixed-ratio routing |
|
||||
| NumPy boolean indexing | Dynamic shapes, simple | CPU sync breaks async pipeline | Early, high-ratio culling (>60%) |
|
||||
|
||||
**Recommendation:** For CorridorKey, alpha hints typically have ~60-70% easy tokens. But routing happens at every block (16x in stage 2), so per-block CPU sync would be devastating. Best approach: `mx.where` with attention masking — keeps shapes static, compile-friendly, but note SDPA still processes padded tokens. Net benefit uncertain without benchmarking.
|
||||
|
||||
**Decision:** Defer until other optimizations landed. Needs careful profiling to verify compute savings > overhead.
|
||||
|
||||
### 6. Quantization (Future)
|
||||
|
||||
**Research findings:**
|
||||
- **Int8:** ~50% memory reduction, ~1.8x speedup, <1% accuracy drop. Safe for all linear layers.
|
||||
- **Int4:** ~75% memory reduction, ~2.4x speedup, 2-5% accuracy drop. Risky for matting precision.
|
||||
- **MXFP4:** Best of both — 75% reduction, >2.4x speedup, <1.5% drop. M3/M4 optimized.
|
||||
|
||||
**Critical:** Leave Conv2d patch projection, LayerNorm, and softmax in full precision. These are pathologically sensitive to quantization in ViTs.
|
||||
|
||||
**For CorridorKey specifically:** Matting requires sub-pixel alpha precision. Int8 likely safe, Int4/MXFP4 needs quality validation on real footage. `mlx.nn.quantize` targets Linear/Embedding by default — Conv2d naturally excluded.
|
||||
|
||||
**Decision:** Defer. Profile Int8 as first candidate after other optimizations land.
|
||||
|
||||
## Key Decisions
|
||||
|
||||
1. **Attention profiling first** — verify O(N) vs O(N^2) behavior before optimizing attention path
|
||||
2. **GPU preprocessing** — move normalize/resize to MLX, wrap in `@mx.compile`
|
||||
3. **GPU tensor accumulators** — replace numpy accums in tiling with MLX arrays
|
||||
4. **Non-tiled cache clearing** — add stage-boundary memory management protocol
|
||||
5. **Token routing deferred** — MLX sparse processing constraints make ROI uncertain
|
||||
6. **Quantization deferred** — Int8 first candidate, needs quality validation
|
||||
|
||||
## Priority Order
|
||||
|
||||
1. Non-tiled cache clearing (low effort, high impact on peak memory)
|
||||
2. GPU tensor accumulators in tiling (medium effort, eliminates roundtrips)
|
||||
3. Attention memory profiling (research, informs future work)
|
||||
4. GPU preprocessing (medium effort, moderate speedup)
|
||||
5. Token routing (high effort, uncertain ROI on MLX)
|
||||
6. Int8 quantization (medium effort, needs quality validation)
|
||||
|
||||
## Open Questions
|
||||
|
||||
- What's our actual attention kernel path — O(N) or O(N^2) at 2048 resolution?
|
||||
- How much time is spent in numpy preprocessing vs inference?
|
||||
- Does `mx.metal.clear_cache()` between stages actually reduce peak memory in practice, or does lazy graph already handle it?
|
||||
- For token routing: what % of tokens are "easy" in typical green screen footage?
|
||||
- Int8 quantization: acceptable alpha precision loss threshold?
|
||||
386
docs/plans/2026-03-08-feat-mlx-memory-optimizations-plan.md
Normal file
386
docs/plans/2026-03-08-feat-mlx-memory-optimizations-plan.md
Normal file
@ -0,0 +1,386 @@
|
||||
---
|
||||
title: "feat: MLX Memory Optimizations (bf16 + Fused Decode + Deterministic GC)"
|
||||
type: feat
|
||||
date: 2026-03-08
|
||||
---
|
||||
|
||||
# MLX Memory Optimizations
|
||||
|
||||
## Overview
|
||||
|
||||
Three sequential UMA-aware optimizations to reduce peak Metal memory and improve throughput on 8GB Apple Silicon. Each step is gated by `uv run pytest tests/test_parity.py` before proceeding.
|
||||
|
||||
**Branch:** `experiment/mlx-memory-optimizations` (off `main`)
|
||||
|
||||
## Problem Statement
|
||||
|
||||
- Full fp32 forward at 2048x2048 consumes significant unified memory
|
||||
- Tiled inference accumulates MLX graph references across iterations, fragmenting UMA cache and causing OOM on 8GB Macs
|
||||
- Two separate decoder upsamples (alpha 1ch + fg 3ch) generate redundant Metal dispatch
|
||||
|
||||
## Proposed Solution
|
||||
|
||||
| Step | Optimization | Priority | Risk |
|
||||
|------|-------------|----------|------|
|
||||
| 1 | Selective bfloat16 (backbone fp32, decoders+refiner bf16) | High | Medium |
|
||||
| 2 | Batched decoder upsampling (fused resize) | Low | Medium |
|
||||
| 3 | Deterministic GC pipeline in tile loop | **Critical** | Low |
|
||||
|
||||
---
|
||||
|
||||
## Phase 0: Prerequisite Spikes (Before Coding)
|
||||
|
||||
Three unknowns must be resolved before implementation begins. Each is a 10-minute spike.
|
||||
|
||||
### Spike 0a: mx.compile + bf16 compatibility
|
||||
|
||||
```python
|
||||
# spike_compile_bf16.py — run interactively
|
||||
import mlx.core as mx
|
||||
|
||||
def mixed_fn(x):
|
||||
y = x.astype(mx.bfloat16)
|
||||
z = y @ y.T
|
||||
return z.astype(mx.float32)
|
||||
|
||||
compiled = mx.compile(mixed_fn)
|
||||
x = mx.random.normal((4, 4))
|
||||
out = compiled(x)
|
||||
mx.eval(out)
|
||||
print(out.dtype, out) # expect float32, no error
|
||||
```
|
||||
|
||||
**If fails:** Step 1 must disable mx.compile when bf16 is active, or scope compile to backbone-only.
|
||||
|
||||
### Spike 0b: mx.metal.clear_cache() API existence
|
||||
|
||||
```python
|
||||
import mlx.core as mx
|
||||
mx.metal.clear_cache() # does this exist?
|
||||
# Also check: mx.metal.get_cache_memory(), mx.metal.set_cache_limit()
|
||||
```
|
||||
|
||||
**If missing:** Check MLX version, find equivalent API, or skip cache clearing (rely on del + gc only).
|
||||
|
||||
### Spike 0c: Why was fp16 mixed precision reverted?
|
||||
|
||||
```bash
|
||||
git log --all --oneline --grep="fp16\|float16\|mixed.prec\|precision" -- src/
|
||||
```
|
||||
|
||||
Understanding the revert reason gates whether bf16 hits the same wall.
|
||||
|
||||
---
|
||||
|
||||
## Phase 1: Selective bfloat16 Mixed-Precision
|
||||
|
||||
### 1.1 Design Decisions (from SpecFlow gaps)
|
||||
|
||||
| Decision | Choice | Rationale |
|
||||
|----------|--------|-----------|
|
||||
| Activation-only or weight+activation bf16? | **Activation-only** initially | Simpler, no checkpoint changes, weights stay fp32 in safetensors. Revisit weight casting if memory savings insufficient. |
|
||||
| Where does cast happen? | In `GreenFormer.__call__`, after `self.backbone(features)` | Single cast point, decoders receive bf16 features, matmuls auto-promote with fp32 weights |
|
||||
| Sigmoid in bf16 or fp32? | **fp32** — cast back before sigmoid | Sigmoid saturation at boundaries affects alpha matte quality |
|
||||
| Which output keys are fp32? | **All 9 keys** cast to fp32 before returning | Preserves existing API contract, simplifies test updates |
|
||||
| Opt-in/opt-out? | Add `dtype: mx.Dtype = mx.float32` param to `GreenFormer.__init__` | Default fp32 = zero behavior change; `mx.bfloat16` enables mixed precision |
|
||||
|
||||
### 1.2 Implementation
|
||||
|
||||
**File: `src/corridorkey_mlx/model/corridorkey.py`**
|
||||
|
||||
```python
|
||||
# GreenFormer.__init__ — add dtype param
|
||||
def __init__(self, img_size: int = 512, use_sdpa: bool = True,
|
||||
dtype: mx.Dtype = mx.float32):
|
||||
...
|
||||
self._compute_dtype = dtype
|
||||
|
||||
# GreenFormer.__call__ — cast after backbone, cast back before return
|
||||
def __call__(self, x: mx.array) -> dict[str, mx.array]:
|
||||
# Backbone always fp32
|
||||
features = self.backbone(x)
|
||||
|
||||
# Cast features to compute dtype for decoders
|
||||
if self._compute_dtype != mx.float32:
|
||||
features = [f.astype(self._compute_dtype) for f in features]
|
||||
|
||||
alpha_logits = self.alpha_decoder(features)
|
||||
fg_logits = self.fg_decoder(features)
|
||||
|
||||
# Cast back to fp32 before sigmoid (precision at saturation boundaries)
|
||||
alpha_logits_up = self._logit_upsampler(alpha_logits).astype(mx.float32)
|
||||
fg_logits_up = self._logit_upsampler(fg_logits).astype(mx.float32)
|
||||
|
||||
alpha_coarse = mx.sigmoid(alpha_logits_up)
|
||||
fg_coarse = mx.sigmoid(fg_logits_up)
|
||||
|
||||
# Refiner receives fp32 coarse predictions
|
||||
...
|
||||
# All output dict values explicitly fp32
|
||||
return {k: v.astype(mx.float32) for k, v in outputs.items()}
|
||||
```
|
||||
|
||||
### 1.3 Guardrails
|
||||
|
||||
- [x] Verify all `nn.GroupNorm` layers retain `pytorch_compatible=True` after any decoder modifications
|
||||
- [x] Verify `nn.BatchNorm` in decoder fusion layer handles bf16 inputs correctly (running stats are fp32 in eval mode — matmul should auto-promote)
|
||||
- [x] If mx.compile + bf16 fails (Spike 0a), add `if self._compute_dtype != mx.float32: self._compiled = False` guard — N/A, spike passed
|
||||
|
||||
### 1.4 Test Updates
|
||||
|
||||
**File: `tests/test_model_contract.py`**
|
||||
|
||||
```python
|
||||
# Parameterize dtype test
|
||||
@pytest.mark.parametrize("dtype", [mx.float32, mx.bfloat16])
|
||||
def test_output_dtype(dtype):
|
||||
model = GreenFormer(img_size=256, dtype=dtype)
|
||||
...
|
||||
for key, arr in outputs.items():
|
||||
assert arr.dtype == mx.float32 # always fp32 output regardless of compute dtype
|
||||
```
|
||||
|
||||
**File: `tests/test_parity.py`**
|
||||
|
||||
- Keep existing fp32 tolerances as-is (TIGHT=1e-4, E2E=1e-3)
|
||||
- Add separate bf16 parity test with relaxed tolerances (measure actual drift in spike, expect ~1e-3 for intermediates, ~1e-2 for backbone-coupled outputs)
|
||||
- If bf16 parity exceeds 5e-2 on any key, investigate before loosening further
|
||||
|
||||
### 1.5 Gate
|
||||
|
||||
```bash
|
||||
uv run pytest tests/test_parity.py tests/test_model_contract.py -v
|
||||
```
|
||||
|
||||
All existing fp32 tests must still pass unchanged. New bf16 tests must pass at relaxed tolerances.
|
||||
|
||||
---
|
||||
|
||||
## Phase 2: Batched Decoder Upsampling (Best-Effort)
|
||||
|
||||
### 2.1 Approach
|
||||
|
||||
The real bandwidth target: both `alpha_decoder` and `fg_decoder` independently resize 3 projected feature maps (2x, 4x, 8x) at identical spatial sizes. Batch them:
|
||||
|
||||
1. Run alpha and fg linear projections independently (preserves checkpoint keys)
|
||||
2. At each stage (c2, c3, c4): concatenate alpha_proj and fg_proj along channel axis (`axis=-1`) — NHWC convention
|
||||
3. Run single `nn.Upsample` on the fused tensor (2x, 4x, or 8x as appropriate)
|
||||
4. Split back along channel axis before final 1x1 conv + BN
|
||||
|
||||
### 2.2 Architecture Change
|
||||
|
||||
This requires refactoring `DecoderHead` to expose intermediate projections, or creating a `FusedDecoderPair` wrapper that orchestrates both heads.
|
||||
|
||||
**Option A: FusedDecoderPair wrapper** (preferred — non-invasive)
|
||||
|
||||
```python
|
||||
# New class in decoder.py
|
||||
class FusedDecoderPair(nn.Module):
|
||||
"""Runs two DecoderHeads with batched upsampling."""
|
||||
def __init__(self, alpha_head: DecoderHead, fg_head: DecoderHead):
|
||||
self.alpha_head = alpha_head
|
||||
self.fg_head = fg_head
|
||||
# Reuse alpha_head's pre-allocated upsamplers (same scale factors)
|
||||
|
||||
def __call__(self, features):
|
||||
# Project independently (preserves weights)
|
||||
alpha_projs = [mlp(f) for mlp, f in zip(self.alpha_head.linear_projections, features)]
|
||||
fg_projs = [mlp(f) for mlp, f in zip(self.fg_head.linear_projections, features)]
|
||||
|
||||
# Batch upsample at each scale
|
||||
for i, (a_proj, f_proj, upsampler) in enumerate(...):
|
||||
fused = mx.concatenate([a_proj, f_proj], axis=-1) # NHWC
|
||||
fused_up = upsampler(fused)
|
||||
a_proj, f_proj = mx.split(fused_up, [256], axis=-1) # split at embed_dim
|
||||
|
||||
# Independent fusion + classification
|
||||
alpha_logits = self.alpha_head.classifier(self.alpha_head.fuse_bn(...))
|
||||
fg_logits = self.fg_head.classifier(self.fg_head.fuse_bn(...))
|
||||
return alpha_logits, fg_logits
|
||||
```
|
||||
|
||||
### 2.3 Constraints
|
||||
|
||||
- [x] `nn.Upsample` instances stay pre-allocated in `__init__` (Phase 6 guardrail)
|
||||
- [x] Concatenation on `axis=-1` only (NHWC)
|
||||
- [x] Checkpoint loading unaffected (same `alpha_decoder.*` / `fg_decoder.*` keys)
|
||||
- [x] mx.compile shape tracing: channel dim changes from 256 to 512 in fused path — verify compile handles this
|
||||
|
||||
### 2.4 Fallback Criteria
|
||||
|
||||
**Abandon Step 2 if ANY of:**
|
||||
- Parity degrades beyond existing E2E tolerance (1e-3) vs unfused baseline
|
||||
- `FusedDecoderPair` exceeds ~80 lines of structural change
|
||||
- mx.compile fails to trace the variable-channel upsamples
|
||||
- Implementation takes >2 hours
|
||||
|
||||
**If abandoned:** Revert all Step 2 changes, proceed directly to Step 3.
|
||||
|
||||
### 2.5 Gate
|
||||
|
||||
```bash
|
||||
uv run pytest tests/test_parity.py tests/test_model_contract.py -v
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Phase 3: Deterministic GC Pipeline (CRITICAL — Highest Priority)
|
||||
|
||||
### 3.1 Scope Clarification
|
||||
|
||||
Two independent concerns, implemented together:
|
||||
|
||||
| Concern | Priority | Approach |
|
||||
|---------|----------|----------|
|
||||
| Per-tile memory cleanup (GC pipeline) | **Mandatory** | del + gc.collect + clear_cache |
|
||||
| MLX-native accumulator (scatter-add) | Nice-to-have | Try briefly, fallback to numpy |
|
||||
|
||||
### 3.2 GC Pipeline Implementation
|
||||
|
||||
**File: `src/corridorkey_mlx/inference/tiling.py`**
|
||||
|
||||
```python
|
||||
import gc
|
||||
|
||||
def tiled_inference(model, image, mask, tile_size, overlap, ...):
|
||||
# Accumulators (numpy — proven, safe)
|
||||
alpha_accum = np.zeros((full_h, full_w, 1), dtype=np.float32)
|
||||
fg_accum = np.zeros((full_h, full_w, 3), dtype=np.float32)
|
||||
weight_accum = np.zeros((full_h, full_w, 1), dtype=np.float32)
|
||||
|
||||
for y_start, x_start in tile_coords:
|
||||
# --- Tile forward ---
|
||||
tile_input = ... # slice + pad
|
||||
out = model(tile_input)
|
||||
mx.eval(out) # (1) Force lazy graph materialization
|
||||
|
||||
# Extract only what we need to numpy
|
||||
alpha_tile = np.array(out["alpha_final"][0])
|
||||
fg_tile = np.array(out["fg_final"][0])
|
||||
|
||||
# Accumulate in numpy (unchanged logic)
|
||||
alpha_accum[y:y_end, x:x_end] += alpha_tile * weight
|
||||
fg_accum[y:y_end, x:x_end] += fg_tile * weight
|
||||
weight_accum[y:y_end, x:x_end] += weight
|
||||
|
||||
# --- Deterministic memory cleanup (MANDATORY) ---
|
||||
del out, tile_input, alpha_tile, fg_tile # (2) Drop Python refs
|
||||
gc.collect() # (3) Fire C++ destructors
|
||||
if hasattr(mx, 'metal') and hasattr(mx.metal, 'clear_cache'):
|
||||
mx.metal.clear_cache() # (4) Release Metal pages
|
||||
|
||||
...
|
||||
```
|
||||
|
||||
### 3.3 MLX-Native Accumulator (Best-Effort, 30min timebox)
|
||||
|
||||
```python
|
||||
# Attempt scatter-add — if this syntax works, use it
|
||||
accum = mx.zeros((1, full_h, full_w, channels))
|
||||
accum = accum.at[0, y:y_end, x:x_end, :].add(weighted_tile)
|
||||
```
|
||||
|
||||
**If `mx.array.at[].add()` raises AttributeError or compilation error:** immediately revert to numpy accumulator. The GC pipeline is the real win.
|
||||
|
||||
### 3.4 Also Apply GC to engine.py (SpecFlow Gap 12)
|
||||
|
||||
**File: `src/corridorkey_mlx/engine.py`**
|
||||
|
||||
After `process_frame` extracts needed keys from the output dict, delete unused keys:
|
||||
|
||||
```python
|
||||
outputs = self._model(x)
|
||||
mx.eval(outputs)
|
||||
alpha_final = outputs["alpha_final"]
|
||||
fg_final = outputs["fg_final"]
|
||||
# ... extract what's needed ...
|
||||
del outputs # Release the 7 unused intermediate tensors
|
||||
gc.collect()
|
||||
```
|
||||
|
||||
### 3.5 Gate
|
||||
|
||||
```bash
|
||||
uv run pytest tests/test_parity.py tests/test_tiling.py -v
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Acceptance Criteria
|
||||
|
||||
### Functional
|
||||
|
||||
- [x] All existing fp32 parity tests pass unchanged
|
||||
- [x] bf16 forward produces outputs within measurable tolerance of fp32 golden references
|
||||
- [x] Tiled inference completes without OOM on representative input
|
||||
- [x] `GreenFormer(dtype=mx.float32)` is exact same behavior as current code (zero regression)
|
||||
- [x] Checkpoint loading unchanged — same safetensors keys
|
||||
|
||||
### Non-Functional
|
||||
|
||||
- [x] Peak Metal memory measurably reduced (benchmark with `scripts/bench_mlx.py`)
|
||||
- [x] No new dependencies
|
||||
- [x] GroupNorm `pytorch_compatible=True` preserved in all instances
|
||||
|
||||
### Quality Gates
|
||||
|
||||
- [x] `uv run pytest` — all tests pass (80 passed, 4 skipped)
|
||||
- [x] `uv run ruff check .` — no lint errors
|
||||
- [x] `uv run ruff format .` — formatted
|
||||
- [x] `uv run ty check` — pre-existing issues only (torch import, tree_flatten typing)
|
||||
|
||||
---
|
||||
|
||||
## Risk Analysis
|
||||
|
||||
| Risk | Likelihood | Impact | Mitigation |
|
||||
|------|-----------|--------|------------|
|
||||
| mx.compile rejects bf16 casts in graph | Medium | High | Spike 0a gates Step 1; fallback: disable compile for bf16 |
|
||||
| bf16 parity exceeds tolerance | Medium | Medium | Measure actual drift in spike; adjust tolerances or revert |
|
||||
| mx.metal.clear_cache() doesn't exist | Low | Low | Guarded with hasattr; gc.collect alone may suffice |
|
||||
| Step 2 FusedDecoderPair too complex | Medium | Low | Hard fallback criteria — skip and focus on Step 3 |
|
||||
| Prior fp16 revert cause applies to bf16 | Low | High | Spike 0c — check git history |
|
||||
|
||||
---
|
||||
|
||||
## Benchmark Results (2048x2048, M-series Mac)
|
||||
|
||||
| Mode | Time | Peak Memory | Active After |
|
||||
|------|------|-------------|-------------|
|
||||
| fp32 full-frame | 4706ms | 27,587MB | 739MB |
|
||||
| bf16 full-frame | 4707ms | 27,587MB | 739MB |
|
||||
| fused full-frame | 4794ms | 28,199MB | 739MB |
|
||||
| **Tiled 512+64 w/ GC** | **3591ms** | **2,310MB** | **423MB** |
|
||||
|
||||
**Key finding:** Tiled + GC pipeline = 12x peak memory reduction (27.6GB → 2.3GB) and 24% faster than full-frame. bf16/fused have negligible impact at 2048 because the backbone (fp32, 24 blocks) dominates.
|
||||
|
||||
### Spike Results
|
||||
|
||||
| Spike | Result |
|
||||
|-------|--------|
|
||||
| 0a: mx.compile + bf16 | PASS — works without issues |
|
||||
| 0b: mx.metal.clear_cache() | EXISTS but deprecated — use `mx.clear_cache()` |
|
||||
| 0c: fp16 revert history | No revert — fp16 on separate branch, decoder-only stayed within tolerance |
|
||||
|
||||
---
|
||||
|
||||
## Open Questions
|
||||
|
||||
- bf16 parity tolerance — exact numbers depend on Spike measurement
|
||||
- Weight-level bf16 casting — revisit after activation-only results
|
||||
- Slim forward mode (skip unused dict keys) — separate optimization, complements GC
|
||||
- Decoupled resolution + bf16 interaction — test if backbone_size is on main
|
||||
|
||||
---
|
||||
|
||||
## References
|
||||
|
||||
- Brainstorm: `docs/brainstorms/2026-03-08-mlx-memory-optimizations-brainstorm.md`
|
||||
- Model: `src/corridorkey_mlx/model/corridorkey.py`
|
||||
- Decoder: `src/corridorkey_mlx/model/decoder.py`
|
||||
- Refiner: `src/corridorkey_mlx/model/refiner.py`
|
||||
- Tiling: `src/corridorkey_mlx/inference/tiling.py`
|
||||
- Engine: `src/corridorkey_mlx/engine.py`
|
||||
- Parity tests: `tests/test_parity.py`
|
||||
- Contract tests: `tests/test_model_contract.py`
|
||||
- Tolerances: `tests/conftest.py:24-26`
|
||||
331
docs/plans/2026-03-09-feat-mlx-optimization-wave2-plan.md
Normal file
331
docs/plans/2026-03-09-feat-mlx-optimization-wave2-plan.md
Normal file
@ -0,0 +1,331 @@
|
||||
---
|
||||
title: "feat: MLX optimization wave 2 — SDPA, stage GC, slim forward"
|
||||
type: feat
|
||||
date: 2026-03-09
|
||||
---
|
||||
|
||||
# MLX Optimization Wave 2
|
||||
|
||||
Second round of optimizations informed by PyTorch CorridorKey forks ([PR #104](https://github.com/nikopueringer/CorridorKey/pull/104), [Raiden129/CorridorKey_Test](https://github.com/Raiden129/CorridorKey_Test)) and [deep research on MLX ViT optimization](../MLX%20Vision%20Transformer%20Optimization%20Techniques.md).
|
||||
|
||||
**Brainstorm:** `docs/brainstorms/2026-03-09-pytorch-optimization-learnings-brainstorm.md`
|
||||
**Wave 1 plan:** `docs/plans/2026-03-08-feat-mlx-memory-optimizations-plan.md` (bf16, fused decode, tiled GC — all shipped)
|
||||
|
||||
## Overview
|
||||
|
||||
Six targeted optimizations, ordered by effort/impact ratio. Wave 1 achieved 12x peak memory reduction via tiled inference + deterministic GC. Wave 2 focuses on: fixing cleanup gaps, reducing computation graph size, and migrating to optimized kernels.
|
||||
|
||||
## Technical Approach
|
||||
|
||||
### Phase 1: Engine Cleanup Fixes (trivial, 15 min)
|
||||
|
||||
**Item 6: Add `mx.clear_cache()` to engine cleanup**
|
||||
|
||||
`engine.py:210` does `gc.collect()` but never calls `mx.clear_cache()`. The tiling module already does this (`tiling.py:164`), showing the pattern is known but wasn't applied to the non-tiled cleanup.
|
||||
|
||||
**Critical placement detail:** `mx.clear_cache()` must go **after postprocessing**, not at line 210. At line 210, `alpha_out`/`fg_out` are still live MLX arrays used by `postprocess_alpha/foreground` (which calls `np.array()` to transfer to CPU). Clearing cache while those references exist is pointless.
|
||||
|
||||
```python
|
||||
# engine.py — after postprocessing extracts numpy arrays (~line 225)
|
||||
del alpha_out, fg_out
|
||||
gc.collect()
|
||||
mx.clear_cache() # NEW: release Metal buffer cache
|
||||
```
|
||||
|
||||
**Acceptance criteria:**
|
||||
- [x] `mx.clear_cache()` called after all MLX arrays are consumed by postprocessing
|
||||
- [x] `del` references to MLX arrays before cache clear
|
||||
- [x] All existing tests pass (`uv run pytest`)
|
||||
|
||||
**Files:** `engine.py`
|
||||
|
||||
---
|
||||
|
||||
### Phase 2: Slim Forward Mode (low effort, 30 min)
|
||||
|
||||
**Item 3: Skip returning unused intermediate tensors**
|
||||
|
||||
`GreenFormer.forward()` (`corridorkey.py:103-113`) returns 9 dict keys. The engine uses at most 4 (`alpha_coarse`, `fg_coarse`, `alpha_final`, `fg_final`). The unused 5 (`alpha_logits`, `fg_logits`, `alpha_logits_up`, `fg_logits_up`, `delta_logits`) hold references that prevent MLX from freeing underlying buffers.
|
||||
|
||||
**Important clarification:** This is about **reference lifetime**, not computation skipping. All intermediates must still be computed to produce the final outputs. The savings come from not keeping references in the returned dict, allowing MLX to reclaim buffers sooner.
|
||||
|
||||
**Design decision:** Add `slim: bool = False` parameter to `GreenFormer.__init__`, not `__call__`. This avoids mx.compile recompilation from a runtime boolean changing the output dict shape. The `pipeline.infer()` API always returns the full dict for debugging. Only the engine sets `slim=True`.
|
||||
|
||||
```python
|
||||
# corridorkey.py
|
||||
class GreenFormer(nn.Module):
|
||||
def __init__(self, ..., slim: bool = False):
|
||||
self.slim = slim
|
||||
...
|
||||
|
||||
def __call__(self, x):
|
||||
...
|
||||
if self.slim:
|
||||
return {
|
||||
"alpha_coarse": alpha_coarse,
|
||||
"fg_coarse": fg_coarse,
|
||||
"alpha_final": alpha_final,
|
||||
"fg_final": fg_final,
|
||||
}
|
||||
return { # full dict for debugging/testing
|
||||
"alpha_logits": ...,
|
||||
...
|
||||
}
|
||||
```
|
||||
|
||||
**Acceptance criteria:**
|
||||
- [x] `slim=True` returns 4-key dict, `slim=False` returns full 9-key dict
|
||||
- [x] Engine passes `slim=True` via `load_model()`
|
||||
- [x] `pipeline.infer()` keeps `slim=False` (preserves debugging API)
|
||||
- [x] Parity tests run with `slim=False` (validate intermediate keys unchanged)
|
||||
- [x] Add test: slim output matches corresponding keys from full output
|
||||
|
||||
**Files:** `corridorkey.py`, `engine.py`, `pipeline.py`, new test in `tests/`
|
||||
|
||||
---
|
||||
|
||||
### Phase 3: Non-Tiled Stage-Boundary Memory Management (medium effort, 1 hr)
|
||||
|
||||
**Item 1: Add intermediate graph materialization between encoder/decoder/refiner**
|
||||
|
||||
The non-tiled path builds one massive computation graph across all 24 Hiera blocks + 2 decoders + refiner before any materialization. Inserting `mx.eval()` at stage boundaries lets MLX free intermediate graph nodes.
|
||||
|
||||
**mx.compile interaction (CRITICAL):**
|
||||
`mx.compile` wraps the entire `GreenFormer.__call__`. Inserting `mx.eval()` inside a compiled function breaks the compilation graph. Solution: add a `_compiled: bool` flag set by `load_model()` when compile is enabled. Skip stage-boundary GC when compiled.
|
||||
|
||||
```python
|
||||
# corridorkey.py — inside __call__
|
||||
features = self.backbone(x)
|
||||
|
||||
if not self._compiled:
|
||||
mx.eval(features) # materialize backbone output
|
||||
gc.collect()
|
||||
mx.clear_cache() # free backbone intermediate graph
|
||||
|
||||
# ... decoders ...
|
||||
coarse = {...}
|
||||
|
||||
if not self._compiled:
|
||||
mx.eval(coarse)
|
||||
gc.collect()
|
||||
mx.clear_cache()
|
||||
|
||||
# ... refiner ...
|
||||
```
|
||||
|
||||
**Key question answered:** Is this worth it for full-frame? The benchmark shows 27.6 GB peak for full-frame. The backbone's intermediate computation graph (24 blocks of attention + MLP, accumulated lazily) likely dominates. Breaking it into 3 smaller graphs (backbone, decoders, refiner) should reduce peak significantly by allowing MLX to reclaim backbone intermediates before running decoders.
|
||||
|
||||
**Acceptance criteria:**
|
||||
- [x] `_compiled` flag set by `compile_model()` when compile=True
|
||||
- [x] Stage-boundary GC only runs when `_compiled=False`
|
||||
- [ ] Full-frame peak memory measurably decreases (measure with `mx.metal.get_peak_memory()`)
|
||||
- [x] Compiled path behavior unchanged (no mx.eval calls inside graph)
|
||||
- [x] All parity tests pass
|
||||
|
||||
**Files:** `corridorkey.py`, `pipeline.py`
|
||||
|
||||
---
|
||||
|
||||
### Phase 4: SDPA Migration (medium effort, 1-2 hr)
|
||||
|
||||
**Item 2: Replace manual attention with `mx.fast.scaled_dot_product_attention`**
|
||||
|
||||
Current (`hiera.py:288-290`):
|
||||
```python
|
||||
attn = (q * self.scale) @ mx.transpose(k, ...)
|
||||
attn = mx.softmax(attn, axis=-1)
|
||||
x = attn @ v
|
||||
```
|
||||
|
||||
This materializes the full `(B, heads, windows, tokens, tokens)` attention matrix. `mx.fast.scaled_dot_product_attention` fuses this into a single Metal kernel.
|
||||
|
||||
**Pre-requisites (spike required, 15 min):**
|
||||
|
||||
Before implementing, verify with a quick spike:
|
||||
|
||||
1. **5D input handling:** Current Q/K/V are 5D `(B, heads, num_windows, tokens_per_window, head_dim)`. SDPA expects 4D `(batch, heads, seq_len, head_dim)`. Must fold `num_windows` into batch dim: reshape to `(B * num_windows, heads, tokens_per_window, head_dim)`, call SDPA, reshape back.
|
||||
|
||||
2. **Asymmetric Q/K/V lengths:** At stride blocks (indices 2, 5, 21), Q is max-pooled to fewer tokens than K/V. Verify `mx.fast.scaled_dot_product_attention` supports `q_len != kv_len` (standard for cross-attention, should work).
|
||||
|
||||
3. **Scale factor application:** Current code does `(q * scale) @ k.T`. SDPA does `(q @ k.T) * scale` internally. Mathematically equivalent in fp32 but different rounding in bf16. Measure parity impact.
|
||||
|
||||
**Implementation:**
|
||||
|
||||
```python
|
||||
# hiera.py — MaskUnitAttention.__call__
|
||||
# Fold windows into batch for SDPA
|
||||
B, H, W, T, D = q.shape
|
||||
q_4d = q.reshape(B * W, H, -1, D) # (B*windows, heads, tokens, head_dim)
|
||||
k_4d = k.reshape(B * W, H, T, D)
|
||||
v_4d = v.reshape(B * W, H, T, D)
|
||||
|
||||
x = mx.fast.scaled_dot_product_attention(
|
||||
q_4d, k_4d, v_4d, scale=self.scale
|
||||
)
|
||||
|
||||
x = x.reshape(B, H, W, -1, D) # unfold windows
|
||||
```
|
||||
|
||||
**Attention matrix sizes (reference):**
|
||||
|
||||
| Stage | Blocks | mask_unit_attn | Tokens per window | Attention matrix |
|
||||
|---|---|---|---|---|
|
||||
| 0 | 0-1 | True | 64 | 64x64 |
|
||||
| 1 | 2-4 | True | 16 | 16x16 |
|
||||
| 2 | 5-20 | False | 256 | 256x256 |
|
||||
| 3 | 21-23 | False | 64 | 64x64 |
|
||||
|
||||
Sizes are modest thanks to Hiera's unrolling. Memory savings per SDPA call are small (~4MB for 256x256). Primary benefit is **kernel fusion** — one Metal dispatch vs three (matmul + softmax + matmul).
|
||||
|
||||
**Acceptance criteria:**
|
||||
- [x] Spike confirms SDPA supports: 4D input, asymmetric Q/K, bf16
|
||||
- [x] All attention calls use `mx.fast.scaled_dot_product_attention`
|
||||
- [x] Window dim folded into batch before SDPA, unfolded after
|
||||
- [x] Stride blocks (Q pooling) produce correct output with asymmetric lengths
|
||||
- [x] E2E parity within `1e-3` tolerance (existing `PARITY_TOL_E2E`)
|
||||
- [x] Backbone stage parity regressions documented if any
|
||||
|
||||
**Files:** `hiera.py`
|
||||
|
||||
---
|
||||
|
||||
### Phase 5: GPU Preprocessing (medium effort, 1-2 hr)
|
||||
|
||||
**Item 5: Move ImageNet normalize + resize from numpy/PIL to MLX**
|
||||
|
||||
Current flow (`io/image.py:48-68`, `engine.py:153-181`):
|
||||
1. `image.astype(np.float32) / 255.0` — numpy
|
||||
2. PIL bicubic resize — numpy/PIL
|
||||
3. ImageNet normalize + concat — numpy
|
||||
4. Single `mx.array()` call — boundary
|
||||
|
||||
**Scope limitation:** Full-frame path only. For tiled inference, the full-res image must stay accessible for tile slicing. Moving full-res preprocessing to GPU would allocate the entire image on GPU before tiling begins — counterproductive.
|
||||
|
||||
**Implementation approach:**
|
||||
|
||||
```python
|
||||
# New: io/preprocess_mlx.py
|
||||
def preprocess_mlx(image_uint8: np.ndarray, mask: np.ndarray, img_size: int) -> mx.array:
|
||||
"""GPU-side preprocessing for full-frame inference."""
|
||||
# Convert to MLX early
|
||||
img = mx.array(image_uint8).astype(mx.float32) / 255.0 # (H, W, 3)
|
||||
mask = mx.array(mask).astype(mx.float32) # (H, W, 1)
|
||||
|
||||
# Resize using nn.Upsample (bilinear, not bicubic)
|
||||
# Note: scale_factor = img_size / current_size
|
||||
img = mx.expand_dims(img, axis=0) # (1, H, W, 3) NHWC
|
||||
img = upsample(img, target_size=img_size)
|
||||
mask = mx.expand_dims(mask, axis=0)
|
||||
mask = upsample(mask, target_size=img_size)
|
||||
|
||||
# ImageNet normalize
|
||||
mean = mx.array([0.485, 0.456, 0.406]) # (3,)
|
||||
std = mx.array([0.229, 0.224, 0.225])
|
||||
img = (img - mean) / std
|
||||
|
||||
# Concat
|
||||
return mx.concatenate([img, mask], axis=-1) # (1, H, W, 4) NHWC
|
||||
```
|
||||
|
||||
**Caveat — resize quality:** PIL `BICUBIC` and MLX `nn.Upsample(mode="cubic")` use different cubic kernels. May introduce parity differences. Run a measurement spike: compare PIL bicubic vs MLX cubic on the golden test image, measure max absolute difference.
|
||||
|
||||
**Acceptance criteria:**
|
||||
- [x] Full-frame normalize+concat runs on GPU (resize stays CPU/PIL)
|
||||
- [x] Tiled path unchanged (keeps numpy preprocessing)
|
||||
- [x] Resize quality spike: PIL bicubic vs MLX cubic max diff = 0.22 (upscale), 0.59 (downscale)
|
||||
- [x] Resize diff >> 1e-3: kept PIL resize, moved only normalize+concat to GPU
|
||||
- [x] E2E parity within tolerance (preprocess_mlx vs preprocess: 0.0 diff)
|
||||
- [ ] Optional: wrap in `@mx.compile` for static input sizes
|
||||
|
||||
**Files:** New `io/preprocess_mlx.py`, `engine.py`
|
||||
|
||||
---
|
||||
|
||||
### Phase 6: GPU Tensor Accumulators in Tiling (nice-to-have, deferred)
|
||||
|
||||
**Item 4: Replace numpy accumulators with MLX arrays**
|
||||
|
||||
Current (`tiling.py:120-122`): Three numpy float32 arrays at full resolution. Per-tile results are transferred via `np.array(out["alpha_final"][0])`.
|
||||
|
||||
**Why deferred:**
|
||||
1. MLX lacks in-place slice assignment (`arr[y:y_end, x:x_end] += tile`)
|
||||
2. `mx.scatter_add` or equivalent needs API verification
|
||||
3. GPU->CPU transfer per tile is ~1MB at 512x512 — microseconds on unified memory
|
||||
4. Wave 1 plan already flagged this as "best-effort, 30min timebox" with numpy fallback
|
||||
5. Total transfer for 16 tiles (2048x2048) is ~16MB — negligible vs compute time
|
||||
|
||||
**If pursued later:** Investigate `mx.scatter_add` API, or accumulate using full-size zero tensors with padded tiles (wasteful but avoids slice assignment).
|
||||
|
||||
---
|
||||
|
||||
## Alternative Approaches Considered
|
||||
|
||||
### Token routing (from Raiden129)
|
||||
Route "easy" tokens to lightweight LTRM module, skip expensive attention. **Rejected for now:** MLX lacks dynamic boolean indexing. `mx.where` + padding doesn't reduce FLOPs. NumPy CPU fallback per-block (16x in stage 2) would destroy async pipeline. Net benefit uncertain. Revisit if profiling shows attention is the bottleneck.
|
||||
|
||||
### Int8/MXFP4 quantization
|
||||
50-75% weight memory reduction. **Deferred:** Matting requires sub-pixel alpha precision. Int8 likely safe (<1% drop) but needs quality validation on real footage. Leave Conv2d, LayerNorm, softmax in full precision. Profile as separate future work.
|
||||
|
||||
### Full weight-level bf16 casting
|
||||
Cast all model weights to bf16 at load time (halves parameter memory). **Not yet explored:** Currently only activations use bf16. Would need to verify backbone parity in bf16 weights (16 blocks = most drift risk). Lower priority than computation graph optimizations.
|
||||
|
||||
## Acceptance Criteria
|
||||
|
||||
### Functional Requirements
|
||||
- [ ] All 94 existing tests pass
|
||||
- [ ] E2E parity within `1e-3` tolerance
|
||||
- [ ] No regression in tiled inference behavior
|
||||
- [ ] Engine cleanup releases all Metal buffers after postprocessing
|
||||
|
||||
### Non-Functional Requirements
|
||||
- [ ] Full-frame peak memory measurably decreases with stage-boundary GC (measure with `mx.metal.get_peak_memory()`)
|
||||
- [ ] No latency regression > 5% in any mode
|
||||
- [ ] `uv run ruff check .` passes
|
||||
- [ ] `uv run ruff format --check .` passes
|
||||
- [ ] `uv run ty check` passes
|
||||
|
||||
### Quality Gates
|
||||
- [ ] SDPA spike completed before Phase 4 implementation
|
||||
- [ ] Resize quality spike completed before Phase 5 implementation
|
||||
- [ ] Each phase merged separately with passing CI
|
||||
|
||||
## Implementation Order & Dependencies
|
||||
|
||||
```
|
||||
Phase 1 (engine cleanup) ─── no deps ──────────────────── 15 min
|
||||
Phase 2 (slim forward) ─── no deps ──────────────────── 30 min
|
||||
Phase 3 (stage GC) ─── depends on Phase 2 (slim changes what's in dict) ── 1 hr
|
||||
Phase 4 (SDPA) ─── no deps (backbone only) ──── spike 15m + impl 1-2 hr
|
||||
Phase 5 (GPU preprocess) ── no deps (I/O only) ────────── spike 15m + impl 1-2 hr
|
||||
Phase 6 (GPU accum) ─── deferred ─────────────────── (skip)
|
||||
```
|
||||
|
||||
Phases 1-2 and Phase 4 can run in parallel. Phase 3 should follow Phase 2.
|
||||
|
||||
## References
|
||||
|
||||
### Internal
|
||||
- Wave 1 plan: `docs/plans/2026-03-08-feat-mlx-memory-optimizations-plan.md`
|
||||
- Brainstorm: `docs/brainstorms/2026-03-09-pytorch-optimization-learnings-brainstorm.md`
|
||||
- Deep research: `docs/MLX Vision Transformer Optimization Techniques.md`
|
||||
- Attention impl: `src/corridorkey_mlx/model/hiera.py:267-297`
|
||||
- Forward pass: `src/corridorkey_mlx/model/corridorkey.py:56-113`
|
||||
- Engine cleanup: `src/corridorkey_mlx/engine.py:187-210`
|
||||
- Tiling accumulators: `src/corridorkey_mlx/inference/tiling.py:120-171`
|
||||
- Preprocessing: `src/corridorkey_mlx/io/image.py:48-68`
|
||||
- Pipeline defaults: `src/corridorkey_mlx/inference/pipeline.py:31-58`
|
||||
|
||||
### External
|
||||
- PR #104: https://github.com/nikopueringer/CorridorKey/pull/104
|
||||
- Raiden129 fork: https://github.com/Raiden129/CorridorKey_Test
|
||||
- MLX SDPA docs: https://ml-explore.github.io/mlx/build/html/python/_autosummary/mlx.core.fast.scaled_dot_product_attention.html
|
||||
- MLX FlashAttention issue: https://github.com/ml-explore/mlx/issues/129
|
||||
- MLX memory management: https://github.com/ml-explore/mlx/issues/742
|
||||
|
||||
## Open Questions
|
||||
|
||||
- Stage-boundary GC: measurable peak memory reduction in full-frame, or negligible?
|
||||
- SDPA spike: does MLX SDPA handle 5D->4D reshape + asymmetric Q/K correctly?
|
||||
- Resize quality: PIL bicubic vs MLX cubic — within parity tolerance?
|
||||
- Weight-level bf16: safe for Hiera backbone (16 blocks of drift)?
|
||||
- Int8 quantization: acceptable alpha precision threshold for matting?
|
||||
450
scripts/bench_optimizations.py
Normal file
450
scripts/bench_optimizations.py
Normal file
@ -0,0 +1,450 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Benchmark matrix for wave 2 optimizations.
|
||||
|
||||
Runs an exhaustive matrix of optimization toggles and reports latency + peak
|
||||
memory for each configuration. Helps isolate the impact of each optimization.
|
||||
|
||||
Usage:
|
||||
uv run python scripts/bench_optimizations.py
|
||||
uv run python scripts/bench_optimizations.py --resolution 512 --bench-runs 5
|
||||
uv run python scripts/bench_optimizations.py --checkpoint ckpt.safetensors
|
||||
|
||||
All mx.eval() calls are MLX array materialization, NOT Python eval().
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import gc
|
||||
import itertools
|
||||
import sys
|
||||
import time
|
||||
from dataclasses import dataclass, field
|
||||
from pathlib import Path
|
||||
|
||||
import mlx.core as mx
|
||||
from rich.console import Console
|
||||
from rich.table import Table
|
||||
|
||||
# Ensure package is importable when running from repo root
|
||||
sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "src"))
|
||||
|
||||
from corridorkey_mlx.io.image import preprocess as preprocess_numpy
|
||||
from corridorkey_mlx.io.preprocess_mlx import preprocess_mlx
|
||||
from corridorkey_mlx.model.corridorkey import GreenFormer
|
||||
|
||||
console = Console()
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Optimization flags
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
TOGGLE_NAMES = ["slim", "stage_gc", "sdpa", "bf16", "fused_decode", "gpu_preprocess"]
|
||||
|
||||
|
||||
@dataclass
|
||||
class OptConfig:
|
||||
"""Single optimization configuration to benchmark."""
|
||||
|
||||
slim: bool = True
|
||||
stage_gc: bool = True
|
||||
sdpa: bool = True
|
||||
bf16: bool = True
|
||||
fused_decode: bool = True
|
||||
gpu_preprocess: bool = True
|
||||
|
||||
@property
|
||||
def label(self) -> str:
|
||||
flags = []
|
||||
for name in TOGGLE_NAMES:
|
||||
val = getattr(self, name)
|
||||
flags.append(f"{name}={'on' if val else 'off'}")
|
||||
return " | ".join(flags)
|
||||
|
||||
@property
|
||||
def short_label(self) -> str:
|
||||
parts = []
|
||||
for name in TOGGLE_NAMES:
|
||||
if getattr(self, name):
|
||||
parts.append(name)
|
||||
return "+".join(parts) if parts else "baseline"
|
||||
|
||||
|
||||
@dataclass
|
||||
class BenchResult:
|
||||
"""Result from a single benchmark configuration."""
|
||||
|
||||
config: OptConfig
|
||||
resolution: int
|
||||
warmup_ms: float = 0.0
|
||||
median_ms: float = 0.0
|
||||
min_ms: float = 0.0
|
||||
peak_memory_mb: float = 0.0
|
||||
all_times: list[float] = field(default_factory=list)
|
||||
error: str | None = None
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Benchmark logic
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def build_model(
|
||||
config: OptConfig,
|
||||
img_size: int,
|
||||
checkpoint: Path | None,
|
||||
) -> GreenFormer:
|
||||
"""Build model with the given optimization config."""
|
||||
dtype = mx.bfloat16 if config.bf16 else mx.float32
|
||||
model = GreenFormer(
|
||||
img_size=img_size,
|
||||
dtype=dtype,
|
||||
fused_decode=config.fused_decode,
|
||||
slim=config.slim,
|
||||
use_sdpa=config.sdpa,
|
||||
stage_gc=config.stage_gc,
|
||||
)
|
||||
if checkpoint and checkpoint.exists():
|
||||
model.load_checkpoint(checkpoint)
|
||||
else:
|
||||
model.eval()
|
||||
mx.eval(model.parameters()) # noqa: S307
|
||||
return model
|
||||
|
||||
|
||||
def make_input(img_size: int, gpu_preprocess: bool) -> mx.array:
|
||||
"""Create input tensor, optionally via GPU preprocessing path."""
|
||||
import numpy as np
|
||||
|
||||
np.random.seed(42)
|
||||
rgb_f32 = np.random.rand(img_size, img_size, 3).astype(np.float32)
|
||||
mask_f32 = np.random.rand(img_size, img_size, 1).astype(np.float32)
|
||||
|
||||
if gpu_preprocess:
|
||||
return preprocess_mlx(rgb_f32, mask_f32)
|
||||
return preprocess_numpy(rgb_f32, mask_f32)
|
||||
|
||||
|
||||
def time_inference(
|
||||
model: GreenFormer,
|
||||
x: mx.array,
|
||||
warmup_runs: int,
|
||||
bench_runs: int,
|
||||
) -> tuple[float, float, float, list[float]]:
|
||||
"""Run warmup + benchmark, return (warmup_ms, median_ms, min_ms, all_times)."""
|
||||
|
||||
def run() -> None:
|
||||
out = model(x)
|
||||
mx.eval(out) # noqa: S307
|
||||
|
||||
# Warmup
|
||||
start = time.perf_counter()
|
||||
run()
|
||||
warmup_ms = (time.perf_counter() - start) * 1000.0
|
||||
|
||||
for _ in range(warmup_runs - 1):
|
||||
run()
|
||||
|
||||
# Benchmark
|
||||
times: list[float] = []
|
||||
for _ in range(bench_runs):
|
||||
start = time.perf_counter()
|
||||
run()
|
||||
times.append((time.perf_counter() - start) * 1000.0)
|
||||
|
||||
times.sort()
|
||||
median_ms = times[len(times) // 2]
|
||||
min_ms = times[0]
|
||||
return warmup_ms, median_ms, min_ms, times
|
||||
|
||||
|
||||
def measure_peak_memory(model: GreenFormer, x: mx.array) -> float:
|
||||
"""Measure peak Metal memory for a single forward pass (MB)."""
|
||||
gc.collect()
|
||||
mx.clear_cache()
|
||||
|
||||
# Use non-deprecated API if available, fall back to mx.metal.*
|
||||
reset_fn = getattr(mx, "reset_peak_memory", None) or getattr(
|
||||
mx.metal, "reset_peak_memory", None
|
||||
)
|
||||
get_fn = getattr(mx, "get_peak_memory", None) or getattr(mx.metal, "get_peak_memory", None)
|
||||
if not reset_fn or not get_fn:
|
||||
return 0.0
|
||||
|
||||
reset_fn()
|
||||
|
||||
out = model(x)
|
||||
mx.eval(out) # noqa: S307
|
||||
|
||||
peak_bytes = get_fn()
|
||||
|
||||
del out
|
||||
gc.collect()
|
||||
mx.clear_cache()
|
||||
|
||||
return peak_bytes / (1024 * 1024)
|
||||
|
||||
|
||||
def bench_config(
|
||||
config: OptConfig,
|
||||
resolution: int,
|
||||
checkpoint: Path | None,
|
||||
warmup_runs: int,
|
||||
bench_runs: int,
|
||||
) -> BenchResult:
|
||||
"""Benchmark a single optimization configuration."""
|
||||
result = BenchResult(config=config, resolution=resolution)
|
||||
|
||||
try:
|
||||
model = build_model(config, resolution, checkpoint)
|
||||
x = make_input(resolution, config.gpu_preprocess)
|
||||
mx.eval(x) # noqa: S307
|
||||
|
||||
warmup_ms, median_ms, min_ms, all_times = time_inference(model, x, warmup_runs, bench_runs)
|
||||
result.warmup_ms = warmup_ms
|
||||
result.median_ms = median_ms
|
||||
result.min_ms = min_ms
|
||||
result.all_times = all_times
|
||||
|
||||
# Measure peak memory with a fresh model to avoid cached state
|
||||
model_fresh = build_model(config, resolution, checkpoint)
|
||||
result.peak_memory_mb = measure_peak_memory(model_fresh, x)
|
||||
|
||||
# Cleanup
|
||||
del model, model_fresh, x
|
||||
gc.collect()
|
||||
mx.clear_cache()
|
||||
|
||||
except Exception as e:
|
||||
result.error = str(e)
|
||||
|
||||
return result
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Matrix generation
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def generate_matrix(sweep: str) -> list[OptConfig]:
|
||||
"""Generate benchmark configurations.
|
||||
|
||||
Args:
|
||||
sweep: "full" for exhaustive 2^6 matrix, "ablation" for toggle-one-off,
|
||||
"key" for important combinations only.
|
||||
"""
|
||||
all_on = OptConfig()
|
||||
|
||||
if sweep == "ablation":
|
||||
# All-on baseline + toggle each flag off one at a time
|
||||
configs = [all_on]
|
||||
for name in TOGGLE_NAMES:
|
||||
off = OptConfig(**{**vars(all_on), name: False})
|
||||
configs.append(off)
|
||||
# Also add all-off baseline
|
||||
configs.append(
|
||||
OptConfig(
|
||||
slim=False,
|
||||
stage_gc=False,
|
||||
sdpa=False,
|
||||
bf16=False,
|
||||
fused_decode=False,
|
||||
gpu_preprocess=False,
|
||||
)
|
||||
)
|
||||
return configs
|
||||
|
||||
if sweep == "full":
|
||||
# Exhaustive 2^6 = 64 configurations
|
||||
configs = []
|
||||
for combo in itertools.product([True, False], repeat=len(TOGGLE_NAMES)):
|
||||
kwargs = dict(zip(TOGGLE_NAMES, combo, strict=True))
|
||||
configs.append(OptConfig(**kwargs))
|
||||
return configs
|
||||
|
||||
# "key" — important configurations
|
||||
return [
|
||||
# All off (wave 1 baseline)
|
||||
OptConfig(
|
||||
slim=False,
|
||||
stage_gc=False,
|
||||
sdpa=False,
|
||||
bf16=False,
|
||||
fused_decode=False,
|
||||
gpu_preprocess=False,
|
||||
),
|
||||
# Wave 1 only (bf16 + fused_decode)
|
||||
OptConfig(
|
||||
slim=False,
|
||||
stage_gc=False,
|
||||
sdpa=False,
|
||||
bf16=True,
|
||||
fused_decode=True,
|
||||
gpu_preprocess=False,
|
||||
),
|
||||
# Wave 2 only (slim + stage_gc + sdpa + gpu_preprocess)
|
||||
OptConfig(
|
||||
slim=True,
|
||||
stage_gc=True,
|
||||
sdpa=True,
|
||||
bf16=False,
|
||||
fused_decode=False,
|
||||
gpu_preprocess=True,
|
||||
),
|
||||
# All on
|
||||
all_on,
|
||||
# SDPA only
|
||||
OptConfig(
|
||||
slim=False,
|
||||
stage_gc=False,
|
||||
sdpa=True,
|
||||
bf16=False,
|
||||
fused_decode=False,
|
||||
gpu_preprocess=False,
|
||||
),
|
||||
# Stage GC only
|
||||
OptConfig(
|
||||
slim=False,
|
||||
stage_gc=True,
|
||||
sdpa=False,
|
||||
bf16=False,
|
||||
fused_decode=False,
|
||||
gpu_preprocess=False,
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Output
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def print_results(results: list[BenchResult]) -> None:
|
||||
"""Print results as a rich table."""
|
||||
table = Table(title="Optimization Benchmark Matrix")
|
||||
|
||||
table.add_column("Config", max_width=50)
|
||||
table.add_column("Res", justify="right")
|
||||
table.add_column("Warmup (ms)", justify="right")
|
||||
table.add_column("Median (ms)", justify="right")
|
||||
table.add_column("Min (ms)", justify="right")
|
||||
table.add_column("Peak Mem (MB)", justify="right")
|
||||
table.add_column("Status", justify="center")
|
||||
|
||||
# Sort by resolution then median time
|
||||
results.sort(key=lambda r: (r.resolution, r.median_ms))
|
||||
|
||||
for r in results:
|
||||
if r.error:
|
||||
table.add_row(
|
||||
r.config.short_label,
|
||||
str(r.resolution),
|
||||
"-",
|
||||
"-",
|
||||
"-",
|
||||
"-",
|
||||
f"[red]FAIL: {r.error[:30]}[/red]",
|
||||
)
|
||||
else:
|
||||
table.add_row(
|
||||
r.config.short_label,
|
||||
str(r.resolution),
|
||||
f"{r.warmup_ms:.1f}",
|
||||
f"{r.median_ms:.1f}",
|
||||
f"{r.min_ms:.1f}",
|
||||
f"{r.peak_memory_mb:.0f}" if r.peak_memory_mb > 0 else "N/A",
|
||||
"[green]OK[/green]",
|
||||
)
|
||||
|
||||
console.print(table)
|
||||
|
||||
# Print speedup summary vs baseline
|
||||
baselines = {r.resolution: r for r in results if r.config.short_label == "baseline"}
|
||||
if baselines:
|
||||
console.print("\n[bold]Speedup vs baseline (all-off):[/bold]")
|
||||
for r in results:
|
||||
if r.error or r.config.short_label == "baseline":
|
||||
continue
|
||||
baseline = baselines.get(r.resolution)
|
||||
if baseline and baseline.median_ms > 0:
|
||||
speedup = baseline.median_ms / r.median_ms
|
||||
mem_ratio = (
|
||||
f"{r.peak_memory_mb / baseline.peak_memory_mb:.2f}x"
|
||||
if baseline.peak_memory_mb > 0 and r.peak_memory_mb > 0
|
||||
else "N/A"
|
||||
)
|
||||
console.print(
|
||||
f" {r.config.short_label:40s} "
|
||||
f"{speedup:.2f}x speed, {mem_ratio} memory "
|
||||
f"({r.resolution})"
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Main
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def main() -> None:
|
||||
parser = argparse.ArgumentParser(description="Benchmark matrix for wave 2 optimizations")
|
||||
parser.add_argument(
|
||||
"--checkpoint",
|
||||
type=Path,
|
||||
default=Path("checkpoints/corridorkey_mlx.safetensors"),
|
||||
)
|
||||
parser.add_argument(
|
||||
"--resolution",
|
||||
type=int,
|
||||
nargs="+",
|
||||
default=[512],
|
||||
help="Resolution(s) to benchmark",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--sweep",
|
||||
choices=["key", "ablation", "full"],
|
||||
default="ablation",
|
||||
help="key=6 configs, ablation=8 configs (toggle one off), full=64 configs",
|
||||
)
|
||||
parser.add_argument("--warmup-runs", type=int, default=3)
|
||||
parser.add_argument("--bench-runs", type=int, default=10)
|
||||
args = parser.parse_args()
|
||||
|
||||
ckpt = args.checkpoint if args.checkpoint.exists() else None
|
||||
configs = generate_matrix(args.sweep)
|
||||
|
||||
console.print("[bold]CorridorKey MLX — Optimization Benchmark Matrix[/bold]")
|
||||
console.print(f" Weights: {args.checkpoint if ckpt else 'random (no checkpoint)'}")
|
||||
console.print(f" Resolutions: {args.resolution}")
|
||||
console.print(f" Sweep: {args.sweep} ({len(configs)} configs)")
|
||||
console.print(f" Warmup: {args.warmup_runs}, Bench: {args.bench_runs}")
|
||||
console.print()
|
||||
|
||||
all_results: list[BenchResult] = []
|
||||
total = len(configs) * len(args.resolution)
|
||||
idx = 0
|
||||
|
||||
for res in args.resolution:
|
||||
for config in configs:
|
||||
idx += 1
|
||||
console.print(
|
||||
f" [{idx}/{total}] {res}x{res} — {config.short_label}...",
|
||||
end=" ",
|
||||
)
|
||||
result = bench_config(config, res, ckpt, args.warmup_runs, args.bench_runs)
|
||||
if result.error:
|
||||
console.print("[red]FAIL[/red]")
|
||||
else:
|
||||
console.print(
|
||||
f"[green]{result.median_ms:.1f}ms[/green] "
|
||||
f"(peak: {result.peak_memory_mb:.0f}MB)"
|
||||
if result.peak_memory_mb > 0
|
||||
else f"[green]{result.median_ms:.1f}ms[/green]"
|
||||
)
|
||||
all_results.append(result)
|
||||
|
||||
console.print()
|
||||
print_results(all_results)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@ -15,10 +15,10 @@ from pathlib import Path
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from PIL import Image
|
||||
|
||||
# Reuse the GreenFormer and loading logic from dump script
|
||||
from dump_pytorch_reference import IMG_SIZE, GreenFormer, load_checkpoint
|
||||
from PIL import Image
|
||||
|
||||
IMAGENET_MEAN = torch.tensor([0.485, 0.456, 0.406]).reshape(1, 3, 1, 1)
|
||||
IMAGENET_STD = torch.tensor([0.229, 0.224, 0.225]).reshape(1, 3, 1, 1)
|
||||
@ -94,7 +94,7 @@ def main() -> None:
|
||||
# Composite over green
|
||||
alpha_f = alpha_np.astype(np.float32) / 255.0
|
||||
green = np.array([0, 177, 64], dtype=np.float32)
|
||||
comp_np = (fg_np.astype(np.float32) * alpha_f[..., None] + green * (1 - alpha_f[..., None]))
|
||||
comp_np = fg_np.astype(np.float32) * alpha_f[..., None] + green * (1 - alpha_f[..., None])
|
||||
comp_np = comp_np.clip(0, 255).astype(np.uint8)
|
||||
|
||||
# Save
|
||||
|
||||
@ -6,6 +6,7 @@ mirrors the Torch engine contract expected by the main CorridorKey repository.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import gc
|
||||
import logging
|
||||
import warnings
|
||||
from pathlib import Path
|
||||
@ -16,11 +17,16 @@ import numpy as np
|
||||
from PIL import Image
|
||||
|
||||
from corridorkey_mlx.inference.pipeline import load_model
|
||||
from corridorkey_mlx.inference.tiling import (
|
||||
DEFAULT_OVERLAP,
|
||||
tiled_inference,
|
||||
)
|
||||
from corridorkey_mlx.io.image import (
|
||||
postprocess_alpha,
|
||||
postprocess_foreground,
|
||||
preprocess,
|
||||
)
|
||||
from corridorkey_mlx.io.preprocess_mlx import preprocess_mlx
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from corridorkey_mlx.model.corridorkey import GreenFormer
|
||||
@ -39,17 +45,27 @@ class CorridorKeyMLXEngine:
|
||||
device: Ignored on MLX (Apple Silicon uses unified memory).
|
||||
Accepted for API compatibility with the Torch engine.
|
||||
img_size: Internal model resolution (square). The model was trained
|
||||
at 2048. Use 512 for fast dev iteration.
|
||||
at 2048. Use 512 for fast dev iteration. When tiling is enabled,
|
||||
this is ignored and the model runs at ``tile_size``.
|
||||
use_refiner: If True, return refined alpha/fg. If False, return
|
||||
coarse predictions (skips refiner in postprocessing, not forward pass).
|
||||
compile: If True, wrap model forward with ``mx.compile`` for faster
|
||||
repeated inference at the same resolution.
|
||||
tile_size: If set, enable tiled inference — split full-res input into
|
||||
overlapping tiles of this size. Model loads at tile_size instead of
|
||||
img_size. Set to 0 or None to disable.
|
||||
overlap: Overlap in pixels between adjacent tiles (default 64).
|
||||
|
||||
Example::
|
||||
|
||||
from corridorkey_mlx import CorridorKeyMLXEngine
|
||||
|
||||
# Full-frame (default)
|
||||
engine = CorridorKeyMLXEngine("/path/to/corridorkey_mlx.safetensors")
|
||||
|
||||
# Tiled — 12x less memory at 2048x2048
|
||||
engine = CorridorKeyMLXEngine("/path/to/ckpt.safetensors", tile_size=512, overlap=64)
|
||||
|
||||
result = engine.process_frame(rgb_uint8, mask_uint8)
|
||||
# result["alpha"] — (H, W) uint8
|
||||
# result["fg"] — (H, W, 3) uint8
|
||||
@ -67,6 +83,8 @@ class CorridorKeyMLXEngine:
|
||||
img_size: int = PRODUCTION_IMG_SIZE,
|
||||
use_refiner: bool = True,
|
||||
compile: bool = True,
|
||||
tile_size: int | None = None,
|
||||
overlap: int = DEFAULT_OVERLAP,
|
||||
) -> None:
|
||||
checkpoint = Path(checkpoint_path)
|
||||
if not checkpoint.exists():
|
||||
@ -76,11 +94,24 @@ class CorridorKeyMLXEngine:
|
||||
if device is not None:
|
||||
logger.info("device=%r ignored on MLX (unified memory)", device)
|
||||
|
||||
self._img_size = img_size
|
||||
self._use_refiner = use_refiner
|
||||
self._model: GreenFormer = load_model(
|
||||
checkpoint, img_size=img_size, compile=compile
|
||||
)
|
||||
self._tiled = bool(tile_size)
|
||||
self._tile_size = tile_size or img_size
|
||||
self._overlap = overlap
|
||||
|
||||
if self._tiled:
|
||||
# Tiled: model runs at tile_size, input stays full-res
|
||||
self._img_size = self._tile_size
|
||||
self._model: GreenFormer = load_model(
|
||||
checkpoint, img_size=self._tile_size, compile=False, slim=True
|
||||
)
|
||||
logger.info(
|
||||
"Tiled inference: tile_size=%d, overlap=%d", self._tile_size, self._overlap
|
||||
)
|
||||
else:
|
||||
# Full-frame: resize input to img_size
|
||||
self._img_size = img_size
|
||||
self._model = load_model(checkpoint, img_size=img_size, compile=compile, slim=True)
|
||||
|
||||
def process_frame(
|
||||
self,
|
||||
@ -126,48 +157,68 @@ class CorridorKeyMLXEngine:
|
||||
if mask_f32.ndim == 2:
|
||||
mask_f32 = mask_f32[:, :, np.newaxis]
|
||||
|
||||
# -- resize to model resolution --
|
||||
if rgb_f32.shape[0] != self._img_size or rgb_f32.shape[1] != self._img_size:
|
||||
rgb_pil = Image.fromarray(image).resize(
|
||||
(self._img_size, self._img_size), Image.Resampling.BICUBIC
|
||||
if self._tiled:
|
||||
# Tiled: keep full-res, preprocess, run tiled_inference
|
||||
x = preprocess(rgb_f32, mask_f32)
|
||||
outputs = tiled_inference(
|
||||
self._model, x, tile_size=self._tile_size, overlap=self._overlap
|
||||
)
|
||||
rgb_f32 = np.asarray(rgb_pil, dtype=np.float32) / 255.0
|
||||
# NOTE: mx.eval is MLX array materialization, not Python eval()
|
||||
mx.eval(outputs) # noqa: S307
|
||||
else:
|
||||
# Full-frame: resize to model resolution
|
||||
if rgb_f32.shape[0] != self._img_size or rgb_f32.shape[1] != self._img_size:
|
||||
rgb_pil = Image.fromarray(image).resize(
|
||||
(self._img_size, self._img_size), Image.Resampling.BICUBIC
|
||||
)
|
||||
rgb_f32 = np.asarray(rgb_pil, dtype=np.float32) / 255.0
|
||||
|
||||
mask_u8 = mask_linear if mask_linear.ndim == 2 else mask_linear[:, :, 0]
|
||||
mask_pil = Image.fromarray(mask_u8, mode="L").resize(
|
||||
(self._img_size, self._img_size), Image.Resampling.BICUBIC
|
||||
)
|
||||
mask_f32 = np.asarray(mask_pil, dtype=np.float32)[:, :, np.newaxis] / 255.0
|
||||
mask_u8 = mask_linear if mask_linear.ndim == 2 else mask_linear[:, :, 0]
|
||||
mask_pil = Image.fromarray(mask_u8, mode="L").resize(
|
||||
(self._img_size, self._img_size), Image.Resampling.BICUBIC
|
||||
)
|
||||
mask_f32 = np.asarray(mask_pil, dtype=np.float32)[:, :, np.newaxis] / 255.0
|
||||
|
||||
# -- preprocess (ImageNet norm + concat) -> (1, H, W, 4) NHWC --
|
||||
x = preprocess(rgb_f32, mask_f32)
|
||||
|
||||
# -- forward --
|
||||
outputs = self._model(x)
|
||||
mx.eval(outputs) # noqa: S307 — mx.eval materializes lazy MLX arrays, not Python eval
|
||||
x = preprocess_mlx(rgb_f32, mask_f32)
|
||||
outputs = self._model(x)
|
||||
# NOTE: mx.eval is MLX array materialization, not Python eval()
|
||||
mx.eval(outputs) # noqa: S307
|
||||
|
||||
# -- select coarse vs refined --
|
||||
alpha_coarse = outputs["alpha_coarse"]
|
||||
fg_coarse = outputs["fg_coarse"]
|
||||
alpha_refined = outputs["alpha_final"]
|
||||
fg_refined = outputs["fg_final"]
|
||||
|
||||
if not self._use_refiner or refiner_scale == 0.0:
|
||||
alpha_out = alpha_coarse
|
||||
fg_out = fg_coarse
|
||||
elif refiner_scale == 1.0:
|
||||
alpha_out = alpha_refined
|
||||
fg_out = fg_refined
|
||||
if self._tiled:
|
||||
# tiled_inference only returns final (refined) outputs
|
||||
alpha_out = outputs["alpha_final"]
|
||||
fg_out = outputs["fg_final"]
|
||||
else:
|
||||
# output-space lerp
|
||||
s = refiner_scale
|
||||
alpha_out = alpha_coarse * (1.0 - s) + alpha_refined * s
|
||||
fg_out = fg_coarse * (1.0 - s) + fg_refined * s
|
||||
alpha_coarse = outputs["alpha_coarse"]
|
||||
fg_coarse = outputs["fg_coarse"]
|
||||
alpha_refined = outputs["alpha_final"]
|
||||
fg_refined = outputs["fg_final"]
|
||||
|
||||
if not self._use_refiner or refiner_scale == 0.0:
|
||||
alpha_out = alpha_coarse
|
||||
fg_out = fg_coarse
|
||||
elif refiner_scale == 1.0:
|
||||
alpha_out = alpha_refined
|
||||
fg_out = fg_refined
|
||||
else:
|
||||
s = refiner_scale
|
||||
alpha_out = alpha_coarse * (1.0 - s) + alpha_refined * s
|
||||
fg_out = fg_coarse * (1.0 - s) + fg_refined * s
|
||||
|
||||
# Release unused intermediate tensors
|
||||
del outputs
|
||||
gc.collect()
|
||||
|
||||
# -- postprocess to uint8 --
|
||||
alpha_u8 = postprocess_alpha(alpha_out)
|
||||
fg_u8 = postprocess_foreground(fg_out)
|
||||
|
||||
# Release all MLX array references, then free Metal buffer cache
|
||||
del alpha_out, fg_out
|
||||
gc.collect()
|
||||
mx.clear_cache()
|
||||
|
||||
# -- resize back to original --
|
||||
if alpha_u8.shape[0] != original_h or alpha_u8.shape[1] != original_w:
|
||||
target = (original_w, original_h)
|
||||
@ -198,11 +249,7 @@ class CorridorKeyMLXEngine:
|
||||
# -- composite over black --
|
||||
alpha_3ch = alpha_u8[:, :, np.newaxis].astype(np.float32) / 255.0
|
||||
fg_float = fg_u8.astype(np.float32)
|
||||
comp = (
|
||||
(fg_float * alpha_3ch).astype(np.uint8)
|
||||
if fg_is_straight
|
||||
else fg_u8.copy()
|
||||
)
|
||||
comp = (fg_float * alpha_3ch).astype(np.uint8) if fg_is_straight else fg_u8.copy()
|
||||
|
||||
return {
|
||||
"alpha": alpha_u8,
|
||||
|
||||
@ -33,6 +33,11 @@ def load_model(
|
||||
img_size: int = DEFAULT_IMG_SIZE,
|
||||
compile: bool = False,
|
||||
shapeless: bool = False,
|
||||
dtype: mx.Dtype = mx.bfloat16,
|
||||
fused_decode: bool = True,
|
||||
slim: bool = False,
|
||||
use_sdpa: bool = True,
|
||||
stage_gc: bool = True,
|
||||
) -> GreenFormer:
|
||||
"""Build GreenFormer and load weights from safetensors checkpoint.
|
||||
|
||||
@ -44,8 +49,23 @@ def load_model(
|
||||
no shape-dependent logic varies across calls. The Hiera backbone
|
||||
uses shape-dependent reshapes, so shapeless is NOT recommended
|
||||
unless all inputs share the same spatial dimensions.
|
||||
dtype: Compute dtype for decoder activations. bfloat16 reduces memory;
|
||||
backbone and sigmoid always stay fp32. All outputs are fp32.
|
||||
fused_decode: If True, batch alpha+fg decoder upsamples to reduce
|
||||
Metal dispatch calls. Bit-exact with unfused path.
|
||||
slim: If True, forward returns only 4 final keys (drops intermediates).
|
||||
Reduces reference lifetime so MLX can reclaim buffers sooner.
|
||||
use_sdpa: If True, use mx.fast.scaled_dot_product_attention in backbone.
|
||||
stage_gc: If True, materialize + GC at backbone/decoder/refiner boundaries.
|
||||
"""
|
||||
model = GreenFormer(img_size=img_size)
|
||||
model = GreenFormer(
|
||||
img_size=img_size,
|
||||
dtype=dtype,
|
||||
fused_decode=fused_decode,
|
||||
slim=slim,
|
||||
use_sdpa=use_sdpa,
|
||||
stage_gc=stage_gc,
|
||||
)
|
||||
model.load_checkpoint(checkpoint)
|
||||
if compile:
|
||||
model = compile_model(model, shapeless=shapeless)
|
||||
@ -59,6 +79,7 @@ def compile_model(model: GreenFormer, shapeless: bool = False) -> GreenFormer:
|
||||
same resolution. Shapeless compile is experimental — the backbone uses
|
||||
shape-dependent reshapes that may trigger recompilation.
|
||||
"""
|
||||
model._compiled = True
|
||||
model.__call__ = mx.compile(model.__call__, shapeless=shapeless) # type: ignore[method-assign]
|
||||
return model
|
||||
|
||||
|
||||
@ -6,6 +6,7 @@ then blends results using linear ramp weights in the overlap region.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import gc
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import mlx.core as mx
|
||||
@ -157,6 +158,11 @@ def tiled_inference(
|
||||
fg_accum[y_start:y_end, x_start:x_end, :] += fg_tile * w3d
|
||||
weight_accum[y_start:y_end, x_start:x_end, :] += w3d
|
||||
|
||||
# Deterministic memory cleanup — release MLX graph refs between tiles
|
||||
del out, tile, alpha_tile, fg_tile
|
||||
gc.collect()
|
||||
mx.clear_cache()
|
||||
|
||||
# Normalize by accumulated weights
|
||||
weight_accum = np.maximum(weight_accum, 1e-8)
|
||||
alpha_final = mx.array(alpha_accum / weight_accum)[None] # (1, H, W, 1)
|
||||
|
||||
41
src/corridorkey_mlx/io/preprocess_mlx.py
Normal file
41
src/corridorkey_mlx/io/preprocess_mlx.py
Normal file
@ -0,0 +1,41 @@
|
||||
"""GPU-side preprocessing for full-frame inference.
|
||||
|
||||
Moves ImageNet normalization and channel concatenation to MLX while
|
||||
keeping PIL bicubic resize on CPU (MLX cubic kernel differs too much).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import mlx.core as mx
|
||||
|
||||
if TYPE_CHECKING:
|
||||
import numpy as np
|
||||
|
||||
# ImageNet normalization constants (GPU-resident)
|
||||
_IMAGENET_MEAN = mx.array([0.485, 0.456, 0.406])
|
||||
_IMAGENET_STD = mx.array([0.229, 0.224, 0.225])
|
||||
|
||||
|
||||
def preprocess_mlx(rgb_f32: np.ndarray, mask_f32: np.ndarray) -> mx.array:
|
||||
"""GPU-side normalize + concat for full-frame inference.
|
||||
|
||||
Assumes resize already happened on CPU via PIL bicubic.
|
||||
|
||||
Args:
|
||||
rgb_f32: (H, W, 3) float32 in [0, 1], already resized.
|
||||
mask_f32: (H, W, 1) float32 in [0, 1], already resized.
|
||||
|
||||
Returns:
|
||||
(1, H, W, 4) mx.array — ImageNet-normalized RGB + raw alpha hint.
|
||||
"""
|
||||
img = mx.array(rgb_f32)
|
||||
mask = mx.array(mask_f32)
|
||||
|
||||
# ImageNet normalize on GPU
|
||||
img = (img - _IMAGENET_MEAN) / _IMAGENET_STD
|
||||
|
||||
# Concat and add batch dim
|
||||
combined = mx.concatenate([img, mask], axis=-1) # (H, W, 4)
|
||||
return mx.expand_dims(combined, axis=0) # (1, H, W, 4)
|
||||
@ -6,6 +6,7 @@ All internal operations use NHWC layout.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import gc
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import mlx.core as mx
|
||||
@ -13,7 +14,7 @@ import mlx.nn as nn
|
||||
from safetensors import safe_open
|
||||
|
||||
from corridorkey_mlx.model.backbone import HieraBackbone
|
||||
from corridorkey_mlx.model.decoder import DecoderHead
|
||||
from corridorkey_mlx.model.decoder import DecoderHead, FusedDecoderPair
|
||||
from corridorkey_mlx.model.hiera import ENCODER_KEY_PREFIX, _interpolate_pos_embed, _prod
|
||||
from corridorkey_mlx.model.refiner import CNNRefinerModule
|
||||
|
||||
@ -31,13 +32,29 @@ class GreenFormer(nn.Module):
|
||||
Output: dict with coarse/final alpha and foreground maps in NHWC.
|
||||
"""
|
||||
|
||||
def __init__(self, img_size: int = 512) -> None:
|
||||
def __init__(
|
||||
self,
|
||||
img_size: int = 512,
|
||||
dtype: mx.Dtype = mx.float32,
|
||||
fused_decode: bool = False,
|
||||
slim: bool = False,
|
||||
use_sdpa: bool = True,
|
||||
stage_gc: bool = True,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
self.backbone = HieraBackbone(img_size=img_size)
|
||||
self._compute_dtype = dtype
|
||||
self._fused_decode = fused_decode
|
||||
self._slim = slim
|
||||
self._stage_gc = stage_gc
|
||||
self._compiled = False
|
||||
self.backbone = HieraBackbone(img_size=img_size, use_sdpa=use_sdpa)
|
||||
self.alpha_decoder = DecoderHead(BACKBONE_CHANNELS, EMBED_DIM, output_dim=1)
|
||||
self.fg_decoder = DecoderHead(BACKBONE_CHANNELS, EMBED_DIM, output_dim=3)
|
||||
self.refiner = CNNRefinerModule()
|
||||
|
||||
if fused_decode:
|
||||
self._fused_pair = FusedDecoderPair(self.alpha_decoder, self.fg_decoder)
|
||||
|
||||
# Decoder outputs at stride-4 (H/4, W/4); upsampler is always 4x
|
||||
self._logit_upsampler = nn.Upsample(
|
||||
scale_factor=(4.0, 4.0), mode="linear", align_corners=False
|
||||
@ -54,22 +71,48 @@ class GreenFormer(nn.Module):
|
||||
alpha_coarse, fg_coarse, delta_logits, alpha_final, fg_final.
|
||||
All tensors in NHWC format.
|
||||
"""
|
||||
# Backbone -> 4 multiscale feature maps in NHWC
|
||||
# Backbone always runs in fp32
|
||||
features = self.backbone(x)
|
||||
|
||||
# Materialize backbone output so MLX can free intermediate graph nodes.
|
||||
# NOTE: mx.eval is MLX array materialization, not Python eval()
|
||||
if self._stage_gc and not self._compiled:
|
||||
mx.eval(features) # noqa: S307
|
||||
gc.collect()
|
||||
mx.clear_cache()
|
||||
|
||||
# Cast features to compute dtype for decoders (bf16 saves memory)
|
||||
if self._compute_dtype != mx.float32:
|
||||
features = [f.astype(self._compute_dtype) for f in features]
|
||||
|
||||
# Decoder heads -> logits at H/4 resolution
|
||||
alpha_logits = self.alpha_decoder(features) # (B, H/4, W/4, 1)
|
||||
fg_logits = self.fg_decoder(features) # (B, H/4, W/4, 3)
|
||||
if self._fused_decode:
|
||||
alpha_logits, fg_logits = self._fused_pair(features)
|
||||
else:
|
||||
alpha_logits = self.alpha_decoder(features) # (B, H/4, W/4, 1)
|
||||
fg_logits = self.fg_decoder(features) # (B, H/4, W/4, 3)
|
||||
|
||||
# Upsample logits to full input resolution (4x from stride-4 decoder)
|
||||
alpha_logits_up = self._logit_upsampler(alpha_logits) # (B, H, W, 1)
|
||||
fg_logits_up = self._logit_upsampler(fg_logits) # (B, H, W, 3)
|
||||
|
||||
# Coarse predictions via sigmoid
|
||||
# Cast back to fp32 before sigmoid (precision at saturation boundaries)
|
||||
if self._compute_dtype != mx.float32:
|
||||
alpha_logits_up = alpha_logits_up.astype(mx.float32)
|
||||
fg_logits_up = fg_logits_up.astype(mx.float32)
|
||||
|
||||
# Coarse predictions via sigmoid (always fp32)
|
||||
alpha_coarse = mx.sigmoid(alpha_logits_up)
|
||||
fg_coarse = mx.sigmoid(fg_logits_up)
|
||||
|
||||
# Refiner: RGB + coarse predictions -> delta logits
|
||||
# Materialize decoder output so MLX can free decoder graph nodes.
|
||||
# NOTE: mx.eval is MLX array materialization, not Python eval()
|
||||
if self._stage_gc and not self._compiled:
|
||||
mx.eval(alpha_coarse, fg_coarse, alpha_logits_up, fg_logits_up) # noqa: S307
|
||||
gc.collect()
|
||||
mx.clear_cache()
|
||||
|
||||
# Refiner receives fp32 coarse predictions
|
||||
rgb = x[:, :, :, :3] # (B, H, W, 3)
|
||||
coarse_pred = mx.concatenate([alpha_coarse, fg_coarse], axis=-1) # (B, H, W, 4)
|
||||
delta_logits = self.refiner(rgb, coarse_pred) # (B, H, W, 4)
|
||||
@ -78,9 +121,17 @@ class GreenFormer(nn.Module):
|
||||
alpha_final = mx.sigmoid(alpha_logits_up + delta_logits[:, :, :, 0:1])
|
||||
fg_final = mx.sigmoid(fg_logits_up + delta_logits[:, :, :, 1:4])
|
||||
|
||||
if self._slim:
|
||||
return {
|
||||
"alpha_coarse": alpha_coarse,
|
||||
"fg_coarse": fg_coarse,
|
||||
"alpha_final": alpha_final,
|
||||
"fg_final": fg_final,
|
||||
}
|
||||
|
||||
return {
|
||||
"alpha_logits": alpha_logits,
|
||||
"fg_logits": fg_logits,
|
||||
"alpha_logits": alpha_logits.astype(mx.float32),
|
||||
"fg_logits": fg_logits.astype(mx.float32),
|
||||
"alpha_logits_up": alpha_logits_up,
|
||||
"fg_logits_up": fg_logits_up,
|
||||
"alpha_coarse": alpha_coarse,
|
||||
|
||||
@ -8,9 +8,14 @@ Architecture mirrors nikopueringer/CorridorKey DecoderHead (SegFormer-style).
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import mlx.core as mx
|
||||
import mlx.nn as nn
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from collections.abc import Sequence
|
||||
|
||||
|
||||
class MLP(nn.Module):
|
||||
"""Single linear projection: input_dim -> embed_dim."""
|
||||
@ -98,3 +103,69 @@ class DecoderHead(nn.Module):
|
||||
fused = nn.relu(fused)
|
||||
# No dropout at inference (eval mode)
|
||||
return self.classifier(fused)
|
||||
|
||||
def _project_features(self, features: list[mx.array]) -> list[mx.array]:
|
||||
"""Project features without upsampling or fusion."""
|
||||
linears = [self.linear_c1, self.linear_c2, self.linear_c3, self.linear_c4]
|
||||
projected = []
|
||||
for feat, linear in zip(features, linears, strict=True):
|
||||
b, h, w, _c = feat.shape
|
||||
x = feat.reshape(b, h * w, _c)
|
||||
x = linear(x)
|
||||
x = x.reshape(b, h, w, -1)
|
||||
projected.append(x)
|
||||
return projected
|
||||
|
||||
def _fuse_and_classify(self, projected: list[mx.array]) -> mx.array:
|
||||
"""Fuse pre-upsampled projections and classify."""
|
||||
fused = mx.concatenate(projected[::-1], axis=-1)
|
||||
fused = self.linear_fuse(fused)
|
||||
fused = self.bn(fused)
|
||||
fused = nn.relu(fused)
|
||||
return self.classifier(fused)
|
||||
|
||||
|
||||
class FusedDecoderPair(nn.Module):
|
||||
"""Runs two DecoderHeads with batched upsampling.
|
||||
|
||||
Concatenates alpha+fg projections along channel axis before each upsample,
|
||||
reducing 6 Metal dispatch calls to 3.
|
||||
"""
|
||||
|
||||
def __init__(self, alpha_head: DecoderHead, fg_head: DecoderHead) -> None:
|
||||
super().__init__()
|
||||
self.alpha_head = alpha_head
|
||||
self.fg_head = fg_head
|
||||
# Reuse alpha_head's pre-allocated upsamplers
|
||||
self._upsamplers: Sequence[nn.Upsample | None] = [
|
||||
None,
|
||||
alpha_head._upsampler_2x,
|
||||
alpha_head._upsampler_4x,
|
||||
alpha_head._upsampler_8x,
|
||||
]
|
||||
|
||||
def __call__(self, features: list[mx.array]) -> tuple[mx.array, mx.array]:
|
||||
"""Forward pass with batched upsampling.
|
||||
|
||||
Returns:
|
||||
(alpha_logits, fg_logits) both in NHWC.
|
||||
"""
|
||||
alpha_projs = self.alpha_head._project_features(features)
|
||||
fg_projs = self.fg_head._project_features(features)
|
||||
|
||||
alpha_up = []
|
||||
fg_up = []
|
||||
for a_proj, f_proj, up in zip(alpha_projs, fg_projs, self._upsamplers, strict=True):
|
||||
if up is not None:
|
||||
# Batch upsample: concat along channel axis (NHWC)
|
||||
fused = mx.concatenate([a_proj, f_proj], axis=-1)
|
||||
fused = up(fused)
|
||||
embed_dim = a_proj.shape[-1]
|
||||
a_proj = fused[:, :, :, :embed_dim]
|
||||
f_proj = fused[:, :, :, embed_dim:]
|
||||
alpha_up.append(a_proj)
|
||||
fg_up.append(f_proj)
|
||||
|
||||
alpha_logits = self.alpha_head._fuse_and_classify(alpha_up)
|
||||
fg_logits = self.fg_head._fuse_and_classify(fg_up)
|
||||
return alpha_logits, fg_logits
|
||||
|
||||
@ -251,6 +251,7 @@ class MaskUnitAttention(nn.Module):
|
||||
q_stride: int = 1,
|
||||
window_size: int = 0,
|
||||
use_mask_unit_attn: bool = False,
|
||||
use_sdpa: bool = True,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
self.dim_out = dim_out
|
||||
@ -260,6 +261,7 @@ class MaskUnitAttention(nn.Module):
|
||||
self.scale = self.head_dim**-0.5
|
||||
self.window_size = window_size
|
||||
self.use_mask_unit_attn = use_mask_unit_attn
|
||||
self._use_sdpa = use_sdpa
|
||||
|
||||
self.qkv = nn.Linear(dim, 3 * dim_out)
|
||||
self.proj = nn.Linear(dim_out, dim_out)
|
||||
@ -284,10 +286,32 @@ class MaskUnitAttention(nn.Module):
|
||||
q = q.reshape(batch_size, self.heads, num_windows, self.q_stride, -1, self.head_dim)
|
||||
q = mx.max(q, axis=3)
|
||||
|
||||
# Scaled dot-product attention (manual implementation)
|
||||
attn = (q * self.scale) @ mx.transpose(k, axes=(0, 1, 2, 4, 3))
|
||||
attn = mx.softmax(attn, axis=-1)
|
||||
x = attn @ v
|
||||
if self._use_sdpa:
|
||||
# Fold window dim into batch for SDPA (expects 4D).
|
||||
# Q/K/V are (B, heads, windows, tokens, head_dim) — transpose windows
|
||||
# next to batch before reshape so (B, W, H, T, D) -> (B*W, H, T, D).
|
||||
q_tokens = q.shape[3]
|
||||
kv_tokens = k.shape[3]
|
||||
q_4d = mx.transpose(q, axes=(0, 2, 1, 3, 4)).reshape(
|
||||
batch_size * num_windows, self.heads, q_tokens, self.head_dim
|
||||
)
|
||||
k_4d = mx.transpose(k, axes=(0, 2, 1, 3, 4)).reshape(
|
||||
batch_size * num_windows, self.heads, kv_tokens, self.head_dim
|
||||
)
|
||||
v_4d = mx.transpose(v, axes=(0, 2, 1, 3, 4)).reshape(
|
||||
batch_size * num_windows, self.heads, kv_tokens, self.head_dim
|
||||
)
|
||||
|
||||
x = mx.fast.scaled_dot_product_attention(q_4d, k_4d, v_4d, scale=self.scale)
|
||||
|
||||
# Unfold: (B*W, H, T, D) -> (B, W, H, T, D) -> (B, H, W, T, D)
|
||||
x = x.reshape(batch_size, num_windows, self.heads, q_tokens, self.head_dim)
|
||||
x = mx.transpose(x, axes=(0, 2, 1, 3, 4))
|
||||
else:
|
||||
# Manual scaled dot-product attention (original implementation)
|
||||
attn = (q * self.scale) @ mx.transpose(k, axes=(0, 1, 2, 4, 3))
|
||||
attn = mx.softmax(attn, axis=-1)
|
||||
x = attn @ v
|
||||
|
||||
# [B, heads, num_windows, tokens, head_dim] -> [B, tokens, num_windows, heads, head_dim]
|
||||
# -> [B, N', dim_out] (matches PyTorch transpose(1, 3))
|
||||
@ -312,6 +336,7 @@ class HieraBlock(nn.Module):
|
||||
q_stride: int = 1,
|
||||
window_size: int = 0,
|
||||
use_mask_unit_attn: bool = False,
|
||||
use_sdpa: bool = True,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
self.dim = dim
|
||||
@ -324,7 +349,7 @@ class HieraBlock(nn.Module):
|
||||
self.proj = nn.Linear(dim, dim_out)
|
||||
|
||||
self.attn = MaskUnitAttention(
|
||||
dim, dim_out, heads, q_stride, window_size, use_mask_unit_attn
|
||||
dim, dim_out, heads, q_stride, window_size, use_mask_unit_attn, use_sdpa=use_sdpa
|
||||
)
|
||||
|
||||
self.norm2 = nn.LayerNorm(dim_out)
|
||||
@ -353,7 +378,7 @@ class HieraBackbone(nn.Module):
|
||||
Hardcoded for hiera_base_plus_224 config with 4-channel input.
|
||||
"""
|
||||
|
||||
def __init__(self, img_size: int = 512) -> None:
|
||||
def __init__(self, img_size: int = 512, use_sdpa: bool = True) -> None:
|
||||
super().__init__()
|
||||
self.img_size = img_size
|
||||
|
||||
@ -418,6 +443,7 @@ class HieraBackbone(nn.Module):
|
||||
q_stride=(flat_q_stride if i in q_pool_blocks else 1),
|
||||
window_size=flat_mu_size,
|
||||
use_mask_unit_attn=use_mu_attn,
|
||||
use_sdpa=use_sdpa,
|
||||
)
|
||||
self.blocks.append(block)
|
||||
embed_dim = dim_out
|
||||
|
||||
@ -116,9 +116,7 @@ def list_assets(release: dict[str, Any]) -> list[dict[str, Any]]:
|
||||
return release.get("assets", [])
|
||||
|
||||
|
||||
def find_asset(
|
||||
release: dict[str, Any], asset_name: str | None = None
|
||||
) -> dict[str, Any]:
|
||||
def find_asset(release: dict[str, Any], asset_name: str | None = None) -> dict[str, Any]:
|
||||
"""Find a specific asset in a release, or infer from platform."""
|
||||
assets = list_assets(release)
|
||||
target = asset_name or default_asset_name()
|
||||
@ -275,9 +273,7 @@ def download_asset(
|
||||
)
|
||||
_console.print("[green]Checksum OK[/green]")
|
||||
else:
|
||||
_console.print(
|
||||
"[yellow]No checksum file found — skipping verification[/yellow]"
|
||||
)
|
||||
_console.print("[yellow]No checksum file found — skipping verification[/yellow]")
|
||||
|
||||
# atomic rename
|
||||
partial_path.rename(final_path)
|
||||
|
||||
@ -5,6 +5,7 @@ from __future__ import annotations
|
||||
import mlx.core as mx
|
||||
import pytest
|
||||
|
||||
from corridorkey_mlx.inference.pipeline import compile_model
|
||||
from corridorkey_mlx.model.corridorkey import GreenFormer
|
||||
|
||||
IMG_SIZE = 256
|
||||
@ -32,11 +33,11 @@ def dummy_input() -> mx.array:
|
||||
def test_compiled_matches_eager(model: GreenFormer, dummy_input: mx.array) -> None:
|
||||
"""Fixed-shape compiled output matches eager output within tolerance."""
|
||||
eager_out = model(dummy_input)
|
||||
mx.eval(eager_out) # noqa: S307
|
||||
mx.eval(eager_out) # noqa: S307 # NOTE: MLX materialization, not Python eval
|
||||
|
||||
compiled_fn = mx.compile(model.__call__)
|
||||
compiled_out = compiled_fn(dummy_input)
|
||||
mx.eval(compiled_out) # noqa: S307
|
||||
compiled_model = compile_model(model)
|
||||
compiled_out = compiled_model(dummy_input)
|
||||
mx.eval(compiled_out) # noqa: S307 # NOTE: MLX materialization
|
||||
|
||||
for key in OUTPUT_KEYS:
|
||||
diff = float(mx.max(mx.abs(eager_out[key] - compiled_out[key])))
|
||||
@ -45,12 +46,12 @@ def test_compiled_matches_eager(model: GreenFormer, dummy_input: mx.array) -> No
|
||||
|
||||
def test_compiled_deterministic(model: GreenFormer, dummy_input: mx.array) -> None:
|
||||
"""Compiled model produces identical results across consecutive calls."""
|
||||
compiled_fn = mx.compile(model.__call__)
|
||||
compiled_model = compile_model(model)
|
||||
|
||||
out1 = compiled_fn(dummy_input)
|
||||
mx.eval(out1) # noqa: S307
|
||||
out2 = compiled_fn(dummy_input)
|
||||
mx.eval(out2) # noqa: S307
|
||||
out1 = compiled_model(dummy_input)
|
||||
mx.eval(out1) # noqa: S307 # NOTE: MLX materialization
|
||||
out2 = compiled_model(dummy_input)
|
||||
mx.eval(out2) # noqa: S307 # NOTE: MLX materialization
|
||||
|
||||
for key in OUTPUT_KEYS:
|
||||
diff = float(mx.max(mx.abs(out1[key] - out2[key])))
|
||||
|
||||
@ -23,7 +23,7 @@ EXPECTED_KEY_COUNT = 365
|
||||
|
||||
def _torch_available() -> bool:
|
||||
try:
|
||||
import torch # noqa: F401
|
||||
import torch # noqa: F401 # type: ignore[import-not-found]
|
||||
|
||||
return True
|
||||
except ImportError:
|
||||
|
||||
@ -192,3 +192,147 @@ def test_no_nan_inf() -> None:
|
||||
arr = np.array(out[key])
|
||||
assert not np.isnan(arr).any(), f"{key} contains NaN"
|
||||
assert not np.isinf(arr).any(), f"{key} contains Inf"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# bf16 mixed precision contract
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def bf16_model_output() -> dict[str, mx.array]:
|
||||
"""Forward pass with bf16 compute dtype, random weights."""
|
||||
model = GreenFormer(img_size=SMALL_IMG_SIZE, dtype=mx.bfloat16)
|
||||
mx.random.seed(42)
|
||||
x = mx.random.normal((1, SMALL_IMG_SIZE, SMALL_IMG_SIZE, 4))
|
||||
# mx.eval is MLX array materialization, not Python eval()
|
||||
mx.eval(x) # noqa: S307
|
||||
out = model(x)
|
||||
mx.eval(out) # noqa: S307
|
||||
return out
|
||||
|
||||
|
||||
def test_bf16_output_dtype_always_fp32(bf16_model_output: dict[str, mx.array]) -> None:
|
||||
"""All outputs must be fp32 regardless of compute dtype."""
|
||||
for key, arr in bf16_model_output.items():
|
||||
assert arr.dtype == mx.float32, f"{key}: expected float32, got {arr.dtype}"
|
||||
|
||||
|
||||
def test_bf16_sigmoid_outputs_in_range(bf16_model_output: dict[str, mx.array]) -> None:
|
||||
"""Post-sigmoid outputs must be in [0, 1] even with bf16 compute."""
|
||||
for key in ("alpha_coarse", "fg_coarse", "alpha_final", "fg_final"):
|
||||
arr = bf16_model_output[key]
|
||||
assert float(mx.min(arr)) >= 0.0, f"{key} has values < 0"
|
||||
assert float(mx.max(arr)) <= 1.0, f"{key} has values > 1"
|
||||
|
||||
|
||||
def test_bf16_no_nan_inf(bf16_model_output: dict[str, mx.array]) -> None:
|
||||
"""bf16 forward produces finite outputs."""
|
||||
for key in ("alpha_final", "fg_final"):
|
||||
arr = np.array(bf16_model_output[key])
|
||||
assert not np.isnan(arr).any(), f"{key} contains NaN"
|
||||
assert not np.isinf(arr).any(), f"{key} contains Inf"
|
||||
|
||||
|
||||
def test_fused_decode_matches_unfused() -> None:
|
||||
"""Fused decoder pair produces bit-exact output vs independent decoders."""
|
||||
from mlx.utils import tree_flatten
|
||||
|
||||
mx.random.seed(77)
|
||||
x = mx.random.normal((1, SMALL_IMG_SIZE, SMALL_IMG_SIZE, 4))
|
||||
# mx.eval is MLX array materialization, not Python eval()
|
||||
mx.eval(x) # noqa: S307
|
||||
|
||||
model_unfused = GreenFormer(img_size=SMALL_IMG_SIZE, fused_decode=False)
|
||||
mx.eval(model_unfused.parameters()) # noqa: S307
|
||||
|
||||
model_fused = GreenFormer(img_size=SMALL_IMG_SIZE, fused_decode=True)
|
||||
model_fused.load_weights(tree_flatten(model_unfused.parameters())) # type: ignore[arg-type]
|
||||
mx.eval(model_fused.parameters()) # noqa: S307
|
||||
|
||||
out_u = model_unfused(x)
|
||||
out_f = model_fused(x)
|
||||
mx.eval(out_u) # noqa: S307
|
||||
mx.eval(out_f) # noqa: S307
|
||||
|
||||
for key in out_u:
|
||||
np.testing.assert_array_equal(
|
||||
np.array(out_u[key]),
|
||||
np.array(out_f[key]),
|
||||
err_msg=f"{key} differs between fused and unfused decode",
|
||||
)
|
||||
|
||||
|
||||
def test_fp32_default_unchanged() -> None:
|
||||
"""GreenFormer(dtype=mx.float32) behaves identically to GreenFormer()."""
|
||||
mx.random.seed(99)
|
||||
x = mx.random.normal((1, SMALL_IMG_SIZE, SMALL_IMG_SIZE, 4))
|
||||
# mx.eval is MLX array materialization, not Python eval()
|
||||
mx.eval(x) # noqa: S307
|
||||
|
||||
model_default = GreenFormer(img_size=SMALL_IMG_SIZE)
|
||||
model_explicit = GreenFormer(img_size=SMALL_IMG_SIZE, dtype=mx.float32)
|
||||
|
||||
# Same weights via flattened tree
|
||||
from mlx.utils import tree_flatten
|
||||
|
||||
model_explicit.load_weights(tree_flatten(model_default.parameters())) # type: ignore[arg-type]
|
||||
mx.eval(model_explicit.parameters()) # noqa: S307
|
||||
|
||||
out_default = model_default(x)
|
||||
out_explicit = model_explicit(x)
|
||||
mx.eval(out_default) # noqa: S307
|
||||
mx.eval(out_explicit) # noqa: S307
|
||||
|
||||
for key in out_default:
|
||||
np.testing.assert_array_equal(
|
||||
np.array(out_default[key]),
|
||||
np.array(out_explicit[key]),
|
||||
err_msg=f"{key} differs between default and explicit fp32",
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Slim forward mode
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
SLIM_KEYS = {"alpha_coarse", "fg_coarse", "alpha_final", "fg_final"}
|
||||
|
||||
|
||||
def test_slim_returns_four_keys() -> None:
|
||||
"""slim=True returns only the 4 final/coarse keys."""
|
||||
model = GreenFormer(img_size=SMALL_IMG_SIZE, slim=True)
|
||||
x = mx.random.normal((1, SMALL_IMG_SIZE, SMALL_IMG_SIZE, 4))
|
||||
out = model(x)
|
||||
# NOTE: mx.eval is MLX array materialization, not Python eval()
|
||||
mx.eval(out) # noqa: S307
|
||||
assert set(out.keys()) == SLIM_KEYS
|
||||
|
||||
|
||||
def test_slim_matches_full_output() -> None:
|
||||
"""slim=True values match the corresponding keys from full output."""
|
||||
from mlx.utils import tree_flatten
|
||||
|
||||
mx.random.seed(55)
|
||||
x = mx.random.normal((1, SMALL_IMG_SIZE, SMALL_IMG_SIZE, 4))
|
||||
# NOTE: mx.eval is MLX array materialization, not Python eval()
|
||||
mx.eval(x) # noqa: S307
|
||||
|
||||
model_full = GreenFormer(img_size=SMALL_IMG_SIZE, slim=False)
|
||||
mx.eval(model_full.parameters()) # noqa: S307
|
||||
|
||||
model_slim = GreenFormer(img_size=SMALL_IMG_SIZE, slim=True)
|
||||
model_slim.load_weights(tree_flatten(model_full.parameters())) # type: ignore[arg-type]
|
||||
mx.eval(model_slim.parameters()) # noqa: S307
|
||||
|
||||
out_full = model_full(x)
|
||||
out_slim = model_slim(x)
|
||||
mx.eval(out_full) # noqa: S307
|
||||
mx.eval(out_slim) # noqa: S307
|
||||
|
||||
for key in SLIM_KEYS:
|
||||
np.testing.assert_array_equal(
|
||||
np.array(out_full[key]),
|
||||
np.array(out_slim[key]),
|
||||
err_msg=f"{key} differs between slim and full forward",
|
||||
)
|
||||
|
||||
Loading…
Reference in New Issue
Block a user