torch-directml: F.cross_entropy(reduction='none') silently produces zero gradient on backward (0.2.5.dev240914 / torch 2.4.1+cpu) #740

Summary

On torch-directml==0.2.5.dev240914 (torch 2.4.1+cpu), calling
torch.nn.functional.cross_entropy(logits, targets, reduction='none')
and then combining/reducing the per-element losses manually (e.g. to
apply a custom per-token weighting before backward) produces a normal
forward pass, a normal-looking loss value, and a normal-looking autograd
graph — but .backward() on the resulting loss yields an all-zero
gradient on every upstream parameter. No error, warning, or NaN is
raised. Training appears to proceed normally (loss values printed each
step look plausible) while no parameter is actually being updated.

We hit this in a real LoRA fine-tuning run: after ~450 optimizer steps
with a per-token-weighted cross-entropy loss built on
reduction='none', none of the LoRA adapter's trainable tensors had
moved from their initial values, despite the run completing without
error and reporting a steadily-changing loss curve throughout.

Environment

• torch==2.4.1+cpu
• torch-directml==0.2.5.dev240914
• transformers==4.46.3, peft==0.20.0, accelerate==1.15.0
• Windows, AMD Radeon 8060S-class integrated GPU, DirectML device

Minimal repro

The following isolates the bug to specifically the fused
reduction='none' kernel, by comparing three mathematically-equivalent
ways of computing the same per-token loss on the same model, inputs,
and device, and checking the gradient norm on a known trainable tensor
after a single backward() call:

import torch
import torch.nn.functional as F
import torch_directml

device = torch_directml.device()

# ... set up a small model with a trainable parameter (e.g. a linear
# layer or a LoRA adapter tensor) on `device`, and a batch of
# logits/targets on the same device ...

# 1. Fused kernel, reduction='none', then manual mean:
loss_a = F.cross_entropy(logits, targets, reduction='none').mean()
loss_a.backward()
print("fused reduction='none':", param.grad.norm().item())  # -> 0.0

# reset gradients, recompute logits fresh, then:

# 2. Manual unfused equivalent:
log_probs = F.log_softmax(logits, dim=-1)
per_token_loss = -log_probs.gather(-1, targets.unsqueeze(-1)).squeeze(-1)
loss_b = per_token_loss.mean()
loss_b.backward()
print("manual log_softmax+gather:", param.grad.norm().item())  # -> nonzero

# reset gradients, recompute logits fresh, then:

# 3. Control: built-in mean-reduction path (as used internally by
#    Hugging Face Transformers' default loss computation):
loss_c = F.cross_entropy(logits, targets, reduction='mean')
loss_c.backward()
print("built-in reduction='mean':", param.grad.norm().item())  # -> nonzero

In our real repro (diag_gradflow_repro.py, run against a live LoRA
adapter tensor on a real training batch), the three gradient norms were:

Only the fused reduction='none' path returns a zero gradient. The
manual unfused equivalent and the built-in reduction='mean' control
both produce a real, nonzero gradient on the identical batch, model, and
device — which rules out the model, the data, or the optimizer as the
cause and isolates the problem to the fused kernel's backward pass under
reduction='none' specifically.

We're happy to share a fully self-contained, minimal script (no
transformers/peft dependency) reproducing this on a plain
nn.Linear if that's more useful for triage — let us know.

Why this is worth flagging as high-severity

There is no error, warning, or NaN anywhere in this failure. The loss
value printed at each step looks plausible and can even appear to
decrease over a run (if the reduction happens elsewhere in the
pipeline), so this is not something a user would discover by watching
training logs — only by explicitly checking that trainable parameters
actually changed after training. We only found it because we added a
pre-flight gradient-flow assertion as a general safety practice. We'd
guess this could silently affect anyone using a custom/per-token loss
built on reduction='none' on this backend, not just our specific
LoRA use case.

What we're NOT reporting here

This is unrelated to #702 (masked_fill uint8 overflow during
gpt-neo inference) — we hit that one separately during the same
project and worked around it, but it's a different bug in a different
op with a different (loud, exception-raising) failure mode. Filing this
as a new issue rather than a comment on #702 since the root cause,
affected op, and failure signature are all different.