Make optional dropout explicit and replayable

Reuse multiplication and keep mode separate from autograd

DropoutProbability::new rejects nonfinite probabilities, negative values and values at least one. It preserves the accepted probability’s exact bits, including an accepted signed zero, and stores the keep probability and scale. Configuration validation still applies in Eval, even though Eval will not use a mask.

Active dropout creates an untracked constant multiplier with the same shape as the input. Each coefficient is zero or the saved training scale. Existing TensorValue::mul_with_context applies it, and the existing Multiply VJP retains the coefficients for backward. The mask is not a parameter, and no new gradient rule is needed.

/// Train samples even under no-grad; Eval and zero probability preserve the handle.
#[allow(clippy::too_many_arguments)]
pub fn apply_dropout_with_context(
    context: AutogradContext,
    input: &TensorValue,
    mode: DropoutMode,
    policy: DropoutPolicy,
    site: DropoutSiteId,
    binding: &DropoutBinding,
    counter: &mut DropoutCounterState,
    limits: DropoutLimits,
) -> Result<DropoutApplication, DropoutError> {
    if !policy.active(mode) {
        return Ok(DropoutApplication {
            output: input.clone(),
            record: None,
        });
    }
    let n = input_entries(input)?;
    let end = checked_range(counter.next, n)?;
    limits.admit(n, 1)?;
    if n == 0 {
        return Ok(DropoutApplication {
            output: input.clone(),
            record: None,
        });
    }
    let (key, coefficients) =
        coefficients_from_validated_range(counter.next, end, n, policy.probability, || {
            counter.key.unwrap_or_else(|| derive_key(counter.seed))
        })?;
    let (multiplier, output) = multiply(context, input, coefficients)?;
    let shape = input.shape();
    let ticket = DropoutReplayTicket {
        version: DROPOUT_COUNTER_VERSION,
        placement: DROPOUT_PLACEMENT,
        purpose: DROPOUT_PURPOSE,
        seed: counter.seed,
        key,
        start: counter.next,
        end,
        probability_bits: policy.probability.bits(),
        scale_bits: policy.probability.scale().to_bits(),
        site,
        shape: [shape[0], shape[1], shape[2]],
        binding: binding.clone(),
    };
    counter.next = end;
    counter.key = Some(key);
    Ok(DropoutApplication {
        output: output.clone(),
        record: Some(DropoutSiteRecord {
            ticket,
            input: input.clone(),
            multiplier,
            output,
        }),
    })
}

DropoutMode chooses Train or Eval. AutogradContext separately permits or disables new graph recording. The following table uses an already tracked input, a nonempty site and a valid positive probability; zero probability returns the input handle in every row.

Eval and zero probability preserve every input bit, including signed zero. They allocate no dropout mask, derive no dropout key and consume no random words. no_grad() controls new operations; an identity return does not retroactively detach an already tracked input handle. If an Eval identity result receives upstream derivative $g_i$, that connection passes $g_i$ to the input without training rescaling.

Place dropout after each complete branch projection

The course policy residual-output-element-dropout-v1 has two sites per block: the complete attention output after $W_o$, just before the attention residual addition, and the complete SwiGLU output after $W_{\mathrm{down}}$, just before the feed-forward residual addition. The untouched residual input is still added. Embeddings, normalization gains, attention probabilities and vocabulary logits are not dropout sites in this policy.

Each active site draws for every element of its full $[B,T,D]$ rectangle, including PAD cells. Blocks run in ascending order, attention before feed-forward, then row-major batch, token, feature order. Padding, causal attention and valid-target masks keep their existing meanings. Repacking or changing microbatch boundaries can assign different counters to valid coordinates, even with the same seed.

Legacy model and run records remain version one with unchanged bytes and digests. DropoutRunConfig is an explicit version-two wrapper containing that legacy run and a version-one dropout policy; its identity domain is decoder-dropout-run-v2. The policy records probability bits, seed, placement, counter mapping and purpose. Changing it changes replay compatibility, while the learned parameter count stays unchanged.

Specify the words and check the cursor

The course registry assigns initialization purpose 1, corpus order 2, packing 3, optional dropout 4, sampling 5 and evaluation bootstrap 6. Dropout owns purpose 4 and its own seed/key and next-word counter. Calls to other mutable generators cannot advance that counter.

course-dropout-counter-v1 reuses SplitMix64. Let $s$ be the dropout seed, $G$ the existing increment 0x9e3779b97f4a7c15, and $\operatorname{next}(a)$ the result of one next_u64() call starting from SplitMix64::from_state(a). The key and counter-indexed word are

