Yeah I would suggest taking a look at PyTorch on AMD before saying stuff like "scaled_dot_product_attention is an NVIDIA CUDA kernel exposed as a PyTorch function", because that is demonstrably false.
Also, FWIW, I would suggest getting a small Llama 3.1 model training fast before trying to do a big 405B model -- faster to iterate and almost everything you'll learn on the small models will scale to the 405B.
Any direct comparisons to 8xH100? 2 toks/sec seems very slow!
I haven't done any LoRA training on MI300x myself, but I have done LLama 3.1 full training on 8xMI300x and got pretty close to 8xH100 performance with my own kernels (ROCm is just too slow).
Does JAX have its own implementations of matmul, flash attention etc? Or does it use the ROCm implementations like PyTorch does? (e.g,. hipblaslt, Composable Kernel FA etc)
Not too familiar with JAX, but the abysmal PyTorch training perf on MI300x is in large part attributable to the slow perf of the ROCm libraries it is using under the hood.
Well yeah I reckon you render a full custom 4004 w/ koru patterned transistors into about 4m x 4m stained glass panel. Would look good as the foyer panel for the CS dept at the University of Waikato
OK, so in the case of llm.c, if you're just including the HIP headers, using hipblasLt, etc, what would be the benefit of using scale instead of hipify?
I ported Karparthy's llm.c repo to AMD devices [1], and have trained GPT2 from scratch with 10B tokens of fineweb-edu on a 4x 7900XTX machine in just a few hours (about $2 worth of electricity) [2].
I've also trained the larger GPT2-XL model from scratch on bigger CDNA machines.
Hi, why do you believe that bfloat16 is not supported? Can you please provide some references (specifically the part about the hardware "doesn't do it")?
For the hardware you are focussing on (gfx11), the reference manual [2] and the list of LLVM gfx11 instructions supported [1] describe the bfloat16 vdot & WMMA operations, and these are in fact implemented and working in various software such as composable kernels and rocBLAS, which I have used (and can guarantee they are not simply being run as float). I've also used these in the AMD fork of llm.c [3]
Outside of gfx11, I have also used bfloat16 in CDNA2 & 3 devices, and they are working and being supported.
Regarding cublasLt, what is your plan for support there? Pass everything through to hipblasLt (hipify style) or something else?
I have not been impressed by the perf. Slower than PyTorch for LLMs, and PyTorch is actually stable on AMD (I've trained 7B/13B models).. so the stability issues seem to be more of a tinygrad problem and less of an AMD problem, despite George's ramblings [0][1]
Final loss from that fineweb-10B run (since then I'm up to ~100k toks/sec/GPU):
step 18865/18865 | train loss 3.280550 | norm 0.4362 | lr 0.00e+00 | 1669.06 ms | 55.4% A100 fp16 MFU | 314058 tok/s
Writing state to log124M/state_00018865_00003.bin
val loss 3.296179
You can buy these GPUs on Amazon for under $1k. I heard the MI300X may be available in Azure now or at least very soon.
I think the matmul issue is symptomatic of a much deeper issue.
It would be nice to see less whining and blaming AMD (PyTorch and llm.c actually work on 7900 XTX, and blow tiny grad out of the water in terms of perf!), and more just getting stuff to work.
Nah, I reproduced on 4x 7900 XTX machine in 8.75 hours, so a single 7900 XTX (costs less than $1k) could do it in under 24 hours. Was hitting 55.4% MFU.