Chapter 04 · Content revision 2

Pad variable-length batches without adding predictions

Extend the Rust batch and decoder APIs with right padding, valid attention prefixes and a loss averaged over real targets.

The problem: equal storage widths, unequal evidence

Tokenizing documents separately preserves their boundaries, but does not make their lengths equal. A rectangular batch needs equally wide rows. When a shorter example shares storage with a longer one, the extra cells have no source token and no next-token target. Treating those cells as ordinary predictions changes the training objective.

The practical crate already has complete causal windows, deterministic batching, causal attention and a mean next-token loss. We will extend those APIs with explicit right padding and validity metadata. A padding marker will occupy storage without becoming a vocabulary token, an attended key or a loss target.

Shift two documents, then align their rows

Use two controlled encoded documents: [0,99,1] and [0,100,101,102,1]. The existing tokenizer layout assigns 0 to BOS, 1 to EOS, and 99, 100, 101, 102 to the bytes a, b, c, d. These are supplied valid sequences; the chapter’s small vocabulary fixture has $266$ entries. We are studying batching, rather than learning another tokenizer policy.

The short document has $3$ encoded tokens. Shifting by one produces inputs [0,99] and targets [99,1]: $2$ predictions. The long document has $5$ encoded tokens, giving inputs [0,100,101,102] and targets [100,101,102,1]: $4$ predictions. EOS is a real next-token target. No target joins one document to the other.

Preserve the supplied order and right-pad to the longest prediction length. The batch has $B=2$ rows and $T=4$ prediction columns. Its input rows are [0,99,-1,-1] and [0,100,101,102]; its target rows are [99,1,-1,-1] and [100,101,102,1].

The value -1 is PAD_CELL, a batch-domain marker. The tokenizer has no PAD ID or PAD embedding row. BOS and EOS retain their meanings; neither substitutes for padding. Branch on the marker before converting to an unsigned token ID or performing an embedding or target lookup.

The short row’s positions are [0,1,0,0] and segment IDs are [0,0,-1,-1]. The long row’s positions are [0,1,2,3] and segment IDs are [1,1,1,1]. Here a segment identifies one example within the batch; NO_SEGMENT is -1. Padding’s stored position 0 is unused. For an explicitly selected source range, positions restart at zero within that example.

Keep three fields distinct. Token validity says an input exists; target validity says its immediate successor exists; loss eligibility says that target is supervised. We supervise every real target in this chapter, so all three short-row fields are [true,true,false,false], and all three long-row fields are [true,true,true,true]. The rectangle has $8$ storage cells and $6$ valid targets.

Attention also needs validity

A valid query can attend only to valid keys in its own row, at positions no later than the query. The short row’s allowed key lists are [0], [0,1], [], []. Its last two queries are inactive. The long row’s lists are [0], [0,1], [0,1,2], [0,1,2,3].

An inactive query skips softmax and produces a zero attention mixture. Applying softmax to an entirely masked row would leave no finite normalization reference. Carry validity through every attention layer so padding never becomes an attended key. A later bias or residual could produce a nonzero padded hidden state; the required invariant is its lack of influence on valid predictions.

Average over the targets that exist

L=-\frac{\sum_{b,t}m_{b,t}\log p(y_{b,t}\mid x_{b,\le t})}{\sum_{b,t}m_{b,t}}

Use the natural logarithm. The numerator sums negative log likelihood over eligible next-token targets; the denominator counts those targets. For our batch the denominator is $6$ , rather than the $8$ storage cells or the $2$ documents. This is a token-weighted mean: each real target gets the same weight.

The notation expresses the intended objective. The implementation excludes a cell before converting its target or reading its logits. Computing an invalid loss and then multiplying it by zero does not safely exclude it: a NaN value would still contaminate the sum. Zero valid targets is an error, not meaningful zero loss, and the loss adapter performs no optimizer update.

What the loss symbols mean

