# Lab Specification — Module FT01: The VRAM Calculator

**Course**: Course 3 — LLM Fine-Tuning Masterclass
**Module**: FT01 — VRAM Math: Can I Actually Run This?
**Duration**: 45–60 minutes
**Environment**: Python 3.10+. **No GPU required.** This is a pure-Python estimation lab — it runs on a laptop, Colab T4 (free), or any machine. Optional: run the validation against a real loaded model on a GPU/MPS box to feel the estimate vs. reality.

---

## Learning objectives

By the end of this lab you will have:

1. **Built a bottom-up VRAM estimator** from the three consumers (weights, optimizer states + gradients, activations) rather than trusting a single multiplier.
2. **Validated it against three real jobs** — 7B QLoRA, 70B QLoRA, 7B full fine-tuning — and confirmed each lands within ~20% of the field rules of thumb.
3. **Turned the three-question framework into code** — model size × method × context length → a GPU-class recommendation.
4. **Internalized the numbers** so you never again open a model loader without first knowing whether it will fit.

This lab is deliberately GPU-free. The math is the point. Once you can estimate, the lab's stretch goal checks your estimate against `nvidia-smi` (or Activity Monitor on Apple Silicon) on a real load.

---

## Phase 0 — Environment setup (2 min)

```bash
python3 -m venv ft01-env && source ft01-env/bin/activate
# No ML deps needed for the core lab. Optional GPU validation (Phase 4) needs torch.
pip install -q torch   # optional — only for the real-model check in Phase 4
```

No other dependencies. The estimator is pure Python.

---

## Phase 1 — The estimator spec

Implement this function in a file called `vram_calc.py`:

```python
def vram_estimate(model_params_b, method, context_len, batch_size=1,
                  gradient_checkpointing=True, flash_attention=True):
    """Estimate training VRAM in GB from the three consumers.

    Args:
        model_params_b:    model size in billions of params (e.g. 7, 70)
        method:            "qlora" | "lora16" | "full"
        context_len:       max sequence length in tokens (use your 99th pctile)
        batch_size:        PHYSICAL micro-batch size (effective batch uses grad accum)
        gradient_checkpointing: discard-and-recompute activations (default on)
        flash_attention:   fused attention, linear memory (default on)

    Returns:
        total training VRAM in GB (float).
    """
```

### The model (the rules of thumb, made executable)

The estimator sums three consumers. Implement them exactly as below — the constants are calibrated against the field rules of thumb (Introl Dec 2025, Spheron PEFT guide, QLoRA paper).

**Consumer 1 — Weights** (the frozen/loaded base):

| Method | Base precision | Bytes/param |
| --- | --- | --- |
| `qlora` | 4-bit (NF4) frozen | 0.5 |
| `lora16` | FP16/BF16 frozen | 2.0 |
| `full` | FP16 trainable | 2.0 |

`weights_gb = bytes_per_param × params_b`

**Consumer 2 — Optimizer states + gradients** (only over *trainable* params):

- `qlora` / `lora16`: only the adapter trains. Assume **adapter ≈ 1% of params**. Per trainable param: FP16 weight (2) + FP16 gradient (2) + 8-bit AdamW states (~2) ≈ **6 bytes/trainable-param**.
- `full`: **all** params trainable. Per trainable param: FP16 gradient (2) + FP32 master copy (4) + AdamW `m` (4) + AdamW `v` (4) = **14 bytes/trainable-param** (on top of the 2-byte FP16 weight counted in Consumer 1, for 16 bytes/param total).

`opt_gb = bytes_per_trainable_param × trainable_params_b`

**Consumer 3 — Activations** (scales with your data):

Activations scale with layers, hidden dim, batch, and context. Derive architecture from params using transformer scaling (`params ≈ 12 × layers × hidden²`, with `layers ≈ hidden/120`), giving:

```python
hidden = round((10 * params_n) ** (1/3) / 64) * 64      # snapped to 64
layers = round(hidden / 120)
```

Then:

