add autocast for conv1/2/3d #797
Conversation
| res = conv_function( | ||
| a, |
There was a problem hiding this comment.
nit: looks like conv_function == _convolution_autocast_impl accepts kwargs here, but _convolution_autocast_impl misses kwargs in its definition.
There was a problem hiding this comment.
Yes, but this is to be able to route dtype through without burdening the conv helper to know about dtypes.
thunder/tests/test_autocast.py
Outdated
| with torch.autocast("cpu", torch.bfloat16): | ||
| eager_out = foo(x, w) | ||
| jit_out = jfoo(x, w) |
There was a problem hiding this comment.
I wonder, will the autograd be affected? Shall we test is as well?
There was a problem hiding this comment.
you mean the backward?
There was a problem hiding this comment.
Added it with terrible bounds...
There was a problem hiding this comment.
Hmmm, maybe rand vs randn could do better?
There was a problem hiding this comment.
Well, so I would suspect that there still is some dragon lingering somewhere, because I cannot really think of much reason to have thunder + autocast return wildly different results than torch + autocast even for wild inputs, but TBH, I would postpone investigating that.
There was a problem hiding this comment.
but you are right, switching to rand buys me an order of magnitude, cool, thank you!
There was a problem hiding this comment.
Turns out that Thunder does funny, slow things for the conv backward.
There was a problem hiding this comment.
#799 <-- filed an issue
Fixes: #796