# Training state and activation rematerialization Canonical: https://semiaifoundry.com/learn/atlas/training-memory-state/ Revision: 2026-10-07 Rights: https://semiaifoundry.com/ai-atlas/licenses.html Model–machine co-design · Atlas v1.3.0 # Training state and activation rematerialization Budgets weights, gradients, optimizer state and saved activations across a training step. [Explore in Atlas](https://semiaifoundry.com/ai-atlas/#training-memory-state)One model step to accelerator memory hierarchy Learning objectiveBuild a training memory ledger that distinguishes state type, precision and lifetime. ## The mechanism Training 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. Training 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. 1. List persistent parameter and optimizer tensors with their actual stored precision rather than assuming a single shared format. 2. Trace forward activations and identify when each is last needed by backward computation. 3. Choose which activations to retain and which to recompute under the available memory budget. 4. Partition replicated state where the distributed method supports it and insert required communication. 5. Measure peak live memory and completed step time, including temporary workspaces and exposed synchronization. 6. Count persistent weight and moment arrays under the chosen numerical representation. 7. Separate the current gradient, frozen reference and saved activations from those arrays. 8. Use a lifetime trace or measured heap evidence before claiming peak runtime memory. ## Science and assumptions ### Training capacity ledger $$ M_{\mathrm{peak}}=M_{\mathrm{weights}}+M_{\mathrm{gradients}}+M_{\mathrm{optimizer}}+M_{\mathrm{live\ activations}}+M_{\mathrm{workspace}} $$ Peak memory is the sum of weights, gradients, optimizer state, live activations and workspace at the same device boundary and the same instant. Do not add independently occurring category maxima as though they were simultaneous. Symbols and units M\_peak Combined occupancy at the chosen peak instant and device boundary bytes M\_weights Resident weight storage at that instant bytes M\_gradients Resident gradient storage at that instant bytes M\_optimizer Resident optimizer state at that instant bytes M\_live activations Simultaneously live activation storage bytes M\_workspace Workspace occupied at that instant bytes Each term is bytes at the same device boundary and instant of peak occupancy. Assumptions Sharding and mixed precision alter each term; summing independent category maxima may overestimate the true simultaneous peak. ### Training tensors with actual lifetimes · worked relation $$ B_{payload}=P\sum_k b_k $$ Inventory real trainable arrays separately from activation and runtime overhead. Symbols and units B\_payload Payload of the explicitly inventoried full-length arrays bytes P Parameters represented in every listed array parameters b\_k Bytes per parameter for listed array k bytes per parameter The 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. Assumptions 423 entries per full array, eight bytes per binary64 number; no claim that JavaScript object storage equals this payload. No rematerialization feature is executed. ## Engineering example Synthetic 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. The 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. ## Distinctions and tradeoffs Activation checkpointing recomputes intermediate tensors; a recovery checkpoint saves durable state for restart. They solve different problems. Rematerialization reduces saved activation memory but adds computation; sharding reduces replication while increasing data exchange and scheduling complexity. Training 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. ## Specifications and interfaces Persistent versus transient Step-to-step state versus temporary tensors Optimizer state persists; many activations expire after backward use. Memory budget Bytes per state category Include workspace, fragmentation and communication buffers separately. ### Interfaces Forward and optimizer configuration Input Weights, batch activations, gradients and update rules Updated model state Output New parameter values and retained optimizer history ## Check your understanding ### What cost accompanies discarding a forward activation that backward computation needs? 1. Recomputing it or an equivalent quantity later 2. Eliminating the backward pass 3. Automatically lowering the optimizer state size Answer and reasoning Recomputing it or an equivalent quantity later Rematerialization trades retained memory for computation; persistent optimizer storage is a separate category. ### What does 10152 bytes include? 1. All future checkpoint JSON and activation history. 2. Weights and the two full Adam moment arrays only. 3. The complete browser heap at peak training. Answer and reasoning Weights and the two full Adam moment arrays only. Weights 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. ## Evidence and scope 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. Mode: 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. Enable the optimizer-state option to budget weights and master/moment state; gradients, activations, workspaces and actual tier transfers are omitted. ## Sources - [ZeRO: Memory Optimizations Toward Training Trillion Parameter Models](https://arxiv.org/abs/1910.02054) - [2 OLMo 2 Furious](https://arxiv.org/html/2501.00656v1) [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) ## Continue exploring ### Foundations - [Loss, gradients & optimizer](https://semiaifoundry.com/learn/atlas/loss-optimizer/) - [Full memory hierarchy and input path](https://semiaifoundry.com/learn/atlas/memory-hierarchy-route/) ### Related lessons - [Data, tensor, pipeline, and state parallelism](https://semiaifoundry.com/learn/atlas/parallelism/) - [Parallelism as a memory and communication budget](https://semiaifoundry.com/learn/atlas/parallelism-communication-budget/) - [Next-token supervision from raw sequences](https://semiaifoundry.com/learn/atlas/causal-pretraining-objective/) - [Initialization, learning rate and stable updates](https://semiaifoundry.com/learn/atlas/optimizer-schedule-stability/) - [Gradient accumulation under a fixed global objective](https://semiaifoundry.com/learn/atlas/microbatch-equivalent-updates/)