[xpu-alignment] [26-10-15] [26-10-15] XPU nll_loss2d accumulates fp16 in global atomics: mean cross_entropy returns inf/NaN and all-zero gradients

Summary

F.cross_entropy with reduction="mean" on a spatial (4-D) fp16 input returns inf, then NaN with an all-zero gradient, once the per-work-group partial sums exceed the fp16 range, even though the true loss is ~3.5. The XPU nll_loss2d forward reduces each work group in acc_type_device<scalar_t, kXPU> (fp32) but casts the result back to scalar_t and accumulates it into fp16 global accumulators, and output/total_weight are allocated in the input dtype.

Found by an upstream alignment scan of the window [2026-10-15T00:00:00Z, 2026-10-16T00:00:00Z). Upstream source: pytorch/pytorch#201120 (reported for CUDA, aten/src/ATen/native/cuda/NLLLoss2d.cu, open, no fix landed). XPU has its own kernels, so the CUDA fix will not cover XPU.

Observed XPU behavior

Intel(R) Arc(TM) Pro B60 Graphics, torch 2.16.0.dev20261015+xpu:

pixels (HxW)	XPU fp16 loss	fp64 CPU reference	all-zero grad
128x128 = 16384	3.513671875	3.5134	no
181x181 = 32761	`inf`	3.5003	no
256x256 = 65536	`nan`	3.5044	yes

The signature matches the upstream CUDA report exactly: the loss sum overflows to inf first, then total_weight overflows and inf/inf gives NaN; the backward divides by the overflowed total_weight, so the input gradient is silently all zero. Autocast is unaffected because autocast runs nll_loss in fp32.

Target-path evidence

The reproducer keeps the tensors on XPU for the whole measured op (input.device=xpu:0, dtype=torch.float16, loss.device=xpu:0, grad.device=xpu:0) and compares against an fp64 CPU reference computed from the identical initialized inputs, so only the XPU kernels are measured. The code path is:

src/ATen/native/xpu/LossNLL2d.cpp:58-59 - output and total_weight are allocated with self.options(), i.e. fp16.
src/ATen/native/xpu/sycl/LossNLL2dKernels.cpp:123-128 - each work group reduces in accscalar_t (fp32) and then does atomicAdd(sycl_global_ptr<scalar_t>(total_weight), static_cast<scalar_t>(acc_weight_reduce)) and the same for output, so the cross-work-group accumulation is fp16.
src/ATen/native/xpu/sycl/LossNLL2dKernels.cpp:137 - *output /= *total_weight in fp16.
src/ATen/native/xpu/sycl/LossNLL2dKernels.cpp:324 - backward computes -(size_average ? *grad_output / *total_weight : *grad_output), which is 0 once total_weight is inf.

Reproducer

import torch
import torch.nn.functional as F

torch.manual_seed(0)
C = 21
for H in (128, 181, 256):
    x_cpu = torch.randn(1, C, H, H, dtype=torch.half)
    y_cpu = torch.randint(0, C, (1, H, H))

    x = x_cpu.detach().to("xpu").requires_grad_(True)
    y = y_cpu.to("xpu")

    loss = F.cross_entropy(x, y)
    loss.backward()

    ref = F.cross_entropy(x_cpu.double(), y_cpu).item()
    print(
        f"{H * H} px: loss {loss.item()} (fp64 {ref:.4f}), "
        f"all-zero grad: {bool((x.grad == 0).all())}"
    )

Output:

16384 px: loss 3.513671875 (fp64 3.5134), all-zero grad: False
32761 px: loss inf (fp64 3.5003), all-zero grad: False
65536 px: loss nan (fp64 3.5044), all-zero grad: True

Environment

torch 2.16.0.dev20261015+xpu
python 3.13.13 (/opt/conda/bin/python3.13)
torch path /opt/conda/lib/python3.13/site-packages/torch/__init__.py
xpu device Intel(R) Arc(TM) Pro B60 Graphics

Ownership