Here $\oplus$ is bitwise exclusive-or, and $c$ is the counter of the requested word. Only the specified generator arithmetic wraps modulo $2^{64}$. For a request starting at $c$ with $n$ entries, the end $c+n$ is checked before allocation. Overflow or an exhausted u64::MAX cursor rejects without partial consumption. The generated range is $[c,c+n)$.

fn derive_key(seed: u64) -> u64 {
    SplitMix64::from_state(seed ^ SplitMix64::WORD_INCREMENT.wrapping_mul(DROPOUT_PURPOSE))
        .next_u64()
}
fn word_at(key: u64, counter: u64) -> u64 {
    SplitMix64::from_state(key.wrapping_add(SplitMix64::WORD_INCREMENT.wrapping_mul(counter)))
        .next_u64()
}
fn checked_range(start: u64, count: usize) -> Result<u64, DropoutError> {
    let n = u64::try_from(count).map_err(|_| DropoutError::Arithmetic("counter count"))?;
    if n > 0 && start == u64::MAX {
        return Err(DropoutError::CounterExhausted);
    }
    start.checked_add(n).ok_or(DropoutError::CounterExhausted)
}

The conversion keeps the high $53$ bits:

The values are $j/2^{53}$ for integer $0\le j<2^{53}$, so one is never generated. Equality with the threshold keeps the element. Assuming uniform words, counting this grid gives effective drop probability

At $p=0.5$, the grid represents the threshold exactly. At requested $p=2^{-54}$, grid zero drops while $2^{-53}$ keeps. With the requested scale, even exact real arithmetic would then give

Stored f64 subtraction and division can add rounding as well. The earlier expectation equality is the ideal Bernoulli statement; it is not an exact claim for every requested floating-point threshold. Separate counter ownership also does not prove statistically independent or cryptographically secure streams: with an odd increment, keyed full-period sequences are cyclic shifts.

Inspect ten predetermined seeded masks

The executed primitive fixture uses seeds 0 through 8, then 39, with the same input, upstream gradient, shape [1,1,4], purpose 4, half-drop policy and initial counter 0. Each record ends at counter 4. The summary retains all four generated hexadecimal words and exact output/final-gradient bits. The table writes those f64 bit patterns as their corresponding scalar values:

For seed 39, the generated words are e0f2d7d7dd3d13ac, bbe95911f05f90b5, 210c4710466fdb4c and 88a1d7a436e3ab9b. At the one-half threshold, their high bits give mask 1101. The seeds were fixed before execution; these records check the specified mapping, not pseudorandom frequencies. A local Multiply edge can contribute negative zero, while accumulation into an initially positive-zero adjoint gives the final stored zeros shown here. Identity paths separately preserve input zero signs.

Commit only a successful complete forward

The active text wrapper checks the full counter demand and dropout resource limits before sampling. One candidate counter then covers every block and both sites. The existing private text decoder consumes that candidate; the caller’s live state is published only after the complete forward succeeds.

    /// One candidate counter covers every site and commits after the complete forward.
    #[allow(clippy::too_many_arguments)]
    pub fn forward_with_dropout(
        &self,
        context: AutogradContext,
        token_ids: &[u32],
        token_shape: &[usize],
        mode: DropoutMode,
        policy: &DropoutRunPolicy,
        state: &mut DropoutCounterState,
        limits: DropoutLimits,
    ) -> Result<DropoutModelForward, DecoderModelError> {
        let p = policy
            .policy()
            .map_err(|_| DropoutError::Binding("invalid run policy"))?;
        if !p.active(mode) {
            return Ok(DropoutModelForward {
                forward: self.forward_with_context(context, token_ids, token_shape)?,
                sites: Vec::new(),
                binding: None,
                parameters: Vec::new(),
            });
        }
        self.validate_tokens(token_ids, token_shape)?;
        let shape = [token_shape[0], token_shape[1], self.config().model_width()];
        let execution = self.fresh_dropout_execution(
            context,
            shape,
            policy,
            ordinary_layout(token_ids, token_shape)?,
            state,
            limits,
        )?;
        let input = TextInputEncoder::tokens(self, token_ids, token_shape)
            .encode(context, &self.backbone())
            .map_err(ordinary_error)?;
        self.consume_dropout_text(input, execution, Some(state))
    }

