Nemotron 3 Ultra: replaying a rounded state

Robert Joseph George · TorchLean Verified Examples

After K3, it was time to look at an open-source model from a US team too. NVIDIA's Nemotron 3 Ultra combines recurrent Mamba layers, attention, and routed experts. It has 550 billion parameters, with 55 billion active for each token, and supports a million-token context. NVIDIA released base, post-trained, and quantized checkpoints, along with training datasets and recipes. Nemotron 3 Ultra release ↗

The memory-saving idea I want to work through is keeping an older checkpoint and replaying recent inputs when we need the current state. A recurrent layer repeatedly overwrites its state, and rounding those numbers can affect later calculations. Let's first check when replay gives us the same state, then use FloatLib to ask what changes when we round the checkpoint.

NVIDIA's account of quantizing Ultra makes the practical choices clear. A quantized model uses several formats: the routed experts use NVFP4, while the recurrent state cache uses FP16 with stochastic rounding. Saving memory means choosing what to round, where to round it, and how to recover the state later. We have the report and implementation available, so we can follow those choices into the equations and state the assumptions our proofs need.

A language model carries information forward as it reads. Attention stores keys and values for individual past tokens. Mamba's state-space model (SSM) keeps a tensor called the state, updating its entries with each new input. Nemotron uses both. The recurrent tensor keeps the same shape throughout a conversation, even though its numbers change at every step.

In Nemotron 3 Ultra, those tensors span 48 Mamba blocks. Keeping them in FP32 takes 384 MiB per main-model sequence before other allocations. That cost grows with the number of requests being served together, even though each request’s state has a fixed shape. The report's §4.5 describes keeping an older checkpoint and replaying the inputs since it was saved. I start with the memory calculation, then check when replay returns the same state and how rounding a checkpoint affects later updates. The example below lets you change the storage schedule and see why a smaller error budget need not mean a smaller actual error.

Counting the state and attention cache

The released BF16 configuration lists 108 residual blocks: 48 Mamba blocks, 48 MoE blocks, and 12 attention blocks. The Mamba blocks hold recurrent state; the attention blocks hold keys and values. releasedSchedule records their order, and Lean checks the counts. Each entry names one residual block.

Each Mamba-2 block has 256 heads, 64 value channels per head, and 128 state coordinates. That gives a state tensor of shape [256, 64, 128]. The 48 FP32 state tensors together hold 384 MiB. Each attention block has two KV heads of width 128, so the 12 blocks add 6144 bytes of FP8 keys and values for each token.

Mamba is a state-space layer: it updates a fixed-size state using coefficients chosen from the current input. A mixture-of-experts (MoE) block selects a few learned feed-forward transformations for each token. Attention instead keeps keys, which future queries compare against, and values, which their attention weights combine. These blocks therefore have different storage needs.

Let's count the entries. One Mamba head stores \(64\times128=8{,}192\) entries; 256 heads give 2,097,152 entries in a block. At four bytes per FP32 entry, that is 8 MiB. Multiplying by 48 gives 384 MiB. Here MiB means \(2^{20}\) bytes. The BF16 label on the model describes its weight representation; it does not mean that every working state has the same precision.

The KV payload grows until it reaches the recurrent payload at 65,536 tokens; below that point it is smaller. released_fp32_fp8_crossover checks the exact boundary. Using FP16 for SSM storage halves it to 32,768 tokens. These counts cover the payload for one main-model sequence. Convolution history, scale metadata, allocator padding, MTP, and retained snapshots add to the serving system’s memory use.

The batch size matters as much as the sequence length here. Sixty-four active requests with separate FP32 states already require \(64\times384\,\mathrm{MiB}=24\,\mathrm{GiB}\) for this payload alone. Reducing the state to FP16 halves that arithmetic count. It also changes what happens every time an updated value is saved, which is why the rounding policy belongs in the calculation below.

/-- The schedule has 48 recurrent, 12 attention, and 48 expert blocks. -/
theorem releasedSchedule_counts :
    releasedSchedule.count .mamba = 48 ∧
    releasedSchedule.count .attention = 12 ∧
    releasedSchedule.count .moe = 48 := by decide +kernel

/-- The idealized main-stack payloads cross at precisely 64 Ki tokens. At the boundary they are
equal, not strictly ordered; other caches and overhead are excluded from this calculation. -/
theorem released_fp32_fp8_crossover (tokens : Nat) :
    kvBytes releasedSchedule tokens 2 128 1 ≤ ssmBytes releasedSchedule 256 64 128 4 ↔
      tokens ≤ 65536 := by
  rw [(released_cache_payloads tokens).1, (released_cache_payloads tokens).2]
  omega

