Train with deliberate precision choices

The problem: a small gradient can disappear

The decoder’s matrix products carry activations and gradients through many layers. Using fewer bits for those values can reduce memory traffic and make supported hardware arithmetic faster, but it also changes which numbers survive. A useful backward signal may round to zero before AdamW ever sees it. A large signal may become infinity instead. Moving an operation to a GPU does not solve either numerical problem.

Start with a tiny gradient equal to $2^{-25}$. In FP16, the smallest positive subnormal value is $2^{-24}$. Our value lies halfway between zero and that value; rounding to nearest with ties to even gives zero. Multiplying by a loss scale $S=128$ first gives $2^{-18}$, which FP16 can represent. Dividing the surviving gradient by $128$ afterward recovers the tiny value in higher precision. Scaling after it has rounded to zero would preserve only zero.

This chapter extends the same matrix product, reverse pass and decoder update. The CPU example rounds matrix operands and results to FP16 or BF16 and accumulates matrix products in FP32. It keeps the existing f64 master weights, AdamW moments and non-matrix operations. All tensor buffers still contain f64 values: this is an inspectable numerical reference, with no demonstrated storage saving or speedup. The existing CUDA path remains f64; requesting a mixed CUDA policy returns an explicit error.

Scale backward, then unscale before clipping

Let $L$ be the scalar mean next-token negative log likelihood, $\theta$ the master parameters and $S$ a finite positive scale fixed for one attempt. Ignoring rounding for the moment, the exact-arithmetic chain rule gives

$\widetilde g$ is the scaled parameter gradient; $g$ is its unscaled value, and $g_{\mathrm{clip}}$ is the gradient passed to AdamW. The notation $\nabla_\theta$ means the gradient with respect to all master parameters, and $\lVert g\rVert_2$ is their global Euclidean gradient norm. $C$ is the positive maximum allowed norm. When the norm is zero, leave the zero gradient unchanged. Otherwise the last line limits the norm to $C$. The global norm includes all named parameter gradients, counting the tied embedding/output table once.

For a one-coordinate scaled gradient of $8$, scale $4$ and norm limit $1$, unscale to $2$, measure norm $2$, then clip to $1$. Clipping $8$ to $1$ before division would produce $0.25$ instead. That would change the gradient passed to AdamW. Loss scaling also leaves AdamW’s decoupled weight decay unscaled.

Represented computation does not cancel the scale perfectly. Rounding occurs inside backward, and a larger scale can move a signal away from zero or into overflow. The scale cannot restore a forward activation already lost to rounding. The numerical policy and the update decision must therefore be observable separately.

Choose the representation at each operation

Both formats use sixteen bits, including a sign bit. FP16 has five exponent bits and ten explicit fraction bits; BF16 has eight and seven. The exponent controls range, while fraction bits control spacing within a magnitude range. Around one, the next FP16 value is $1+2^{-10}$ and the next BF16 value is $1+2^{-7}$. Thus BF16 spans a wider range but has coarser spacing there. The halfway inputs $1+2^{-11}$ for FP16 and $1+2^{-8}$ for BF16 both round to $1$ under ties-to-even conversion. BF16 preserves $2^{-25}$ in this fixture, whereas FP16 rounds it to zero. Neither format represents arbitrary real numbers exactly.

For one matrix output cell, let $a_k$ and $w_k$ be its paired input values, and let $K$ count contracted features. For $K>0$, the zero-based index $k$ runs from $0$ through $K-1$. Write $Q_h$ for rounding to the chosen half format $h$ and $\operatorname{RN}_{32}$ for FP32 rounding. Let $p_k$ be the rounded product for input pair $k$, $s_k$ the FP32 running sum after $k$ paired products and $y$ the represented output value. The CPU policy is

The loop visits $k$ in ascending order; when $K=0$, the empty contraction returns zero. The rounded result $y$ is then widened exactly into an f64 tensor slot. FP32 accumulation protects the running sum from repeated half rounding; it cannot repair rounded inputs or prevent the final half result from overflowing. The default f64 policy retains the original contraction.

