Chapter 13 · Content revision 1

Accumulate gradients and recompute activations

Preserve token-weighted loss across smaller batches, replay decoder blocks with the same dropout choices, and measure retained tensor payload at a declared point.

The problem: one training update can retain too much

A training forward pass keeps values that reverse mode will need to calculate gradients. More batch rows, longer sequences and more decoder blocks increase the amount retained. Parameters and AdamW state also occupy memory, but reducing the saved forward values addresses a different part of the problem.

Two choices help at different points. Gradient accumulation processes a chosen logical batch in smaller microbatches, adds their gradients and delays the optimizer until every contribution is ready. Each microbatch’s graph can then be released before the next forward pass. Activation recomputation keeps block boundaries and reconstructs block interiors during backward, exchanging retained values for extra forward calculations. Both choices must preserve the loss whose gradient the optimizer receives.

First consider a small constructed loss example. One microbatch contains $2$ valid next-token targets with mean loss $1$ ; another contains $4$ with mean loss $3$ . Their loss sums are $2$ and $12$ . The combined mean is therefore $14/6=7/3$ . Averaging the two means equally would give $2$ , because it would give the smaller microbatch’s targets more weight. Padded positions do not add targets to that denominator.

The decoder experiment uses three checked text examples containing $2$ , $4$ and $3$ valid targets: $9$ in all. With the default maximum of two rows per microbatch, their token grids have shapes $[2,4]$ and $[1,3]$ , where the axes count rows and supplied positions per row. The first microbatch contains $6$ valid targets and the final one contains $3$ . Its final row count is smaller, but every real target still belongs to the same logical update.

The full-batch and accumulated runs start from matching seed- $39$ models with three decoder blocks, residual width $4$ , one attention head and feed-forward width $8$ . They use the existing text encoder and tied text head, with $29$ named parameter tensors and $1572$ learned scalars. Dropout is disabled for this comparison so that changing the grouping does not change random masks.

Weight each mean by the targets it represents

Let $K$ count the microbatches in one logical update and let $i$ index them from $1$ through $K$ . Microbatch $i$ has $n_i>0$ valid targets and scalar mean negative log likelihood $L_i$ . Let $N$ be the total target count and $L$ the mean over that total. All forwards use the same parameters $\theta$ , held fixed until the update. Hold the finite positive loss scale $S$ fixed as well. Then, in exact arithmetic,

\begin{aligned} N&=\sum_{i=1}^{K}n_i,\\ L&=\sum_{i=1}^{K}\frac{n_i}{N}L_i,\\ \widetilde g&=\sum_{i=1}^{K}\nabla_\theta\!\left(S\frac{n_i}{N}L_i\right),\\ g&=\widetilde g/S. \end{aligned}

Here $\nabla_\theta$ takes the derivative with respect to the model parameters. $\widetilde g$ is the accumulated scaled gradient, and $g$ is its unscaled value before clipping. The target counts do not depend on the parameters, so linearity of differentiation gives the gradient of the combined mean when each microbatch represents the same per-example objective. Padding contributes to neither $n_i$ nor $N$ .

Reverse mode can multiply a scalar loss gradient by choosing its starting seed. With $S=128$ , the recorded f64 seed for the first microbatch’s target share $6/9$ is $85.33333333333333$ ; the final microbatch’s share $3/9$ uses $42.666666666666664$ . Each backward adds its contribution to the original parameter leaves. The optimizer has not changed those leaves’ values between the two passes.

After all contributions succeed, divide by $S$ , measure the global gradient norm, apply the existing norm clipping and call AdamW once. Count the tied input/output table once in that global norm. AdamW’s moments, decoupled weight decay and step advance belong to that one update; clipping each microbatch separately would generally change the resulting gradient. Finally clear the gradients and make one loss-scale growth decision.

In the measured CPU f64 comparison, both groupings have mean loss $5.622831723383019$ and global gradient norm $0.9389921230306614$ . The norm limit is $1$ , so neither run clips. Both commit optimizer step $1$ , clear gradients and change the scale from $128$ to $256$ because this fixture uses a growth interval of one committed update.

The comparison checks every coordinate of all $29$ scaled gradient tensors and the updated weights. For reference coordinate $r$ and compared coordinate $x$ , it requires $|x-r|\le 10^{-11}+10^{-10}|r|$ . The maximum observed absolute scaled-gradient error is $1.7763568394002505\times10^{-15}$ ; the maximum weight error is $3.469446951953614\times10^{-18}$ . These observations support the declared tolerance for this fixture, rather than bitwise equality for arbitrary groupings or backends.

Reconstruct the same block before differentiating it

Accumulation shortens the lifetime of each microbatch’s graph. Within a graph, ordinary training still retains each decoder block’s intermediate values and saved derivative contexts. A block activation checkpoint replaces that interior tape with enough information to run the same block again.