$L$ is mean negative log likelihood, measured in nats per valid target.
$b$ is the zero-based batch row index, selecting one document-local example.
$t$ is the zero-based prediction column within that row.
$m_{b,t}$ is loss eligibility: $1$ for a supervised real target, $0$ for padding.
$y_{b,t}$ is the immediate next source token at an eligible prediction cell.
$x_{b,\le t}$ is that row’s valid input prefix ending at column $t$ .
$p(y_{b,t}\mid x_{b,\le t})$ is the model probability of that real target given its own row’s prefix.

A constructed loss probe exposes two wrong denominators

Assign gold probabilities $[1/2,1/4]$ to the short row and $[1/2,1/2,1/4,1/2]$ to the long row. These probabilities are constructed for the loss probe; they are not trained-decoder predictions. Their negative log likelihood sums are $3\ln2$ and $5\ln2$ nats. The correct combined mean is $8\ln2/6=(4/3)\ln2\approx0.924196$ nats per valid target.

Dividing the numerator by $8$ storage cells gives $\ln2\approx0.693147$ nats per stored cell. Averaging the two document means equally gives $(11/8)\ln2\approx0.953077$ nats per target under a different document weighting. Both calculations differ from our six-target objective. To combine separately computed document means, multiply each by its valid-target count, add those raw sums, and divide once by the total count. Apply the same weighting to their parameter gradients.

The probe uses finite full-vocabulary logits. For vocabulary size $V=266$ and desired gold probability $0<q<1$ , set nongold logits to zero and the gold logit to $z_y=\ln((V-1)q/(1-q))$ . Softmax then assigns the gold class probability $q$ in ideal arithmetic. The Rust calculation uses f64; padded-versus-separate comparisons allow absolute error $10^{-10}$ plus relative error $10^{-9}$ .

Extend the existing owners

SequenceExample validates the full document using the existing tokenizer’s framing and bounded decoding. It then asks CausalWindowConfig for one shifted pair from the explicitly selected range. It retains the example ID, source document, partition, range, tokenizer payload identity and split-receipt identity. A receipt digest identifies caller-accepted evidence; it does not approve the evidence by itself.

Full-document examples reject empty content. Selected ranges need at least two source tokens to produce a prediction. A row exceeding the configured prediction width fails. The adapter does not truncate, split or join documents silently; an upstream caller may select an explicit range using the existing window policy.

Validate a document, retain its selected source range and reuse the existing causal-window iterator. rust/crates/llm-from-scratch-practical/src/training/sequence_example.rs#document-range-pair

    pub fn from_range(
        id: &str,
        document: BatchDocument<'a>,
        source_range: Range<usize>,
        binding: &BatchBinding,
    ) -> Result<Self, VariableBatchError> {
        if id.is_empty() {
            return Err(VariableBatchError::Invalid("example ID must not be empty"));
        }
        if binding
            .tokenizer()
            .decode_document_bounded(document.token_ids(), binding.limits)?
            .is_empty()
        {
            return Err(VariableBatchError::Invalid(
                "an empty-content document supplies no example",
            ));
        }
        let selected =
            document
                .token_ids()
                .get(source_range.clone())
                .ok_or(VariableBatchError::Invalid(
                    "source range must lie inside its document",
                ))?;
        let length = selected
            .len()
            .checked_sub(1)
            .filter(|&length| length > 0)
            .ok_or(VariableBatchError::Invalid(
                "a source range needs at least one prediction",
            ))?;
        let config = CausalWindowConfig::new(length, length)
            .map_err(|_| VariableBatchError::Invalid("source range cannot form a causal window"))?;
        let pair = config
            .windows(selected)
            .next()
            .ok_or(VariableBatchError::Invalid(
                "source range has no complete shifted pair",
            ))?;
        Ok(Self {
            id: id.to_owned(),
            document,
            source_range,
            pair,
            binding: binding.clone(),
        })
    }

VariableBatch derives its private arrays together from those examples. It uses the longest real row in each batch and retains the supplied order. The final batch has its actual row count; no extra all-padding rows are added. Lengths and validity define attention eligibility without allocating a separate dense mask for every batch row and attention head. The reference decoder still computes rectangular score tensors; these changes establish correct padding semantics. This example makes no speedup claim.