Softmax, RMSNorm, general reductions, residual arithmetic and parameter-gradient accumulation retain their existing f64 implementations. A half-precision matrix product inside those larger computations still follows the selected matrix policy. The trace distinguishes this limited policy from a model whose every activation and operation uses a half dtype.

A master weight is the higher-precision value AdamW updates. In the fixture, subtracting $0.0001$ from an f64 master initialized to $1$ produces the displayed value $0.9999$, while its FP16 copy still rounds to $1$. Later updates can accumulate in the master even when an individual rounded copy does not change. We retain f64 masters here; FP32 masters are common in the historical method discussed below.

Extend the existing numerical and update paths

MatmulPrecision delegates format conversion to the half library. TensorExecution::cpu_with_precision selects the matrix policy explicitly. The existing checked matrix plan still owns shapes, broadcasting, views and transposes. No new tensor shape checker or floating-point codec is needed.

/// Rounding policy for the existing matrix product and its local pullbacks.
///
/// Mixed modes round each operand to the named half format, multiply and sum
/// in f32, then round the result to that half format. Tensor buffers, saved
/// primals, accumulated parameter gradients and master weights remain f64.
/// This CPU numerical reference provides no compact-storage or speed claim.
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
pub enum MatmulPrecision {
    #[default]
    F64,
    F16F32,
    Bf16F32,
}

impl MatmulPrecision {
    pub const fn as_str(self) -> &'static str {
        match self {
            Self::F64 => "f64",
            Self::F16F32 => "fp16-fp32",
            Self::Bf16F32 => "bf16-fp32",
        }
    }

    /// Type of the transient represented operands and rounded matrix result.
    pub const fn represented_dtype(self) -> &'static str {
        match self {
            Self::F64 => "f64",
            Self::F16F32 => "fp16",
            Self::Bf16F32 => "bf16",
        }
    }

    pub const fn accumulator_dtype(self) -> &'static str {
        match self {
            Self::F64 => "f64",
            Self::F16F32 | Self::Bf16F32 => "f32",
        }
    }

    /// Uses half's format conversion, including ties, subnormals and overflow.
    /// Widening the represented result back to f64 loses no further bits.
    pub fn round(self, value: f64) -> f64 {
        match self {
            Self::F64 => value,
            Self::F16F32 => f16::from_f64(value).to_f64(),
            Self::Bf16F32 => bf16::from_f64(value).to_f64(),
        }
    }

    pub(crate) fn operand_f32(self, value: f64) -> f32 {
        match self {
            Self::F64 => value as f32,
            Self::F16F32 => f16::from_f64(value).to_f32(),
            Self::Bf16F32 => bf16::from_f64(value).to_f32(),
        }
    }
}

Each matrix operation saves its precision on the existing tape. Its pullback uses the recorded policy, including when ordinary CPU backward() is called. The matrix derivative is evaluated with represented operands; casts are treated as identity paths for training. This is the usual training approximation, not the mathematical derivative of the discontinuous rounding function. The text encoder, shared transformer and tied text head keep their existing composition; precision introduces no tokenizer requirement into the backbone.

decoder_update_with_loss_scaling enters the existing update owner. It supplies a recording context with the selected execution and starts backward with seed $S$. This differentiates $SL$ without adding a separate multiplication node. The shared norm calculation uses gradients divided by $S$, and the existing AdamW gradient transform applies unscale and clipping to moment inputs. Raw stored gradients need no intermediate rewrite.

