3 ms·
This is a really cool effort! I’m curious, where do the bottlenecks live right now?
by brrrrrm 4y ago
This is a really cool effort! I’m curious, where do the bottlenecks live right now?
- tysam_and 4y agoGreetings, Bram! I hope you are well (we don't know each other, I just felt compelled to say that XD :D) I browsed your Twitter, this might be the shortlist that you find the most interesting: > fp16 -> fp8 >>>>> loss(dtype=fp16).sum() -> loss(dtype=fp8).mean() # the summing is for regularization's sake > nn.MaxPool2d's backward is hellaciously slow due to what appears to be indexing. right now we're under 10 seconds without _any_ custom kernels or JITting (!!!, as far as I know!!!). this may need to change, though, because I think max_pooling is now taking longer than some of the convolutions themselves. it's also annoyingly necessary for performance, very much so > BatchNorm2d(bias=True, weight=False) + Conv2d(bias=False) is paradoxically faster than BatchNorm2d(affine=False) + Conv2d(bias=True). This shouldn't be the case, as having a fused kernel like that should be faster. > Information flow engineering. Knowing which layers are more information-stressed during training could help with micro-adjustments of layer depths to an uncomfortable level, but right now we're just sort of guessing in the dark. Being able to visualize and track this would be nice. > Waiting for Pytorch 2.0! That will open up a lot of neat tricks that would otherwise slow down training with tons of small kernel launches. However, that likely will mean we end up losing our JIT card. I'm hoping that's not the Faustian bargain we have to make. And a few others, however, those are under wraps as I think we'd like to publish those if/when they work out and once they're completed. We always have other improvements in the pipeline even as we're finalizing and publishing a release. :D Hope that helps, the library you're working on looks super cool! Thanks for sharing the pictures/notes of that on your Twitter, I truly enjoyed reading through that.