A later-layer failure discards partial results and leaves the live cursor, parameters and pre-existing gradients unchanged. Active nonfinite arithmetic is an error, rather than a value repaired by clipping or replacing it with zero. This forward boundary does not rewind every allocation or internal graph ID, and it does not make backward or the optimizer update transactional. After a successful forward, the cursor has advanced; a failed update needs restoration of the last committed training state before retry.

Cached decoding explicitly selects Eval and its existing no_grad context. The mode and graph choices are separate. An explicit cached Train request returns CachedTrain before cache, work or dropout state changes. Disabled ordinary and cached calls retain the preceding error order and behavior.

Bind replay to the original computation

DropoutReplayTicket is immutable evidence for a site. It binds mapping and placement versions, purpose, seed/key, counter start/end, probability/scale bits, layer/site, actual shape and run/text-layout/model-revision identities. Full-model replay also checks the original live parameter nodes and revisions before using the same private decoder core. A same-shaped different packed or ordinary batch does not gain replay authority.

replay_dropout_with_context regenerates the saved coefficient tensor from the ticket. It accepts no live counter and consumes zero live words. The complete text replay validates all site tickets before decoder work. A ticket for old weights cannot establish recomputation under changed weights.

For the supplied bridge, $B=1,T=3,D=16,L=2$: each site has $1\cdot3\cdot16=48$ entries and a fresh forward uses $2\cdot2\cdot48=192$ words. Inputs are [0,99,100], targets [99,100,1], positions [0,1,2], initialization and dropout seeds 39, and drop probability $0.5$. All three targets belong to the existing token-mean loss.

The full replay compares every logit, token-mean loss and all named parameter gradients, using all four saved tickets. Gradients are deliberately cleared between comparison backward passes. These are local site and same-core recomputation checks; they do not provide Chapter 13’s activation scheduler or demonstrate a memory-saving training schedule.

The executed R1, R2 and R3 records each compare $48$ multiplier, output and input-gradient coordinates exactly and retain live cursor 192. Full-core replay reports equal bits for all $798$ logits, the token-mean loss and all $8304$ parameter-gradient coordinates, plus all four masks. A changed same-shaped text input is rejected.

The continuation fixture runs four shared-trainer updates on that repeated batch, with learning rate $0.001$, AdamW coefficients $0.9$ and $0.999$, epsilon $10^{-8}$, weight decay $0.01$ and global gradient-norm cap $1$. U1, U2 and U3 capture committed zero-gradient states after updates one, two and three, at dropout cursors 192, 384 and 576; four fresh forwards end at 768.

Each in-memory snapshot carries parameter names/order/values and admitted model identity, AdamW moments/step, the zero-gradient and commit phase, dropout mapping/seed/cursor, run/config/dtype/backend bindings, batch identity/cursor, valid-target count and denominator, optimizer settings, schedule and remaining updates. Restoration creates independent live parameter nodes. Comparisons cover every remaining mask, logit, loss, raw and clipped gradient, update, parameter, moment, step, final cursor and next word. Missing or mismatched counter, phase, run, probability or batch cursor rejects.

The executed U1, U2 and U3 continuations report all_remaining_tapes_equal: true for three, two and one remaining updates respectively. All finish at counter 768, with next word a8c1103f72c54933, matching final parameter, optimizer and declared-state bits. In the injected post-forward failure, the counter first reaches 192; restoring the last committed snapshot returns it to 0 and restores the commit phase before an identical retry. The seven missing/mismatched snapshot cases in the report all reject.

This snapshot is scoped to a four-update in-memory fixture. The unchanged version-one component checkpoint is not a complete training-job record. Chapter 15 still owns durable job continuation and restart policy.

Count mask data separately from memory peaks

The multiplier stores f64 coefficients, so one active site owns $8BTD$ mask-data bytes. If both multipliers in all blocks are live, their subtotal is

At nominal full-context shapes, this gives reference $4096$ bytes, bridge $65536$ bytes and laptop $33554432$ bytes. These are data subtotals, not allocator peaks: output tensors, saved graph primals, adjoints, temporary storage, alignment and metadata also matter. The existing planner accounts for their lifetimes. The laptop number is counterfactual active-f64 accounting; mandatory 8gb-gpu-core keeps dropout off, and no GPU execution is established here. Eval and zero probability allocate zero dropout masks.

The recorded compiled CPU summary run had peak resident set size $48050176$ bytes and took about $0.314$ seconds. That process measurement covers this bounded fixture, separate from compilation, container setup and nominal full-context planning. It does not measure the laptop or a GPU allocator peak.