```python
BASE_K = 2.9e-8   # calibrated constant (FlashAttention-on baseline)
act_gb = BASE_K * layers * batch_size * context_len * hidden
if gradient_checkpointing: act_gb *= 0.35     # ~65% activation-memory cut
if not flash_attention:    act_gb *= 1.7      # naive N×N attention penalty
```

`total_gb = weights_gb + opt_gb + act_gb`

> **Why derive `hidden`/`layers` from params?** Because activation memory is `layers × batch × seq × hidden`, and larger models have bigger hidden dims — so activation memory does *not* scale linearly with parameter count. This is why a 70B is not 10× a 7B in activation cost. The `hidden = (10·params)^(1/3)` relation reproduces real architectures to within ~10% (7B→~4096/34, 70B→~8900/74).

---

## Phase 2 — Implement it (15 min)

Write the full file. A reference skeleton:

```python
"""vram_calc.py — FT01 lab: bottom-up training VRAM estimator."""

def _arch(params_b):
    """Derive (hidden, layers) from param count via transformer scaling."""
    params_n = params_b * 1e9
    hidden = round((10 * params_n) ** (1 / 3) / 64) * 64
    layers = max(1, round(hidden / 120))
    return hidden, layers


def vram_breakdown(model_params_b, method, context_len, batch_size=1,
                   gradient_checkpointing=True, flash_attention=True):
    """Return the per-consumer breakdown as a dict."""
    assert method in ("qlora", "lora16", "full"), f"unknown method: {method}"
    hidden, layers = _arch(model_params_b)
    params_n = model_params_b * 1e9

    # --- Consumer 1: weights ---
    base_bytes = {"qlora": 0.5, "lora16": 2.0, "full": 2.0}[method]
    weights_gb = base_bytes * model_params_b

    # --- Consumer 2: optimizer states + gradients (trainable params only) ---
    if method in ("qlora", "lora16"):
        trainable_b = 0.01 * model_params_b          # adapter ~1% of params
        bytes_per_trainable = 6                        # FP16 wt+grad + 8-bit opt
    else:  # full
        trainable_b = model_params_b                   # all params trainable
        bytes_per_trainable = 14                        # grad + FP32 master + m + v
    opt_gb = bytes_per_trainable * trainable_b

    # --- Consumer 3: activations ---
    BASE_K = 2.9e-8
    act_gb = BASE_K * layers * batch_size * context_len * hidden
    if gradient_checkpointing:
        act_gb *= 0.35
    if not flash_attention:
        act_gb *= 1.7

    return {
        "method": method,
        "params_b": model_params_b,
        "context_len": context_len,
        "batch_size": batch_size,
        "arch": {"hidden": hidden, "layers": layers},
        "weights_gb": round(weights_gb, 1),
        "optimizer_gradients_gb": round(opt_gb, 1),
        "activations_gb": round(act_gb, 1),
        "total_gb": round(weights_gb + opt_gb + act_gb, 1),
    }


def vram_estimate(model_params_b, method, context_len, batch_size=1,
                  gradient_checkpointing=True, flash_attention=True):
    """Headline API: return total training VRAM in GB (float)."""
    return vram_breakdown(model_params_b, method, context_len, batch_size,
                          gradient_checkpointing, flash_attention)["total_gb"]


def recommend_gpu(total_gb):
    """Map a VRAM budget to a GPU class (the three-question framework, step 4).

    Boundaries leave ~25% headroom: a 15 GB estimate lands on a 24 GB card,
    not a 16 GB card, because CUDA context + fragmentation + optimizer buffers
    eat into the nominal budget. Buy headroom; don't squeeze onto the ceiling.
    """
    if total_gb <= 12:  return "Laptop / Apple Silicon / Colab T4 (free)"
    if total_gb <= 18:  return "RTX 4090 24GB ($1,500) — 7B QLoRA sweet spot"
    if total_gb <= 30:  return "RTX 4090 24GB with checkpointing, or A100 40GB"
    if total_gb <= 60:  return "1× A100 80GB"
    if total_gb <= 120: return "2× A100 80GB"
    if total_gb <= 160: return "multi-A100 80GB (~$50K of GPUs)"
    if total_gb <= 640: return "8× A100/H100 single node"
    return "multi-node H100 cluster (8–16× H100)"


if __name__ == "__main__":
    import json
    cases = [
        (7,  "qlora", 4096, 2),
        (70, "qlora", 2048, 1),
        (7,  "full",  4096, 1),
    ]
    for params, method, ctx, bs in cases:
        bd = vram_breakdown(params, method, ctx, bs)
        print(f"\n=== {params}B {method} @ {ctx} ctx, bs {bs} ===")
        print(json.dumps(bd, indent=2))
        print(f"  -> recommend: {recommend_gpu(bd['total_gb'])}")
```

