[vforno]

Hi, in nanoeuler I use cuBLAS (NVIDIA's super-optimized library) for all matrix multiplications, with the tensor cores in TF32 mode. It's the same thing PyTorch uses underneath, so it's very fast. What I've optimized (and will improve even more) and written by hand are the kernels for the other parts (like FlashAttention, which gave a nice 3x speedup), while I've delegated the large matrices to cuBLAS. Training the 116M model on a 4070 runs well and in reasonable times. Compared to PyTorch, it's a bit slower (probably 1.5-2.5x), but nothing dramatic, especially considering it's all done from scratch without a framework and there are no other optimizations that would make it faster. I'm working on it.

[novaRom]

Do you have a guess why your code is so much slower than torch? I didn't look, but there must be no reason to have 2x slower code esp. for a simple grid of FMAs.

[vforno]

Yes, because it has many separate kernels instead of aggressive merges like PyTorch (with Torch Compile). Each pass (norm, matmul, residual, RoPE, etc.) launches its own kernel, which increases launch overhead and memory traffic. CuBLAS helps, but it's not enough to compensate.

[AndReics]

i see really cool, where i failed was trying to build my own matrix operations library, it was just too much, but using cuBLAS definitely helps, i'll look into the custom kernels you wrote they seem interesting!

did you build the backprop yourself? it is a really cool project to build and i think you can agree that it teaches you a lot of how LLMS and machine learning in general works.

[vforno]

Absolutely yes! With nanoeuler I learned so much by testing every little detail of the project. Every little part you see has been tested and proven several times so that it could be understood and worked.