The accumulation and the output allocation are XPU code in this repository (LossNLL2d.cpp, LossNLL2dKernels.cpp), so the fix is independent XPU work: allocate output/total_weight in the accumulate type (or accumulate into fp32 scratch) and cast only the final mean back to the input dtype, matching what the upstream issue proposes for NLLLoss2d.cu. Upstream issue 201120 is CUDA-scoped and has no linked fix PR, and no intel/torch-xpu-ops issue tracks this signature (the closest, #5050, is tolerance flakiness in the fused linear_cross_entropy, a different op and root cause). Since cross-work-group accumulation is also the source of nondeterminism here, the fix may be worth coordinating with the upstream direction in pytorch/pytorch#201120.

This issue was produced by an automated upstream-alignment scan with an AI assistant; the evidence above was gathered on an XPU host and independently reviewed.

@torchxpubot

@torchxpubot fix

Automated fix session for this issue. Stages are appended below as they complete.

Issue Triage

Field	Value
Issue type	single-bug
Reproduction missing	no
Scope	torch-xpu-ops
Dependencies	(none)
Handling	agent-fixable

Automated by issue-triage.

Reproduce

Ran the issue's reproducer against nightly torch 2.16.0.dev20261019+xpu on
Intel(R) Arc(TM) Pro B60 Graphics (BMG).

Test case	Verdict	Observed
`F.cross_entropy` fp16 4-D mean, 181x181 px	REPRODUCED	`loss inf` (fp64 reference 3.5003)
`F.cross_entropy` fp16 4-D mean, 256x256 px	REPRODUCED	`loss nan`, `all-zero grad: True`

A direct aten::nll_loss2d_forward comparison shows CPU fp16 saturates
identically, so this is not an XPU-vs-CPU divergence but an fp16
accumulation defect that XPU shares with CPU/CUDA; the XPU kernels are
ours to fix, as the issue notes.

Automated by fix-reproduce.

Root Cause

Analyzed at torch-xpu-ops@76cc8324, pytorch origin/main@9544ec5e.

The reduction path of the XPU nll_loss2d forward reduces each work group in
the accumulate type but then leaves the accumulate type to combine work
groups: src/ATen/native/xpu/sycl/LossNLL2dKernels.cpp:122-129 casts
acc_weight_reduce / input_sum_reduce back to scalar_t and atomically adds
them into the fp16 output / total_weight scalars, and
LossNLL2dKernels.cpp:137 then divides in fp16. For a 4-D fp16 input of
181x181 pixels the loss sum is ~1.1e5, past the fp16 maximum of 65504, so
output saturates to inf; at 256x256 the weight sum (65536) overflows too
and inf/inf yields NaN. The backward at LossNLL2dKernels.cpp:324 divides
by that same total_weight, so once it is inf every gradient is exactly
zero.

The in-repo 1-D sibling already does this correctly: nll_loss reduces in a
single work group and performs the mean division in the accumulate type,
casting only the final values out
(src/ATen/native/xpu/sycl/LossNLLKernel.cpp:152-156). The 2-D kernel needs
the same, adapted to its multi-work-group atomics: accumulate into an
accscalar_t scratch buffer and let a finalize kernel compute
output = static_cast<scalar_t>(sum / weight_sum).

The public dtypes cannot change: CUDA allocates both outputs with
self.options() (aten/src/ATen/native/cuda/NLLLoss2d.cu:502-503) and
aten.nll_loss2d_forward has a registered decomposition
(torch/_decomp/decompositions.py:4932) that fixes the meta dtype of
total_weight to the input dtype. So total_weight stays fp16 and still
saturates above 65504 elements; the loss becomes correct in every case, and
the gradient is correct whenever the total weight is representable in fp16.

Verdict: IMPLEMENTING -- accumulate across work groups in the accumulate type and divide there, mirroring the 1-D kernel.
Target repo: torch-xpu-ops
Domain: xpu-kernel

Automated by fix-root-cause.

Implement Result

Target repo: torch-xpu-ops
Analyzed at: torch-xpu-ops@76cc8324
What I changed: src/ATen/native/xpu/sycl/LossNLL2dKernels.cpp (+41/-25): the forward reduce kernel now atomically accumulates into an accscalar_t scratch buffer, and a new nll_loss2d_forward_finalize_kernel performs the mean division in accscalar_t and casts the loss and total weight back to the input dtype (it replaces the fp16 nll_loss2d_forward_average_kernel). test/regressions/test_nll_loss2d_fp16_overflow.py (+66): new regression test.
Why: cross-work-group accumulation and the division were the only fp16 steps in an otherwise accscalar_t reduction, which is what overflowed; the 1-D nll_loss kernel already keeps both in the accumulate type and casts once at the end (src/ATen/native/xpu/sycl/LossNLLKernel.cpp:152-156).
Skip added: no
Ready for verify: yes

diff --git a/src/ATen/native/xpu/sycl/LossNLL2dKernels.cpp b/src/ATen/native/xpu/sycl/LossNLL2dKernels.cpp
index 60b1dbaa..276c1835 100644
--- a/src/ATen/native/xpu/sycl/LossNLL2dKernels.cpp
+++ b/src/ATen/native/xpu/sycl/LossNLL2dKernels.cpp
@@ -75,8 +75,8 @@ template <typename scalar_t, typename accscalar_t, typename index_t, int SIMD>
 SYCL_EXT_ONEAPI_FUNCTION_PROPERTY((syclexp::nd_range_kernel<1>))
 SYCL_EXT_ONEAPI_FUNCTION_PROPERTY((syclexp::sub_group_size<SIMD>))
 void nll_loss2d_forward_kernel_(
-    scalar_t* output,
-    scalar_t* total_weight,
+    accscalar_t* output_acc,
+    accscalar_t* total_weight_acc,
     const scalar_t* input,
     const int64_t* target,
     const scalar_t* weight,
@@ -121,20 +121,27 @@ void nll_loss2d_forward_kernel_(
 
   if (item.get_local_id(0) == 0) {
     atomicAdd(
-        sycl_global_ptr<scalar_t>(total_weight),
-        static_cast<scalar_t>(acc_weight_reduce));
-    atomicAdd(
-        sycl_global_ptr<scalar_t>(output),
-        static_cast<scalar_t>(input_sum_reduce));
+        sycl_global_ptr<accscalar_t>(total_weight_acc), acc_weight_reduce);
+    atomicAdd(sycl_global_ptr<accscalar_t>(output_acc), input_sum_reduce);
   }
 }
 
-template <typename scalar_t>
+// Casts the accumulators back to the output dtype. The division happens in
+// the accumulate type so that a loss sum beyond the range of scalar_t (e.g.
+// any fp16 spatial input with more than a few thousand pixels) still yields
+// a finite mean.
+template <typename scalar_t, typename accscalar_t>
 SYCL_EXT_ONEAPI_FUNCTION_PROPERTY((syclexp::nd_range_kernel<1>))
-void nll_loss2d_forward_average_kernel(
+void nll_loss2d_forward_finalize_kernel(
     scalar_t* output,
-    const scalar_t* total_weight) {
-  *output /= *total_weight;
+    scalar_t* total_weight,
+    const accscalar_t* output_acc,
+    const accscalar_t* total_weight_acc,
+    bool size_average) {
+  const accscalar_t weight_sum = *total_weight_acc;
+  *total_weight = static_cast<scalar_t>(weight_sum);
+  *output = static_cast<scalar_t>(
+      size_average ? *output_acc / weight_sum : *output_acc);
 }
 
 void nll_loss2d_forward_kernel(
@@ -236,6 +243,15 @@ void nll_loss2d_forward_kernel(
               blocks_per_sample =
                   (blocks_per_sample == 0) ? 1 : blocks_per_sample;
               int total_blocks = blocks_per_sample * batch_size;
+              // Work groups accumulate into the accumulate type: adding the
+              // partial sums straight into a scalar_t scalar overflows for
+              // low-precision dtypes.
+              auto acc_buffer = at::zeros(
+                  {2},
+                  input_.options().dtype(
+                      c10::CppTypeToScalarType<accscalar_t>::value));
+              auto* output_acc = acc_buffer.mutable_data_ptr<accscalar_t>();
+              auto* total_weight_acc = output_acc + 1;
               constexpr auto kfn = nll_loss2d_forward_kernel_<
                   scalar_t,
                   accscalar_t,
@@ -246,8 +262,8 @@ void nll_loss2d_forward_kernel(
                   work_group_size,
                   getCurrentSYCLQueue(),
                   2 * sizeof(accscalar_t) * work_group_size,
-                  output.mutable_data_ptr<scalar_t>(),
-                  total_weight.mutable_data_ptr<scalar_t>(),
+                  output_acc,
+                  total_weight_acc,
                   input_.const_data_ptr<scalar_t>(),
                   target_.const_data_ptr<int64_t>(),
                   optional_data<scalar_t>(weight_),
@@ -256,18 +272,18 @@ void nll_loss2d_forward_kernel(
                   blocks_per_sample,
                   ignore_index,
                   work_group_size);
-              // Divide by total_weight
-              if (reduction == at::Reduction::Mean) {
-                constexpr auto kfn =
-                    nll_loss2d_forward_average_kernel<scalar_t>;
-                sycl_kernel_submit<kfn>(
-                    1,
-                    1,
-                    getCurrentSYCLQueue(),
-                    0,
-                    output.mutable_data_ptr<scalar_t>(),
-                    total_weight.const_data_ptr<scalar_t>());
-              }
+              constexpr auto finalize_kfn =
+                  nll_loss2d_forward_finalize_kernel<scalar_t, accscalar_t>;
+              sycl_kernel_submit<finalize_kfn>(
+                  1,
+                  1,
+                  getCurrentSYCLQueue(),
+                  0,
+                  output.mutable_data_ptr<scalar_t>(),
+                  total_weight.mutable_data_ptr<scalar_t>(),
+                  output_acc,
+                  total_weight_acc,
+                  reduction == at::Reduction::Mean);
             });
       });
 }