Run it:

```bash
python vram_calc.py
```

---

## Phase 3 — Validate against three real jobs (15 min)

Add this validation block to the bottom of `vram_calc.py` (or a separate `test_vram.py` run with `pytest`), then run it. All three asserts must pass.

```python
def validate():
    """Three real jobs. Each must land within ~20% of the field rule of thumb."""
    TOL = 0.20

    # Job 1: 7B QLoRA, 4K context, batch 2. Rule of thumb: 10-16 GB. Fits RTX 4090 24GB.
    j1 = vram_estimate(7, "qlora", 4096, batch_size=2)
    assert 10 <= j1 <= 16, f"7B QLoRA: expected 10-16 GB, got {j1}"
    assert j1 <= 24, f"7B QLoRA must fit RTX 4090 24GB, got {j1}"

    # Job 2: 70B QLoRA, 2K context, batch 1. Rule of thumb: 48-60 GB. Fits 1× A100 80GB.
    j2 = vram_estimate(70, "qlora", 2048, batch_size=1)
    assert 48 <= j2 <= 60, f"70B QLoRA: expected 48-60 GB, got {j2}"
    assert j2 <= 80, f"70B QLoRA must fit 1× A100 80GB, got {j2}"

    # Job 3: 7B full FT, 4K context, batch 1, checkpointing on. Rule of thumb: 100-160 GB.
    j3 = vram_estimate(7, "full", 4096, batch_size=1, gradient_checkpointing=True)
    assert 100 <= j3 <= 160, f"7B full FT: expected 100-160 GB, got {j3}"
    assert j3 > 80, f"7B full FT needs multi-A100 (>80GB), got {j3}"

    print("\nAll three validation jobs PASS.")
    print(f"  7B QLoRA : {j1} GB  -> {recommend_gpu(j1)}")
    print(f"  70B QLoRA: {j2} GB  -> {recommend_gpu(j2)}")
    print(f"  7B Full  : {j3} GB -> {recommend_gpu(j3)}")


if __name__ == "__main__":
    validate()
```

**Expected output** (your numbers should match to within rounding):

```
All three validation jobs PASS.
  7B QLoRA : 15.5 GB  -> RTX 4090 24GB ($1,500) — 7B QLoRA sweet spot
  70B QLoRA: 52.9 GB  -> 1× A100 80GB
  7B Full  : 117.8 GB -> 2× A100 80GB
```

(7B Full at ~118 GB lands on 2× A100 80GB = 160 GB capacity — that *is* the "multi-A100 (~$50K)" class the teaching doc describes; `recommend_gpu` just names the minimum card count.)

If an assert fails, debug *which consumer* is off. The most common student bugs:

- **Off by ~10× on full FT?** You forgot that *all* params are trainable, not 1%. Recheck Consumer 2.
- **Activations tiny for 70B?** You used `params_b` directly instead of deriving `hidden`/`layers`. Activations need the architecture.
- **7B QLoRA too high?** Make sure gradient checkpointing (×0.35) and FlashAttention (×1.7 penalty *off*) are applied.

---

## Phase 4 — Feel it: estimate vs. a real load (optional, needs a GPU/MPS) (10 min)

If you have a GPU (or Apple Silicon), load a small model and compare your estimate to reality. This is the moment the numbers stop being abstract.