/// Extends the same update owner with a scaled backward seed and skip outcome.
/// Start with cleared gradients. Forward receives the explicit execution
/// context; the reverse pass remains transactional on numerical overflow.
#[allow(clippy::too_many_arguments)]
pub fn decoder_update_with_loss_scaling<F, O>(
    model: &DecoderModel,
    optimizer: &mut AdamW,
    config: &TrainerConfig,
    step: usize,
    execution: &TensorExecution,
    scaler: &mut DynamicLossScaler,
    forward: F,
    observer: &mut O,
) -> Result<ScaledDecoderUpdate, TrainerError>
where
    F: FnOnce(AutogradContext<'_>) -> Result<TensorValue, TrainerError>,
    O: DecoderUpdateObserver,
{
    for parameter in model.parameters() {
        let gradient =
            parameter
                .tensor()
                .gradient()
                .ok_or_else(|| TrainerError::MissingGradient {
                    name: parameter.name().to_owned(),
                })?;
        if gradient.as_slice().iter().any(|&value| value != 0.0) {
            return Err(TrainerError::ScaledUpdateRequiresZeroGradient {
                name: parameter.name().to_owned(),
            });
        }
    }
    decoder_update_attempt(
        model,
        optimizer,
        config,
        step,
        execution,
        Some(scaler),
        || forward(AutogradContext::recording().with_execution(execution)),
        observer,
    )
}

The scaled entry requires zero initial gradients. Otherwise a new scale could be applied to previously accumulated work that used another scale. A dirty gradient, stale tape, shape error or unavailable backend is an error to fix; it is not reported as an overflow skip.

The existing backward pass stages all local pullbacks and gradient additions before writing parameter gradients. If one becomes nonfinite, it returns a numerical error without committing those writes. The scaled update turns that specific failure into Skipped, halves the scale and returns before AdamW. Weights, gradients, parameter revisions, moments, beta powers and optimizer step remain unchanged. No weight decay is applied, and no learning-rate-schedule step is consumed. A skip therefore differs from an AdamW step with zero gradients, which can still decay weights and advance optimizer state. This does not undo arbitrary side effects inside a caller’s forward function or observer.

The controller accepts finite power-of-two scales of at least $1$, with a minimum, maximum and positive growth interval. After that many committed finite updates, it doubles the scale up to its maximum. An overflow resets the count and halves the scale down to its minimum. At the minimum, a persistent overflow still returns a skip; the caller must decide whether to retry or stop. There is no unbounded automatic retry loop in the trainer.

/// Course-local dynamic policy: halve on overflow; double after an interval of
/// committed finite updates, within declared bounds. Skips reset the interval.
/// This is run state, independent of learned model identity.
#[derive(Clone, Debug, PartialEq)]
pub struct DynamicLossScaler {
    scale: f64,
    minimum: f64,
    maximum: f64,
    growth_interval: usize,
    finite_updates: usize,
}

impl DynamicLossScaler {
    pub fn new(
        initial: f64,
        minimum: f64,
        maximum: f64,
        growth_interval: usize,
    ) -> Result<Self, InvalidLossScale> {
        let valid = |value: f64| {
            value.is_finite()
                && value >= 1.0
                && value.log2().fract() == 0.0
                && value.log2().exp2() == value
        };
        if !valid(initial)
            || !valid(minimum)
            || !valid(maximum)
            || minimum > initial
            || initial > maximum
            || growth_interval == 0
        {
            return Err(InvalidLossScale);
        }
        Ok(Self {
            scale: initial,
            minimum,
            maximum,
            growth_interval,
            finite_updates: 0,
        })
    }

    pub const fn scale(&self) -> f64 {
        self.scale
    }
    pub const fn finite_updates(&self) -> usize {
        self.finite_updates
    }

    pub(crate) fn overflow(&mut self) {
        self.scale = (self.scale / 2.0).max(self.minimum);
        self.finite_updates = 0;
    }

    pub(crate) fn committed(&mut self) {
        self.finite_updates += 1;
        if self.finite_updates == self.growth_interval {
            // Divide first so the growth itself cannot overflow f64.
            self.scale = if self.scale <= self.maximum / 2.0 {
                self.scale * 2.0
            } else {
                self.maximum
            };
            self.finite_updates = 0;
        }
    }
}

From FP16 range problems to deliberate mixed precision

FP16 saves representation bits but can erase small updates and gradients or overflow; choosing a lower dtype for every operation does not address those failures. Micikevicius and colleagues’ 2017 paper, Mixed Precision Training, describes FP16 training with FP32 master weights, loss scaling and FP32 accumulation, including language-model experiments.

BF16 provides a wider exponent range than FP16, while using fewer fraction bits; the precision choice must still match the numerical operation. Kalamkar and colleagues’ 2019 BF16 study examines training across tasks including language modeling, emphasizing BF16’s FP32-like exponent range.

Our Rust contrast is narrower: run the same decoder under both half-format matrix policies and the retained f64 reference, then inspect an overflowing update. It illustrates numerical mechanisms with existing f64 masters and protected operations. It does not reproduce either paper’s training results or demonstrate accelerated mixed-precision GPU training.

Decoder projections, attention products and their adjoints reuse matrix operations whose representation and accumulation precision affect training.

Follow the values and the update decision

Reproduce and inspect the bounded examples

The existing clap CLI passes the selected precision to the CPU numerical demonstration.

use clap::{Parser, ValueEnum};
use llm_from_scratch_practical::tensor::precision::MatmulPrecision;

#[derive(Clone, Copy, ValueEnum)]
enum Precision {
    Fp16,
    Bf16,
    F64,
}

#[derive(Parser)]
#[command(
    about = "Inspect mixed-precision CPU arithmetic and transactional loss-scaled decoder updates"
)]
struct Args {
    #[arg(long, value_enum, default_value = "fp16")]
    precision: Precision,
}

