Lightning-AI / Lightning-AI/lightning-thunder

Cannot correctly compute cross entropy loss (backward) when the reduction used is "sum" and torch executor is used.

Open
#2,056 2 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

Dominant language
Python
Stars
1.5k
Forks
121
PR merge metrics
No merged PRs in 30d

Description

*Note*: If you have a model or program that is not supported yet but should be, please use the program coverage template.

## 🐛 Bug

With the torch executor, Thunder cannot correctly compute cross entropy loss (backward) when the reduction used is "sum". It works correctly with "mean"

### To Reproduce

Run the code snippet below.

#### Code sample

```
device='cuda'
executor=TorchExecutor

def cross_entropy_fn(logits, labels):
return torch.nn.functional.cross_entropy(logits, labels, reduction="sum")

sequence_length, vocab_size = 8192, 32064
logits = torch.rand((sequence_length, vocab_size), device=device, dtype=thunder_dtype, requires_grad=True)
labels = torch.randint(0, sequence_length, (sequence_length,), requires_grad=False, device=device)

fn = executor.make_callable(cross_entropy_fn)
a = fn(logits, labels)
a.backward()
print(logits.grad)
```

### Expected behavior

Should run correctly to completion with output matching that of torch compile.

### Environment

- PyTorch Version (e.g., 1.0):
- OS (e.g., Linux):
- How you installed PyTorch (`conda`, `pip`, source):
- Build command you used (if compiling from source):
- Python version:
- CUDA/cuDNN version:
- GPU models and configuration:
- Any other relevant information:

### Additional context

The error seen is:
https://gist.github.com/protonu/ff88fbc9508947f45e39a6190acc6508

Contributor guide

No contributing guide indexed for this repository

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start with the cross_entropy_fn reproducer in this issue and inspect the linked error gist, comparing the TorchExecutor behavior for reduction="sum" with reduction="mean" and torch.compile. The work is done when the backward pass completes and its output matches torch.compile for the provided case.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, pytorch
Domain
machine-learning
Issue type
Bug
Difficulty
3/5
Estimated time
1-2 days
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
45/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.