{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "# ZeRO’s Memory Budget, Byte by Byte\n",
        "\n",
        "Parameters, gradients, and optimizer state account for different parts of training memory. ZeRO changes who stores each part.\n",
        "\n",
        "Original dtype-explicit arithmetic; persistent state only, with no framework tensor inspection or peak-memory measurement.\n",
        "\n",
        "The code uses Python 3 and the standard library. Saved output is from the checked CPU run.\n"
      ],
      "id": "zero-training-memory-accountin-0"
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {},
      "source": [
        "parameters = 1_000_000_000\n",
        "workers = 4\n",
        "bytes_per_parameter = {\n",
        "    \"parameters\": 2,\n",
        "    \"gradients\": 2,\n",
        "    \"optimizer\": 4 + 4 + 4,  # Master weights and two moments.\n",
        "}\n",
        "partitioned = {\n",
        "    0: set(),\n",
        "    1: {\"optimizer\"},\n",
        "    2: {\"optimizer\", \"gradients\"},\n",
        "    3: {\"optimizer\", \"gradients\", \"parameters\"},\n",
        "}\n",
        "totals = []\n",
        "for stage in range(4):\n",
        "    parts = {}\n",
        "    for name, width in bytes_per_parameter.items():\n",
        "        divisor = workers if name in partitioned[stage] else 1\n",
        "        parts[name] = parameters * width / divisor / 1_000_000_000\n",
        "    total = sum(parts.values())\n",
        "    totals.append(total)\n",
        "    print(f\"stage={stage}: params={parts['parameters']:.1f} GB, \"\n",
        "          f\"grads={parts['gradients']:.1f} GB, \"\n",
        "          f\"optimizer={parts['optimizer']:.1f} GB, total={total:.1f} GB\")\n",
        "assert totals == [16.0, 7.0, 5.5, 4.0]\n",
        "assert totals[-1] * workers == totals[0]\n"
      ],
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "stage=0: params=2.0 GB, grads=2.0 GB, optimizer=12.0 GB, total=16.0 GB\n",
            "stage=1: params=2.0 GB, grads=2.0 GB, optimizer=3.0 GB, total=7.0 GB\n",
            "stage=2: params=2.0 GB, grads=0.5 GB, optimizer=3.0 GB, total=5.5 GB\n",
            "stage=3: params=0.5 GB, grads=0.5 GB, optimizer=3.0 GB, total=4.0 GB\n"
          ]
        }
      ],
      "id": "zero-training-memory-accountin-1"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Sources\n",
        "\n",
        "- ZeRO: Memory Optimizations Toward Training Trillion Parameter Models: https://arxiv.org/html/1910.02054v3\n"
      ],
      "id": "zero-training-memory-accountin-2"
    }
  ],
  "metadata": {
    "kernelspec": {
      "display_name": "Python 3",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "name": "python",
      "version": "3.12.14"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 5
}