fn main() -> Result<(), Box<dyn std::error::Error>> {
    let precision = match Args::parse().precision {
        Precision::Fp16 => MatmulPrecision::F16F32,
        Precision::Bf16 => MatmulPrecision::Bf16F32,
        Precision::F64 => MatmulPrecision::F64,
    };
    println!(
        "{}",
        serde_json::to_string(&practical_ch12_mixed_precision_training::demonstrate(
            precision
        )?)?
    );
    Ok(())
}

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

COURSE_CORPUS=false ./course run cargo run --quiet --offline --locked -p practical-ch12-mixed-precision-training

The command prints one JSON record. Its dtypes object reports actual host storage and matrix arithmetic separately. representation records the rounding examples. finite_updates contains two committed updates, growing scale from $128$ to $256$ after the second. text_reference compares the existing seed-39 ab model against f64, retaining $20$ named gradient tensors, $1372$ learned scalars and the tied input/output table.

The overflow_probe first commits a warm-up update so AdamW already has moments, then deliberately starts a separate scale controller at $16777216$. Its first skip halves that scale to $8388608$ and keeps optimizer step $1$. The fixture checks all preserved state on each skip and retries scheduled step $2$, allowing at most $32$ attempts. The recorded run skips $8$ attempts, then commits that same step at scale $65536$ and clears gradients. These are numerical examples, not evidence of a trained language model’s quality.

Use the same setup and working directory to inspect BF16, then select the chapter’s four focused reproduction tests:

COURSE_CORPUS=false ./course run cargo run --quiet --offline --locked -p practical-ch12-mixed-precision-training -- --precision bf16
COURSE_CORPUS=false ./course run cargo test --offline --locked -p practical-ch12-mixed-precision-training

The BF16 command prints the same record structure with the selected matrix dtype changed; the deliberately overflowing probe remains explicitly FP16. The test command should report four passing library tests. They cover preserved populated optimizer state with a successful retry, both half formats through the existing text model, agreement with the original clipped update under scaled f64 arithmetic, and refusal of structural errors and dirty gradients.