With graph recording enabled, the initial checkpointed block forward uses a no-grad context under the selected executor. After obtaining the block output, it drops the interior tensors and retains the input boundary, aliases to the original parameters, their revisions, the exact layout and positions, the execution session and precision, and a block recipe. The recipe also retains fixed rotary-position tables and compact tickets for the block’s two dropout sites. A parameter alias refers to the original learned node; it does not copy another set of learned weights. Text lookup and the tied text head remain outside these block checkpoints.

At backward time, the reverse engine first checks the original operand revisions. It recreates the block from a fresh tracked boundary leaf with the same parameter nodes, layout, execution and tickets. It checks the output’s shape and f64 bits against the original block output before calculating the block’s local derivatives. Those derivatives provide a contribution for the input boundary and for each block parameter. The local calculation does not write gradients to the original parameters; the complete outer reverse pass stages and checks all contributions before committing leaf gradients. A successful releasing backward clears the checkpoint recipes and edges.

Dropout replay is essential to reconstructing that same computation. The separate recomputation experiment keeps the layout $[B,T,D]=[3,4,4]$ , with $B$ batch rows, $T$ positions per row and $D$ residual features. At dropout probability $0.5$ and seed $39$ , the three blocks use six sites. Each site consumes $48$ counter words, including padded positions. Their consecutive counter ranges are $[0,48)$ , $[48,96)$ , $[96,144)$ , $[144,192)$ , $[192,240)$ and $[240,288)$ .

The compact tickets reproduce all six original masks. Reverse replay receives no live random counter, so it leaves the live counter at $288$ after backward. A seed alone would not provide this guarantee after repartitioning: changing microbatch shape or padding changes the existing counter assignment to sites. That is why the dropout replay comparison keeps the layout fixed and remains separate from the accumulation comparison.

The retained and recomputed dropout runs have bit-identical forward logits. Their measured maximum scaled-gradient and updated-weight errors are both $0$ in this fixture. Both have mean loss $5.535604128733965$ and global norm $1.4853618224483591$ , clipped to $1$ before their one AdamW update. Each clears gradients and advances the scale from $128$ to $256$ . Reconstructing the three blocks performs three extra block forwards during backward. These are bounded observations, with no universal bitwise or speed claim.

Measure what remains after forward

The observation point is completion of the loss forward, before backward. For the tape reachable from the scalar mean-loss node, retained_tape_storage counts unique nonparameter node primals, the tensor payload of each edge’s saved context, and the fixed rotary-position tensors captured by checkpoint recipes. A primal is a stored forward value; a saved context holds values needed by a local derivative. Multiplying their f64 value counts by $8$ gives the reported bytes.

For the fixed-layout dropout experiment, ordinary retention keeps $16808$ nonparameter primal values and $9736$ saved-context values: $212352$ bytes. Recomputation keeps $10241$ primal values, $3751$ saved-context values and $102$ fixed recipe values: $112752$ bytes. The last count includes, per block, $2$ inverse frequencies, $16$ cosines and $16$ sines, for $816$ recipe bytes across three blocks. The reachable tape has $210$ nodes ordinarily and $48$ with recomputation, including its three checkpoint nodes.

Saved contexts count any copied parameter primals they own. The observation excludes the original master-parameter buffers, gradient and optimizer buffers, Rust metadata, layout vectors, tickets, allocation capacity and transient forward or reverse work. It therefore measures retained tensor payload at this point, not total training memory, process memory, a peak or GPU storage. The added backward work can have a different temporary lifetime from the values counted here.

Trade retained activations for repeated computation

Keeping every intermediate available for backward consumes more retained tensor storage as decoder batches, sequences and depth grow. Chen and colleagues’ 2016 paper, Training Deep Nets with Sublinear Memory Cost, presents a computation-versus-memory tradeoff by recomputing intermediate activations. The useful intuition is to retain selected boundaries and reconstruct what backward needs between them. Repeating whole blocks also repeats forward work.

Selective recomputation chooses which intermediate calculations to repeat instead of recomputing every operation in a block. Korthikanti and colleagues’ 2022 paper, Reducing Activation Recomputation in Large Transformer Models, studies selective recomputation and sequence parallelism to reduce unnecessary recomputation in large transformer language-model training. Transformer training must preserve each block’s parameters, layout and random choices when reconstructing the values needed for its gradients.

We retain block boundaries and repeat the existing block computation during backward, while accumulating token-weighted microbatch gradients before one shared optimizer update. The Rust example implements bounded whole-block recomputation. Its retained-payload observation and extra-forward count do not reproduce either paper’s schedules, parallelism or performance results.

Extend the existing update and reverse owners

decoder_update_with_accumulation accepts a checked VariableBatchEpoch as the explicit group for one logical update. The caller chooses that group; the type’s name does not add whole-dataset scheduling. The entry requires training batches, positive checked target counts and cleared model gradients. The total must be at most $2^{53}$ so its integer count converts exactly to f64. Each callback returns the existing MaskedLoss with the checked batch’s target count.

