Lightning-AI / Lightning-AI/lightning-thunder
Cannot correctly compute cross entropy loss (backward) when the reduction used is "sum" and torch executor is used.
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
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- 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