Training Considerations: Objectives, Scaling Laws, Precision, Memory

Causal language modelling

The objective behind every model in module 15. Predict the next token:

Its virtue is that it is self-supervised: any text is training data, no labels required. And because of the causal mask (module 08), all T positions are trained in one forward pass — position t predicts t+1 for every t simultaneously.

logits = model(input_ids)                    # (B, T, V)
loss = F.cross_entropy(
    logits[:, :-1].reshape(-1, V),           # predictions for positions 0..T-2
    input_ids[:, 1:].reshape(-1),            # targets are the NEXT tokens
)

That one-position shift is the entire objective.

Masked language modelling

BERT's objective: mask ~15% of tokens and predict them from both directions.

MLM is more sample-efficient per token of context but trains on far fewer positions per pass. Causal LM won because generation is what people want, and because it scales more cleanly.

Multi-token prediction (MTP)

A 2024–26 addition Raschka flags across several models.

The richer signal speeds up training. Then it pays a second dividend: the extra heads become a draft model for speculative decoding (module 14) at no extra cost. Qwen3-Next "introduces a native Multi-Token Prediction (MTP) mechanism, which not only yields an MTP module with a high acceptance rate for Speculative Decoding but also enhances the overall performance."

Used by DeepSeek V3/V3.2, GLM-4.5, MiniMax-M2, Qwen3-Next, Xiaomi MiMo-V2-Flash, and Nemotron 3 Super.

Nemotron 3 Super goes furthest: it uses MTP at inference too, where "the shared-weight MTP head acts as an internal draft model for native speculative decoding." Raschka suggests calling this "shared-weight MTP for speculative decoding" rather than plain MTP, since standard MTP is training-only.

Kaplan et al. (2020)

Loss falls as a power law in model size, dataset size, and compute:

The finding that shaped 2020–22: performance improves smoothly and predictably with scale, so you can extrapolate from small runs. Kaplan's analysis suggested spending most of a compute increase on parameters, which is why GPT-3 (175B) was trained on only ~300B tokens.

Chinchilla (Hoffmann et al., 2022)

Then Kaplan's conclusion turned out to be wrong, because of a learning-rate schedule artefact. Chinchilla's corrected finding:

Chinchilla-70B beat GPT-3-175B while being 2.5× smaller. Same compute, better allocation.

Why modern models blow past Chinchilla

Look at Llama 3 8B: 1875 tokens per parameter, roughly 90× past compute-optimal. Deliberately.

Chinchilla optimises training compute. It ignores inference entirely. If you serve a model to millions of users, inference cost dominates lifetime cost by orders of magnitude — so it is rational to overtrain a smaller model, paying more once to pay less forever.

The practical lesson: scaling laws are about compute allocation, not about what is best to deploy. A "compute-optimal" model is rarely the one you want to serve.

They also say nothing about data quality, architecture, or post-training — all of which matter enormously in 2026. Raschka's repeated caution applies: "datasets, training techniques, and hyperparameters vary widely and are often not well documented."

The formats

The key comparison is FP16 versus BF16. Both are 16 bits, but they spend those bits differently. BF16 keeps FP32's 8-bit exponent — the same dynamic range — and sacrifices mantissa precision.

For deep learning that is the right trade. Gradients span many orders of magnitude, so range matters more than precision, and FP16's narrow range causes underflow. BF16 is the default for training in 2026.

The recipe

from torch.amp import autocast, GradScaler

scaler = GradScaler()                       # only needed for fp16, not bf16

for batch in loader:
    optimizer.zero_grad()
    with autocast(device_type='cuda', dtype=torch.bfloat16):
        loss = model(batch)                 # forward in bf16
    scaler.scale(loss).backward()           # gradients in bf16
    scaler.step(optimizer)                  # master weights updated in fp32
    scaler.update()

Three components:

1. Half-precision forward/backward — 2× less memory, 2× more throughput on tensor cores.
2. FP32 master weights — the optimizer keeps full-precision copies, because accumulating tiny updates into bf16 loses them entirely.
3. Loss scaling — multiply the loss before backward to push small gradients above the representable minimum, then unscale before the step. Required for fp16; unnecessary for bf16, whose range already covers it.

Certain operations always run in fp32 regardless: softmax, layer/RMS norm reductions, and loss computation. Module 06's RMSNorm upcast is an instance of this rule.

FP8 and below

FP8 training is now used at frontier scale — DeepSeek V3 trained in FP8, and FlashAttention-3 supports it. It needs careful per-tensor scaling and is not yet routine. NVFP4 appears in Nemotron 3 deployment.

Training memory for a model with N parameters, bf16 + Adam:

A 7B model needs ~112 GB just for state — before a single activation. This is why training requires clusters and why optimizer sharding (ZeRO / FSDP) exists.

Then add activations, which scale with batch size, sequence length, and depth:

At long context this exceeds the parameter state.

Gradient checkpointing

The standard fix. Backprop needs forward activations; storing all of them is expensive. Instead, store activations only at block boundaries and recompute the rest during the backward pass.

from torch.utils.checkpoint import checkpoint

class CheckpointedStack(nn.Module):
    def forward(self, x, *args):
        for block in self.blocks:
            # store only this block's input; recompute its internals in backward
            x = checkpoint(block, x, *args, use_reentrant=False)
        return x

Roughly 30% more compute for a large memory reduction — and the same principle as FlashAttention's backward recomputation (module 11): recompute is cheaper than remember when you are memory-bound.

Other memory techniques

Large runs combine several — commonly called 3D or 4D parallelism.

AdamW is the default: Adam with decoupled weight decay. Typical settings β₁=0.9, β₂=0.95, ε=1e-8, weight decay 0.1, cosine schedule with warm-up.

The notable 2025 development, per Raschka, is Kimi K2:

Muon orthogonalises the update direction for matrix-shaped parameters (via Newton-Schulz iteration), which conditions the update better than elementwise adaptive scaling.

Raschka is careful about the evidence. On the widely-shared claim that Kimi K2's loss curve was exceptionally smooth:

A good instance of reading a training plot properly: smoothness is not the same as fast convergence, and the second is what matters.

Arcee Trinity Large uses a related MuOpt optimizer.

Large-model training diverges. The architectural defences, cross-referenced:

Raschka's OLMo 2 caveat is worth carrying into any such discussion: their stability figure "shows the results of the reordering together with QK-Norm, which is a separate concept. So, it's hard to tell how much the normalization layer reordering contributed by itself." Stability techniques are usually shipped together and rarely ablated individually.

Pretraining produces a base model that continues text. Turning it into an assistant takes more:

Raschka scopes his article to architecture and repeatedly defers here — and it is worth taking seriously that most 2026 quality differences come from this stage, not from architecture. His comment on Mistral 3 Large adopting DeepSeek V3's architecture wholesale says it directly: "A lot of the secret sauce these days is in the training pipeline as well as the inference scaling strategies."

Two models with identical architectures can differ enormously. Architecture is necessary to understand; it is not where the remaining differentiation lives.