Lightning-AI / Lightning-AI/lightning-thunder

Difference in backward output compared to eager

Open
#969 1 comment 1 reaction 0 assignees View on GitHub
numerical accuracy
Dominant language
Python
Stars
1.5k
Forks
121
PR merge metrics
No merged PRs in 30d

Description

```python
import torch
import thunder
from thunder.benchmarks.targets import LitGPTConfig, LitGPTBenchmark

name = "thunder"

def get_model_and_args(name):
with torch.device("cuda"):
cfg: LitGPTConfig = LitGPTConfig.from_name("Llama-2-7b-hf")
cfg.n_layer = 2
cfg.block_size = 512
b = LitGPTBenchmark(cfg, batchdims=(2,), dtype=thunder.float64)
model = b.fn()
args, kwargs = b.make_batch()

if name == "thunder":
jmodel = thunder.jit(model)
elif name == "tcompile":
jmodel = torch.compile(model)
elif name == "eager":
jmodel = model

return jmodel, model, args, kwargs, cfg

jmodel, model, args, kwargs, cfg = get_model_and_args(name)

a = jmodel(*args, **kwargs)

g = torch.rand_like(a)
actual_grads = torch.autograd.grad(a, model.parameters(), g)

# Sanity Check values vs Eager
e = model(*args, **kwargs)
expected_grads = torch.autograd.grad(e, model.parameters(), g)
try:
torch.testing.assert_close(a, e)
except Exception as exception:
print("Difference in forward")
print(exception)

try:
torch.testing.assert_close(actual_grads, expected_grads)
except Exception as exception:
print("Difference in backward")
print(exception)

```

Output
```
Difference in forward
Tensor-likes are not close!

Mismatched elements: 701 / 32768000 (0.0%)
Greatest absolute difference: 3.4200238951953565e-07 at index (0, 196, 27841) (up to 1e-07 allowed)
Greatest relative difference: 0.0003868698751363598 at index (0, 232, 16553) (up to 1e-07 allowed)
Difference in backward
Tensor-likes are not close!

Mismatched elements: 20801070 / 131072000 (15.9%)
Greatest absolute difference: 5.836758326438485e-06 at index (28794, 902) (up to 1e-07 allowed)
Greatest relative difference: 2.512434422647769 at index (22700, 3441) (up to 1e-07 allowed)

The failure occurred for item [0]
```

Note that in backward - `Greatest relative difference: 2.512434422647769` and this is at `float64`.

This also happens with PyTorch native `torch.compile(model)` - so it seems that it happens in general with compilation. Filing this issue to document this.

cc: @IvanYashchuk

Contributor guide

No contributing guide indexed for this repository

Research direction

Start with the provided LitGPTBenchmark reproduction using thunder.jit and torch.compile, and compare both forward and autograd.grad results against eager at float64. No source file or test is named; done means isolating the compilation-related backward discrepancy and documenting or correcting the behavior.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, pytorch
Domain
compilers, machine-learning
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Stale
Clarity
Needs clarification
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.