burn-optim: FP16 norm clipping can zero finite gradients; FP32 promotion or max scaling?

Describe the bug

GradientClipping::Norm(1.0) turns a finite FP16 gradient [300.0] into [0.0], instead of approximately [1.0].

The current norm implementation computes square().sum().sqrt() in separate tensor operations. The square of 300 is 90,000, which exceeds FP16's maximum finite value, 65,504. The norm becomes infinity, and the clipping coefficient becomes zero. Higher-precision accumulation after the square cannot recover the lost value.

Here 300 is a gradient, not a model weight. For example, w = 0.5, x = 10, y = -10, and loss (w*x - y)^2 produce dL/dw = 300 through ordinary backward.

Reproduction

Confirmed on Burn 0.22.0-pre.3, local revision f168ea1f070fa9814d6dbc2267faf4cfe50e13ab, CPU/Flex. The clipping source matched upstream when checked. Burn GPU behavior has not been tested for this report.

As an integration test in crates/burn-train/tests/fp16_norm_clipping.rs:

use burn_core::tensor::{Device, FloatDType, Tensor};
use burn_optim::grad_clipping::GradientClipping;

#[test]
fn finite_fp16_gradient_should_not_be_zeroed_by_norm_clipping() {
    let device = Device::flex();
    let grad = Tensor::<1>::from_floats([300.0], &device)
        .cast(FloatDType::F16);

    let actual = GradientClipping::Norm(1.0)
        .clip_gradient(grad)
        .into_scalar::<f32>();

    // Currently actual == 0.0.
    assert!((actual - 1.0).abs() < 0.001, "actual={actual}");
}

Run: cargo test -p burn-train --test fp16_norm_clipping -- --nocapture.

PyTorch comparison and scope

Tested PyTorch 2.8.0+cu128 on CPU and an NVIDIA RTX 3060/CUDA, including foreach=None, False, and True. For an FP16 parameter initialized to 0.5, backward through the scalar example above produces gradient 300, and clip_grad_norm_([w], 1.0) correctly produces gradient 1.

PyTorch's norm kernels use FP32 intermediates for FP16/BF16; the L2 reduction casts before squaring. They do not use max scaling for this path.

PyTorch is not universally overflow-safe: FP16 [60000, 60000] overflows the returned FP16 norm, and FP32 [1e20, 1e20] also returns infinity. Default clipping produces zeros; error_if_nonfinite=True rejects both cases. These results were observed on CPU and CUDA.

This Burn report concerns gradients that actually reach clipping in FP16. A standard AMP test with FP32 model parameters and gradients did not encounter this specific failure. The reproduction establishes correctness, not its frequency in production workloads.

Two possible solutions

1. Promote low-precision clipping calculations to FP32

For FP16/BF16 gradients, promote before squaring, keep the norm, coefficient, and gradient scaling in FP32, and cast only the final clipped gradient back. Preserve FP32/FP64 behavior.

• Small, localized change with no public API changes.
• Follows PyTorch's wider-intermediate approach. Keeping the norm in FP32 additionally avoids its FP16 norm-output limitation.
• Requires casts and potentially a temporary FP32 tensor; does not solve extreme FP32/BF16 or FP64 overflow generally.

2. Use max scaling for a robust clipping calculation

For finite gradients and a nonnegative clipping limit, the core calculation can be expressed as:

m = max(abs(g))
if m == 0: return g
u = g / m
r = sqrt(sum(u * u))
clipped = u * min(m, max_norm / r)

Promote FP16/BF16 intermediates to FP32 here as well, so long reductions do not overflow FP16. Avoid reconstructing m * r, which could overflow even when the clipped result is representable. Handle zero and nonfinite inputs explicitly, preserve output dtype, and settle epsilon/unchanged-gradient semantics in the implementation.

• Handles extreme finite gradients beyond the range covered by solution 1.
• Adds a maximum reduction and extra operations/passes; performance and rounding effects need measurement.
• This is a proposal, not an implemented or benchmarked fix.

Both options keep the existing per-parameter clipping scope. Adding global gradient-norm clipping is a separate feature.

Questions for the reviewer

Could a maintainer/reviewer share their viewpoint on these two approaches before implementation?

1. Do you prefer the smaller FP32-promotion fix, or max scaling for broader numerical robustness despite the extra reduction?
2. Should this remain local to gradient clipping, or should Burn expose a reusable norm operation with an explicit accumulation dtype?
3. Should nonfinite norms preserve existing behavior or gain an explicit error option, and which backend/performance checks would you require?

Suggested regression coverage: [300], 1,024 components of 10, [60000, 60000] in FP16; zero and below-threshold gradients; output dtype preservation; and extreme FP32/FP64 values if solution 2 is selected.