Build aligned input, target, position and segment arrays with separate validity and loss fields. rust/crates/llm-from-scratch-practical/src/training/variable_batch.rs#right-padded-batch

    pub fn build(
        partition: Partition,
        examples: &[SequenceExample<'_>],
        expected: &BatchBinding,
        config: VariableBatchConfig,
    ) -> Result<Self, VariableBatchError> {
        validate_examples(partition, examples, expected, config)?;
        Self::build_checked(partition, examples, expected, config)
    }

    fn build_checked(
        partition: Partition,
        examples: &[SequenceExample<'_>],
        expected: &BatchBinding,
        config: VariableBatchConfig,
    ) -> Result<Self, VariableBatchError> {
        if examples.len() > config.batch_size() {
            return Err(VariableBatchError::Invalid("too many rows for this batch"));
        }
        let columns = examples
            .iter()
            .map(SequenceExample::len)
            .max()
            .expect("nonempty checked examples");
        let rows = examples.len();
        let cells = rows
            .checked_mul(columns)
            .ok_or(VariableBatchError::Invalid(
                "batch shape product overflows usize",
            ))?;
        if cells > config.max_cells {
            return Err(VariableBatchError::Capacity {
                cells,
                maximum: config.max_cells,
            });
        }
        if rows > i32::MAX as usize + 1 {
            return Err(VariableBatchError::Invalid("segment IDs must fit i32"));
        }
        let mut placements = reserve(rows)?;
        for row in 0..rows {
            placements.push(BatchPlacement {
                occurrence: row,
                row,
                start: 0,
            });
        }
        Self::build_placed(
            partition,
            examples,
            expected,
            [rows, columns],
            &placements,
            0..rows,
        )
    }

The embedding adapter excludes PAD_CELL before conversion. It adds exclusion bits to the existing row-gather plan, which skips those reads and their backward scatter. Rotary position application likewise skips inactive positions.

Keep padding outside the vocabulary and reuse the existing gather kernel without reading or updating a padding row. rust/crates/llm-from-scratch-practical/src/nn/embedding.rs#sentinel-safe-embedding

    /// Batch padding has no embedding row. Exclude it before converting IDs;
    /// the existing gather kernel skips excluded rows in forward and backward.
    pub fn forward_padded_with_context(
        &self,
        context: AutogradContext,
        token_ids: &[i32],
        token_shape: &[usize],
    ) -> Result<TensorValue, EmbeddingError> {
        let (_, expected) =
            checked_row_major_layout(token_shape).map_err(EmbeddingError::TokenShape)?;
        if expected != token_ids.len() {
            return Err(EmbeddingError::TokenCountMismatch {
                expected,
                actual: token_ids.len(),
            });
        }
        let mut indices = Vec::new();
        let mut active = Vec::new();
        indices
            .try_reserve_exact(expected)
            .map_err(|_| EmbeddingError::IndexAllocationFailed { elements: expected })?;
        active
            .try_reserve_exact(expected)
            .map_err(|_| EmbeddingError::IndexAllocationFailed { elements: expected })?;
        for (position, &id) in token_ids.iter().enumerate() {
            if id == crate::training::variable_batch::PAD_CELL {
                indices.push(0);
                active.push(false);
                continue;
            }
            let unsigned =
                u32::try_from(id).map_err(|_| EmbeddingError::NegativeTokenId { position, id })?;
            let index = usize::try_from(unsigned)
                .ok()
                .filter(|&index| index < self.vocabulary_size)
                .ok_or(EmbeddingError::TokenIdOutOfBounds {
                    position,
                    id: unsigned,
                    vocabulary_size: self.vocabulary_size,
                })?;
            indices.push(index);
            active.push(true);
        }
        self.table
            .tensor()
            .gather_rows_with_plan_and_context(context, move |table| {
                RowGatherPlan::from_validated_indices(table, indices, token_shape.to_vec())?
                    .with_active_rows(active)
            })
            .map_err(EmbeddingError::Autodiff)
    }

The existing causal softmax loop accepts one valid prefix length per score grid. It normalizes only real queries; untouched inactive query rows remain zero. Multi-head attention passes the same row length to each head in every layer.

Reuse causal softmax over real query prefixes and skip inactive-query normalization. rust/crates/llm-from-scratch-practical/src/autograd/model_ops.rs#padding-aware-causal-softmax

pub(crate) fn validate_causal_layout(
    grids: usize,
    queries: usize,
    lengths: Option<&[usize]>,
    starts: Option<&[usize]>,
) -> Result<(), TensorAutodiffError> {
    if let Some(lengths) = lengths {
        if lengths.len() != grids {
            return Err(ModelOpError::CausalLengthCount {
                expected: grids,
                actual: lengths.len(),
            }
            .into());
        }
        for (grid, &length) in lengths.iter().enumerate() {
            if length > queries {
                return Err(ModelOpError::CausalLengthOutOfBounds {
                    grid,
                    length,
                    tokens: queries,
                }
                .into());
            }
        }
    }
    if let Some(starts) = starts {
        let expected = grids
            .checked_mul(queries)
            .ok_or(TensorError::ShapeOverflow)?;
        if starts.len() != expected {
            return Err(ModelOpError::CausalRangeCount {
                expected,
                actual: starts.len(),
            }
            .into());
        }
        for grid in 0..grids {
            let valid = lengths.map_or(queries, |lengths| lengths[grid]);
            for query in 0..valid {
                let start = starts[grid * queries + query];
                if start > query {
                    return Err(ModelOpError::CausalRangeStart { grid, query, start }.into());
                }
            }
        }
    }
    Ok(())
}

fn causal_softmax_forward_with_lengths(
    input: &Tensor,
    lengths: Option<&[usize]>,
    starts: Option<&[usize]>,
) -> Result<Tensor, TensorAutodiffError> {
    if input.rank() < 2 {
        return Err(ModelOpError::CausalSoftmaxRank { rank: input.rank() }.into());
    }
    let queries = input.shape()[input.rank() - 2];
    let keys = input.shape()[input.rank() - 1];
    if queries != keys {
        return Err(ModelOpError::CausalSoftmaxNonSquare { queries, keys }.into());
    }
    if queries == 0 {
        return Err(ModelOpError::CausalSoftmaxEmptyTokens.into());
    }
    let grid_cells = queries
        .checked_mul(keys)
        .ok_or(TensorError::ShapeOverflow)?;
    let grids = input.len() / grid_cells;
    validate_causal_layout(grids, queries, lengths, starts)?;
    let mut probabilities = zeros(input.shape())?;
    for grid in 0..grids {
        let valid = lengths.map_or(queries, |lengths| lengths[grid]);
        // Skip inactive queries before maximum, exponentiation or division.
        for query in 0..valid {
            let row_start = (grid * queries + query) * keys;
            let start = starts.map_or(0, |starts| starts[grid * queries + query]);
            let allowed = &input.as_slice()[row_start + start..=row_start + query];
            let maximum = allowed.iter().copied().fold(f64::NEG_INFINITY, f64::max);
            let mut exponential_tail = 0.0;
            let mut skipped_one_maximum = false;
            for &score in allowed {
                let shifted = score - maximum;
                if shifted == 0.0 && !skipped_one_maximum {
                    skipped_one_maximum = true;
                } else {
                    exponential_tail += shifted.exp();
                }
            }
            debug_assert!(skipped_one_maximum);
            let denominator = 1.0 + exponential_tail;
            for (key, &score) in allowed.iter().enumerate() {
                let probability = (score - maximum).exp() / denominator;
                probabilities.as_mut_slice()[row_start + start + key] =
                    if probability == 0.0 { 0.0 } else { probability };
            }
        }
    }
    Ok(probabilities)
}

The loss adapter gathers eligible logit rows and targets before invoking the existing stable indexed NLL. The existing gather pullback scatters gradients to selected positions only; excluded logit positions receive exactly zero. For a valid position, its logit derivative is $(p_v-\mathbf{1}[v=y])/N_{\mathrm{valid}}$ , where $v$ is a vocabulary class, $p_v$ its probability, $y$ the gold class and $N_{\mathrm{valid}}$ the batch’s valid-target count.

Select real prediction rows before the existing NLL and expose their mean, raw loss sum and count. rust/crates/llm-from-scratch-practical/src/training/variable_batch.rs#valid-target-loss

    /// Gather before NLL; the existing gather pullback scatters only selected rows.
    pub fn loss_with_context(
        &self,
        context: AutogradContext,
        logits: &TensorValue,
    ) -> Result<MaskedLoss, VariableBatchError> {
        let classes = self.check_logits(&logits.shape())?;
        let (positions, targets) = self.active_targets()?;
        let flat = logits.reshape_with_context(context, &[self.inputs.len(), classes])?;
        let selected = flat.gather_rows_with_context(context, &positions, &[targets.len()])?;
        let mean = selected.indexed_mean_nll_with_context(context, 1, &targets)?;
        let count = targets.len();
        let scale = TensorValue::constant(Tensor::from_vec(vec![], vec![count as f64])?)?;
        let loss_sum = mean.mul_with_context(context, &scale)?;
        Ok(MaskedLoss {
            mean,
            loss_sum,
            count,
        })
    }

The new padded entry point calls the same decoder block loop as the old API. Its vocabulary size and exact batch binding must agree with the declared inputs. The padded path does not change the old rectangular loss method’s behavior.

Validate the batch binding and pass valid lengths into the existing decoder assembly. rust/crates/llm-from-scratch-practical/src/models/decoder.rs#padded-decoder-entry

    pub fn forward_padded(
        &self,
        batch: &VariableBatch,
        expected: &BatchBinding,
    ) -> Result<DecoderModelForward, PaddedDecoderError> {
        self.forward_padded_with_context(AutogradContext::recording(), batch, expected)
    }

    /// A checked batch owns validity. Padding never becomes an unsigned ID.
    pub fn forward_padded_with_context(
        &self,
        context: AutogradContext,
        batch: &VariableBatch,
        expected: &BatchBinding,
    ) -> Result<DecoderModelForward, PaddedDecoderError> {
        let input =
            TextInputEncoder::padded(self, batch, expected).encode(context, &self.backbone())?;
        Ok(self.forward_prepared_text(input)?)
    }

    pub(super) fn prepare_padded_text(
        &self,
        context: AutogradContext,
        batch: &VariableBatch,
        expected: &BatchBinding,
    ) -> Result<TensorValue, PaddedDecoderError> {
        self.validate_padded_text(batch, expected)?;
        let embedding = self
            .embedding
            .forward_padded_with_context(context, batch.inputs(), &batch.shape())
            .map_err(DecoderModelError::Embedding)?;
        Ok(embedding)
    }
    fn validate_padded_text(
        &self,
        batch: &VariableBatch,
        expected: &BatchBinding,
    ) -> Result<(), PaddedDecoderError> {
        batch.ensure_binding(expected)?;
        if expected.vocabulary_size() != self.config.vocabulary_size {
            return Err(VariableBatchError::Invalid(
                "batch vocabulary differs from the decoder vocabulary",
            )
            .into());
        }
        if batch.shape()[1] > self.config.max_positions {
            return Err(DecoderModelError::ContextLengthExceeded {
                tokens: batch.shape()[1],
                max_positions: self.config.max_positions,
            }
            .into());
        }
        Ok(())
    }

Ignoring padding does not freeze actual vocabulary parameters. The decoder ties its input embedding table to the output vocabulary projection. Real vocabulary rows can receive output-normalization gradients even when they are not selected as input tokens. There is no PAD vocabulary row to freeze here.

With Docker running, use the repository’s Rust workspace setup to run the supplied CPU example from the repository root:

COURSE_CORPUS=false ./course run cargo run --offline --locked \
  -p practical-ch04-padded-variable-batches

The JSON trace includes both shifted and padded rows, all validity fields, allowed-key lists, six valid targets and the constructed mean 0.924196. It also compares all $1596$ valid vocabulary logits from padded and separate execution with identical untrained parameters. The fixture demonstrates correctness; it makes no claim about trained-model quality, full-corpus throughput or GPU execution.

Print the batch metadata, controlled loss probe and valid-logit comparison. rust/demos/practical-ch04-padded-variable-batches/src/main.rs#padded-variable-batches-example

fn main() -> Result<(), Box<dyn std::error::Error>> {
    let trace = practical_ch04_padded_variable_batches::demo_trace()?;
    println!("{}", serde_json::to_string_pretty(&trace)?);
    Ok(())
}

History: different kinds of wasted work

Rectangular training batches allocate storage for their longest row, so shorter examples leave padding that supplies no training target. Packing work on BERT addressed padding waste with changes to attention and masked-language-model loss; later Orca work addressed unequal generation lifetimes through iteration-level scheduling.

The 2021 paper Packing: Towards 2x NLP BERT Acceleration studied BERT pretraining. The authors identify padding waste in BERT pretraining and describe packing with additional attention masking and a modified masked-language-model loss. Those BERT results do not establish performance for our causal decoder.

The 2022 Orca paper examines autoregressive serving. Orca schedules generation at iteration granularity so the batch of requests can change between model iterations. Unequal request lifetimes in serving and unequal target lengths in training are different problems. Orca’s scheduling does not supply our training loss mask or token denominator.

Preserve each example’s valid causal predictions and token-weighted objective before changing its storage arrangement or execution schedule. This chapter adds correct padded storage; Chapter 5 will address packing. Padding changes storage capacity and must not create targets or attended keys. Packing and request scheduling address different sources of wasted work and require their own correctness rules.

The runnable contrast reuses MiniBatchEpoch with context length $4$ and stride $4$ . That API selects complete windows only: the long document supplies one four-target window and the short document supplies none. Our explicit full-document variable examples supply six targets. The difference comes from example selection, not an alternative loss formula or a newly implemented batching algorithm.

Call the existing complete-window builder and distinguish its four admitted targets from the six explicit full-document predictions. rust/demos/practical-ch04-padded-variable-batches/src/lib.rs#fixed-window-contrast

    // Reuse the old complete-window policy. At context four it admits only the
    // long document; this is a different example-selection policy, not padding.
    let fixed = MiniBatchEpoch::build(
        Partition::Train,
        &documents,
        CausalWindowConfig::new(4, 4)?,
        MiniBatchConfig::new(2, BatchOrder::Sequential)?,
    )?;
    let fixed_targets: usize = fixed
        .batches()
        .iter()
        .map(|batch| batch.token_count())
        .sum();

See the storage and prediction boundaries together

The aligned rows make each storage cell’s input, target, position, segment and eligibility visible. The legal-key lists distinguish inactive queries from real queries, and the loss count keeps the six-target denominator beside the eight-cell rectangle.

Eight storage cells, six next-token targets

The short document [0,99,1] produces inputs [0,99] and targets [99,1]; the long document [0,100,101,102,1] produces inputs [0,100,101,102] and targets [100,101,102,1]. Short-row columns 2 and 3 are padding, with inactive queries and no allowed keys. BOS 0 and EOS 1 remain real controls. Only the six real targets enter the constructed loss sum and denominator; padding adds no vocabulary token.

Source document: `short`

Encoded document tokens: [0,99,1] · $3$
Real next-token predictions: $2$

Prediction columns: `short`
Batch field	$t=0$	$t=1$	$t=2$	$t=3$
Input ID	`0`	`99`	`-1` · PAD storage	`-1` · PAD storage
Target ID	`99`	`1`	`-1` · PAD storage	`-1` · PAD storage
Local position	`0`	`1`	`0`	`0`
Segment ID	`0`	`0`	`-1` · No segment	`-1` · No segment
Input exists	Yes	Yes	No	No
Target exists	Yes	Yes	No	No
Target supervised	Yes	Yes	No	No

Allowed keys by query column

$t=0$ · Real query: [0]
$t=1$ · Real query: [0,1]
$t=2$ · Inactive padding query: []
$t=3$ · Inactive padding query: []

Source document: `long`

Encoded document tokens: [0,100,101,102,1] · $5$
Real next-token predictions: $4$

Prediction columns: `long`
Batch field	$t=0$	$t=1$	$t=2$	$t=3$
Input ID	`0`	`100`	`101`	`102`
Target ID	`100`	`101`	`102`	`1`
Local position	`0`	`1`	`2`	`3`
Segment ID	`1`	`1`	`1`	`1`
Input exists	Yes	Yes	Yes	Yes
Target exists	Yes	Yes	Yes	Yes
Target supervised	Yes	Yes	Yes	Yes

Allowed keys by query column

$t=0$ · Real query: [0]
$t=1$ · Real query: [0,1]
$t=2$ · Real query: [0,1,2]
$t=3$ · Real query: [0,1,2,3]

Constructed loss, counted by real targets

Loss sum in nats: $\approx 5.545177$
Valid-target count: $6$
Mean nats per valid target: $\approx 0.924196$

The gold probabilities are constructed for this probe. The two excluded storage cells contribute neither a target nor a loss.

Padding stays outside the vocabulary

The batch marker -1 is not a token ID. BOS 0 and EOS 1 remain real controls. An inactive query skips normalization and produces a zero attention mixture.

Optional practice: inspect the count and compare execution

Reproduce the two rows and inspect the short row’s three validity fields and allowed-key lists. Explain why its EOS target is included and its padding cells are excluded.
Inspect the constructed loss sum and compare the six-target mean with division by eight cells and with an unweighted average of document means.
Run the seeded comparisons and inspect their full-gradient coverage. They use identical parameters for padded and separate examples; no optimizer update is needed.

From the repository root, with the same Docker setup as the example, run:

COURSE_CORPUS=false ./course run cargo test --offline --locked \
  -p llm-from-scratch-practical --test practical_ch04_variable_batches

The smaller CPU example has $1$ decoder layer, model width $4$ , $2$ heads, feed-forward width $4$ and a maximum prediction length of $4$ . The second has $2$ layers, model width $8$ , $2$ heads, feed-forward width $16$ and a maximum prediction length of $16$ . Both use the $266$ -entry fixture vocabulary. Within each comparison, padded and separate execution use identical parameters.

Expect a nonzero test count and 0 failed. The suite includes exact metadata, the loss denominator, rejected inputs, explicit ranges, future-token isolation and $100$ fixed seeds in each of the two CPU model profiles. Each seeded comparison checks every valid logit, token-summed loss and every parameter-gradient coordinate within the declared f64 tolerance.

Explanations for padding eligibility, normalization and gradient comparison

The short document supplies two immediate next-token pairs: BOS to a, then a to EOS. Both have real targets. Its final two storage cells have no source pair, so all three validity fields are false and their query key lists are empty.
The raw sums are $3\ln2$ and $5\ln2$ nats. Adding and dividing by six gives $(4/3)\ln2$ nats per valid target. Eight-cell division lowers the value by counting padding; an unweighted document mean gives a short document the same total weight as a long one. Neither preserves equal weight per target.
Multiply each separate document mean and parameter-gradient vector by its valid-target count before combining, then divide once by the total count. The padded loss already uses that total. Comparing full vectors checks each parameter’s contribution; a shared gradient norm would not establish it.

Carry example identity into packing

The decoder now consumes checked padded examples with its existing weights and vocabulary. Each example retains its source range, partition and tokenizer and split identities. Its valid positions keep the same causal prefixes and target weighting as separate execution.

Chapter 5 will place multiple examples in one packed row. It must preserve their source ranges, local positions, attention boundaries and valid-target count. Simply concatenating rows would let a later document attend to an earlier one and change the examples we just preserved.

The problem: equal storage widths, unequal evidence

Shift two documents, then align their rows

Attention also needs validity

Average over the targets that exist

What the loss symbols mean

A constructed loss probe exposes two wrong denominators

Extend the existing owners

History: different kinds of wasted work

See the storage and prediction boundaries together

Eight storage cells, six next-token targets

Source document: short

Allowed keys by query column

Source document: long

Allowed keys by query column

Constructed loss, counted by real targets

Padding stays outside the vocabulary

Optional practice: inspect the count and compare execution

Carry example identity into packing

Source document: `short`

Source document: `long`