Float16 support with loss / gradient scaling

[wwwjn]

-- From user feedback

While might sound counter intuitive, but fp16 is back! https://arxiv.org/abs/2510.26788
Supporting fp16 with proper scaling is essential for a more stable RL run. At the moment titan has no support for such a regime from our understanding.

[zpcore]

Paste explanation of https://arxiv.org/abs/2510.26788 from Claude:

The Paper: "Defeating the Training-Inference Mismatch via FP16"

  The paper identifies a critical problem in RL fine-tuning of LLMs: numerical mismatch between the
  training engine and the inference engine caused by floating-point precision.

  The Root Cause: BF16 vs FP16
  ┌────────┬───────────────┬───────────────┬──────────────────────────────────────┐
  │ Format │ Mantissa bits │ Exponent bits │                Trait                 │
  ├────────┼───────────────┼───────────────┼──────────────────────────────────────┤
  │ FP16   │ 10            │ 5             │ Higher precision, limited range      │
  ├────────┼───────────────┼───────────────┼──────────────────────────────────────┤
  │ BF16   │ 7             │ 8             │ Lower precision, wider dynamic range │
  └────────┴───────────────┴───────────────┴──────────────────────────────────────┘
  BF16 has become the default for LLM training because its wide dynamic range avoids overflow/underflow
  issues during pre-training. However, during RL fine-tuning, models operate within already-established
  numerical ranges, so the dynamic range advantage of BF16 becomes less important — and its lower
  precision (only 7 mantissa bits) introduces large rounding errors.

  Why This Matters for RL

  In RL fine-tuning (e.g., GRPO, PPO), there are two policies:
  1. Training policy — runs through the training engine (e.g., FSDP, torch compile)
  2. Inference/rollout policy — runs through a separate inference engine (e.g., vLLM, SGLang)

  Even when both use identical weights, BF16 rounding causes them to produce different probability
  distributions. This leads to:

  1. Biased gradients — the importance sampling ratio between training and inference log-probs is
  corrupted
  2. Training instability — GRPO under BF16 collapses to 73-84% accuracy on a sanity test where the model
   should achieve ~100%
  3. Deployment gap — models optimized under the training engine don't perform optimally when deployed
  with the inference engine

  The Fix: Switch to FP16

  FP16's 10 mantissa bits provide enough precision to absorb the implementation differences between
  training and inference engines. Results:
  - All RL algorithms (GRPO, GSPO, PG variants) achieve ~99% accuracy on the sanity test under FP16
  - More stable optimization, faster convergence, stronger performance
  - Requires only a few lines of code change with proper loss scaling