The entry seeds each mean backward with its scaled target share, releases that graph and accumulates gradients owned by the same parameter leaves. After the last microbatch, finish_decoder_update performs the shared unscale, global clipping, AdamW and gradient clearing. The loss-scale controller advances only after that successful logical update.

Weight each mean loss by its valid-target share, release each graph and finish one logical update. rust/crates/llm-from-scratch-practical/src/training/trainer.rs#accumulated-decoder-update

/// One logical update: token-weighted microbatch backward passes, then one
/// unscale, global clipping, AdamW step, clear and loss-scale growth decision.
/// A failure clears all partial gradients; caller-owned forward/observer side
/// effects (including live RNG advancement) are outside that rollback boundary.
#[allow(clippy::too_many_arguments)]
pub fn decoder_update_with_accumulation<F, O>(
    model: &DecoderModel,
    optimizer: &mut AdamW,
    config: &TrainerConfig,
    step: usize,
    execution: &TensorExecution,
    scaler: &mut DynamicLossScaler,
    microbatches: &VariableBatchEpoch,
    mut forward: F,
    observer: &mut O,
) -> Result<AccumulatedDecoderUpdate, TrainerError>
where
    F: FnMut(AutogradContext<'_>, &VariableBatch) -> Result<MaskedLoss, TrainerError>,
    O: DecoderUpdateObserver,
{
    let (learning_rate, expected_step) = validate_update_step(optimizer, config, step)?;
    let total = microbatches.valid_target_count()?;
    if total == 0 || !u64::try_from(total).is_ok_and(|n| n <= (1_u64 << 53)) {
        return Err(TrainerError::InvalidAccumulation {
            reason: "valid-target total must be positive and exactly representable in f64",
        });
    }
    if microbatches
        .batches()
        .iter()
        .any(|b| b.partition() != Partition::Train || b.valid_target_count() == 0)
    {
        return Err(TrainerError::InvalidAccumulation {
            reason: "every microbatch must have train targets",
        });
    }
    for parameter in model.parameters() {
        let gradient =
            parameter
                .tensor()
                .gradient()
                .ok_or_else(|| TrainerError::MissingGradient {
                    name: parameter.name().into(),
                })?;
        if gradient.as_slice().iter().any(|&g| g != 0.0) {
            return Err(TrainerError::ScaledUpdateRequiresZeroGradient {
                name: parameter.name().into(),
            });
        }
    }
    let scale = scaler.scale();
    let context = AutogradContext::recording().with_execution(execution);
    let mut train_loss = 0.0;
    for (index, batch) in microbatches.batches().iter().enumerate() {
        let attempt = (|| {
            let loss = forward(context, batch)?;
            if loss.count() != batch.valid_target_count() {
                return Err(TrainerError::InvalidAccumulation {
                    reason: "forward loss count differs from the checked microbatch count",
                });
            }
            let weight = loss.count() as f64 / total as f64;
            train_loss += scalar_loss(loss.mean())? * weight;
            if !train_loss.is_finite() {
                return Err(TrainerError::NonFiniteLoss { value: train_loss });
            }
            observer.after_forward(loss.mean())?;
            let seed = Tensor::from_vec(Vec::new(), vec![scale * weight])?;
            loss.mean().backward_with_seed_and_execution(
                &seed.view(),
                GraphRetention::Release,
                execution,
            )?;
            Ok(())
        })();
        if let Err(error) = attempt {
            clear_and_verify_gradients(model)?;
            if let TrainerError::Autodiff(overflow) = error {
                if is_backward_overflow(&overflow) {
                    scaler.overflow();
                    return Ok(AccumulatedDecoderUpdate::Skipped {
                        microbatches_completed: index,
                        loss_scale: scale,
                        next_loss_scale: scaler.scale(),
                        overflow,
                    });
                }
                return Err(TrainerError::Autodiff(overflow));
            }
            return Err(error);
        }
    }
    let update = finish_decoder_update(
        model,
        optimizer,
        config,
        step,
        learning_rate,
        expected_step,
        train_loss,
        scale,
        observer,
    );
    let applied = match update {
        Ok(update) => update,
        Err(error) => {
            clear_and_verify_gradients(model)?;
            return Err(error);
        }
    };
    scaler.committed();
    Ok(AccumulatedDecoderUpdate::Applied {
        execution: applied,
        microbatches: microbatches.batch_count(),
        valid_target_count: total,
        loss_scale: scale,
        next_loss_scale: scaler.scale(),
    })
}

The checkpoint node belongs to the existing tensor tape. It stores boundary and parameter edges plus the explicit execution and recipe; it does not introduce another automatic-differentiation engine.

Keep a block boundary and original parameter aliases with the execution and recipe needed for replay. rust/crates/llm-from-scratch-practical/src/autograd/tensor_core.rs#block-activation-checkpoint

    /// Retain a boundary and parameter aliases, rather than the block's inner tape.
    /// Account explicitly for fixed tensors captured by the crate-owned recipe.
    pub(crate) fn checkpoint_with_context(
        context: AutogradContext,
        input: &Self,
        parameters: &[Self],
        output: Tensor,
        recipe_tensor_values: usize,
        recompute: impl for<'a> Fn(&Self, AutogradContext<'a>) -> Result<Self, TensorAutodiffError>
        + 'static,
    ) -> Result<Self, TensorAutodiffError> {
        if !context.records_graph() {
            return Self::constant(output);
        }
        let mut parents = Vec::with_capacity(parameters.len() + 1);
        parents.push(ParentEdge::capture(input, TensorSavedContext::Checkpoint));
        parents.extend(
            parameters
                .iter()
                .map(|p| ParentEdge::capture(p, TensorSavedContext::Checkpoint)),
        );
        let node = Self::operation_node(context, output, TensorOperation::Checkpoint, parents)?;
        node.node.state.borrow_mut().checkpoint = Some(Checkpoint {
            recompute: Rc::new(recompute),
            execution: context.execution().clone(),
            counter: RecomputationCounter::default(),
            recipe_tensor_values,
        });
        Ok(node)
    }