```python
import torch
from transformers import AutoModelForCausalLM

# Pick a small open-data model (from FT00): MiniCPM5-1B
MODEL_ID = "openbmb/MiniCPM5-1B"

# Step 1: PREDICT before loading (inference floor = 4-bit weights for QLoRA-style load)
params_b = 1.0
pred_weights_gb = 0.5 * params_b            # 4-bit
print(f"Predicted 4-bit weight footprint: {pred_weights_gb} GB")

# Step 2: LOAD and measure
model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID, torch_dtype=torch.float16, device_map="auto", trust_remote_code=True,
)
if torch.cuda.is_available():
    actual = torch.cuda.memory_allocated() / 1e9
    print(f"Actual FP16 loaded VRAM: {actual:.2f} GB")
    print(f"Predicted FP16: {2.0 * params_b} GB  (factor ~{actual/(2.0*params_b):.2f}x = overhead)")
elif torch.backends.mps.is_available():
    print("On Apple Silicon: watch Activity Monitor GPU memory while loading.")
```

The *inference* footprint (Consumer 1 only) will be ~1.5–2× the bare weight footprint because of CUDA context, KV cache, and activation buffers. For *training* you add Consumers 2 and 3 on top — which is exactly why training costs 3–10× inference for the same model. Your estimator captures this; `nvidia-smi` confirms it.

---

## Deliverables

Submit `ft01-lab-report.md` containing:

- [ ] Your `vram_calc.py` (the function + validation).
- [ ] The output of `python vram_calc.py` showing all three validation jobs PASS, with the per-consumer breakdown printed.
- [ ] A one-paragraph answer: **for each of the three validation jobs, which consumer dominates, and what is the single cheapest knob that would shrink it?** (e.g., for 7B full FT: optimizer states dominate → the knob is "switch to QLoRA so only 1% of params carry optimizer states.")
- [ ] (If you did Phase 4) the predicted-vs-actual inference footprint and the overhead factor.

---

## Solution key

- **Job 1 (7B QLoRA, bs 2, 4K):** ~15.5 GB. Dominant consumer: activations (~11.6 GB). Cheapest knob: lower physical batch to 1 (halves activations) and use gradient accumulation to keep effective batch 2. GPU: RTX 4090 24GB.
- **Job 2 (70B QLoRA, bs 1, 2K):** ~52.9 GB. Dominant consumer: 4-bit weights (35 GB). Cheapest knob: none that's cheap — the weights are the floor; you cannot shrink frozen 4-bit weights further without going to 2-bit (quality cost). GPU: 1× A100 80GB.
- **Job 3 (7B full FT, bs 1, 4K, ckpt on):** ~117.8 GB. Dominant consumer: optimizer states + gradients (~98 GB). Cheapest knob: switch from full FT to QLoRA — trainable params drop from 100% to 1%, and optimizer states collapse from ~98 GB to <0.5 GB. Total falls from ~118 GB to ~10–16 GB. GPU without the knob: multi-A100 (~$50K). GPU with the knob: RTX 4090 ($1,500).

The lesson, stated bluntly: **for Job 3, the single decision "full FT vs QLoRA" is worth ~$48,500 and ~10× the hardware.** That is the entire economic argument for PEFT, and it falls out of the math.

---

## Stretch goals

1. **Add LoRA (16-bit base).** Extend the estimator to `lora16`. Validate: 7B LoRA @ 4K should land ~18–30 GB (2–3× the FP16 model). Which GPU class? (A100 40GB, or 4090 with aggressive checkpointing.)
2. **Model the KV cache at inference.** Add an `inference_estimate(params_b, quant_bits, context_len)` that returns weights + KV cache. Validate: Q4_K_M 7B at 4K context runs in ~6 GB; Q4 70B at 4K in ~40 GB.
3. **Build a CLI.** `python vram_calc.py --params 7 --method qlora --ctx 4096 --batch 2` prints the breakdown and GPU recommendation. Make it a tool you actually use before renting GPUs.
4. **Reconcile the shorthand.** The field says "QLoRA ≈ 1.5–2× the 4-bit model size." Show in code that this holds *if* you use the deployed 4-bit footprint (~6 GB for 7B, including runtime overhead) rather than the bare NF4 weights (3.5 GB) — and explain in a comment why the bottom-up estimator is more trustworthy than any single multiplier.