[johnb-code]

This paper is from DeepMind, and I want people to have the most efficient training code.

Codex prepared the code and the quoted text below at my request.

Description

Add an opt-in loss reference, a gradient runner, CPU tests, and instructions under examples/numerical_stability/.
The example follows arXiv:2610.31642v1.
Codex authored this example and helped prepare this draft.

Related Issue

Closes #1022.

Motivation and Context

On CPU, FP32 logits [9, -9, -9] with target 0 produce zero stock loss and target gradient.
The reference preserves the small residual.
This motivates gradient checks at the same represented inputs.
No Neuron hardware defect is claimed.
Searches for the paper ID, softmax collapse, and numerical stability found no matching open PR.
The issue-linked PR search also found no existing contribution for this proposal.

The example leaves the dispatcher and default primitives unchanged.
It supports replicated, class-last logits, unweighted integer targets, and fixed teacher logits.
It excludes ignored labels, weights, smoothing, sharding, and higher-order distillation derivatives.
FP16 distillation requires explicit FP32 or FP64 computation.
This prevents an FP16 0 * inf case in the distillation reference.

The reference is a correctness oracle for a later native implementation.
That implementation must preserve fused reductions and the vocabulary memory contract.
Compare loss forward and backward time, peak memory, training throughput, and time to target quality.
Report overhead together with any stability benefit.
This example does not establish a Neuron efficiency benefit.

Type of Change

• Bug fix
• New feature
• Breaking change
• Documentation update
• Refactoring

Testing

Ran locally (no Neuron hardware needed):

The CPU environment used Python 3.12.13, PyTorch 2.14.1, and pytest 9.1.1.
The following repository-root commands used the isolated CPU interpreter.
Set RESULTS to a writable receipt directory.
The commands below abbreviate recorded output paths with RESULTS.

PYTHONPATH=examples/numerical_stability \
python -m pytest examples/numerical_stability/test_stable_softmax.py -q \
  -o addopts='' --confcutdir=examples/numerical_stability

python examples/numerical_stability/check_loss_gradients.py \
  --device cpu --mode both --dtype float32 \
  --output "$RESULTS/cpu-fp32.json"

python examples/numerical_stability/check_loss_gradients.py \
  --device cpu --mode both --dtype bfloat16 \
  --output "$RESULTS/cpu-bf16.json"

python examples/numerical_stability/check_loss_gradients.py \
  --device cpu --mode both --dtype bfloat16 --autocast --loss-fp32 \
  --output "$RESULTS/cpu-bf16-amp-fp32.json"

pre-commit run --all-files

All 41 CPU tests passed.
The suite includes 10 fullgraph aot_eager value and gradient comparisons.
The FP16 tests check both validation modes and explicit FP32 and FP64 computation.
All three CPU runner configurations passed their candidate gates.
All nine pre-commit hooks passed or skipped files outside their scope.
One PyTorch torch.jit.script deprecation warning came from the forward-mode rejection test.
CPU checks do not verify Neuron lowering, fusion, convergence, memory use, or speed.

Ran on Neuron hardware:

No Neuron host was available.
A maintainer must start hardware CI after review.

Checklist

• Added or updated tests for this change
• I've performed a self-review for my code
• Ran the tests, or explained under Testing why I could not
• Updated the corresponding documentation
• My changes generate no new warnings
• Said under Description if an AI assistant helped with this change

[tdalmia-aws]

Thanks for the PR, @johnb-code. Closing for the same reason as this issue: #1022 (comment).