Recovering the state by replay

Before replaying several tokens, let's run one prepared SSM update in Lean. We retain [10, 0], write [2, 0], and read the first coordinate. The output is 12. Hover over the tensors and step to see how their dimensions appear in the types.

import Nemotron.SSM
import NN.Tensor

open TorchLean

-- A rank-one write updates a two-coordinate state, then reads its first coordinate.
def ssmMemory : Nemotron.SSM.State Float32 1 2 := [[10, 0]]
def ssmToken : Nemotron.SSM.Input Float32 1 2 :=
  { retention := 1
    value := [2]
    write := [1, 0]
    read := [1, 0] }

#eval (Nemotron.SSM.step ssmMemory ssmToken).2.getScalar ⟨0, by decide⟩
#check Nemotron.SSM.step

12.000000

Nemotron.SSM.step {α : Type} [Storage α] [Add α] [Mul α] [Zero α] {valueDim stateDim : ℕ}
  (state : Nemotron.SSM.State α valueDim stateDim) (input : Nemotron.SSM.Input α valueDim stateDim) :
  Nemotron.SSM.State α valueDim stateDim × Tensor α [valueDim]

Instead of writing the state after every token, save it at an earlier point and retain the inputs needed for the updates since then. Starting from that checkpoint, run those updates again in order to recover the current state. Here \(s\) is the saved state and \(u_1,\ldots,u_k\) are the prepared inputs: the values and coefficients consumed by the recurrent update, after the earlier processing of each token. These are not raw token IDs. Replaying a token ID through a changed projection or convolution would not necessarily recover the prepared values used the first time.

In §4.5, the authors describe how they catch up from the older state:

we cache the input activations and apply activation replay to forward the state on-the-fly.

Let's make that replay step precise. We have to retain the values the update consumed, not resample them or regenerate them from a different state. With those values fixed, we can prove the replay identity for a general transition before looking at rounding. I'll use “scan” for applying that transition in order, carrying each updated state into the next step.

Replaying the old buffer and one new input gives the same result as recovering the current state and taking one more step. In Lean, this is observe_push. run_from_replay extends the identity to a whole suffix, including every emitted output. Both use the causal scan from the Kimi K3 example.

If a step uses randomness, replay must retain its random choices as well. Once those choices are fixed, the same deterministic identity applies. Drawing again on each replay changes the question to one about distributions, which needs a separate theorem.

