Metadata-Version: 2.4
Name: ml-memguard
Version: 0.5.0
Summary: Stop OOM crashes in vLLM, SGLang, Unsloth, and HuggingFace. Proactive memory estimation and runtime KV-cache monitoring for LLM inference serving and fine-tuning on Apple Silicon, CUDA, and CPU.
Project-URL: Homepage, https://github.com/vgpprasad91/ml-memguard
Project-URL: Repository, https://github.com/vgpprasad91/ml-memguard
Project-URL: Issues, https://github.com/vgpprasad91/ml-memguard/issues
Project-URL: Changelog, https://github.com/vgpprasad91/ml-memguard/blob/main/CHANGELOG.md
Author-email: Guru Prasad Venkataraghavan <vguruprasad91@gmail.com>
License-Expression: Apache-2.0
License-File: LICENSE
Keywords: apple-silicon,cuda,fine-tuning,gpu-memory,inference,inference-server,kv-cache,llm,llm-serving,lora,machine-learning,memory,moe,oom,oom-prevention,serving,sglang,training,vllm
Classifier: Development Status :: 4 - Beta
Classifier: Intended Audience :: Developers
Classifier: Intended Audience :: Science/Research
Classifier: License :: OSI Approved :: Apache Software License
Classifier: Operating System :: MacOS
Classifier: Operating System :: Microsoft :: Windows
Classifier: Operating System :: POSIX :: Linux
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.9
Classifier: Programming Language :: Python :: 3.10
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Programming Language :: Python :: 3.13
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Classifier: Typing :: Typed
Requires-Python: >=3.9
Provides-Extra: all
Requires-Dist: httpx>=0.24; extra == 'all'
Requires-Dist: mlx>=0.10; extra == 'all'
Requires-Dist: psutil>=5.9; extra == 'all'
Requires-Dist: sglang>=0.3; extra == 'all'
Requires-Dist: torch>=2.0; extra == 'all'
Requires-Dist: transformers>=4.30; extra == 'all'
Requires-Dist: vllm>=0.4; extra == 'all'
Provides-Extra: apple
Requires-Dist: mlx>=0.10; extra == 'apple'
Provides-Extra: cloud
Requires-Dist: httpx>=0.24; extra == 'cloud'
Provides-Extra: cuda
Requires-Dist: torch>=2.0; extra == 'cuda'
Provides-Extra: dev
Requires-Dist: mypy>=1.0; extra == 'dev'
Requires-Dist: pytest-cov>=4.0; extra == 'dev'
Requires-Dist: pytest>=7.0; extra == 'dev'
Requires-Dist: ruff>=0.1; extra == 'dev'
Provides-Extra: hf
Requires-Dist: accelerate>=1.1.0; extra == 'hf'
Requires-Dist: torch>=2.0; extra == 'hf'
Requires-Dist: transformers>=4.30; extra == 'hf'
Provides-Extra: sglang
Requires-Dist: sglang>=0.3; extra == 'sglang'
Provides-Extra: unsloth
Requires-Dist: accelerate>=1.1.0; extra == 'unsloth'
Requires-Dist: torch>=2.0; extra == 'unsloth'
Requires-Dist: transformers>=4.30; extra == 'unsloth'
Requires-Dist: unsloth; extra == 'unsloth'
Provides-Extra: vllm
Requires-Dist: vllm>=0.4; extra == 'vllm'
Description-Content-Type: text/markdown

# memory-guard

**Stop `No available memory for cache blocks` in vLLM. Stop `CUDA out of memory` in Unsloth. Stop frozen Macs in mlx_lm. Works across inference serving and fine-tuning — and learns optimal configs from experience.**

No more `gpu_memory_utilization` trial-and-error. No more KV cache crashes at 3 AM. No more wasted GPU-hours on jobs that OOM in the first minute.

