{ "cells": [ { "cell_type": "markdown", "metadata": {}, "source": [ "# Fine-tune an 8B model on 4 GB of VRAM — run it yourself\n", "\n", "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/MakazhanAlpamys/Soup/blob/main/notebooks/proof-4gb.ipynb)\n", "\n", "[Soup](https://github.com/MakazhanAlpamys/Soup) trains a model whose weights do not fit\n", "in your GPU. The frozen base stays in host RAM and is streamed to the GPU one decoder\n", "layer at a time, so peak VRAM is bounded by **one layer** instead of by the model.\n", "\n", "This notebook does not ask you to believe that. It caps this process to **4 GB** on\n", "Colab's free T4 and then measures what actually happens.\n", "\n", "| Section | What it proves | Time |\n", "|---|---|---|\n", "| 1–3 | The cap is real, and this GPU has no bf16 | ~2 min |\n", "| 4 | A streamed model and a normal one produce **bit-identical** logits | ~3 min |\n", "| 5 | Llama-3.1-8B trains with a measured peak under 4 GB | ~20 min |\n", "\n", "Sections 1–4 are the argument. Section 5 is the headline and is optional.\n", "\n", "**Runtime → Change runtime type → T4 GPU** before you start.\n" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## 1. Install\n", "\n", "**From git, not PyPI, and the reason is the point of section 2.** The fix that makes\n", "this pick the right precision on a T4\n", "([#385](https://github.com/MakazhanAlpamys/Soup/issues/385),\n", "[#387](https://github.com/MakazhanAlpamys/Soup/issues/387)) is on `main` and is not in a\n", "release yet, so on the published 0.73.0 the next cell raises `ImportError`. Switch this\n", "line back to `soup-cli[train]` once the next version ships.\n", "\n", "The `torchao` line is not incidental either. Colab preinstalls **torchao 0.10.0**, and\n", "`peft` does not merely decline to use a version it considers too old — it *raises*\n", "`ImportError` from `is_torchao_available()`, several frames inside `get_peft_model`.\n", "Nothing here needs torchao, so it is removed rather than upgraded (upgrading risks\n", "pulling a wheel built against a different torch).\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "%pip uninstall -q -y torchao\n", "%pip install -q \"soup-cli[train] @ git+https://github.com/MakazhanAlpamys/Soup.git\"\n", "\n", "import importlib.util\n", "\n", "import soup_cli\n", "import soup_cli.utils.gpu\n", "\n", "print(\"soup\", soup_cli.__version__)\n", "print(\"pre-Ampere fix present:\", hasattr(soup_cli.utils.gpu, \"bf16_fp16_flags\"))\n", "print(\"torchao gone (peft raises on an old one):\",\n", " importlib.util.find_spec(\"torchao\") is None)\n" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## 2. What card did we get, and does it have bf16?\n", "\n", "Colab's free tier is a **T4** — Turing, sm_75. bf16 hardware arrived with Ampere, so a\n", "T4 has none.\n", "\n", "**Read the two lines below carefully, because they disagree, and the disagreement is\n", "the point.** `torch.cuda.is_bf16_supported()` defaults to `including_emulation=True`:\n", "when the compute-capability check fails it falls through to *constructing* a bf16\n", "tensor, which software emulation satisfies. So a T4 answers **True** to the question\n", "everyone asks, and False only to `is_bf16_supported(including_emulation=False)`.\n", "\n", "Soup asked the permissive question and therefore handed bf16 to a card with no bf16\n", "units. The first version of this fix asked it too, and was a no-op on exactly the\n", "hardware it was written for — caught by running this notebook on a real T4, not before\n", "([#385](https://github.com/MakazhanAlpamys/Soup/issues/385),\n", "[#387](https://github.com/MakazhanAlpamys/Soup/issues/387)).\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "import torch\n", "\n", "from soup_cli.utils.gpu import bf16_fp16_flags\n", "\n", "assert torch.cuda.is_available(), \"No GPU. Runtime -> Change runtime type -> T4 GPU.\"\n", "\n", "name = torch.cuda.get_device_name(0)\n", "major, minor = torch.cuda.get_device_capability(0)\n", "total = torch.cuda.get_device_properties(0).total_memory\n", "\n", "print(f\"GPU {name} (sm_{major}{minor})\")\n", "print(f\"VRAM {total / 1e9:.1f} GB\")\n", "print(f\"bf16, incl. emulation {torch.cuda.is_bf16_supported()}\")\n", "print(f\"bf16 IN HARDWARE \"\n", " f\"{torch.cuda.is_bf16_supported(including_emulation=False)}\")\n", "\n", "bf16, fp16 = bf16_fp16_flags(\"cuda\")\n", "print(f\"Soup will train in {'bf16' if bf16 else 'fp16' if fp16 else 'fp32'}\")\n" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## 3. Cap this process to 4 GB\n", "\n", "`set_per_process_memory_fraction` caps PyTorch's allocator. Everything after this cell\n", "runs as if the card were a 4 GB laptop GPU — an allocation past the cap raises, exactly\n", "as it would on the real thing.\n", "\n", "**One honest caveat.** The cap is enforced by the allocator, not by the driver, so\n", "`torch.cuda.mem_get_info()` keeps reporting the *whole* card. Soup's pre-flight reads\n", "that, so the \"free VRAM\" line it prints later belongs to the host card and not to this\n", "capped process — it will happily allow a configuration that the allocator then refuses.\n", "That is [#347](https://github.com/MakazhanAlpamys/Soup/issues/347), it is open, and it\n", "does not affect anything measured here: the proof below is the **peak VRAM torch\n", "actually reports**, not what the pre-flight predicted.\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "BUDGET_BYTES = 4 * 1000**3 # 4 GB, the card this method was developed on\n", "\n", "fraction = BUDGET_BYTES / total\n", "torch.cuda.set_per_process_memory_fraction(fraction)\n", "print(f\"capped at {BUDGET_BYTES / 1e9:.2f} GB (fraction {fraction:.3f} of this card)\")\n", "\n", "# Prove the cap bites: ask for 15% more than the budget and expect a refusal.\n", "try:\n", " _ = torch.empty(int(BUDGET_BYTES * 1.15), dtype=torch.uint8, device=\"cuda\")\n", " print(\"WARNING: the allocation succeeded — the cap is NOT in force\")\n", "except RuntimeError as exc:\n", " print(\"refused, as it should be:\", str(exc).splitlines()[0][:90])\n" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## 4. The claim that matters: streamed == resident, bit for bit\n", "\n", "A streaming bug is silent. If the base were substituted wrongly, or the autograd path\n", "severed, the loss would still fall — the upper layers keep learning — and you would ship\n", "a damaged model without an error anywhere.\n", "\n", "So the check is not \"does it train\". It is: **the same weights, through the same kernels,\n", "must produce the same numbers.** Below, one model is streamed layer-by-layer and the\n", "other is an ordinary resident model, both carrying identical adapter weights.\n", "`torch.equal` is exact equality, not a tolerance.\n", "\n", "This runs on a small model because the reference has to fit in memory next to the\n", "streamed copy — that is the whole reason the headline size cannot be checked this way.\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "import tempfile\n", "from pathlib import Path\n", "\n", "from peft import LoraConfig, TaskType, get_peft_model\n", "from transformers import AutoModelForCausalLM\n", "\n", "from soup_cli.utils.layer_shard import shard_checkpoint\n", "from soup_cli.utils.layer_stream import resolve_stream_dtype\n", "from soup_cli.utils.layer_stream_runtime import build_streamed_model\n", "from soup_cli.utils.spectrum_scan import resolve_model_weights\n", "\n", "MODEL = \"HuggingFaceTB/SmolLM2-135M-Instruct\"\n", "DTYPE = resolve_stream_dtype(\"cuda\") # fp16 on a T4, bf16 on an Ampere card\n", "LORA = LoraConfig(\n", " r=8, lora_alpha=16, lora_dropout=0.0, bias=\"none\",\n", " target_modules=[\"q_proj\", \"v_proj\"], task_type=TaskType.CAUSAL_LM,\n", ")\n", "\n", "workdir = Path(tempfile.mkdtemp())\n", "weights = resolve_model_weights(MODEL) # downloads on first use\n", "index = shard_checkpoint(weights, str(workdir / \"shards\"), dtype=DTYPE, arch=\"llama\")\n", "streamed, runtime = build_streamed_model(\n", " model_id=weights, shard_dir=str(workdir / \"shards\"), index=index,\n", " lora_config=LORA, device=\"cuda\", dtype=DTYPE, buffers=2, pin=True, seed=0,\n", ")\n", "print(f\"streamed: {index.n_layers} layers, dtype={DTYPE}\")\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# PEFT initialises lora_B to zero, so an untrained adapter contributes NOTHING and any\n", "# comparison would silently be about the base model alone. Make it load-bearing first.\n", "gen = torch.Generator().manual_seed(7)\n", "with torch.no_grad():\n", " for pname, param in streamed.named_parameters():\n", " if \"lora_B\" in pname:\n", " param.copy_(torch.randn(param.shape, generator=gen).to(param.device, param.dtype))\n", "\n", "resident = AutoModelForCausalLM.from_pretrained(\n", " MODEL, dtype=getattr(torch, DTYPE), device_map={\"\": \"cuda\"}\n", ")\n", "resident = get_peft_model(resident, LORA)\n", "\n", "# Copy the adapter across. The streamed wrapper inserts an '.inner.' segment in its keys.\n", "src = {k.replace(\".inner.\", \".\"): v for k, v in streamed.state_dict().items() if \"lora_\" in k}\n", "dst = {k.replace(\".inner.\", \".\"): v for k, v in resident.state_dict().items() if \"lora_\" in k}\n", "assert src and set(src) == set(dst)\n", "with torch.no_grad():\n", " for key, val in src.items():\n", " dst[key].copy_(val.to(dst[key].dtype))\n", "\n", "ids = torch.randint(0, 4096, (1, 32), device=\"cuda\")\n", "with torch.no_grad():\n", " a = streamed(input_ids=ids).logits\n", " b = resident(input_ids=ids).logits\n", "\n", "print(\"max |streamed - resident| =\", (a.float() - b.float()).abs().max().item())\n", "print(\"torch.equal =\", torch.equal(a, b))\n", "assert torch.equal(a, b), \"NOT bit-exact — please open an issue with this output\"\n", "print(\"\\nBit-exact. The streamed model is the same model.\")\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "import gc\n", "\n", "# Free the reference before the headline run.\n", "runtime.close()\n", "del streamed, resident, a, b\n", "gc.collect()\n", "torch.cuda.empty_cache()\n", "torch.cuda.reset_peak_memory_stats()\n", "print(\"reset\")\n" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## 5. The headline: Llama-3.1-8B, trained under the 4 GB cap\n", "\n", "8B parameters. In NF4 the weights alone are about **4.5 GB** — more than the budget this\n", "process is allowed, before activations, gradients or the optimizer. It trains anyway,\n", "because at any moment only a couple of decoder layers are resident.\n", "\n", "Expect ~20 minutes, most of it the download. The number to watch is the **peak VRAM**\n", "at the end — not the loss. This is a handful of steps on 32 toy rows; over that distance\n", "the loss can go up as easily as down, and it would prove nothing either way. What is\n", "being demonstrated is that the run *happens at all* inside the budget.\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "import json\n", "\n", "# Varied on both sides on purpose: 32 copies of one answer drive the loss to 0.000\n", "# immediately and the printed curve stops meaning anything.\n", "TOPICS = [\n", " (\"streaming\", \"Only a couple of decoder layers are resident at any moment.\"),\n", " (\"NF4\", \"Four-bit weights make the host-side store about four times smaller.\"),\n", " (\"LoRA\", \"The base is frozen, so it is read and never written.\"),\n", " (\"VRAM\", \"Peak memory is bounded by one layer instead of by the model.\"),\n", "]\n", "rows = [\n", " {\"messages\": [\n", " {\"role\": \"user\", \"content\": f\"Question {i}: tell me about {topic}.\"},\n", " {\"role\": \"assistant\", \"content\": answer},\n", " ]}\n", " for i in range(8)\n", " for topic, answer in TOPICS\n", "]\n", "Path(\"train.jsonl\").write_text(\"\\n\".join(json.dumps(r) for r in rows), encoding=\"utf-8\")\n", "\n", "config = \"\"\"\n", "base: NousResearch/Meta-Llama-3.1-8B-Instruct\n", "task: sft\n", "data:\n", " train: train.jsonl\n", " max_length: 256\n", "training:\n", " epochs: 1\n", " batch_size: 1\n", " lr: 0.0002\n", " logging_steps: 1 # short run — without this nothing gets logged\n", " quantization: 4bit # NF4 — the store is ~4x smaller than bf16\n", " stream_layers: true # the feature\n", " stream_buffers: 2\n", " lora:\n", " r: 8\n", " alpha: 16\n", "output: ./out-8b\n", "\"\"\"\n", "Path(\"soup.yaml\").write_text(config, encoding=\"utf-8\")\n", "print(config)\n" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "The cell below runs the trainer **in this process** rather than shelling out to\n", "`soup train`. That is deliberate and it is the only reason: `max_memory_allocated()`\n", "reports the peak of *the process that calls it*, so a subprocess would train fine and\n", "leave us measuring nothing. It is the same code path the CLI runs on the same\n", "`soup.yaml` above — the CLI adds argument parsing and the pre-flight panel, neither of\n", "which changes what the GPU does.\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "from soup_cli.config.loader import load_config_from_string\n", "from soup_cli.data.loader import load_dataset\n", "from soup_cli.trainer.sft import SFTTrainerWrapper\n", "\n", "cfg = load_config_from_string(config)\n", "dataset = load_dataset(cfg.data)\n", "\n", "torch.cuda.empty_cache()\n", "torch.cuda.reset_peak_memory_stats()\n", "\n", "wrapper = SFTTrainerWrapper(cfg)\n", "wrapper.setup(dataset) # downloads, shards to NF4, builds the streamed model\n", "result = wrapper.train()\n", "\n", "print(f\"\\nsteps: {result['total_steps']} loss: {result['initial_loss']:.3f}\"\n", " f\" -> {result['final_loss']:.3f}\")\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "peak = torch.cuda.max_memory_allocated()\n", "print(f\"peak VRAM allocated by this process: {peak / 1e9:.2f} GB\")\n", "print(f\"budget this process was capped to: {BUDGET_BYTES / 1e9:.2f} GB\")\n", "print(\"model weights in NF4, for scale: ~4.5 GB\")\n", "\n", "adapter = Path(\"out-8b/adapter_model.safetensors\")\n", "print(f\"\\nadapter written: {adapter.exists()}\")\n", "if adapter.exists():\n", " from safetensors.torch import load_file\n", "\n", " tensors = load_file(str(adapter))\n", " live = sum(1 for v in tensors.values() if v.abs().max().item() > 0)\n", " print(f\"adapter tensors: {len(tensors)}, non-zero: {live}\")\n" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## What this proved, and what it did not\n", "\n", "**Proved, on your hardware:**\n", "\n", "- A streamed model returns **bit-identical** logits to an ordinary one (§4).\n", "- An 8B model trained with a measured peak below a cap smaller than its own weights (§5).\n", "- Both on a GPU with **no bf16** — the case that was broken until recently.\n", "\n", "**Not proved:**\n", "\n", "- *Backward* exactness at this size. §4 compares the forward. Gradient exactness is\n", " verified up to 14B against resident references on hardware that can hold them, and\n", " a defect **above** that size was found, named upstream and repaired — see\n", " [`benchmarks/`](https://github.com/MakazhanAlpamys/Soup/tree/main/benchmarks).\n", "- Speed. A T4 under an artificial cap is not a throughput benchmark, and this notebook\n", " deliberately does not quote tok/s.\n", "\n", "Layer streaming is **BETA** and opt-in (`stream_layers: true`).\n", "\n", "**If any assertion above failed, that is worth an issue** — with the cell output. A\n", "reproduction on hardware we do not own is more useful to this project than a star.\n", "\n", "The method, the correctness protocol and every measurement:\n", "[10.5281/zenodo.21771064](https://doi.org/10.5281/zenodo.21771064) ·\n", "[measurement records](https://github.com/MakazhanAlpamys/Soup/tree/main/benchmarks)\n" ] } ], "metadata": { "accelerator": "GPU", "colab": { "gpuType": "T4", "provenance": [] }, "kernelspec": { "display_name": "Python 3", "name": "python3" }, "language_info": { "name": "python" } }, "nbformat": 4, "nbformat_minor": 0 }