Automated by fix-implement.

Verify

Source build of pytorch origin/main@9544ec5e with torch-xpu-ops at the fix
commit, TORCH_XPU_ARCH_LIST=bmg,pvc, on Intel(R) Arc(TM) Pro B60 Graphics.
The baseline and fixed builds differ only by
src/ATen/native/xpu/sycl/LossNLL2dKernels.cpp: the kernel hunk was reverted
in the same tree and rebuilt to get the Before column, so this is a true A/B
on one build rather than a wheel-vs-source comparison.

Test case	Before	After
`test_cross_entropy_fp16_spatial_mean_H_181`	FAIL	PASS
`test_cross_entropy_fp16_spatial_mean_H_256`	FAIL	PASS
`test_cross_entropy_fp16_spatial_mean_backward`	PASS	PASS
issue reproducer, 181x181 px	`loss inf`	`loss 3.5` (fp64 3.5003)
issue reproducer, 256x256 px	`loss nan`, all-zero grad	`loss 3.50390625` (fp64 3.5044), all-zero grad

Before was loss inf at 32761 px and loss nan with an all-zero gradient at
65536 px. The backward case passes on both builds: it pins the gradient
semantics but is not the regression guard. The 256x256 gradient is still
all-zero, which this fix cannot change -- total_weight is returned in the
input dtype (CPU and CUDA do the same, and the registered
aten.nll_loss2d_forward decomposition fixes its meta dtype), so a weight sum
of 65536 saturates to inf in fp16 whatever the kernel does. CPU fp16 returns
nan and all-zero gradients for both sizes, so XPU is now strictly better
than the CPU reference here.