[![PyPI version](https://img.shields.io/pypi/v/ml-memguard.svg)](https://pypi.org/project/ml-memguard/)
[![License](https://img.shields.io/badge/license-Apache%202.0-blue.svg)](LICENSE)
[![Python](https://img.shields.io/pypi/pyversions/ml-memguard.svg)](https://pypi.org/project/ml-memguard/)

**→ [How much is your team wasting on OOM crashes? Find your number.](https://memguard-calculator.pages.dev)**

```bash
pip install ml-memguard                  # core (zero dependencies)
pip install ml-memguard[hf]             # + HuggingFace Transformers adapter
pip install ml-memguard[unsloth]        # + Unsloth adapter
pip install ml-memguard[apple]          # + MLX Metal ground-truth monitoring
pip install ml-memguard[cuda]           # + CUDA OOM recovery
pip install ml-memguard[vllm]           # + vLLM inference serving adapter
pip install ml-memguard[sglang]         # + SGLang inference serving adapter
```

## The Problem

### Inference serving (vLLM / SGLang / Ollama)

`No available memory for the cache blocks. Try increasing gpu_memory_utilization` is vLLM's most-filed error. It appears when `--gpu-memory-utilization`, `--max-num-seqs`, and `--max-num-batched-tokens` are misconfigured — which they almost always are on first deploy. There is no formula; the official advice is to tune until it stops crashing. When it does crash mid-serving, it takes live user traffic down with it.

- vLLM has 30+ distinct numbered OOM issues. SGLang has a dedicated OOM tracking issue. Ollama makes the host unresponsive.
- KV cache grows linearly with context length × batch size × layers. A 128k-context Llama 3 70B needs ~40 GB of KV cache on top of ~140 GB for weights.
- There is no built-in tool that tells you the right `max_num_seqs` before you launch.

### Fine-tuning (Unsloth / HuggingFace / mlx_lm)

- **Apple Silicon**: No OOM exception exists. When you exceed memory, macOS silently swaps to disk, your Mac freezes for minutes, and eventually the OS kills your process.
- **CUDA**: `torch.cuda.OutOfMemoryError` crashes your training run. You restart, guess a smaller batch size, and pray.
- **Containers**: cgroups silently kill your process with no warning when you hit the memory limit.

Existing solutions (PyTorch Lightning BatchSizeFinder, HuggingFace `accelerate`) are CUDA-only and reactive — they catch OOM exceptions that don't exist on Apple Silicon and do nothing for inference servers.

## The Solution

`memory-guard` is **proactive, not reactive**. For inference: it calculates the safe `max_num_seqs` before you launch the server and monitors KV cache utilization at runtime. For fine-tuning: it estimates peak memory before training starts and auto-adjusts batch size, LoRA rank, and sequence length to fit. Both paths use an RL optimizer that learns your specific device over time.

**Inference serving (vLLM)**

```python
from memory_guard import guard_vllm
from vllm import LLM

llm = LLM(model="meta-llama/Llama-3.1-8B-Instruct", gpu_memory_utilization=0.9)

safe = guard_vllm(llm)
# InferenceSafeConfig:
#   max_num_seqs:     32         <- largest concurrent batch that fits
#   max_seq_len:      4096
#   estimated memory: 18,240 MB
#   budget:           19,456 MB

# Wire the KV cache monitor — fires on_shed_load at 92% utilization
safe.monitor.on_shed_load = lambda u: load_balancer.reduce_weight("primary", 0)
safe.monitor.on_warning   = lambda u: logger.warning("KV cache at %.0f%%", u * 100)

with safe.monitor.session():
    server.serve_forever()   # monitor runs in background thread
```

**Fine-tuning (Unsloth / HuggingFace / mlx_lm)**

```python
from memory_guard import MemoryGuard

guard = MemoryGuard.auto()

# Pre-flight: estimate memory and auto-downgrade config
safe = guard.preflight(
    model_params=9_000_000_000,  # 9B parameter model
    model_bits=4,                # 4-bit quantized
    hidden_dim=4096,
    num_heads=32,
    num_layers=32,
    batch_size=4,
    seq_length=2048,
    lora_rank=32,
    lora_layers=16,
)

print(safe)
# SafeConfig (FITS):
#   batch_size:       2          <- auto-reduced from 4
#   grad_checkpoint:  True       <- auto-enabled
#   grad_accumulation:4          <- compensates for smaller batch
#   estimated memory: 3835 MB
#   budget:           4643 MB

# Runtime monitoring: polls memory pressure every 5s
with guard.monitor(safe.batch_size) as mon:
    for step in range(1000):
        # Batch size may decrease mid-training if pressure rises
        train_step(batch_size=mon.current_batch_size)
```

## Features

### Inference serving

| Feature | vLLM | SGLang | Ollama / custom |
|---------|:---:|:---:|:---:|
| Safe `max_num_seqs` pre-flight | Yes | Yes | Yes (via `preflight_inference`) |
| KV cache utilization monitoring | Yes | Yes | Yes (`KVCacheMonitor`) |
| Load-shed signal at 92% KV utilization | Yes | Yes | Yes |
| Warning signal at 80% KV utilization | Yes | Yes | Yes |
| RL optimizer (learns per device/model) | Yes | Yes | Yes |
| Architecture auto-introspection | Yes (hf_config) | Yes (server_args) | Manual |

### Fine-tuning

| Feature | Apple Silicon | CUDA | Linux CPU | Windows |
|---------|:---:|:---:|:---:|:---:|
| Proactive memory estimation | Yes | Yes | Yes | Yes |
| Auto-downgrade config | Yes | Yes | Yes | Yes |
| RL optimizer (learns per device) | Yes | Yes | Yes | Yes |
| Runtime pressure monitoring | Yes (Mach kernel + MLX Metal) | Yes (torch.cuda) | Yes (PSI, cgroups) | Yes (kernel32) |
| MLX Metal ground-truth | Yes (mx.metal.get_active_memory) | N/A | N/A | N/A |
| OOM catch & retry | N/A (no OOM on Metal) | Yes | N/A | N/A |
| Container-aware (cgroups v1/v2) | N/A | Yes | Yes | N/A |
| Auto-calibration | Yes | Yes | Yes | Yes |
| FlashAttention-aware | Yes | Yes | Yes | Yes |
| GQA / MoE / Multi-modal | Yes | Yes | Yes | Yes |

## How It Works

### Inference serving path

`preflight_inference()` computes the memory footprint of model weights + KV cache at a given `max_num_seqs` and `max_seq_len`, then binary-searches for the largest concurrent batch that fits within your GPU budget (default: 80% of available VRAM). The RL optimizer learns which `max_num_seqs` value worked well on your device and model over time, replacing binary search with a confident recommendation after a handful of runs.

`KVCacheMonitor` runs a background thread polling the live KV cache token counts from vLLM or SGLang. At 80% utilization it fires `on_warning`; at 92% it fires `on_shed_load`. Neither callback does anything by default — they are signals. Your load balancer or health endpoint decides what to do (reduce upstream weight, return 503, etc.). The engine is never mutated while serving.

### Fine-tuning path

### 1. Proactive Estimation

Calculates peak memory from model architecture, accounting for:
- Per-projection LoRA input buffers (Q, K, V, O)
- FlashAttention O(n) vs standard O(n^2) attention scores
- GQA-aware KV cache (uses `num_kv_heads`, not `num_heads`)
- MoE routing buffers and active expert activations
- Optimizer states (Adam 3x, SGD 2x, Adafactor 1.5x)
- MLX lazy evaluation discount (20% reduction on Apple Silicon)
- Framework overhead (25% proportional + 400MB fixed runtime cost)

With gradient checkpointing, activation memory drops to `sqrt(layers)`.

### 2. Auto-Downgrade (quality-preserving order)

When estimate exceeds budget (available × 80%):

1. Enable gradient checkpointing (free quality, ~40% activation savings)
2. Halve batch size (compensate with gradient accumulation)
3. Halve sequence length
4. Halve LoRA rank
5. Halve LoRA layers

### 3. Runtime Monitoring

Background thread polls memory pressure every 5 seconds:

- **Apple Silicon**: `mx.metal.get_active_memory()` (ground-truth from Metal allocator), with `kern.memorystatus_level` as fallback. Detects the monotonic memory growth pattern from [mlx-examples#1262](https://github.com/ml-explore/mlx-examples/issues/1262).
- **CUDA**: `torch.cuda.memory_allocated()` vs total VRAM
- **Linux**: `/proc/pressure/memory` (PSI), cgroup-aware (`memory.high` preferred over `memory.max`)
- **Windows**: `GlobalMemoryStatusEx`

When pressure exceeds 85%, batch size is halved mid-training.

### 4. Auto-Calibration

After each training run, the actual peak memory (from `mx.metal.get_peak_memory()` or `torch.cuda.max_memory_allocated()`) is recorded alongside the formula estimate. After 3+ runs, a median correction factor is applied to future estimates, narrowing the gap between predicted and actual memory usage over time.

### 5. RL Optimizer (v0.4)

A contextual bandit that learns which `(batch_size, lora_rank)` combination works best on your specific device and model. On cold start it falls back to the binary-search path from step 2. After a handful of runs it starts recommending configs it has learned are safe and efficient — and still falls back to binary search on the 5 % exploration floor so novel model architectures always get probed.

```python
guard = MemoryGuard.auto()   # loads ~/.memory-guard/rl_policy.json on disk

safe = guard.preflight(...)  # bandit recommends once it has learned; binary search until then

# ... training loop ...

guard.record_result(
    actual_peak_mb=get_peak_memory(),
    oom_occurred=False,          # set True if training crashed with OOM
)
# → updates the Q-table and saves the policy file atomically
```

The policy is a plain JSON file — human-readable, editable, and deletable.  See [`docs/rl_optimizer.md`](docs/rl_optimizer.md) for the full reference and [`docs/decisions/004-rl-contextual-bandit.md`](docs/decisions/004-rl-contextual-bandit.md) for the design rationale.

## Framework Integration

### With vLLM

```python
from memory_guard import guard_vllm
from vllm import LLM

llm = LLM(model="meta-llama/Llama-3.1-8B-Instruct", gpu_memory_utilization=0.9)
safe = guard_vllm(llm)

# safe.max_num_seqs  → pass to --max-num-seqs
# safe.monitor       → KVCacheMonitor, ready to start

safe.monitor.on_shed_load = lambda u: load_balancer.reduce_weight("primary", 0)
safe.monitor.on_warning   = lambda u: logger.warning("KV cache %.0f%%", u * 100)

with safe.monitor.session():
    server.serve_forever()
```

`guard_vllm` reads architecture directly from `model_config.hf_config` — no manual `hidden_dim` or `num_layers` required.

### With SGLang

```python
from memory_guard import guard_sglang
from sglang import Runtime

runtime = Runtime(model_path="meta-llama/Llama-3.1-8B-Instruct")
safe = guard_sglang(runtime)

safe.monitor.on_shed_load = lambda u: nginx.upstream_weight("primary", 0)

with safe.monitor.session():
    runtime.wait()
```

Polls `token_to_kv_pool` (preferred) or `scheduler.get_stats()` as fallback. Rolling-max smoothing suppresses false-recovery signals from RadixAttention prefix-cache evictions.

### With mlx_lm (Apple Silicon)

```python
import mlx.optimizers as optim
from memory_guard import MemoryGuard
from mlx_lm import load
from mlx_lm.tuner.trainer import train, TrainingArgs
from mlx_lm.tuner.utils import linear_to_lora_layers

guard = MemoryGuard.auto()
model, tokenizer = load("mlx-community/Qwen3.5-9B-MLX-4bit")

safe = guard.preflight(
    model_params=9e9, model_bits=4,
    hidden_dim=4096, num_heads=32, num_layers=32,
    batch_size=4, seq_length=2048,
    lora_rank=32, lora_layers=16,
)

model.freeze()
linear_to_lora_layers(
    model, safe.lora_layers,
    {"rank": safe.lora_rank, "scale": 20.0, "dropout": 0.0},
)
optimizer = optim.Adam(learning_rate=1e-4)

# The monitor runs in the background and logs if memory pressure rises.
# Note: mlx_lm's train() uses a fixed batch size. For dynamic adjustment,
# use a custom training loop that reads mon.current_batch_size each step.
with guard.monitor(safe.batch_size) as mon:
    train(
        model=model, optimizer=optimizer, train_dataset=train_set,
        args=TrainingArgs(
            batch_size=safe.batch_size,
            iters=1000,
            max_seq_length=safe.seq_length,
            grad_checkpoint=safe.grad_checkpoint,
            adapter_file="adapters.safetensors",
        ),
    )
```

### With HuggingFace Transformers (CUDA / CPU)

One call reads the model's architecture, runs preflight, patches `trainer.args`,
and attaches `MemoryGuardCallback` for mid-training batch-size downgrade.

```python
from memory_guard import guard_trainer
from transformers import AutoModelForCausalLM, Trainer, TrainingArguments

model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-3.1-8B")
trainer = Trainer(
    model=model,
    args=TrainingArguments(output_dir="./output", max_steps=1000),
    train_dataset=train_set,
)

guard_trainer(trainer)   # reads model, runs preflight, patches args + callback
trainer.train()
```

Pass `preflight_overrides` to lock in specific values the adapter can't infer
(e.g. `guard_trainer(trainer, batch_size=8, seq_length=4096, lora_rank=32)`).

### With Unsloth

Three lines — no manual architecture spelunking required.  `guard_unsloth_model`
introspects the loaded model, runs preflight, and returns a `SafeConfig` you
thread directly into `get_peft_model`.

```python
from memory_guard import guard_unsloth_model, guard_sft_trainer
from unsloth import FastLanguageModel
from trl import SFTTrainer

# 1. Load model (before LoRA)
model, tokenizer = FastLanguageModel.from_pretrained(
    "unsloth/Meta-Llama-3.1-8B-bnb-4bit",
    max_seq_length=2048,
    load_in_4bit=True,
)

# 2. Preflight — reads architecture automatically, auto-downgrades if needed
safe = guard_unsloth_model(model)   # ← the one line that replaces all the math

# 3. Attach LoRA using the safe values
model = FastLanguageModel.get_peft_model(
    model,
    r=safe.lora_rank,
    lora_alpha=safe.lora_rank * 2,
    max_seq_length=safe.seq_length,
    target_modules=["q_proj", "k_proj", "v_proj", "o_proj"],
)

# 4. Train with mid-training downgrade protection
trainer = SFTTrainer(model=model, tokenizer=tokenizer, ...)
guard_sft_trainer(trainer)   # patches trainer.args + adds MemoryGuardCallback
trainer.train()
```

> **BnB double-quantization**: Unsloth loads with `bnb_4bit_use_double_quant=True`
> by default.  memory-guard detects this and applies a 5 % correction to the
> weight-memory estimate.  Auto-calibration refines the correction after 3+
> training runs.

## Framework Adapters

*New in v0.2.0* — adapters read the model's architecture automatically so you
don't have to look up `hidden_size`, `num_heads`, or `num_layers`.
Inference serving adapters added in v0.3.0; RL optimizer integrated in v0.4.0.
Full reference: [`docs/adapters.md`](docs/adapters.md).

### How model introspection works

`introspect_model(model)` reads directly from `model.config` and
`model.parameters()` without importing torch or transformers at the call site:

| Field read | Source |
|---|---|
| `hidden_size` | `model.config.hidden_size` |
| `num_attention_heads` | `model.config.num_attention_heads` |
| `num_hidden_layers` | `model.config.num_hidden_layers` |
| `num_key_value_heads` | `model.config.num_key_value_heads` (falls back to `num_attention_heads` for MHA) |
| `model_bits` | `quantization_config.load_in_4bit / load_in_8bit`, else `model.dtype` (fp16/bf16 → 16, fp32 → 32) |
| `num_parameters` | `sum(p.numel() for p in model.parameters())` |

### When to pass `preflight_overrides`

Introspected values cover most cases.  Override when:

| Scenario | Pass |
|---|---|
| Specific training batch size | `batch_size=8` |
| Non-default sequence length | `seq_length=4096` |
| Fixed LoRA config | `lora_rank=32, lora_layers=24` |
| Model loaded at different precision | `model_bits=16` |

```python
# Override batch_size and lora_rank; everything else is introspected
guard_trainer(trainer, batch_size=8, lora_rank=32)
safe = guard_unsloth_model(model, seq_length=4096, lora_rank=16)
```

### QLoRA with BnB double-quantization

When Unsloth (or any HF model) loads with `bnb_4bit_use_double_quant=True`,
`guard_unsloth_model` automatically applies a 5 % correction to the weight-memory
estimate.  `model_bits` stays 4; only `num_parameters` is scaled down to proxy
the reduced quantization-constant footprint.  After 3+ runs, auto-calibration
refines the correction further.

## API Reference

### Inference Serving (v0.3+)

#### `guard.preflight_inference(...) -> InferenceSafeConfig`
Find the largest `max_num_seqs` that fits in your GPU budget. Binary-searches from your requested max down to 1; uses the RL optimizer if it has learned this device/model combination.

```python
safe = guard.preflight_inference(
    model_params=8e9, model_bits=4,
    hidden_dim=4096, num_kv_heads=8, num_layers=32,
    max_seq_len=4096, max_num_seqs=256,
)
# safe.max_num_seqs  → int, largest safe concurrent batch
# safe.monitor       → KVCacheMonitor (not yet started)
```

#### `guard_vllm(llm, ...) -> InferenceSafeConfig`
One call. Reads architecture from `model_config.hf_config`, runs preflight, returns `InferenceSafeConfig`. No manual config spelunking.

#### `guard_sglang(runtime, ...) -> InferenceSafeConfig`
Same as `guard_vllm` for SGLang. Reads `server_args.context_length` and `server_args.max_running_requests`.

#### `KVCacheMonitor`
Background-thread monitor. Fires `on_warning` at 80% KV utilization, `on_shed_load` at 92%. Start with `safe.monitor.session()` context manager or `safe.monitor.start()` / `safe.monitor.stop()` directly.

### Fine-Tuning

#### `MemoryGuard.auto(safety_ratio=0.80)`
Create with auto-detected platform. `safety_ratio` controls headroom (0.80 = use 80% of available).

#### `guard.preflight(**config) -> SafeConfig`
Estimate memory and auto-downgrade. Returns safe config.

#### `guard.monitor(batch_size) -> RuntimeMonitor`
Context manager for runtime monitoring. Use `mon.current_batch_size` in training loop.

#### `guard.estimate(**config) -> MemoryEstimate`
Pure estimation without auto-downgrade.

#### `estimate_training_memory(**config) -> MemoryEstimate`
Standalone estimation function.

#### `auto_downgrade(budget_mb, **config) -> DowngradeResult`
Standalone downgrade function.

#### `CUDAOOMRecovery(initial_batch_size)`
CUDA-specific OOM catch-and-retry wrapper.

### RL Optimizer (v0.4)

#### `guard.record_result(actual_peak_mb=None, oom_occurred=False, policy_update=True, model_name="")`
Call after each training run to update the calibration store and the RL
policy.  `actual_peak_mb` is auto-detected from MLX/CUDA if not supplied.
Set `oom_occurred=True` if the run ended with OOM — the policy learns a
negative reward and avoids that config in future.

#### `BanditPolicy.load(path=None) -> BanditPolicy`
Load the policy from disk (defaults to `~/.memory-guard/rl_policy.json`).
Returns a fresh cold-start policy silently if the file is absent or corrupt.

#### `BanditPolicy.q_value(state_key, action) -> float`
Read the current Q-value for a `(StateKey, ConfigAction)` pair (0.0 for
unseen entries).

#### `StateKey.from_values(available_mb, backend, model_params, model_bits)`
Convenience constructor for use with `BanditPolicy.q_value()` and
`BanditPolicy.update()` directly.

Full reference: [`docs/rl_optimizer.md`](docs/rl_optimizer.md).

### Fine-Tuning Adapters (v0.2, `pip install ml-memguard[hf]`)

#### `guard_trainer(trainer, guard=None, **preflight_overrides) -> SafeConfig`
Attach memory-guard to a HuggingFace `Trainer` in one call.  Introspects the
model, runs preflight, writes safe values to `trainer.args`, and appends
`MemoryGuardCallback`.

#### `MemoryGuardCallback(guard)`
`TrainerCallback` subclass.  `on_train_begin` starts the monitor;
`on_step_begin` records a pending batch-size downgrade when the monitor signals
pressure; `on_epoch_begin` applies it (scales `gradient_accumulation_steps` to
preserve effective batch); `on_train_end` stops the monitor and records
calibration data.

#### `guard_unsloth_model(model, guard=None, **preflight_overrides) -> SafeConfig`
Run preflight on an Unsloth model before `FastLanguageModel.get_peft_model` is
called.  Thread `safe.lora_rank`, `safe.lora_layers`, `safe.seq_length` into
`get_peft_model`.  Detects BnB double-quantization and applies a 5 % correction.

#### `guard_sft_trainer(trainer, guard=None, **preflight_overrides) -> SafeConfig`
Identical to `guard_trainer` but named for TRL `SFTTrainer` workflows.

**Design constraint** (ADR 003): `guard_vllm` and `guard_sglang` emit signals only — they never mutate a running engine. Load-shedding requires a load balancer or health endpoint in front of the engine. See [`docs/adapters.md`](docs/adapters.md) for the full reference.

## Estimation Accuracy

Measured accuracy on real training runs. **We need your help expanding this table** — see [Contributing](#help-us-benchmark) below.

| Model | Device | Batch | Seq | Rank | Estimated | Actual | Error |
|-------|--------|------:|----:|-----:|----------:|-------:|------:|
| Qwen3.5-9B-4bit | M4 Max 36GB | 1 | 512 | 8 | 6,193 MB | 7,048 MB | 12.1% under |
| Qwen3.5-9B-4bit | M4 Max 36GB | 1 | 128 | 16 | 9,522 MB | 8,879 MB | 7.2% over |

**What's tested**:
- LoRA fine-tuning with mlx_lm on Apple Silicon (M4 Max)
- HuggingFace `Trainer` + `MemoryGuardCallback` end-to-end (distilgpt2, CPU, fp32) — integration smoke test passes in CI
- BnB 4-bit + double-quantization detection and 5 % correction (unit-tested against mock models)

**What's NOT tested yet**:
- CUDA GPUs (RTX 3060/4090, A100, H100)
- AMD ROCm (RX 7900, MI300X)
- Smaller devices (M1/M2 MacBook Air 8-16GB)
- Models below 7B or above 13B
- MoE architectures (Mixtral, DeepSeek-MoE)
- Multi-modal models (LLaVA, Qwen-VL)
- DoRA, full fine-tuning
- PyTorch Lightning, Axolotl, LitGPT

The estimation formula is based on published research (FlashAttention, HyC-LoRA, LoRA-FA) and verified on one configuration. Auto-calibration improves accuracy after 3+ runs on any given setup.

## Known Limitations

- **Single validation point**: Estimation accuracy is verified on one model/device combination. Your results may differ significantly — please report them.
- **Inference monitoring**: `KVCacheMonitor` (v0.3.0) emits signals only — it never mutates a running vLLM or SGLang engine. Load-shedding requires a load balancer or health-endpoint pattern in front of the engine.
- **Calibration cold start**: Auto-calibration needs 3+ training runs on a given device before corrections kick in.
- **Custom kernels**: Frameworks with heavily fused kernels (Unsloth) use less memory than the formula predicts. Calibration corrects this over time.
- **MLX Metal thread safety**: `mx.metal.get_active_memory()` is called from a background thread. MLX's Metal backend has [known thread safety limitations](https://github.com/ml-explore/mlx/issues/2133). Memory counter reads work in practice but aren't guaranteed thread-safe by the MLX API.
- **Windows**: CUDA path uses well-tested `torch.cuda` APIs. The CPU-only fallback (`GlobalMemoryStatusEx`) hasn't been validated across Windows versions.

## Contributing

### Help Us Benchmark

The single most valuable contribution right now is **running the benchmark on your hardware** and sharing the results. This directly improves estimation accuracy for everyone.

```bash
# Install
pip install ml-memguard mlx-lm

# Run with default small model (fast, ~2 minutes)
python bench/bench_accuracy.py

# Run with a specific model
python bench/bench_accuracy.py --model mlx-community/Mistral-7B-Instruct-v0.3-4bit

# Generate a pre-formatted GitHub issue with your results
python bench/bench_accuracy.py --model mlx-community/Qwen3.5-9B-MLX-4bit --submit
```

Then open a [GitHub issue](https://github.com/vgpprasad91/ml-memguard/issues/new) with the output. We'll add your results to the accuracy table above.

**Devices we especially need data from**:
- M1/M2 MacBook Air (8GB, 16GB)
- M3/M4 MacBook Pro (18GB, 36GB)
- RTX 3060/3090, RTX 4070/4090
- A100, H100
- AMD Radeon RX 7900 / MI300X
- Docker/Kubernetes containers with memory limits

### Other Contributions

- **Framework adapters**: PyTorch Lightning, Axolotl, LitGPT wrappers (HF Transformers and Unsloth ship in v0.2.0; vLLM and SGLang in v0.3.0; RL optimizer in v0.4.0)
- **Accuracy data**: Real training runs on CUDA or non-Apple hardware — see the table above
- **Bug reports**: If the estimate was off by >30%, that's a bug — please report it with your config

## License

Apache 2.0