Its local pullback uses the same reverse engine in a mode that returns selected derivatives without committing original parameter gradients. That keeps shared parameter contributions under the outer transaction’s control.

Verify the reconstructed block output, then return local boundary and parameter derivatives to the outer reverse pass. rust/crates/llm-from-scratch-practical/src/autograd/tensor_core.rs#recomputed-block-vjp

    fn checkpoint_vjps(
        &self,
        checkpoint: &Checkpoint,
        parents: &[ParentEdge],
        upstream: &Tensor,
        execution: &TensorExecution,
    ) -> Result<Vec<Tensor>, TensorAutodiffError> {
        if execution != &checkpoint.execution {
            return Err(TensorAutodiffError::Checkpoint {
                reason: "backward must use the forward execution session and precision",
            });
        }
        let boundary = Self::parameter(parents[0].parent.value_snapshot())?;
        let output = (checkpoint.recompute)(
            &boundary,
            AutogradContext::recording().with_execution(execution),
        )?;
        if output.shape() != self.shape()
            || !output
                .value()
                .as_slice()
                .iter()
                .map(|x| x.to_bits())
                .eq(self.value().as_slice().iter().map(|x| x.to_bits()))
        {
            return Err(TensorAutodiffError::Checkpoint {
                reason: "recomputed output differs from the original block output",
            });
        }
        checkpoint
            .counter
            .0
            .set(checkpoint.counter.0.get().checked_add(1).ok_or(
                TensorAutodiffError::Checkpoint {
                    reason: "recomputation counter overflows usize",
                },
            )?);
        let selected = std::iter::once(boundary)
            .chain(parents.iter().skip(1).map(|e| e.parent.clone()))
            .collect::<Vec<_>>();
        // The existing reverse engine computes local VJPs without committing
        // parameter gradients. Only the complete outer pass may commit them.
        let adjoints = output.reverse_with_observer(
            &upstream.view(),
            GraphRetention::Release,
            RecordSelectedAdjoints {
                selected: &selected,
            },
            execution,
            false,
        )?;
        adjoints
            .into_iter()
            .zip(selected)
            .map(|(adjoint, parent)| adjoint.map_or_else(|| zeros(&parent.shape()), Ok))
            .collect()
    }

The checked text path still composes TextInputEncoder, the source-independent backbone and the tied head. Compact replay tickets change what the block tape retains; tokenization and text loss keep their existing owners. The backbone continues to consume vectors, positions and attention layout without requiring a tokenizer or vocabulary.

Use the checked text encoder and existing backbone with compact dropout tickets for block recomputation. rust/crates/llm-from-scratch-practical/src/models/decoder_dropout.rs#recomputed-text-forward

    /// Extend the checked text path with per-block recomputation. Train mode
    /// draws only in the initial forward; reverse replay accepts no live counter.
    #[allow(clippy::too_many_arguments)]
    pub fn forward_padded_recomputed_with_dropout(
        &self,
        context: AutogradContext,
        batch: &VariableBatch,
        expected: &BatchBinding,
        mode: DropoutMode,
        policy: &DropoutRunPolicy,
        state: &mut DropoutCounterState,
        limits: DropoutLimits,
    ) -> Result<RecomputedModelForward, PaddedDecoderError> {
        self.validate_padded_text(batch, expected)?;
        let p = policy
            .policy()
            .map_err(|_| DecoderModelError::Dropout(DropoutError::Binding("invalid run policy")))?;
        let shape = [
            batch.shape()[0],
            batch.shape()[1],
            self.config().model_width(),
        ];
        let mut dropout = if p.active(mode) {
            self.fresh_dropout_execution(
                context,
                shape,
                policy,
                checked_layout(batch).map_err(DecoderModelError::Dropout)?,
                state,
                limits,
            )
            .and_then(DropoutExecution::into_compact)
            .map_err(DecoderModelError::Dropout)?
        } else {
            DropoutExecution::Identity
        };
        let input =
            TextInputEncoder::padded(self, batch, expected).encode(context, &self.backbone())?;
        let forward =
            self.forward_prepared_text_with_activation_policy(input, &mut dropout, true)?;
        let tickets = dropout
            .finish_compact(Some(state))
            .map_err(DecoderModelError::Dropout)?;
        Ok(RecomputedModelForward { forward, tickets })
    }

