{
  "schemaVersion": 1,
  "id": "atlas-lesson-training-memory-state",
  "title": "Training state and activation rematerialization",
  "category": "learning",
  "format": "Atlas lesson",
  "description": "Budgets weights, gradients, optimizer state and saved activations across a training step.",
  "canonicalUrl": "https://semiaifoundry.com/learn/atlas/training-memory-state/",
  "publishedDate": null,
  "revisionDate": "2026-10-07",
  "author": null,
  "topic": "Model–machine co-design",
  "series": null,
  "license": null,
  "rightsUrl": "https://semiaifoundry.com/ai-atlas/licenses.html",
  "version": "1.3.0",
  "status": "Published educational lesson",
  "evidenceScope": "This is an explicit synthetic format choice, not a universal training multiplier. ZeRO stages and implementation buffers must be identified before applying a sharding estimate.\n\nMode: executable arrays; payload arithmetic. 423 entries per full array, eight bytes per binary64 number; no claim that JavaScript object storage equals this payload. No rematerialization feature is executed. Numerical examples are authored calculations, except explicitly identified automated core fixtures. Primary sources support the scoped mechanism; no human specialist review or learner outcome is implied.",
  "contentCoverage": "published-lesson",
  "artifactUrl": "https://semiaifoundry.com/ai-atlas/lesson-data/training-memory-state-0e965c34b3054246.json",
  "interactiveUrl": "https://semiaifoundry.com/ai-atlas/#training-memory-state",
  "related": [],
  "supersedes": [],
  "supersededBy": [],
  "contentHash": "3fabff49b5c87b0d4c56644b4887147b63239c6659285dd3a20b5f565e501862",
  "contentLength": 9061,
  "jsonUrl": "https://semiaifoundry.com/agents/publications/atlas-lesson-training-memory-state.json",
  "markdownUrl": "https://semiaifoundry.com/agents/publications/atlas-lesson-training-memory-state.md",
  "sections": [
    {
      "id": "section-1",
      "title": "Training state and activation rematerialization",
      "start": 40,
      "end": 412
    },
    {
      "id": "section-2",
      "title": "The mechanism",
      "start": 412,
      "end": 2551
    },
    {
      "id": "section-3",
      "title": "Science and assumptions",
      "start": 2551,
      "end": 2579
    },
    {
      "id": "section-4",
      "title": "Training capacity ledger",
      "start": 2579,
      "end": 3619
    },
    {
      "id": "section-5",
      "title": "Training tensors with actual lifetimes · worked relation",
      "start": 3619,
      "end": 4548
    },
    {
      "id": "section-6",
      "title": "Engineering example",
      "start": 4548,
      "end": 5153
    },
    {
      "id": "section-7",
      "title": "Distinctions and tradeoffs",
      "start": 5153,
      "end": 5643
    },
    {
      "id": "section-8",
      "title": "Specifications and interfaces",
      "start": 5643,
      "end": 5935
    },
    {
      "id": "section-9",
      "title": "Interfaces",
      "start": 5935,
      "end": 6131
    },
    {
      "id": "section-10",
      "title": "Check your understanding",
      "start": 6131,
      "end": 6160
    },
    {
      "id": "section-11",
      "title": "What cost accompanies discarding a forward activation that backward computation needs?",
      "start": 6160,
      "end": 6572
    },
    {
      "id": "section-12",
      "title": "What does 10152 bytes include?",
      "start": 6572,
      "end": 7054
    },
    {
      "id": "section-13",
      "title": "Evidence and scope",
      "start": 7054,
      "end": 7821
    },
    {
      "id": "section-14",
      "title": "Sources",
      "start": 7821,
      "end": 8203
    },
    {
      "id": "section-15",
      "title": "Continue exploring",
      "start": 8203,
      "end": 8226
    },
    {
      "id": "section-16",
      "title": "Foundations",
      "start": 8226,
      "end": 8439
    },
    {
      "id": "section-17",
      "title": "Related lessons",
      "start": 8439,
      "end": 9061
    }
  ],
  "citations": [
    {
      "title": "Explore in Atlas",
      "url": "https://semiaifoundry.com/ai-atlas/#training-memory-state"
    },
    {
      "title": "ZeRO: Memory Optimizations Toward Training Trillion Parameter Models",
      "url": "https://arxiv.org/abs/1910.02054"
    },
    {
      "title": "2 OLMo 2 Furious",
      "url": "https://arxiv.org/html/2501.00656v1"
    },
    {
      "title": "Published lesson data",
      "url": "https://semiaifoundry.com/ai-atlas/lesson-data/training-memory-state-0e965c34b3054246.json"
    },
    {
      "title": "Atlas releases and licensing",
      "url": "https://semiaifoundry.com/ai-atlas/licenses.html"
    },
    {
      "title": "Loss, gradients & optimizer",
      "url": "https://semiaifoundry.com/learn/atlas/loss-optimizer/"
    },
    {
      "title": "Full memory hierarchy and input path",
      "url": "https://semiaifoundry.com/learn/atlas/memory-hierarchy-route/"
    },
    {
      "title": "Data, tensor, pipeline, and state parallelism",
      "url": "https://semiaifoundry.com/learn/atlas/parallelism/"
    },
    {
      "title": "Parallelism as a memory and communication budget",
      "url": "https://semiaifoundry.com/learn/atlas/parallelism-communication-budget/"
    },
    {
      "title": "Next-token supervision from raw sequences",
      "url": "https://semiaifoundry.com/learn/atlas/causal-pretraining-objective/"
    },
    {
      "title": "Initialization, learning rate and stable updates",
      "url": "https://semiaifoundry.com/learn/atlas/optimizer-schedule-stability/"
    },
    {
      "title": "Gradient accumulation under a fixed global objective",
      "url": "https://semiaifoundry.com/learn/atlas/microbatch-equivalent-updates/"
    }
  ],
  "recordHash": "f14d6af6ec970aa34598ac769b5aa839492d0f2444bab154805f7b8a74741f17",
  "contentFormat": "text/markdown",
  "content": "Model–machine co-design · Atlas v1.3.0\n\n# Training state and activation rematerialization\n\nBudgets weights, gradients, optimizer state and saved activations across a training step.\n\n[Explore in Atlas](https://semiaifoundry.com/ai-atlas/#training-memory-state)One model step to accelerator memory hierarchy\n\nLearning objectiveBuild a training memory ledger that distinguishes state type, precision and lifetime.\n\n## The mechanism\n\nTraining keeps more state than inference because it must calculate gradients and update parameters. Weights, gradients and optimizer state have different lifetimes and may use different numerical formats. Forward activations can be saved for the backward pass or recomputed later, exchanging extra arithmetic for reduced storage. Distributed sharding can partition optimizer state, gradients and parameters, while communication supplies the pieces required by each operation. ZeRO is a disclosed method for these partitions. A memory estimate must state its optimizer, precision and sharding configuration; multiplying parameter count by one universal bytes-per-parameter constant conceals the choices that determine whether training fits.\n\nTraining memory includes persistent parameters and optimizer arrays plus temporary gradients and saved forward intermediates. Their lifetimes differ. A coefficient payload calculation is useful only when it names precision and excluded structures. Activation rematerialization trades recomputation for retained intermediates, but is not implemented merely because a UI shows an estimated memory total. The lifecycle core exposes actual scalar tape-node counts and parameter arrays; those nodes are not hardware FLOPs or a byte-accurate JavaScript heap measurement.\n\n1.  List persistent parameter and optimizer tensors with their actual stored precision rather than assuming a single shared format.\n2.  Trace forward activations and identify when each is last needed by backward computation.\n3.  Choose which activations to retain and which to recompute under the available memory budget.\n4.  Partition replicated state where the distributed method supports it and insert required communication.\n5.  Measure peak live memory and completed step time, including temporary workspaces and exposed synchronization.\n6.  Count persistent weight and moment arrays under the chosen numerical representation.\n7.  Separate the current gradient, frozen reference and saved activations from those arrays.\n8.  Use a lifetime trace or measured heap evidence before claiming peak runtime memory.\n\n## Science and assumptions\n\n### Training capacity ledger\n\n$$\nM_{\\mathrm{peak}}=M_{\\mathrm{weights}}+M_{\\mathrm{gradients}}+M_{\\mathrm{optimizer}}+M_{\\mathrm{live\\ activations}}+M_{\\mathrm{workspace}}\n$$\n\nPeak memory is the sum of weights, gradients, optimizer state, live activations and workspace at the same device boundary and the same instant.\n\nDo not add independently occurring category maxima as though they were simultaneous.\n\nSymbols and units\n\nM\\_peak\n\nCombined occupancy at the chosen peak instant and device boundary\n\nbytes\n\nM\\_weights\n\nResident weight storage at that instant\n\nbytes\n\nM\\_gradients\n\nResident gradient storage at that instant\n\nbytes\n\nM\\_optimizer\n\nResident optimizer state at that instant\n\nbytes\n\nM\\_live activations\n\nSimultaneously live activation storage\n\nbytes\n\nM\\_workspace\n\nWorkspace occupied at that instant\n\nbytes\n\nEach term is bytes at the same device boundary and instant of peak occupancy.\n\nAssumptions Sharding and mixed precision alter each term; summing independent category maxima may overestimate the true simultaneous peak.\n\n### Training tensors with actual lifetimes · worked relation\n\n$$\nB_{payload}=P\\sum_k b_k\n$$\n\nInventory real trainable arrays separately from activation and runtime overhead.\n\nSymbols and units\n\nB\\_payload\n\nPayload of the explicitly inventoried full-length arrays\n\nbytes\n\nP\n\nParameters represented in every listed array\n\nparameters\n\nb\\_k\n\nBytes per parameter for listed array k\n\nbytes per parameter\n\nThe core has 423 binary64 weights, requiring 3384 raw bytes. Two Adam moment arrays add 6768 bytes, so weights plus moments total 10152. Adding one equally sized gradient vector yields 13536 raw bytes. A frozen reference adds another 3384. These sums exclude the scalar tape’s graph objects, history, JSON encoding, React state and allocator overhead.\n\nAssumptions 423 entries per full array, eight bytes per binary64 number; no claim that JavaScript object storage equals this payload. No rematerialization feature is executed.\n\n## Engineering example\n\nSynthetic parameter ledger: 1 billion parameters with 2-byte weights, 2-byte gradients and two 4-byte optimizer tensors occupy 12 GB before activations, master copies or workspace. Adding a 4-byte master weight would add 4 GB.\n\nThe core has 423 binary64 weights, requiring 3384 raw bytes. Two Adam moment arrays add 6768 bytes, so weights plus moments total 10152. Adding one equally sized gradient vector yields 13536 raw bytes. A frozen reference adds another 3384. These sums exclude the scalar tape’s graph objects, history, JSON encoding, React state and allocator overhead.\n\n## Distinctions and tradeoffs\n\nActivation checkpointing recomputes intermediate tensors; a recovery checkpoint saves durable state for restart. They solve different problems.\n\nRematerialization reduces saved activation memory but adds computation; sharding reduces replication while increasing data exchange and scheduling complexity.\n\nTraining can be compute-heavy and still fail a capacity limit or stall on collectives, so fitting the full step is separate from feeding its arithmetic.\n\n## Specifications and interfaces\n\nPersistent versus transient\n\nStep-to-step state versus temporary tensors\n\nOptimizer state persists; many activations expire after backward use.\n\nMemory budget\n\nBytes per state category\n\nInclude workspace, fragmentation and communication buffers separately.\n\n### Interfaces\n\nForward and optimizer configuration Input\n\nWeights, batch activations, gradients and update rules\n\nUpdated model state Output\n\nNew parameter values and retained optimizer history\n\n## Check your understanding\n\n### What cost accompanies discarding a forward activation that backward computation needs?\n\n1.  Recomputing it or an equivalent quantity later\n2.  Eliminating the backward pass\n3.  Automatically lowering the optimizer state size\n\nAnswer and reasoning\n\nRecomputing it or an equivalent quantity later\n\nRematerialization trades retained memory for computation; persistent optimizer storage is a separate category.\n\n### What does 10152 bytes include?\n\n1.  All future checkpoint JSON and activation history.\n2.  Weights and the two full Adam moment arrays only.\n3.  The complete browser heap at peak training.\n\nAnswer and reasoning\n\nWeights and the two full Adam moment arrays only.\n\nWeights and the two full Adam moment arrays only. 423 entries per full array, eight bytes per binary64 number; no claim that JavaScript object storage equals this payload. No rematerialization feature is executed.\n\n## Evidence and scope\n\nThis is an explicit synthetic format choice, not a universal training multiplier. ZeRO stages and implementation buffers must be identified before applying a sharding estimate.\n\nMode: executable arrays; payload arithmetic. 423 entries per full array, eight bytes per binary64 number; no claim that JavaScript object storage equals this payload. No rematerialization feature is executed. Numerical examples are authored calculations, except explicitly identified automated core fixtures. Primary sources support the scoped mechanism; no human specialist review or learner outcome is implied.\n\nEnable the optimizer-state option to budget weights and master/moment state; gradients, activations, workspaces and actual tier transfers are omitted.\n\n## Sources\n\n-   [ZeRO: Memory Optimizations Toward Training Trillion Parameter Models](https://arxiv.org/abs/1910.02054)\n-   [2 OLMo 2 Furious](https://arxiv.org/html/2501.00656v1)\n\n[Published lesson data](https://semiaifoundry.com/ai-atlas/lesson-data/training-memory-state-0e965c34b3054246.json) · [Atlas releases and licensing](https://semiaifoundry.com/ai-atlas/licenses.html)\n\n## Continue exploring\n\n### Foundations\n\n-   [Loss, gradients & optimizer](https://semiaifoundry.com/learn/atlas/loss-optimizer/)\n-   [Full memory hierarchy and input path](https://semiaifoundry.com/learn/atlas/memory-hierarchy-route/)\n\n### Related lessons\n\n-   [Data, tensor, pipeline, and state parallelism](https://semiaifoundry.com/learn/atlas/parallelism/)\n-   [Parallelism as a memory and communication budget](https://semiaifoundry.com/learn/atlas/parallelism-communication-budget/)\n-   [Next-token supervision from raw sequences](https://semiaifoundry.com/learn/atlas/causal-pretraining-objective/)\n-   [Initialization, learning rate and stable updates](https://semiaifoundry.com/learn/atlas/optimizer-schedule-stability/)\n-   [Gradient accumulation under a fixed global objective](https://semiaifoundry.com/learn/atlas/microbatch-equivalent-updates/)"
}