Regression checks on the fixed build:

Suite	Result
`test/regressions/test_nll_loss2d_fp16_overflow.py`	3 passed
`test/xpu/test_nn_xpu.py -k "nll_loss"`	25 passed, 6 skipped
`test/xpu/test_ops_xpu.py -k "nll_loss or cross_entropy"`	213 passed, 12 skipped, 7 xfailed

Verdict: PASSED -- the mean-reduced fp16 loss is finite and matches an fp64 reference at every size tested.
Lint: clean (lintrunner -a on both changed files, including CLANGFORMAT and CLANGTIDY)

Automated by fix-verify.

Summary

Outcome: IMPLEMENTING(fix_verified)

The XPU nll_loss2d forward was the only step of an otherwise fp32 reduction
that fell back to the input dtype: it cast each work group's partial sums back
to scalar_t, atomically added them into the fp16 output / total_weight
scalars, and divided in fp16, so a mean-reduced fp16 F.cross_entropy over a
few tens of thousands of pixels overflowed to inf and then NaN. The forward
now accumulates into an accumulate-type scratch buffer and a finalize kernel
does the division there, casting only the final loss and total weight back --
the same shape as the 1-D nll_loss kernel in this repo. On a source build the
reported inf / NaN losses become 3.5 and 3.504 against fp64 references of
3.5003 and 3.5044, the new regression test fails with the kernel hunk reverted
in the same tree, and the nll_loss / cross_entropy suites stay green. One
reported symptom is out of reach: total_weight is returned in the input dtype
by schema, so beyond 65504 total weight it still saturates to inf and that
gradient is still all-zero -- CPU and CUDA behave the same way, and
pytorch/pytorch#201120 tracks the shared defect upstream.