Discard a failed partial logical update

An earlier microbatch may have finished backward when a later one fails. The accumulation owner clears those partial model gradients before returning a forward, observer or reverse failure. A numerical overflow in local backward or gradient addition returns Skipped and backs off the loss scale; structural errors propagate as errors. Neither case applies AdamW before the shared finish phase. Arbitrary effects inside caller callbacks and advancement of a live RNG are outside this cleanup; a retry that needs identical randomness must restore caller-owned RNG state or use replay.

The overflow example first commits a warm-up so AdamW already has moments. For the second microbatch of the next attempted update, it substitutes a constructed zero-logit branch multiplied by f64::MAX. Forward remains finite, but backward becomes nonfinite. One microbatch’s model gradients have already accumulated. The skip clears them and preserves the original zero-gradient bits, weight bits and revisions, AdamW moments and beta powers, and optimizer step $1$ . Scale changes from $128$ to $64$ . Scheduled step $2$ remains available to the caller; this example records that eligibility without performing a retry. An AdamW step with zero gradients can still apply weight decay and advance optimizer state; the skipped attempt does neither. The injection is a failure probe rather than ordinary model logits.

Follow one update and two ways to retain block values

Accumulate one update; reconstruct block interiors

Microbatches contribute their valid-target shares with model parameters and loss scale fixed, then combine gradients before one unscale, global clipping decision and AdamW step. Separately, fixed-layout backward uses retained block interiors or reconstructs them from boundaries, parameter aliases and tickets that reproduce the original dropout choices. Recomputing adds block forwards and reduces the tensor payload retained at loss-forward completion under the exclusions stated with the byte counts; this observation does not measure total or peak training memory.

Count targets, accumulate contributions, update once

Valid targets in the logical update: $N=9$ . Loss scale fixed across microbatches: $S=128$ .

First microbatch

Valid next-token targets: $6$

Share of the logical mean loss: $\frac{6}{9}$

Scaled backward seed: $85.33333333333333$
Final microbatch

Valid next-token targets: $3$

Share of the logical mean loss: $\frac{3}{9}$

Scaled backward seed: $42.666666666666664$
After every microbatch succeeds

Unscale the combined gradient, measure its global norm, clip if needed and apply AdamW once. Then clear gradients and count one committed update toward loss-scale growth.

Optimizer step after the accumulated update: $1$

Loss scale before and after the committed update: $128\rightarrow256$

Keep layout and dropout choices fixed while changing retention

Ordinary tape: retain block interiors

Backward uses the block primals and saved derivative contexts retained by the original forward.

Tensor payload at loss-forward completion: $212352\,\mathrm{bytes}$

Checkpointed tape: reconstruct block interiors

Retain boundaries, original parameter aliases, fixed rotary tables and compact dropout tickets. Backward repeats each block and regenerates its original masks before returning local derivatives to the outer reverse pass.

Tensor payload at loss-forward completion: $112752\,\mathrm{bytes}$

Additional block forwards during backward: $3$

Both paths use the same batch layout and dropout-site tickets. Replay consumes the recorded counter ranges without advancing the live random counter.

Batch rows, positions and residual features: $[B,T,D]=[3,4,4]$ . Live random counter after forward and backward: $288$ .

Count the tape reachable from the scalar mean-loss node: nonparameter primals, owned edge saved-context tensors and fixed recipe tensors. Include copied parameter primals in saved contexts; exclude original master-parameter, gradient and optimizer buffers, Rust metadata, layout and tickets, allocation capacity and transient forward or reverse work. These bytes are retained tensor payload at one observation point, not total training memory, process memory, peak memory or GPU storage.

Reproduce the comparisons and explain their boundaries

The existing clap CLI sets the maximum rows per microbatch to $1$ , $2$ or $3$ , defaulting to $2$ , and prints the bounded demonstration’s JSON record.

Use the existing clap CLI to choose a microbatch row limit from one to three, defaulting to two, and print the bounded demonstration JSON. rust/demos/practical-ch13-gradient-accumulation-recomputation/src/main.rs

