- I don't think it hurts to learn PyTorch (and having learned JAX is good, too). I don't know if JAX + triton is as impossible as you make it out, but it seems that PyTorch integration is quite good for many things.
- For pallas, triton and CUDA/C++, you probably want to know a bit about how GPU works. There is the GPU-Mode discord / lectures / ressources if you are looking for material https://github.com/gpu-mode/ .
- In my experience how well Triton works varies depending on what you want to do (depending on the how well the programming model fits the task). If it does, it is quite nice to get something reasonably fast reasonably fast. PyTorch (in the inductor torch.compile backend) has made many things work well, so you could check that out if you run out of examples elsewhere).
Note that the NVIDIA container uses CUDA+cuBLAS 13.0.2 which cites "Improved performance on NVIDIA DGX Spark for FP16/BF16 and FP8 GEMMs", which seems to be your use-case.
In general, I would suspect that it mostly comes to versions of the libs.
Interestingly, there is a cuBLAS 13.1 whl on PyPI, not sure what that does.
It seems to me that in 2016 people did (have to) play a lot more tricks with the backpropagation than today. Back then it was common to meddle with gradients in between the gradient propagation.
For example, Alex Graves's (great! with attention) 2013 paper "Sequence Generation with Recurrent Neural Networks" has this line:
One difficulty when training LSTM with the full gradient is that the derivatives sometimes become excessively large, leading to numerical problems. To prevent this, all the experiments in this paper clipped the derivative of the loss with respect to the network inputs to the
LSTM layers (before the sigmoid and tanh functions are applied) to lie within
a predefined range.
with this footnote:
In fact this technique was used in all my previous papers on LSTM, and in my publicly available LSTM code, but I forgot to mention it anywhere—mea culpa.
That said, backpropagation seems important enough to me that I once did a specialized videocourse just about PyTorch (1.x) autograd.
The adapter and LoRa have a drastically fewer parameters, so one might expect that forward + backward is roughly 2x the cost of forward.
Then (as far as I know), in contrast to generation, training is done on the entire output of the transformer (so all tokens of the full input) rather than serially token-by-token (in the RNN days, this was called teacher-forcing), so that may give you a significant boost in the tokens per second rate over generation.
https://landscape.pytorch.org/