Branch: agent/fix-issue-5722
Base: 76cc832470c07a51e9d3595317d4b459a0ae5677
Changed files: src/ATen/native/xpu/sycl/LossNLL2dKernels.cpp, test/regressions/test_nll_loss2d_fp16_overflow.py
Nothing was pushed. The patch is exported as a workflow artifact for a human to review and apply.

Generated by fix job.

Review

Copy into a new comment, keep ONE verdict, add a note. Two verdicts in
one comment are dropped as ambiguous.

**strong-accepted** <!-- review: strong-accepted -->
**weak-accepted** <!-- review: weak-accepted -->
**weak-rejected** <!-- review: weak-rejected -->
**strong-rejected** <!-- review: strong-rejected -->
**issue-drop** <!-- review: issue-drop -->

notes:

On a batch run add item=<pytorch issue> to say which sub-item the
verdict is about, in the bold text and in the marker -- the marker
is invisible, so the bold line is all a reader sees. One comment per
sub-item, since a bare verdict on a batch does not say which test it
means and is dropped. run=<id> only when judging a run other than the
newest, in both places too.

A CI DISABLED sub-item, item= from the fan-out table:

**strong-accepted item=197334** <!-- review: strong-accepted item=197334 -->
notes: same TensorAssign overload as upstream CUDA, fp8 dispatch list matches

Another sub-item of the same batch, judged on an earlier run:

**weak-rejected run=35171979678 item=197335** <!-- review: weak-rejected run=35171979678 item=197335 -->
notes: green only because the test is skipped, the kernel still lacks fp8

An [xpu-alignment] issue that should not have been filed -- one issue
per difference, so no item=:

**issue-drop** <!-- review: issue-drop -->
notes: accepted XPU/CUDA numerics difference, nothing to repair

weak-rejected
notes: Local testing confirms the issue, but the proposed algorithm only fixes forward loss while all-zero gradients persist at 256×256.

label-issue: #5722

Root cause: nll_loss2d forward casts fp32 work-group sums to fp16 before global atomicAdd (src/ATen/native/xpu/sycl/LossNLL2dKernels.cpp:122-129)
into fp16 output/total_weight (src/ATen/native/xpu/LossNLL2d.cpp:58-59), overflowing to inf/NaN; owner torch-xpu-ops.

axis	value	reason
`type`	`Bug`	`issue_type` empty; mean cross_entropy returns inf/NaN instead of ~3.5.
`test`	`test: ut`	No e2e/oob signal; op-level reproducer, default UT surface.
`module`	`module: reduction`	Defect is the cross-work-group sum/mean accumulation in LossNLL2dKernels.cpp:122-129.
`priority`	`Medium`	Wrong result (inf/NaN, zero grad) without crash; no prior passing version cited.
`dtype`	`dtype: float16`	Only fp16 overflows (`dtype=torch.float16`; fp16 global atomics).
symptom	`Accuracy`	Loss inf/NaN vs fp64 reference 3.50; numerical overflow.
symptom	`training`	Backward yields all-zero input gradient (LossNLL2dKernels.cpp:324 divides by inf total_weight).

Applied to this issue. Generated by a local label-issue run.

Navigation Menu

Repository navigation

[xpu-alignment] [26-10-15] [26-10-15] XPU nll_loss2d accumulates fp16 in global atomics: mean cross_entropy returns inf/NaN and all-zero gradients #5722

Summary

Observed XPU behavior

Target-path evidence

Reproducer

Environment

Ownership

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

[xpu-alignment] [26-10-15] [26-10-15] XPU nll_loss2d accumulates fp16 in global atomics: mean cross_entropy returns inf/NaN and all-zero gradients #5722

Description

Summary

Observed XPU behavior

Target-path evidence

Reproducer

Environment

Ownership

Activity

laifenxiawucha commented on Oct 20, 2026

torchxpubot commented on Oct 20, 2026

laifenxiawucha commented on Oct 21, 2026

daisyden commented on Nov 10, 2026

Metadata

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Issue actions