use clap::Parser;

#[derive(Parser)]
#[command(
    about = "Compare token-weighted microbatches and per-block activation recomputation on a bounded text decoder"
)]
struct Args {
    #[arg(long, default_value_t = 2, value_parser = clap::value_parser!(u8).range(1..=3))]
    microbatch_size: u8,
}

fn main() -> Result<(), Box<dyn std::error::Error>> {
    println!(
        "{}",
        serde_json::to_string(
            &practical_ch13_gradient_accumulation_recomputation::demonstrate(usize::from(
                Args::parse().microbatch_size
            ))?
        )?
    );
    Ok(())
}

From the repository root, with the existing course Docker setup available, run the default CPU example. It needs no GPU, external corpus or downloaded model:

COURSE_CORPUS=false ./course run cargo run --quiet --offline --locked -p practical-ch13-gradient-accumulation-recomputation

The command prints one JSON record with accumulation, recomputation and late_overflow objects. accumulation.microbatch_valid_targets is [6,3]; the corresponding shapes are [[2,4],[1,3]]. Both accumulation updates reach optimizer step $1$ . The recomputation object reports all six ticket and mask matches, the retained-payload counts and three extra block forwards.

Compare full and microbatch updates from matching initial models with dropout disabled and the declared f64 tolerance. rust/demos/practical-ch13-gradient-accumulation-recomputation/src/lib.rs#logical-batch-comparison

/// Disable dropout for this comparison: repartitioning changes counter assignment.
pub fn accumulation_probe(microbatch_size: usize) -> DemoResult<Value> {
    let full = batches(3)?;
    let micro = batches(microbatch_size)?;
    let execution = TensorExecution::cpu();
    let full_model = model()?;
    let micro_model = model()?;
    let mut full_optimizer = optimizer()?;
    let mut micro_optimizer = optimizer()?;
    let mut full_scale = DynamicLossScaler::new(128., 1., 256., 1)?;
    let mut micro_scale = full_scale.clone();
    let mut reference = Observe::default();
    let mut accumulated = Observe::default();
    let full_update = applied(decoder_update_with_accumulation(
        &full_model,
        &mut full_optimizer,
        &plan()?,
        1,
        &execution,
        &mut full_scale,
        &full,
        |context, batch| {
            Ok(full_model.loss_padded_with_context(context, batch, batch.binding())?)
        },
        &mut reference,
    )?)?;
    let micro_update = applied(decoder_update_with_accumulation(
        &micro_model,
        &mut micro_optimizer,
        &plan()?,
        1,
        &execution,
        &mut micro_scale,
        &micro,
        |context, batch| {
            Ok(micro_model.loss_padded_with_context(context, batch, batch.binding())?)
        },
        &mut accumulated,
    )?)?;
    let gradient_error = maximum_error(&reference.gradients, &accumulated.gradients)?;
    let weight_error = maximum_error(&weights(&full_model), &weights(&micro_model))?;
    Ok(json!({
        "dropout":"disabled","precision":"f64","target_count":micro.valid_target_count()?,
        "microbatch_shapes":micro.batches().iter().map(|b|b.shape()).collect::<Vec<_>>(),
        "microbatch_valid_targets":micro.batches().iter().map(|b|b.valid_target_count()).collect::<Vec<_>>(),
        "backward_seeds":micro.batches().iter().map(|b|128.*(b.valid_target_count() as f64/9.)).collect::<Vec<_>>(),
        "full_update":full_update,"accumulated_update":micro_update,
        "gradient_tensors_compared":reference.gradients.len(),"gradient_values_compared":reference.gradients.iter().map(Tensor::len).sum::<usize>(),
        "maximum_absolute_scaled_gradient_error":gradient_error,"maximum_absolute_weight_error":weight_error,
        "absolute_tolerance":practical_ch11_accelerated_tensor_execution::ABSOLUTE_TOLERANCE,
        "relative_tolerance":practical_ch11_accelerated_tensor_execution::RELATIVE_TOLERANCE,
        "updates_observed":[reference.before_updates,accumulated.before_updates],
        "maximum_retained_tape_bytes_full":reference.storage.iter().map(|s|s.tensor_payload_bytes()).collect::<Result<Vec<_>,_>>()?.into_iter().max(),
        "maximum_retained_tape_bytes_micro":accumulated.storage.iter().map(|s|s.tensor_payload_bytes()).collect::<Result<Vec<_>,_>>()?.into_iter().max(),
    }))
}

Keep the same working directory and Docker setup to run the other supported row limits and the chapter’s seven library tests:

COURSE_CORPUS=false ./course run cargo run --quiet --offline --locked -p practical-ch13-gradient-accumulation-recomputation -- --microbatch-size 1
COURSE_CORPUS=false ./course run cargo run --quiet --offline --locked -p practical-ch13-gradient-accumulation-recomputation -- --microbatch-size 3
COURSE_CORPUS=false ./course run cargo test --offline --locked -p practical-ch13-gradient-accumulation-recomputation

