{"id":"training-memory-state","level":"software","name":"Training state and activation rematerialization","shortName":"Training state and activation rematerialization","role":"Budgets weights, gradients, optimizer state and saved activations across a training step.","description":"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.\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.","specs":[{"label":"Persistent versus transient","value":"Step-to-step state versus temporary tensors","context":"Optimizer state persists; many activations expire after backward use."},{"label":"Memory budget","value":"Bytes per state category","context":"Include workspace, fragmentation and communication buffers separately."}],"interfaces":[{"name":"Forward and optimizer configuration","direction":"in","signal":"Weights, batch activations, gradients and update rules"},{"name":"Updated model state","direction":"out","signal":"New parameter values and retained optimizer history"}],"tradeoff":"Rematerialization reduces saved activation memory but adds computation; sharding reduces replication while increasing data exchange and scheduling complexity.","whyAI":"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.","related":["loss-optimizer","parallelism","parallelism-communication-budget","causal-pretraining-objective","optimizer-schedule-stability","microbatch-equivalent-updates"],"sources":[{"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"}],"scale":"One model step to accelerator memory hierarchy","kind":"workload","branch":"Model–machine co-design","subbranch":"Training, prefill and decode","mechanism":["List persistent parameter and optimizer tensors with their actual stored precision rather than assuming a single shared format.","Trace forward activations and identify when each is last needed by backward computation.","Choose which activations to retain and which to recompute under the available memory budget.","Partition replicated state where the distributed method supports it and insert required communication.","Measure peak live memory and completed step time, including temporary workspaces and exposed synchronization.","Count persistent weight and moment arrays under the chosen numerical representation.","Separate the current gradient, frozen reference and saved activations from those arrays.","Use a lifetime trace or measured heap evidence before claiming peak runtime memory."],"misconception":"Activation checkpointing recomputes intermediate tensors; a recovery checkpoint saves durable state for restart. They solve different problems.","check":{"question":"What cost accompanies discarding a forward activation that backward computation needs?","options":["Recomputing it or an equivalent quantity later","Eliminating the backward pass","Automatically lowering the optimizer state size"],"answer":0,"explanation":"Rematerialization trades retained memory for computation; persistent optimizer storage is a separate category."},"science":[{"title":"Training capacity ledger","equation":"Mpeak = Mweights + Mgradients + Moptimizer + Mlive_activations + Mworkspace","explanation":"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."},{"title":"Training tensors with actual lifetimes · worked relation","equation":"B_{payload}=P\\sum_k b_k","explanation":"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."}],"prerequisites":["loss-optimizer","memory-hierarchy-route"],"learningObjective":"Build a training memory ledger that distinguishes state type, precision and lifetime.","workstream":"W11","visualFamily":"model-workload","engineeringExample":"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.\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.","evidenceNotes":"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.","checks":[{"question":"What does 10152 bytes include?","options":["All future checkpoint JSON and activation history.","Weights and the two full Adam moment arrays only.","The complete browser heap at peak training."],"answer":1,"explanation":"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."}],"labId":"engineering","scope":"Enable the optimizer-state option to budget weights and master/moment state; gradients, activations, workspaces and actual tier transfers are omitted.","deepLabId":"footprint","investigationId":"execution","labScope":"Enable the optimizer-state option to budget weights and master/moment state; gradients, activations, workspaces and actual tier transfers are omitted."}