mlx-knife/docs/ADR/ADR-016-Memory-Aware-Model-Loading.md

# ADR-016: Memory-Aware Model Loading

**Status:** Accepted (Phase 1-2 Complete)
**Created:** 2025-12-05
**Context:** Vision models crash with Metal OOM without warning

## Problem

Vision models (e.g., Llama-3.2-90B-Vision-4bit) cause hard Metal crashes when they exceed available unified memory. The crash happens deep in the Metal layer with no graceful error handling:

```
[METAL] Command buffer execution failed: Insufficient Memory
Abort trap: 6
```

Additionally, leaked semaphore warnings appear due to ungraceful shutdown:
```
UserWarning: resource_tracker: There appear to be 1 leaked semaphore objects to clean up at shutdown
```

### Key Observation

Vision models have **additional overhead** beyond their weight size:
- Vision Encoder (ViT): ~2-5 GB
- Projection Layer: ~0.5-1 GB
- Image Tensors
- Larger KV-Cache (image tokens)

A 49.9GB Vision model on a 68.7GB system (73%) crashed, while similar-sized text-only models run with just a warning.

## Decision

### JSON-API 0.1.6

Add `system.memory_total_bytes` to API responses:

```json
{
  "system": {
    "memory_total_bytes": 68719476736
  },
  "models": [...]
}
```

This is a **hardware fact** (from `sysctl -n hw.memsize`), not a heuristic.

### Memory Thresholds (mlx-knife internal)

| Model Type | >70% of total | Action |
|------------|---------------|--------|
| Vision     | **ERROR + abort** | CLI: stderr error, exit 1<br>Server: HTTP 507 Insufficient Storage |
| Text-only  | **No user-facing action** | Server: `logger.warning()` only (internal) |

**Rationale:**

**Vision (ERROR + abort):**
- **Empirically confirmed:** Llama-3.2-90B-Vision-4bit @ 73% → Metal OOM crash + semaphore leak
- Vision Encoder cannot swap (Metal GPU operations)
- Additional overhead (~2-5GB) beyond weight size
- Hard limit necessary to prevent system crash
- **User-facing error justified:** Prevents crash

**Text (no user-facing action):**
- **Empirically confirmed:** Qwen2.5-Coder-32B-bf16 @ 95-97% → swap (10-20 tok/6h), no crash
- Text inference swaps gracefully (extremely slow but stable)
- No crash risk → no crash prevention needed
- **CLI:** No warning (backwards compatible, user notices slowdown anyway)
- **Server:** Internal log only (operator awareness, no client-facing change)

**70% threshold chosen because:**
- mlx-lm's own WARNING triggers at ~75% ("49152 MB maximum recommended" on 64GB)
- Provides safety margin for OS + KV-Cache + activations
- Vision models crash above this threshold, text models do not

### Implementation Location

- `run.py`: Pre-load check before `VisionRunner` or `MLXRunner` instantiation
- Uses `size_bytes` from `build_model_object()` (already available)
- Compares against `memory_total_bytes * 0.70`

**Messages:**

**CLI:**
- Vision >70%: `[ERROR] Model size (XX.X GB) exceeds 70% of system memory (YY.Y GB). Vision models crash with Metal OOM due to Vision Encoder overhead. Aborting.` → stderr, exit 1
- Text >70%: No user-facing message (backwards compatible)

**Server:**
- Vision >70%: HTTP 507 Insufficient Storage + JSON error response
- Text >70%: `logger.warning("Model size XX.X GB exceeds 70% of YY.Y GB system memory. Expect extreme slowness due to swapping.")` → visible via `--log-level warning` (default) and `--log-json` if enabled

## Status

**Phase 1+2:** ✅ Complete (2.0.4-beta.1) - See CHANGELOG.md

**Phase 2b:** ✅ Complete (2.0.4-beta.9) - Model Switching Memory Gate

**Phase 3 (Future):** Issue #46

### Phase 3a: User-Configurable Limits
- [ ] `MLXK_MAX_IMAGES` env var
- [ ] `--max-images N` CLI flag
- [ ] `MLXK_MEMORY_THRESHOLD` env var (override 70%)

### Phase 3b: Benchmark-Based Heuristics

**Approach:** Empirische Daten statt theoretischer Berechnung.

**Daten sammeln:**
```
Für jedes Vision-Modell:
  - Config: image_size, patch_size, hidden_size, vision_config
  - Test: 1, 2, 4, 8, 16 Images (verschiedene Größen)
  - Messen: Peak Memory (memmon.py), OOM-Punkt
  - Hardware: M1/M2/M3, 16/32/64/96 GB
```