The row-limit variants retain the nine-target logical update and print the same record structure with changed accumulation grouping. The test command reports seven passing library tests. They cover row limits $1$ , $2$ and $3$ , partial final batches, six-site replay, late overflow, structural-error cleanup, stale revisions before replay, packed positions with recorded f64, FP16 and BF16 execution, and refusal of dirty gradients or non-training data. The default measurements described in this lesson use CPU f64 arithmetic.

Hold layout and dropout sites fixed while observing numerical agreement, retained tensor payload and repeated block forwards. rust/demos/practical-ch13-gradient-accumulation-recomputation/src/lib.rs#recomputation-comparison

/// Keep layout, site ranges and scale fixed while changing activation retention.
pub fn recomputation_probe() -> DemoResult<Value> {
    let epoch = batches(3)?;
    let execution = TensorExecution::cpu();
    let policy = DropoutRunPolicy::new(0.5, 39)?;
    let retained_model = model()?;
    let recomputed_model = model()?;
    let mut retained_optimizer = optimizer()?;
    let mut recomputed_optimizer = optimizer()?;
    let mut retained_scale = DynamicLossScaler::new(128., 1., 256., 1)?;
    let mut recomputed_scale = retained_scale.clone();
    let mut retained_state = DropoutCounterState::from_seed(39);
    let mut replay_state = retained_state.clone();
    let mut retained = Observe::default();
    let mut recomputed = Observe::default();
    let mut reference_tickets = Vec::new();
    let mut reference_masks = Vec::new();
    let mut reference_logits = None;
    let retained_update = applied(decoder_update_with_accumulation(
        &retained_model,
        &mut retained_optimizer,
        &plan()?,
        1,
        &execution,
        &mut retained_scale,
        &epoch,
        |context, batch| {
            let forward = retained_model.forward_padded_with_dropout(
                context,
                batch,
                batch.binding(),
                DropoutMode::Train,
                &policy,
                &mut retained_state,
                DropoutLimits::default(),
            )?;
            reference_tickets = forward.sites().iter().map(|s| s.ticket().clone()).collect();
            reference_masks = forward.sites().iter().map(|s| s.mask()).collect();
            reference_logits = Some(forward.logits().value_snapshot());
            Ok(batch.loss_with_context(context, forward.logits())?)
        },
        &mut retained,
    )?)?;
    let mut tickets = Vec::new();
    let mut logits = None;
    let recomputed_update = applied(decoder_update_with_accumulation(
        &recomputed_model,
        &mut recomputed_optimizer,
        &plan()?,
        1,
        &execution,
        &mut recomputed_scale,
        &epoch,
        |context, batch| {
            let forward = recomputed_model.forward_padded_recomputed_with_dropout(
                context,
                batch,
                batch.binding(),
                DropoutMode::Train,
                &policy,
                &mut replay_state,
                DropoutLimits::default(),
            )?;
            tickets = forward.tickets().to_vec();
            logits = Some(forward.logits().value_snapshot());
            Ok(batch.loss_with_context(context, forward.logits())?)
        },
        &mut recomputed,
    )?)?;
    let masks_equal = tickets.iter().zip(&reference_masks).all(|(ticket, mask)| {
        let p = DropoutProbability::from_bits(ticket.probability_bits())
            .expect("validated ticket probability");
        ticket
            .words()
            .into_iter()
            .map(|w| u8::from(p.keeps_word(w)))
            .eq(mask.iter().copied())
    }) && tickets.len() == reference_masks.len();
    if tickets != reference_tickets || !masks_equal || replay_state != retained_state {
        return Err("site-local dropout replay changed".into());
    }
    let logits = logits.unwrap();
    let reference_logits = reference_logits.unwrap();
    let logit_bits_equal = logits
        .as_slice()
        .iter()
        .map(|x| x.to_bits())
        .eq(reference_logits.as_slice().iter().map(|x| x.to_bits()));
    if !logit_bits_equal {
        return Err("checkpointed forward logits changed".into());
    }
    let gradient_error = maximum_error(&retained.gradients, &recomputed.gradients)?;
    let weight_error = maximum_error(&weights(&retained_model), &weights(&recomputed_model))?;
    let kept = retained.storage[0];
    let compact = recomputed.storage[0];
    if compact.tensor_payload_bytes()? >= kept.tensor_payload_bytes()?
        || compact.checkpoint_nodes != 3
        || recomputed
            .counters
            .iter()
            .any(|c| c.completed_forwards() != 1)
    {
        return Err("block recomputation did not reduce retained tape payload".into());
    }
    Ok(json!({
        "dropout_probability":0.5,"seed":39,"layout":[3,4,4],"site_tickets_equal":tickets==reference_tickets,"site_masks_equal":masks_equal,
        "logit_bits_equal":logit_bits_equal,"maximum_absolute_scaled_gradient_error":gradient_error,"maximum_absolute_weight_error":weight_error,
        "retained_update":retained_update,"recomputed_update":recomputed_update,
        "live_counter_after_forward_and_backward":replay_state.next_counter(),
        "site_ranges":tickets.iter().map(|t|json!({"layer":t.site().layer,"start":t.start(),"end":t.end()})).collect::<Vec<_>>(),
        "retained_tape":{"ordinary_bytes":kept.tensor_payload_bytes()?,"recomputed_bytes":compact.tensor_payload_bytes()?,
            "ordinary_nodes":kept.nodes,"recomputed_nodes":compact.nodes,"checkpoint_nodes":compact.checkpoint_nodes,
            "ordinary_nonparameter_primal_values":kept.nonparameter_primal_values,"ordinary_saved_context_values":kept.saved_context_values,
            "recomputed_nonparameter_primal_values":compact.nonparameter_primal_values,"recomputed_saved_context_values":compact.saved_context_values,
            "recomputed_recipe_values":compact.checkpoint_recipe_values},
        "recomputed_block_forwards":recomputed.counters.iter().map(RecomputationCounter::completed_forwards).sum::<usize>(),
        "measurement":"reachable nonparameter primal and saved-context f64 tensor lengths plus fixed RoPE tensors captured by checkpoint recipes, at loss-forward completion; excludes parameter/gradient/optimizer buffers, metadata, allocator capacity and reverse temporaries; not RSS or peak memory",
    }))
}

