Lightning-AI / Lightning-AI/lightning-thunder
conv backward in thunder
- Dominant language
- Python
- Stars
- 1.5k
- Forks
- 121
- PR merge metrics
- No merged PRs in 30d
Description
while testing #797 , it seems that Thunder's backward might lead to things not being optimal (accurary, speed?):
```
def foo(x, w, b=None):
return torch.nn.functional.conv2d(x, w, b)
x = torch.randn(1, 2, 8, 8, requires_grad=True)
w = torch.randn(3, 2, 4, 4, requires_grad=True)
b = torch.randn(3, requires_grad=True)
go = torch.randn(1, 3, 5, 5)
jfoo = thunder.jit(foo)
x64 = x.to(torch.float64)
w64 = w.to(torch.float64)
b64 = b.to(torch.float64)
ref_eager_out = foo(x64, w64, b64)
ref_eager_grads = torch.autograd.grad(ref_eager_out, [x64, w64, b64], go.to(torch.float64))
with torch.autocast("cpu", torch.float16):
print("eager")
with torch.profiler.profile() as prof:
eager_out = foo(x, w, b)
eager_grads = torch.autograd.grad(eager_out, [x, w, b], go)
print(prof.key_averages().table())
print("thunder")
with torch.profiler.profile() as prof:
jit_out = jfoo(x, w, b)
jit_grads = torch.autograd.grad(jit_out, [x, w, b], go)
print(prof.key_averages().table())
torch.testing.assert_close(eager_out, jit_out)
for eg, jg, rg in zip(eager_grads, jit_grads, ref_eager_grads):
# TODO: tighten check?
print(f"ref - eager {(eg - rg).abs().max().item():.4f} ref - thunder {(jg - rg).abs().max().item():.4f}")
torch.testing.assert_close(eg, jg, atol=1e-2, rtol=1e-2)
```
gives (note the backward running conv2d forward twice again(?))
```
eager
------------------------------------------------------- ------------ ------------ ------------ ------------ ------------ ------------
Name Self CPU % Self CPU CPU total % CPU total CPU time avg # of Calls
------------------------------------------------------- ------------ ------------ ------------ ------------ ------------ ------------
aten::conv2d 4.64% 16.030us 99.37% 343.382us 171.691us 2
aten::to 2.59% 8.950us 20.87% 72.130us 10.304us 7
aten::_to_copy 10.49% 36.240us 18.28% 63.180us 9.026us 7
aten::empty_strided 4.68% 16.190us 4.68% 16.190us 2.313us 7
aten::copy_ 6.26% 21.620us 6.26% 21.620us 2.702us 8
aten::convolution 4.14% 14.290us 40.00% 138.241us 138.241us 1
aten::_convolution 3.99% 13.790us 35.87% 123.951us 123.951us 1
aten::_nnpack_available 0.25% 0.850us 0.25% 0.850us 0.425us 2
aten::thnn_conv2d 0.74% 2.570us 31.70% 109.561us 109.561us 1
aten::_slow_conv2d_forward 24.09% 83.260us 30.96% 106.991us 106.991us 1
aten::empty 2.14% 7.390us 2.14% 7.390us 1.232us 6
aten::view 2.17% 7.491us 2.17% 7.491us 1.873us 4
aten::resize_ 1.17% 4.060us 1.17% 4.060us 1.353us 3
aten::reshape 0.41% 1.400us 0.62% 2.150us 2.150us 1
autograd::engine::evaluate_function: ConvolutionBack... 1.84% 6.350us 26.54% 91.701us 91.701us 1
ConvolutionBackward0 1.86% 6.430us 24.70% 85.351us 85.351us 1
aten::convolution_backward 3.16% 10.911us 22.84% 78.921us 78.921us 1
aten::_slow_conv2d_backward 10.20% 35.240us 19.61% 67.760us 67.760us 1
aten::resize_as_ 0.46% 1.600us 0.86% 2.980us 2.980us 1
aten::zero_ 0.40% 1.390us 0.40% 1.390us 0.695us 2
aten::sum 5.02% 17.360us 6.38% 22.050us 22.050us 1
aten::as_strided 0.52% 1.790us 0.52% 1.790us 1.790us 1
aten::fill_ 0.84% 2.900us 0.84% 2.900us 2.900us 1
autograd::engine::evaluate_function: ToCopyBackward0... 1.85% 6.410us 7.86% 27.150us 9.050us 3
ToCopyBackward0 1.22% 4.230us 6.00% 20.740us 6.913us 3
autograd::engine::evaluate_function: torch::autograd... 0.54% 1.870us 0.54% 1.870us 0.623us 3
cudaDeviceSynchronize 4.33% 14.960us 4.33% 14.960us 14.960us 1
------------------------------------------------------- ------------ ------------ ------------ ------------ ------------ ------------
Self CPU time total: 345.572us
thunder
------------------------------------------------------- ------------ ------------ ------------ ------------ ------------ ------------
Name Self CPU % Self CPU CPU total % CPU total CPU time avg # of Calls
------------------------------------------------------- ------------ ------------ ------------ ------------ ------------ ------------
aten::to 0.67% 4.030us 4.20% 25.310us 3.616us 7
aten::_to_copy 1.67% 10.060us 3.53% 21.280us 3.040us 7
aten::empty_strided 1.63% 9.790us 1.63% 9.790us 0.890us 11
aten::copy_ 2.20% 13.260us 2.20% 13.260us 1.105us 12
aten::convolution 0.73% 4.390us 65.19% 392.634us 130.878us 3
aten::_convolution 1.44% 8.660us 64.46% 388.244us 129.415us 3
aten::_nnpack_available 0.04% 0.260us 0.04% 0.260us 0.087us 3
aten::thnn_conv2d 0.29% 1.730us 62.19% 374.564us 124.855us 3
aten::_slow_conv2d_forward 59.96% 361.134us 61.90% 372.834us 124.278us 3
aten::empty 0.69% 4.150us 0.69% 4.150us 0.593us 7
aten::view 0.84% 5.050us 0.84% 5.050us 0.842us 6
aten::resize_ 0.34% 2.020us 0.34% 2.020us 0.673us 3
aten::reshape 1.22% 7.320us 2.41% 14.530us 2.906us 5
ThunderFunction 2.41% 14.510us 2.41% 14.510us 14.510us 1
autograd::engine::evaluate_function: ThunderFunction... 0.64% 3.870us 34.79% 209.572us 209.572us 1
ThunderFunctionBackward 17.19% 103.561us 34.15% 205.702us 205.702us 1
aten::permute 1.46% 8.770us 1.99% 11.980us 1.997us 6
aten::as_strided 0.64% 3.880us 0.64% 3.880us 0.485us 8
aten::sum 1.31% 7.920us 1.57% 9.480us 9.480us 1
aten::fill_ 0.21% 1.240us 0.21% 1.240us 1.240us 1
aten::pad 0.59% 3.570us 2.54% 15.310us 5.103us 3
aten::constant_pad_nd 0.54% 3.230us 1.95% 11.740us 3.913us 3
aten::clone 0.62% 3.740us 2.08% 12.540us 3.135us 4
aten::_reshape_alias 0.63% 3.790us 0.63% 3.790us 1.895us 2
aten::flip 0.97% 5.830us 1.31% 7.890us 7.890us 1
aten::empty_like 0.26% 1.550us 0.50% 3.020us 1.510us 2
aten::contiguous 0.12% 0.730us 0.79% 4.760us 4.760us 1
autograd::engine::evaluate_function: torch::autograd... 0.16% 0.960us 0.16% 0.960us 0.320us 3
cudaDeviceSynchronize 0.55% 3.330us 0.55% 3.330us 3.330us 1
------------------------------------------------------- ------------ ------------ ------------ ------------ ------------ ------------
Self CPU time total: 602.335us
```
Edit: I don't have the output for the accuracy that I edited into the script here, but my impression is that the accuracy of the Thunder backward is not worse than the eager one in this example and this example is not terribly relevant for perf. It's just that we should develop insight into what's going on because we will bump into the question.
I wonder if #655 is related, maybe the same method could provide a 0th-order analysis of what's going on.
cc @tfogal
Contributor guide
No contributing guide indexed for this repository
Research direction
Start by running the convolution example in the issue with profiling enabled, then compare eager and Thunder call counts, timings, and gradient accuracy. Read issue #655 for the suggested 0th-order analysis approach; done means establishing whether the extra convolution calls or accuracy difference is a real Thunder problem and documenting evidence for the next fix.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, pytorch
- Domain
- compilers, performance
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Stale
- Clarity
- Needs clarification
- Newbie friendliness
- 35/100