Add automatic mixed precision training with BF16 autocast and FP16 gradient scaling

Feature description

Add automatic mixed precision (AMP) training with operation-specific autocasting and optional dynamic gradient scaling. Include CPU BF16 where backend kernels support it, alongside GPU BF16/FP16.

Burn already provides low-precision dtypes, explicit differentiable casts, and checkpoint dtype conversion. The proposed feature adds automatic runtime precision selection and training-loop integration. BF16 dtype support already exists; the missing piece is its integration into AMP.

Related discussion: https://github.com/tracel-ai/burn/issues/4332. Maintainers note that some reductions and matrix multiplications already use higher-precision accumulators; the policy audit should build on those existing safeguards.

Feature motivation

Allow users to reduce activation memory and improve training throughput without manually casting every operation. Keep numerically sensitive computation, model parameters, and optimizer state in FP32. CPU speedups depend on hardware and backend kernels and must be measured.

Suggested implementation plan

1. Backend capabilities and API: audit BF16/FP16 operation support and define a scoped, device-aware autocast API with explicit behavior for unsupported operations. Keep normal FP32 training as the default.
2. BF16 autocast first: start with one supported backend and a small, documented operation policy. Cast eligible matrix multiplications/convolutions to BF16, retain FP32 for sensitive operations, and preserve gradients through casts. Validate nesting, thread isolation, and interaction with fusion. BF16 should default to no gradient scaling.
3. FP16 and dynamic scaling: scale loss before backward, unscale gradients before clipping and optimizer updates, detect nonfinite gradients, skip invalid updates, and adapt the scale.
4. Training integration: support gradient accumulation with one fixed scale per accumulation window and one scale update per optimizer step attempt. Coordinate overflow decisions across distributed workers, advance update-based schedulers only after successful optimizer updates, and checkpoint/restore scaler state.
5. Validation and rollout: compare loss/gradients and convergence against FP32; test overflow recovery, accumulation, clipping, distributed skipped updates, and checkpoint resume. Benchmark representative CPU/GPU workloads and publish a backend support matrix and usage example. Deliver this in incremental PRs.

QAT is a separate feature and is outside this proposal.

Questions for the reviewer

• Does this scope and incremental rollout fit Burn's direction, or is there an existing AMP effort we should coordinate with?
• Where should autocast policies live in Burn's architecture so they compose correctly with autodiff, dispatch, and fusion?
• Which backend should be the first implementation target? Should CPU BF16 be part of the initial milestone or follow GPU support?
• Is keeping parameters and optimizer state in FP32, with BF16 autocast first and FP16 scaling next, the preferred starting point?
• What numerical tolerances, convergence checks, and performance criteria should gate acceptance?

Please review the proposed approach and recommend any changes before implementation begins.