For a small example, let the transition be \(s'=s+u\), the stored checkpoint be 10, and the pending inputs be 2 and −3. Observing the buffer gives state 9. Appending input 4 and replaying from 10 gives 13, exactly as starting from the observed state 9 and applying that one new update. The identity uses the order and contents of the buffer. It does not permit us to discard the −3 or replay a newly sampled substitute for it.

/-- Exact replay before a suffix returns the same final state and suffix outputs as the full scan.
This includes every output, rather than merely equality of final states. -/
theorem run_from_replay (step : State → Input → State × Output)
    (cache : Replay State Input) (suffix : List Input) :
    CausalScan.run step cache.checkpoint (cache.pending ++ suffix) =
      let prior := cache.observe step
      let next := CausalScan.run step prior.1 suffix
      (next.1, prior.2 ++ next.2) :=
  CausalScan.run_append step cache.checkpoint cache.pending suffix

An error budget for checkpoint writes

A rounded checkpoint can differ from the reference state in every entry. To track that difference, suppose one update bounds each coordinate error by \(L\) times the previous allowance, with the prepared input held fixed. Suppose each storage write adds at most \(\delta\) to that allowance. Running \(n\) updates before writing multiplies it by \(L^n\) and then adds one \(\delta\).

Writing \(E_j\) for the budget after the \(j\)th checkpoint gives the recurrence:

Here \(E_0\) is the error in the initial state, and \(n_{j+1}\) is the number of tokens in the next chunk. If each update shrinks or preserves a difference between states, then \(0\le L\le1\). After \(m\) writes, we get:

For \(E_0\ge0\), \(\delta\ge0\), and \(0\le L\le1\), we pay at most one \(\delta\) per write. If updates can amplify a difference, the weighted recurrence still applies but this linear bound can fail. Even under the linear bound, moving the writes changes where their errors enter the computation.

The factor \(L^n\) can make the weighted budget much smaller than the count-only bound. With no initial error, \(L=1/2\), a write allowance \(\delta=0.1\), and two chunks of two tokens, the first budget is 0.1. The second is \(\frac14\times0.1+0.1=0.125\). The older allowance has shrunk through two updates before the new write contributes its allowance. These are calculations of a worst-case budget, not measurements of an actual rounding error.

/-- A bound at every actual storage write gives a bound for any checkpoint partition. Sensitivity
depends on token count, while the local storage allowance is paid only at checkpoint sites. -/
theorem quantizedChunks_within_of_writes {shape : Shape} {Input Output : Type}
    (step : Tensor ℝ shape → Input → Tensor ℝ shape × Output)
    (quantize : Tensor ℝ shape → Tensor ℝ shape) (L δ : ℝ)
    (hstep : ∀ input e x y, Within e x y → Within (L * e) (step x input).1 (step y input).1)
    (chunks : List (List Input)) (x y : Tensor ℝ shape) (e : ℝ) (hxy : Within e x y)
    (hwrites : WritesWithin step quantize δ x chunks) :
    Within (checkpointBudget L δ e chunks)
      (quantizedChunks step quantize x chunks) (CausalScan.run step y chunks.flatten).1 

by
  induction chunks generalizing x y e with
  | nil => simpa [quantizedChunks, checkpointBudget, CausalScan.run] using hxy
  | cons chunk rest ih =>
      simp only [WritesWithin] at hwrites
      -- First propagate the old error through the chunk, then pay for its one storage write.
      have hscan := scan_state_within step L hstep chunk x y e hxy
      have hwrite := within_triangle hwrites.1 hscan
      have hwrite' : Within (L ^ chunk.length * e + δ)
          (quantize (CausalScan.run step x chunk).1) (CausalScan.run step y chunk).1 := by
        simpa [add_comm] using hwrite
      -- Both executions resume from their own resulting state with the same remaining inputs.
      have hrest := ih (quantize (CausalScan.run step x chunk).1)
        (CausalScan.run step y chunk).1 (L ^ chunk.length * e + δ) hwrite' hwrites.2
      simpa only [quantizedChunks, checkpointBudget, List.flatten_cons,
        CausalScan.run_append] using hrest

/-- Nonexpansive transitions give a storage-error budget linear in the number of writes.
This bounds worst-case error; it does not compare the actual errors of two schedules. -/
theorem checkpointBudget_le_count {Input : Type} (L δ e : ℝ)
    (hL : 0 ≤ L) (hLone : L ≤ 1) (hδ : 0 ≤ δ) (he : 0 ≤ e)
    (chunks : List (List Input)) :
    checkpointBudget L δ e chunks ≤ e + chunks.length * δ 

by
  induction chunks generalizing e with
  | nil => simp [checkpointBudget]
  | cons chunk rest ih =>
      have hpower : 0 ≤ L ^ chunk.length := pow_nonneg hL _
      have hpowerOne : L ^ chunk.length ≤ 1 := pow_le_one₀ hL hLone
      have hnext : 0 ≤ L ^ chunk.length * e + δ :=
        add_nonneg (mul_nonneg hpower he) hδ
      have hbound := ih (L ^ chunk.length * e + δ) hnext
      simp only [checkpointBudget, List.length_cons, Nat.cast_add, Nat.cast_one]
      nlinarith [mul_le_mul_of_nonneg_right hpowerOne he]

Fewer writes can leave a larger error

Table 16 reports emulated checkpointing experiments on Nemotron 3 Super. The released Ultra cache uses FP16 with stochastic rounding, described on page 46. For the example here, I use an integer grid so we can follow every rounding decision. The grid isolates the effect of the storage schedule; reproducing the report’s measurements would require its model and FP16 encoder.

The table is worth reading by rounding policy. With FP16 round-to-nearest, its reported mean accuracy drop changes from 1.07% without checkpointing to 0.03% with period eight. With FP16 stochastic rounding, the corresponding values are −0.26% and 0.19%. Checkpointing helps several rows, but the table does not show improvement for every policy. Our example asks the related mathematical question: does reducing the number of writes force a smaller terminal error for every input sequence?

This example uses \(s'=as+u\) and nearest-even rounding to the integer grid. Every write has absolute error at most ½. Select the cancellation example to see why fewer writes need not decrease the actual terminal error, even when \(a=1\).

The cancellation sequence is −0.9, −0.9, −0.5, −0.1. Its exact sum is −2.4. Writing each token finishes at −2, with error 0.4. Writing every two tokens finishes at −3, with error 0.6. Both runs satisfy their error budgets. The signs of the rounding errors explain the difference: some cancel, and changing the schedule changes that cancellation. Fewer writes therefore need not reduce the actual final error. The report’s measurements are compatible with this example.

/-- Four nearest-even writes can have less actual error than two writes, through cancellation.
This refutes a universal monotonic-error reading of "fewer sequential quantization steps".
It is a fixed-grid arithmetic example, not a counterexample to the report's measured results. -/
theorem fewer_writes_can_increase_error :
    |quantizedChunks addStep quantize 0
        [[-9 / 10], [-9 / 10], [-1 / 2], [-1 / 10]] - (-12 / 5)| <
      |quantizedChunks addStep quantize 0
        [[-9 / 10, -9 / 10], [-1 / 2, -1 / 10]] - (-12 / 5)| := by
  decide +kernel

The Mamba state update

Inside one recurrent head, the state is a matrix with one row per value channel and one column per state coordinate. The tensor type records this as [valueDim, stateDim]. For each token, the prepared input tells the head how much old state to retain (\(a\)), what values to write (\(x\) and \(B\)), and how to read an output (\(C\)).

The scalar \(a\) determines how much of the old state we retain. The outer product \(xB^{\mathsf T}\) writes new information into the state. Multiplying by \(C\) reads an output vector.

Now start with two different states, \(S\) and \(T\), and apply the same prepared input. Their new input writes cancel when we subtract:

Each coordinate difference is multiplied by \(a\), so \(|a|\le L\) gives the sensitivity condition used in the error budget. BoundedInput pairs a prepared input with evidence of that inequality. The dimensions remain parameters throughout.

The model source gives us a reason for the retention bound. Its rate is \(A=-\exp(\mathrm{logA})\); a mathematical softplus gives a positive timestep \(\Delta\). Consequently \(a=\exp(\Delta A)\) lies between zero and one. The parameterization gives this bound for every real log-rate, so we can check it without inspecting learned values.

I use the pinned model implementation for these equations because the report’s cache section does not give the complete selective SSM update. The adapter begins after projection and convolution. Connecting it to the full decoder still requires proofs for gating, grouped normalization, and output projection, among other operations. Its mathematical softplus also needs a refinement argument for PyTorch’s finite-precision threshold branch.

From state error to output error

A small state error can still give a large output error if the read coefficients are large. If each state coordinate differs by at most \(e\) and both runs use the same read vector \(C\), then every output channel satisfies:

The absolute values are necessary. Suppose two read coefficients are 10 and −10. Their sum is zero, but coordinate errors of +0.01 and −0.01 produce an output difference of \(10\times0.01+(-10)\times(-0.01)=0.2\). The two products have the same sign, so their errors add. The bound uses \(0.01\times(|10|+|-10|)=0.2\) and covers that case.

SSM.read_within follows from the tensor library's matrix-vector coordinate theorem and mathlib’s absolute-sum inequality. step_output_within also accounts for retention in the next update. Both hold at arbitrary dimensions, including an empty state axis. To turn this formula into a small numerical bound, we need a bound on the sum of absolute read coefficients. I have not established one for learned weights. A logit bound would also need to follow the error through gating, normalization, skip connections, and later blocks.

/-- Discretization is nonexpansive in the prior state for every nonnegative time step and every
real learned log-rate. No boundedness assumption on the learned log-rate is required. -/
theorem retention_bounds (delta logA : ℝ) (hdelta : 0 ≤ delta) :
    0 < retention delta logA ∧ retention delta logA ≤ 1 := by
  constructor
  · exact Real.exp_pos _
  · apply Real.exp_le_one_iff.mpr
    exact neg_nonpos.mpr (mul_nonneg hdelta (Real.exp_nonneg _))

/-- The coordinate update has the source's scalar Mamba-2 equation, not a generic affine
transition mislabeled as the entire selective mixer. -/
theorem updateHead_apply {groups headsPerGroup headDim stateDim : Nat}
    (hHeadsPerGroup : 0 < headsPerGroup)
    (parameters : Parameters (groups * headsPerGroup))
    (input : Token groups headsPerGroup headDim stateDim)
    (head : Fin (groups * headsPerGroup)) (state : SSM.State ℝ headDim stateDim)
    (valueIndex : Fin headDim) (stateIndex : Fin stateDim) :
    SSM.update state (prepareHead hHeadsPerGroup parameters input head)
        (valueIndex, stateIndex, PUnit.unit) =
      let delta := timeStep (input.rawDt.getScalar head) (parameters.dtBias.getScalar head)
      state (valueIndex, stateIndex, PUnit.unit) * retention delta (parameters.logA.getScalar head) +
        (Tensor.unstack input.value head).getScalar valueIndex *
          ((Tensor.unstack input.B (groupIndex hHeadsPerGroup head)).getScalar stateIndex * delta) := by
  simp [SSM.update, prepareHead, Tensor.addSpec, Tensor.scaleSpec, Tensor.outerProductSpec,
    Tensor.map2Spec, Tensor.map, Tensor.getScalar_eq_apply]

/-- Grouped coefficient preparation discharges the retention hypothesis for the whole tensor:
prior-state coordinate errors do not increase in one recurrent update. This does not bound
differences in the supplied values/B/C/time steps or in the surrounding residual network. -/
theorem updateHeads_within {groups headsPerGroup headDim stateDim : Nat}
    (hHeadsPerGroup : 0 < headsPerGroup)
    (parameters : Parameters (groups * headsPerGroup))
    (input : Token groups headsPerGroup headDim stateDim)
    (x y : Tensor ℝ [groups * headsPerGroup, headDim, stateDim])
    (e : ℝ) (hxy : Cache.Within e x y) :
    Cache.Within e (updateHeads hHeadsPerGroup parameters input x)
      (updateHeads hHeadsPerGroup parameters input y) 

by
  intro coordinate
  have hhead : Cache.Within e (Tensor.unstack x coordinate.1)
      (Tensor.unstack y coordinate.1) := by
    intro inner
    simpa [Tensor.unstack] using hxy (coordinate.1, inner)
  have hnext := SSM.update_within (Tensor.unstack x coordinate.1)
    (Tensor.unstack y coordinate.1)
    (prepareHead hHeadsPerGroup parameters input coordinate.1) 1 e (by norm_num)
    (prepareHead_bounded hHeadsPerGroup parameters input coordinate.1) hhead
  simpa [updateHeads, Tensor.stackLeading, Tensor.dim] using hnext coordinate.2

/-- The rank-one write cancels when the same prepared input is applied to two states. -/
theorem update_sub_apply {valueDim stateDim : Nat}
    (x y : State ℝ valueDim stateDim) (input : Input ℝ valueDim stateDim)
    (coordinate : Shape.Coord [valueDim, stateDim]) :
    update x input coordinate - update y input coordinate =
      input.retention * (x coordinate - y coordinate) := by
  simp [update, Tensor.addSpec, Tensor.scaleSpec, Tensor.map2Spec, Tensor.map]
  ring

/-- A fixed read vector turns a coordinatewise state error into an output error. The multiplier
is its absolute coefficient sum, so a stable state alone does not imply a small readout error. -/
theorem read_within {valueDim stateDim : Nat}
    (x y : State ℝ valueDim stateDim) (read : Tensor ℝ [stateDim])
    (e : ℝ) (hxy : Cache.Within e x y) :
    Cache.Within (e * ∑ j : Fin stateDim, |read.getScalar j|)
      (Tensor.matVecMulSpec x read) (Tensor.matVecMulSpec y read) 

by
  rintro ⟨i, rest⟩
  cases rest
  change |(Tensor.matVecMulSpec x read).getScalar i -
    (Tensor.matVecMulSpec y read).getScalar i| ≤ _
  rw [Proofs.TensorAlgebra.getScalar_mat_vec_mul_spec,
    Proofs.TensorAlgebra.getScalar_mat_vec_mul_spec, ← Finset.sum_sub_distrib,
    Finset.mul_sum]
  calc
    |∑ j : Fin stateDim, (get2 x i j * read.getScalar j - get2 y i j * read.getScalar j)|
        ≤ ∑ j : Fin stateDim,
            |get2 x i j * read.getScalar j - get2 y i j * read.getScalar j| :=
      Finset.abs_sum_le_sum_abs _ _
    _ ≤ ∑ j : Fin stateDim, e * |read.getScalar j| := by
      apply Finset.sum_le_sum
      intro j _
      rw [← sub_mul, abs_mul]
      apply mul_le_mul_of_nonneg_right _ (abs_nonneg _)
      simpa [get2, Spec.get, Tensor.getScalar, Tensor.unstack, Tensor.item] using
        hxy ⟨i, j, ()⟩

Rounding theory and stored bits

For nearest-even rounding, FloatLib bounds the error by half a unit in the last place (ULP). roundTensor_within applies that theorem to each tensor entry. If we bound all those half-ULPs by \(\delta\), we get the storage allowance needed above. Near zero, the theorem uses an absolute error bound that includes gradual underflow.

FloatLib’s real-valued grid has no upper exponent limit. To connect it to stored bits, its refinement theorems require finiteness premises: an overflowing result has no real value. The released cache uses stochastic rounding, so I also examine the local random rounding choice and the assumptions needed to preserve its mean.

/-- A supplied upper bound on half-ULP values discharges the tensor storage-error contract.
The premise includes subnormal values and does not assume a relative error bound near zero. -/
theorem roundTensor_within {shape : Shape} (fmt : FloatFormat)
    (x : Tensor ℝ shape) (δ : ℝ)
    (hδ : ∀ coordinate, Model.epsilonAt fmt (x coordinate) ≤ δ) :
    Within δ (roundTensor fmt x) x := by
  intro coordinate
  simpa [roundTensor, Tensor.map] using
    (Model.abs_roundAt_sub_le fmt (x coordinate)).trans (hδ coordinate)

Stochastic rounding and the mean state

Consider a value 3.25 between grid points 3 and 4. Nearest rounding always chooses 3. Stochastic rounding chooses 4 one quarter of the time and 3 the rest of the time. The mean is 3.25, although each individual write changes the value. Repeated state writes make us care about both the average and the error along an individual run.

FloatLib's roundQuot chooses the upper integer when the entropy residue is smaller than the discarded remainder. roundQuot_sum counts the choices made by this executable rounder. Averaging them over uniform residues recovers the original quotient.

The prepared SSM update multiplies the old state and adds a fixed write, so updating a nonempty ensemble of states and then averaging gives the update of their mean. update_mean checks this identity. scan_mean carries it through a sequence when the actual ensemble’s storage perturbations have mean zero at every step.

The mean-zero condition must hold at the states reached during execution. Conditionally unbiased rounding can establish it; checking the encoder’s mean at one isolated input is insufficient. The inputs also stay fixed throughout this argument. Nonlinear decoder layers and changes in sampled tokens need their own analysis.

For the example 3.25, an unbiased write can leave an error of −0.25 or +0.75. Averaging those errors with probabilities three quarters and one quarter gives zero, but neither individual write is exact. To describe one run we still need its error bound; to describe the spread across runs we need a variance bound. A zero mean error gives us neither, nor does it make the stored states identical.

/-- Each coordinate is unbiased under a uniform residue. Sharing a residue across coordinates
is allowed here; this theorem makes no claim that the coordinate errors are independent. -/
theorem roundTensor_mean {shape : Shape} (quantum : ℝ) (sign : Bool) (den : Nat)
    (numerators : Tensor Nat shape) (hden : 0 < den) :
    mean (Tensor.stackLeading fun e : Fin den =>
      roundTensor quantum sign den numerators (Tensor.full shape e.val)) =
      Tensor.map (fun (n : Nat) => (if sign then -1 else 1) * quantum * ((n : ℝ) / den))
        numerators := by
  apply TorchLean.Tensor.Internal.Rep.ext
  intro coordinate
  simpa [mean_apply, roundTensor, Tensor.map2Spec, Tensor.map] using
    round_mean quantum sign (numerators coordinate) den hden

/-- Coordinatewise error for any tensor shape, independent of how entropy is correlated. -/
theorem roundTensor_within {shape : Shape} (quantum : ℝ) (sign : Bool) (den : Nat)
    (numerators entropy : Tensor Nat shape) (hden : 0 < den) :
    Within |quantum| (roundTensor quantum sign den numerators entropy)
      (Tensor.map (fun (n : Nat) => (if sign then -1 else 1) * quantum * ((n : ℝ) / den))
        numerators) := by
  intro coordinate
  simpa [roundTensor, Tensor.map2Spec, Tensor.map] using
    round_error quantum sign (entropy coordinate) (numerators coordinate) den hden

/-- A fixed prepared SSM update is affine, so it commutes with the mean of a nonempty ensemble.
This is the step needed to compose conditionally unbiased state writes. Input-dependent
token selection, normalization, and nonlinear decoder layers are outside this identity. -/
theorem update_mean {samples valueDim stateDim : Nat}
    (states : Tensor ℝ [samples, valueDim, stateDim])
    (input : SSM.Input ℝ valueDim stateDim) (hsamples : 0 < samples) :
    SSM.update (mean states) input =
      mean (Tensor.stackLeading fun i => SSM.update (Tensor.unstack states i) input) 

by
  apply TorchLean.Tensor.Internal.Rep.ext
  intro coordinate
  have hn : (samples : ℝ) ≠ 0 := by exact_mod_cast Nat.ne_of_gt hsamples
  simp only [SSM.update, Tensor.addSpec, Tensor.scaleSpec, Tensor.map2Spec_apply,
    Tensor.map, TorchLean.Tensor.Internal.Rep.map_apply,
    mean_apply, Tensor.unstack_stackLeading]
  rw [Finset.sum_add_distrib, ← Finset.sum_mul]
  simp only [Finset.sum_const, Finset.card_univ, Fintype.card_fin, nsmul_eq_mul]
  field_simp

/-- An actual ensemble recurrence with mean-zero storage perturbations preserves the reference
mean through the full sequence. Conditional write unbiasedness can establish `herrors` on
finite histories; neither independence nor zero error on individual histories is assumed. -/
theorem scan_mean {samples valueDim stateDim : Nat}
    (inputs : List (Input ℝ valueDim stateDim))
    (states errors : Nat → Tensor ℝ [samples, valueDim, stateDim]) (hsamples : 0 < samples)
    (hstep : ∀ i : Fin inputs.length, states (i.val + 1) =
      Tensor.addSpec (Tensor.stackLeading fun sample =>
        update (Tensor.unstack (states i.val) sample) inputs[i.val]) (errors i.val))
    (herrors : ∀ i : Fin inputs.length,
      mean (errors i.val) = Tensor.full [valueDim, stateDim] 0) :
    mean (states inputs.length) = (CausalScan.run step (mean (states 0)) inputs).1 

by
  apply scan_mean_of_step inputs states
  intro i
  rw [hstep i, Cache.Stochastic.mean_add, herrors i]
  rw [← update_mean (states i.val) inputs[i.val] hsamples]
  apply TorchLean.Tensor.Internal.Rep.ext
  intro coordinate
  simp [Tensor.addSpec, Tensor.map2Spec]

Why the mean doesn't survive every nonlinear map

Let's carry the 3.25 example one step further. Suppose a later calculation squares the stored value. Three quarters of the writes return three and one quarter return four. The average stored value is 3.25, but the average square is

The difference is 3/16, the variance of this rounded value. The mean-zero write error hasn't disappeared; a nonlinear calculation can turn its spread into a change in the mean output. This hypothetical square illustrates the issue, rather than reproducing one of Nemotron's decoder operations.

For the prepared SSM transition, fixed input coefficients make the update affine in the state, so we can carry the mean through it. We cannot use that argument unchanged for nonlinear gates, normalization, or token selection. There we'd need to bound what the next operation does to the rounding error, much as squaring did in this example.

Losing small gradient contributions in BF16

The report's first training-divergence account describes a divergence after changing the output layer's local gradient accumulation from FP32 to BF16. Each multi-token prediction head's loss had a weight of 0.05, and the small contributions to the shared output layer were lost. Returning to the FP32 recipe re-stabilized training. Let's look at how a representable contribution can disappear during accumulation.

Multi-token prediction asks additional heads to predict farther-ahead targets during training. Their losses contribute gradients to shared parameters, including the output layer. Multiplying each head's loss by 0.05 reduces its gradient contribution before accumulation. The accumulator therefore has to preserve small additions beside larger values; choosing a precision only from the size of each contribution is insufficient.

Take a running sum of 256 and add 1. The next BF16 value above 256 is 258, so 257 falls halfway between them. Nearest-even rounding returns 256; the addition has lost the contribution. Binary32 represents 257 exactly. FloatLib’s finite-format arithmetic checks both equalities in Lean.

The contribution 1 is itself representable in BF16. The loss happens when it is added to the larger running sum. These numbers illustrate that mechanism; they are not gradients recorded from NVIDIA’s run. The report also describes a second training divergence whose cause remained undetermined.

/-- Reuse FloatLib's local-error sum for gradient tensors of any shape and scalar carrier.
The reference sums the interpreted leaves exactly; errors in producing those leaves are separate.
-/
theorem accumulate_within {α : Type} [Storage α] {shape : Shape}
    (combine : α → α → α) (interpret : α → ℝ)
    (schedule : ReductionTree (Tensor α shape))
    (allowance : Tensor α shape → Tensor α shape → ℝ)
    (h : schedule.AllNodes (Tensor.map2Spec combine) id (fun x y =>
      Cache.Within (allowance x y) (Tensor.map interpret (Tensor.map2Spec combine x y))
        (Tensor.addSpec (Tensor.map interpret x) (Tensor.map interpret y)))) :
    Cache.Within (schedule.errorBudget (Tensor.map2Spec combine) id allowance)
      (Tensor.map interpret (accumulate combine schedule))
      (accumulate (· + ·) (schedule.map (Tensor.map interpret))) 

by
  intro coordinate
  have hError := schedule.abs_eval_sub_exact_le_errorBudget (Tensor.map2Spec combine) id
    (fun x => interpret (x coordinate)) allowance
    (by
      induction schedule with
      | leaf x => trivial
      | node left right hl hr =>
          -- Project the tensor's local bound to this coordinate; keep the actual tree order.
          exact ⟨hl h.1, hr h.2.1, by
            simpa [Tensor.map, Tensor.addSpec, Tensor.map2Spec] using h.2.2 coordinate⟩)
  change |interpret (accumulate combine schedule coordinate) -
    schedule.eval (· + ·) (fun x => interpret (x coordinate))| ≤ _ at hError
  rw [accumulate_apply] at hError
  simpa [accumulate_apply, ReductionTree.eval_map, Function.comp_def, Tensor.map] using hError

/-- BF16 loses this representable contribution at an even halfway boundary.
The numbers illustrate the mechanism; they are not recorded gradients from NVIDIA's run. -/
theorem bf16_add_one_eq_self :
    Model.add (Model.roundRat .bfloat16 false 256 1) (Model.roundRat .bfloat16 false 1 1) =
      Model.roundRat .bfloat16 false 256 1 := by
  decide +kernel

/-- The same contribution survives binary32 accumulation. -/
theorem binary32_add_one_eq_sum :
    Model.add (Model.roundRat .binary32 false 256 1) (Model.roundRat .binary32 false 1 1) =
      Model.roundRat .binary32 false 257 1 := by
  decide +kernel

Rolling back rejected draft tokens

Suppose a draft proposes three tokens but the main model accepts only the first. The next computation must start from the state after that first token. A state incorporating the two rejected tokens would describe the wrong prefix.

Attention can discard the keys and values belonging to a rejected suffix. A Mamba update overwrites the fixed-size recurrent state, so discarding token records alone cannot restore it. NVIDIA retains a snapshot at every draft step. I include the initial state in the snapshot sequence too, so rejecting the entire draft selects that state.

These rollback laws use the same causal scan and apply to any deterministic transition. A random transition must retain its choices in the state or prepared inputs for exact replay. The probability law of the speculative acceptance sampler is a separate obligation.

/-- The selected stored state is exactly the state obtained by running the accepted prefix. -/
theorem restore_eq_prefix (step : State → Input → State × Output)
    (initial : State) (draft : List Input) (accepted : Fin (draft.length + 1)) :
    restore step initial draft accepted =
      (CausalScan.run step initial (draft.take accepted.val)).1 

by
  induction draft generalizing initial with
  | nil =>
      simp [restore, snapshots, CausalScan.run]
  | cons input rest ih =>
      cases accepted using Fin.cases with
      | zero => simp [restore, snapshots, CausalScan.run]
      | succ accepted =>
          simpa [restore, snapshots, CausalScan.run] using ih (step initial input).1 accepted

/-- Continuing from the selected snapshot preserves both final state and all committed outputs.
Rejected draft inputs contribute neither recurrent state nor outputs to the committed run. -/
theorem run_after_restore (step : State → Input → State × Output)
    (initial : State) (draft : List Input) (accepted : Fin (draft.length + 1))
    (continuation : List Input) :
    CausalScan.run step initial (draft.take accepted.val ++ continuation) =
      let prior := CausalScan.run step initial (draft.take accepted.val)
      let next := CausalScan.run step (restore step initial draft accepted) continuation
      (next.1, prior.2 ++ next.2) := by
  rw [restore_eq_prefix]
  exact CausalScan.run_append step initial (draft.take accepted.val) continuation

/-- Selecting an accepted-length row restores precisely that prefix's recurrent state. -/
theorem snapshots_unstack {α : Type} [Storage α] [Add α] [Mul α] [Zero α]
    {tokens valueDim stateDim : Nat} (inputs : Sequence α tokens valueDim stateDim)
    (initial : State α valueDim stateDim) (accepted : Fin (tokens + 1)) :
    Tensor.unstack (inputs.snapshots initial) accepted =
      (CausalScan.run step initial (inputs.toList.take accepted.val)).1 := by
  simp only [snapshots, Tensor.unstack_stackLeading, Cache.Speculation.restore_eq_prefix]

Conclusion

I'd love to extend this to the token scores of the released model. We need to follow the error from the actual storage writes through replay and the rest of the decoder, then compare it with the gap between the winning token and its competitors. The cancellation example is a reason to measure the actual error as well as carry a worst-case bound: reducing that bound does not guarantee a smaller error on every input.

The replay proof takes the transition and checkpoint schedule as parameters. If you're trying a different recurrent model or storage format, we can reuse that argument and work on the numerical assumptions specific to your update. I'd be happy to compare examples :)

@misc{george2026nemotron,
  author = {Robert Joseph George},
  title = {{Nemotron 3 Ultra: replaying a rounded state}},
  year = {2026},
  howpublished = {TorchLean Verified Examples},
  url = {https://robertboy18.github.io/TorchLean-Verified-Examples/week-05-nemotron-specification/site/}
}