**Korrelation finden:**
```
Peak_Memory ≈ f(image_size², patch_size, num_images, model_size)
```

**Heuristik-Formel ableiten:**
```python
def estimate_vision_memory(config, num_images):
    base = model_size_bytes
    image_size = config.get("vision_config", {}).get("image_size", 384)
    patch_size = config.get("vision_config", {}).get("patch_size", 14)
    hidden_size = config.get("hidden_size", 4096)

    # Empirisch kalibrierte Konstanten aus Benchmarks
    patches_per_image = (image_size / patch_size) ** 2
    per_image_overhead = patches_per_image * hidden_size * BYTES_PER_ACTIVATION

    return base + (per_image_overhead * num_images * SAFETY_FACTOR)
```

**Infrastruktur:** `memmon.py`, `memplot.py` aus beta.9 existieren bereits.

**Komplexität bei Multimodal:**
- Dynamic Tiling (Qwen2-VL, MiMo-VL) → variable patches pro Bild
- Audio + Vision (Gemma-3n) → zusätzlicher Audio Encoder
- Mixed Input → Kombinations-Overhead

**Pragmatischer Fallback:** Wenn Heuristik unsicher → User-Limit (Phase 3a) verwenden.

---

## Phase 2b: Model Switching Memory Gate

**Problem:** Metal GPU cache is released asynchronously. During model switching, the new model may start loading before memory from the old model is actually freed → OOM / "Broken pipe" crashes.

**Root Cause Analysis:**
- `mx.metal.clear_cache()` releases the cache, but **asynchronously**
- macOS needs time to return memory to the system
- Pre-load check (Phase 1-2) validates `model_size / total_memory` → looks OK
- But **available memory** is still occupied by the previous model

**Solution: Active Polling for Available Memory**

```python
def _wait_for_memory_release(required_bytes, timeout_seconds=10.0):
    """Wait for memory to be released after model unload."""
    while time.time() - start < timeout_seconds:
        available = _get_available_memory_bytes()  # free + speculative
        if available >= required_bytes:
            return True
        time.sleep(0.5)
    return False  # Timeout - continue with warning
```

**Thresholds:**

| Context | Min Available | Timeout | Rationale |
|---------|---------------|---------|-----------|
| Vision Model Switch | 20 GB | 10s | Pixtral-8bit = 13.5 GB + overhead |
| Audio Model Switch | 10 GB | 10s | Whisper ~1.5 GB, Voxtral ~10 GB |
| Test Infrastructure | 20 GB | 15s | Between-test cleanup, larger buffer |

**Implementation:**

| Location | Function |
|----------|----------|
| `server_base.py` | `_get_available_memory_bytes()`, `_wait_for_memory_release()` |
| `server_base.py` | `get_or_load_model()` - 20 GB gate after cleanup |
| `server_base.py` | `get_or_load_audio_model()` - 10 GB gate after cleanup |
| `server_context.py` | `LocalServer` cleanup - 20 GB gate between tests |

**Key Difference from Phase 1-2:**

| Aspect | Phase 1-2 (Pre-load) | Phase 2b (Model Switch) |
|--------|---------------------|------------------------|
| Measures | `total_memory` | `available_memory` |
| Timing | Before first load | After unload, before new load |
| Method | Static check | Active polling with timeout |
| Failure | HTTP 507 (hard block) | Warning + continue (soft) |

**Behavior on Timeout:**
- Log warning with actual available memory
- Continue anyway (probe/policy check will catch real OOM)
- Prevents indefinite blocking on edge cases

## Empirical Data

| Model | Size | System | % Used | Result |
|-------|------|--------|--------|--------|
| **Vision Models** |
| Llama-3.2-90B-Vision-4bit | 49.9GB | 68.7GB | 73% | **CRASH** (Metal OOM + semaphore leak) |
| Llama-3.2-11B-Vision-4bit | 6.0GB | 68.7GB | 9% | OK |
| pixtral-12b-8bit | 13.5GB | 68.7GB | 20% | OK |
| **Text Models** |
| Qwen2.5-Coder-32B-Instruct-bf16 | 61-62GB | 64GB | 95-97% | **Swap, no crash** (10-20 tokens in 6h, Ctrl+C) |

**Key Finding:** Text models swap gracefully (extremely slow but no crash), while Vision models hard-crash at ~73% due to additional Vision Encoder overhead.

## References

- [mlx-vlm Issue #100: High Memory Usage](https://github.com/Blaizzy/mlx-vlm/issues/100)
- [LLaMA 3.2 90B VRAM Requirements](https://blogs.novita.ai/llama-3-2-90b-vram/)
- mlx-lm WARNING message: "maximum recommended size of 49152 MB"