/// Populate AdamW moments first, then deliberately stress the same live model
/// with an oversized scale. Retry the same scheduled step after each skip.
pub fn overflow_fixture() -> DemoResult<Value> {
    let model = model_fixture()?;
    let mut optimizer = optimizer()?;
    let execution = TensorExecution::cpu_with_precision(MatmulPrecision::F16F32);
    let mut warmup = DynamicLossScaler::new(128.0, 1.0, 128.0, 2)?;
    if !matches!(
        update(&model, &mut optimizer, &execution, &mut warmup, 1)?,
        ScaledDecoderUpdate::Applied { .. }
    ) {
        return Err("bounded warm-up unexpectedly overflowed".into());
    }
    let model_before = DecoderModelState::snapshot(&model).bit_pattern();
    let optimizer_before = optimizer.persistence_state();
    let gradients_before = gradients(&model);
    let revisions_before = revisions(&model);
    let mut scaler = DynamicLossScaler::new(16777216.0, 1.0, 16777216.0, 2)?;
    let mut skipped = 0;
    let mut first_skip = None;
    for _ in 0..32 {
        match update(&model, &mut optimizer, &execution, &mut scaler, 2)? {
            ScaledDecoderUpdate::Skipped {
                loss_scale,
                next_loss_scale,
                overflow,
                ..
            } => {
                skipped += 1;
                let preserved = model_before == DecoderModelState::snapshot(&model).bit_pattern()
                    && same_optimizer_bits(&optimizer_before, &optimizer.persistence_state())
                    && gradients_before == gradients(&model)
                    && revisions_before == revisions(&model);
                if !preserved {
                    return Err("overflow changed model, optimizer, gradients or revisions".into());
                }
                first_skip.get_or_insert(
                    json!({"loss_scale":loss_scale,"next_loss_scale":next_loss_scale,
                    "optimizer_step_before":1,"optimizer_step_after":optimizer.step_count(),
                    "weights_moments_powers_gradients_revisions_preserved":preserved,
                    "overflow_operation":overflow.to_string()}),
                );
            }
            ScaledDecoderUpdate::Applied {
                execution: applied,
                loss_scale,
                next_loss_scale,
            } => {
                if skipped == 0 {
                    return Err("overflow fixture did not skip".into());
                }
                return Ok(
                    json!({"precision":"fp16-fp32","first_skip":first_skip,"skipped_attempts":skipped,
                    "retry":{"scheduled_step":2,"optimizer_step":optimizer.step_count(),"loss_scale":loss_scale,
                        "next_loss_scale":next_loss_scale,"clipped":applied.step.clipped(),
                        "cleared_gradients":gradients(&model).iter().all(|&bits| bits == 0)}}),
                );
            }
        }
    }
    Err("bounded retry limit reached".into())
}

1. With the course Docker setup available, run from the repository root:

   COURSE_CORPUS=false ./course run cargo run --quiet --offline --locked -p practical-ch12-mixed-precision-training

   Inspect the dtype record. Explain why half-format matrix arithmetic does not imply that the host tensors use two bytes per element.

2. Work through the explained scaled-gradient example: gradient $8$, scale $4$ and norm limit $1$. State the gradient passed to AdamW and why unscaling must precede clipping.

3. With the course Docker setup available, run from the repository root:

   COURSE_CORPUS=false ./course run cargo run --quiet --offline --locked -p practical-ch12-mixed-precision-training

   Inspect overflow_probe in the JSON output. Identify the state preserved on its first skip and the scheduled step used by its successful retry.

Carry numerical policy into larger batches

The existing text encoder, source-independent transformer and tied head now share a recorded matrix precision policy. Matrix precision and loss scale are run choices. Selecting them leaves the master parameters and logical model architecture unchanged. A committed AdamW update changes the master parameter values. CPU inference remains available, and native mixed-precision CUDA execution requires a later primitive implementation and its own validation.

Chapter 13 will cover gradient accumulation and activation recomputation. Combining gradients from several contributions must preserve their intended normalization and scale. This chapter deliberately requires cleared gradients for each scaled update; it does not yet implement that accumulation policy. Portable model saving and loading remain in Chapter 15.