Explain the constructed two-microbatch example with target counts $2$ and $4$ and mean losses $1$ and $3$ . State the combined mean and why equal weighting of the two means changes the objective.
Use the default Chapter 13 command and setup above. Inspect accumulation; identify the final microbatch’s shape, target count and backward seed, and explain why AdamW runs only once for the logical group.
Use the default Chapter 13 command and setup above. Inspect recomputation; explain what the two byte counts include, which work is repeated, and why the live counter stays unchanged during backward.
Use the default Chapter 13 command and setup above. Inspect late_overflow; identify the cleared partial work, preserved optimizer state and available scheduled step, distinguishing eligibility for a retry from an executed retry.

Checked explanations for normalization, replay, retained payload and partial-update cleanup

The first microbatch has $2$ valid next-token targets and mean loss $1$ ; the second has $4$ valid targets and mean loss $3$ . Their loss sums are $2\cdot1=2$ and $4\cdot3=12$ across $6$ targets, so the combined mean is $14/6=7/3$ . Equal weighting gives $(1+3)/2=2$ and assigns each target in the smaller microbatch twice the weight of each target in the larger one.
The default final token grid has shape $[1,3]$ : one row and three supplied positions. It has $3$ valid targets out of $9$ , so its scale- $128$ backward seed is the recorded $42.666666666666664$ . The first contribution has already accumulated, with weights still fixed. One final unscale, global clipping decision and AdamW step updates the logical group; the measured norm is below the limit, so this run does not clip.
In the fixed-layout three-block dropout experiment, ordinary retention reports $212352$ bytes and recomputation $112752$ bytes at loss-forward completion. These counts cover the tape reachable from the scalar mean-loss node: nonparameter primals, owned saved-context tensors and, for checkpoint recipes, fixed rotary tensors. They exclude original parameter, gradient and optimizer buffers, metadata, layout and tickets, capacity and transient work, so they are neither total nor peak memory. Backward performs three additional block forwards. Each block regenerates its original masks from exact tickets without a live counter, leaving the live counter at $288$ .
One microbatch completed before the constructed second-microbatch backward overflowed. Cleanup discards that partial model gradient and restores the original zero-gradient bits. Weight bits and revisions, populated AdamW moments and beta powers, and optimizer step $1$ remain unchanged; scale backs off from $128$ to $64$ . Scheduled step $2$ remains available to the caller, but this probe performs no retry. Caller callback effects and live RNG advancement are outside the gradient cleanup boundary.

Carry one logical update into schedule reasoning

The existing text encoder, source-independent transformer, tied head, reverse engine and AdamW now support token-weighted logical updates and per-block activation recomputation. Accumulation controls when gradients combine and the optimizer acts; recomputation controls which block values remain between forward and backward. Both preserve the existing component boundaries and use text as the worked input.

Chapter 14 will extend schedule and clipping reasoning around this existing update owner, including warmup and decay. Global clipping already exists and is reused here. Chapter 15 will address portable model saving and loading and the additional state needed for training continuation. An activation checkpoint is a temporary differentiation recipe, not a durable file of learned weights. Full-corpus training, physical-GPU measurements and accelerated mixed-precision execution still need their own evidence.