torch-directml: F.cross_entropy(reduction='none') silently produces a zero backward gradient (loss value correct, gradient vanishes)

[Rook-CodeVolt]

Summary

torch.nn.functional.cross_entropy(..., reduction='none'), when combined with a downstream reduction (e.g. (per_token_loss * weights).sum() / weights.sum()), produces a plausible forward-pass loss value and a normal-looking autograd graph (NllLossBackward0) on torch-directml, but the backward pass silently yields an exactly-zero gradient for all upstream trainable parameters. No exception, warning, or NaN is raised — the run appears to train normally (the loss value itself is correct), but no learning occurs.

This is distinct from the already-tracked masked_fill uint8-overflow issue (#702) — it is a separate failure in the cross_entropy(reduction='none') backward kernel path itself, not in a preceding mask-construction op.

Environment

• torch: 2.4.1+cpu
• torch-directml: 0.2.5.dev240914
• transformers: 4.46.3
• peft: 0.20.0
• OS: Windows-10-10.0.26200-SP0
• GPU: AMD Radeon 8060S (integrated, Ryzen AI Max+ 395 APU), DirectML device (privateuseone:0)
• Model: HuggingFaceTB/SmolLM2-135M + LoRA adapter (PEFT, r=8, alpha=16, target_modules q/k/v/o_proj), attn_implementation="eager"

Observed behaviour

Given identical logits/labels for one real training batch:

Loss computation path	Loss value	Backward grad norm on a LoRA lora_B param
F.cross_entropy(reduction='none') then weighted mean, .backward()	9.8083	0.0 (exactly zero)
Manual log_softmax + gather NLL, same weighting, .backward()	9.8083 (identical)	0.0534
HF-internal model(..., labels=...) default mean-reduction cross_entropy path (separate control run, same batch/model state)	(not directly comparable value; separate forward)	0.0699

The loss value computed by the reduction='none' path is numerically identical to the manual reimplementation (both 9.8083 to 4 dp) — so the forward pass and the reduction math are correct. Only the backward gradient silently vanishes for the fused reduction='none' kernel path. .grad is not None (so requires_grad and graph connectivity are intact) — it is a real tensor containing all zeros.

This was originally discovered during real LoRA fine-tuning: a 6-epoch training run using this loss path completed without any error, logged what looked like a plausible flat loss curve (~9.5–9.6 throughout, no explosion/NaN), and finished normally — but post-hoc inspection showed all 120 lora_B tensors were still exactly PEFT's zero-init value (i.e. the optimizer had received a zero gradient at every one of 444 update steps). A second run using the manual log_softmax+gather reimplementation (only that one line changed) produced real, diversified nonzero lora_B values and a real loss collapse (9.6 → 0.009) on the same data/hyperparameters — confirming the fused reduction='none' backward is the sole cause.

Why this is a high-severity, easy-to-miss bug

cross_entropy(reduction='none') is a completely standard pattern for any form of per-token loss weighting (e.g. cold-start/curriculum weighting, focal loss, token-class-balanced loss). On this DirectML build, a training script written this way will run start-to-finish with no exception and a plausible loss log, giving no obvious signal that gradients never flowed downstream — it only surfaces if a user separately inspects trained-parameter deltas, which most training scripts do not do by default.

Minimal reproduction

Repro script (self-contained, requires torch, torch-directml, transformers, peft, and a local causal-LM checkpoint + a JSONL file of {"input_text": ..., "output_json": ...} records — no proprietary or sensitive data, any small local text works):

#!/usr/bin/env python3
"""
Minimal repro: F.cross_entropy(reduction='none') silently produces a
zero backward gradient on torch-directml, while a mathematically
equivalent manual log_softmax+gather NLL gives a real nonzero gradient
for the identical forward pass/batch.
"""
import argparse
import json


def build_prompt(input_text):
    return f"{input_text}\n\nExtract:\n"


def canonical_target(output_json):
    return json.dumps(output_json, separators=(",", ":"), sort_keys=False)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--model-path", required=True)
    ap.add_argument("--train-file", required=True)
    ap.add_argument("--seq-len", type=int, default=384)
    args = ap.parse_args()

    import torch
    import torch_directml
    from transformers import AutoModelForCausalLM, AutoTokenizer
    from peft import LoraConfig, get_peft_model, TaskType

    dml = torch_directml.device()
    _orig_masked_fill = torch.Tensor.masked_fill

    def _dml_safe_masked_fill(self, mask, value):
        # workaround for #702 (uint8 overflow), unrelated to this bug
        if self.device.type == "privateuseone":
            if not torch.is_tensor(value):
                value = torch.tensor(value, dtype=self.dtype, device=self.device)
            return torch.where(mask, value, self)
        return _orig_masked_fill(self, mask, value)

    torch.Tensor.masked_fill = _dml_safe_masked_fill

    tokenizer = AutoTokenizer.from_pretrained(args.model_path, local_files_only=True)
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token

    records = []
    with open(args.train_file, "r", encoding="utf-8") as f:
        for line in f:
            line = line.strip()
            if line:
                records.append(json.loads(line))
    sample = records[:4]

    model = AutoModelForCausalLM.from_pretrained(
        args.model_path, torch_dtype=torch.float32, local_files_only=True, attn_implementation="eager",
    )
    lora_config = LoraConfig(
        task_type=TaskType.CAUSAL_LM, r=8, lora_alpha=16, lora_dropout=0.0,
        target_modules=["q_proj", "k_proj", "v_proj", "o_proj"], bias="none",
    )
    model = get_peft_model(model, lora_config)
    model.to(dml)
    model.config.use_cache = False
    model.train()

    lora_b_param = None
    for name, p in model.named_parameters():
        if "lora_B" in name and p.requires_grad:
            lora_b_param = p
            break

    def encode(record):
        prompt = build_prompt(record["input_text"])
        target = canonical_target(record["output_json"])
        full_text = prompt + target + tokenizer.eos_token
        prompt_ids = tokenizer(prompt, add_special_tokens=False)["input_ids"]
        full_ids = tokenizer(full_text, add_special_tokens=False, truncation=True, max_length=args.seq_len)["input_ids"]
        labels = list(full_ids)
        prompt_len = min(len(prompt_ids), len(full_ids))
        for i in range(prompt_len):
            labels[i] = -100
        loss_weights = [0.0] * len(full_ids)
        target_pos = 0
        for i in range(len(full_ids)):
            if labels[i] == -100:
                continue
            loss_weights[i] = 15.0 if target_pos == 0 else (5.0 if target_pos < 3 else 1.0)
            target_pos += 1
        pad_len = args.seq_len - len(full_ids)
        attention_mask = [1] * len(full_ids) + [0] * pad_len
        input_ids = full_ids + [tokenizer.pad_token_id] * pad_len
        labels = labels + [-100] * pad_len
        loss_weights = loss_weights + [0.0] * pad_len
        return input_ids, attention_mask, labels, loss_weights

    encoded = [encode(r) for r in sample]
    input_ids = torch.tensor([e[0] for e in encoded], dtype=torch.long).to(dml)
    attention_mask = torch.tensor([e[1] for e in encoded], dtype=torch.long).to(dml).bool()
    labels = torch.tensor([e[2] for e in encoded], dtype=torch.long).to(dml)
    loss_weights = torch.tensor([e[3] for e in encoded], dtype=torch.float32).to(dml)

    out = model(input_ids=input_ids, attention_mask=attention_mask, labels=None)
    logits = out.logits
    shift_logits = logits[:, :-1, :].contiguous()
    shift_labels = labels[:, 1:].contiguous()
    shift_weights = loss_weights[:, 1:].contiguous()

    # --- TEST A: F.cross_entropy(reduction='none') path (BUGGY: zero grad) ---
    per_token_loss_a = torch.nn.functional.cross_entropy(
        shift_logits.view(-1, shift_logits.size(-1)).float(),
        shift_labels.view(-1), ignore_index=-100, reduction="none",
    )
    weighted_loss_a = (per_token_loss_a * shift_weights.view(-1)).sum() / shift_weights.view(-1).sum().clamp(min=1e-8)
    model.zero_grad()
    weighted_loss_a.backward(retain_graph=False)
    grad_a = lora_b_param.grad.norm().item() if lora_b_param.grad is not None else None
    print(f"TEST A (F.cross_entropy reduction='none'): loss={weighted_loss_a.item():.4f} lora_B.grad_norm={grad_a}")

    # --- TEST B: manual log_softmax + gather NLL (avoids fused reduction='none' kernel) ---
    model.zero_grad()
    out2 = model(input_ids=input_ids, attention_mask=attention_mask, labels=None)
    logits2 = out2.logits
    shift_logits2 = logits2[:, :-1, :].contiguous()
    vocab_size = shift_logits2.size(-1)
    log_probs = torch.nn.functional.log_softmax(shift_logits2.view(-1, vocab_size).float(), dim=-1)
    safe_labels = shift_labels.view(-1).clone()
    ignore_mask = (safe_labels == -100)
    safe_labels[ignore_mask] = 0
    gathered = -log_probs.gather(1, safe_labels.unsqueeze(1)).squeeze(1)
    gathered = torch.where(ignore_mask, torch.zeros_like(gathered), gathered)
    weighted_loss_b = (gathered * shift_weights.view(-1)).sum() / shift_weights.view(-1).sum().clamp(min=1e-8)
    weighted_loss_b.backward()
    grad_b = lora_b_param.grad.norm().item() if lora_b_param.grad is not None else None
    print(f"TEST B (manual log_softmax+gather): loss={weighted_loss_b.item():.4f} lora_B.grad_norm={grad_b}")


if __name__ == "__main__":
    main()

Actual output on our hardware/build:

TEST A (F.cross_entropy reduction='none'): loss=9.8083 lora_B.grad_norm=0.0
TEST B (manual log_softmax+gather): loss=9.8083 lora_B.grad_norm=0.05341910570859909

Expected behaviour

F.cross_entropy(reduction='none') backward should produce the same (nonzero) gradient as the mathematically equivalent manual log_softmax+gather computation, given the loss values themselves already agree to within floating-point precision.

Workaround (used in our project)

Replace F.cross_entropy(..., reduction='none') in any per-token/weighted-loss path with a manual log_softmax(...).gather(...) NLL computation, and add a hard pre-flight check (one real batch forward/backward before the full training loop, asserting a trainable parameter's .grad.norm() > 1e-9, hard-aborting the run otherwise) to catch this class of silent-zero-gradient failure mechanically rather than relying on post-hoc inspection.

Additional notes

• This is unrelated to masked_fill uint8-overflow (RuntimeError: value cannot be converted to type uint8_t without overflow with masked_fill in GPT-Neo causal mask creation on DirectML. torch.where works. #702); our repro includes the standard masked_fill workaround for that separate issue so it does not interfere with reproducing this one.
• Searched existing open/closed issues for cross_entropy, reduction, NllLoss, zero gradient, silent, lora gradient, grad is None against this repo before filing; found no existing report of this exact symptom.
• Discovered during real LoRA adapter training work on a small (135M) model; happy to share additional diagnostic output (LoRA weight-tensor inspection before/after training) if useful, with no proprietary data involved (synthetic/internal structured-extraction task only).