Deep Learning from Scratch in 1400 Lines - Neve & Frost Framework

A developer has released Frost, a deep learning framework written in roughly 1,400 lines of code, built on top of the Neve programming language. The framework includes parallel dataloaders and GPU kernels, and ships with a ResNet-18 benchmark that users can run themselves. The developer says the project targets three pain points in existing tooling: optimizer complexity in PyTorch, difficult C++/CUDA interoperability, and limited parallelism.

The title literally means: "I have parallel dataloaders, GPU kernels and a high-performance computing programming language all expressed in a framework with 1400 lines of code". You may run the ResNet-18 benchmark yourself github.com/NoSavedDATA/Neve benchmarks https://github.com/NoSavedDATA/Neve benchmarks I released the Neve programming language a while ago. Now, this is the release of the Frost deep learning framework, alongside with the first benchmark. Other results for Neve: Currently working in a better GPU programming interface, towards the implementation of flash-attention. Neve documentation neve-lang.dev https://neve-lang.dev Neve repo github.com/NoSavedDATA/Neve https://github.com/NoSavedDATA/Neve . Youtube for updates youtube.com/@nosaveddata3994 https://www.youtube.com/@nosaveddata3994 . Discord for extensive talks/suggestions discord.gg/hP5feM7cV https://discord.gg/hP5feM7cV ──────────────────────────────────────── Once day I was reading some papers, and a very interesting paper was published. It was the sophia optimizer. I took a glance in an unnoficial code https://github.com/kyegomez/Sophia/blob/main/Sophia/main.py for it, and I questioned myself why did it have to be so difficult to add new optimizers in PyTorch. I experimented the optimizer, and the results were quite bad with a lot of NaNs. Turns out another paper published later claimed this and other optimizers had overstated claims. Imagine wasting hours studying a 10 pages of a paper, then hardly trying to debug it and asses whether other person discoveries are true. All that code reading complexity makes this a challengeful task. Problem 1: even optmizers are hard to understand in PyTorch. Few weeks later, flash attention https://arxiv.org/pdf/2205.14135 was released, and the algorithm actually achieved a speed-up of 2x. The problem, it was C++ CUDA. Most high-level GPU kernel frameworks were imature to the point the flash attention author chose not to use them. Now take a look what is necessary for adding C++ code in PyTorch python from setuptools import setup, Extension from torch.utils import cpp extension setup name="extension cpp", ext modules= cpp extension.CppExtension "extension cpp", "muladd.cpp" , extra compile args={ "cxx": define Py LIMITED API with min version 3.9 to expose only the stable limited API subset from Python.h "-DPy LIMITED API=0x03090000", define TORCH TARGET VERSION with min version 2.10 to expose only the stable API subset from torch "-DTORCH TARGET VERSION=0x020a000000000000", }, py limited api=True , Build 1 wheel across multiple Python versions cmdclass={'build ext': cpp extension.BuildExtension}, options={"bdist wheel": {"py limited api": "cp39"}} 3.9 is minimum supported Python version That comprehends problem 2: lack of high-level CUDA code and hard interoperability. For my Bachelor's thesis, I implemented the BBF https://github.com/NoSavedDATA/PyTorch-BBF-Bigger-Better-Faster-Atari-100k Reinforcement Learning for Atary. A bit before that, I took a glance code of the Efficient Zero https://github.com/YeWR/EfficientZero/tree/main/core reinforcement learning model. It has a parallelism that PyTorch does not handle, and the implementation required using Cython packages for having threads literaly coding in C, then just calling C functions from Python . Later, I realized PyTorch also needed to implement its data worker threads in C, another workaround over Python Global Interpreter Lock GIL . Not only that, even preprocessing implementations like the BPE are made in C, C++, Rust, etc... That leads us problem 3, lack of parallelism. That is when I decided to create a programming language, a few months before finishing my Bachelor's, which matured to my Master's project Summing up, currently, people must choose between languages like Python for high-level productivity, C and relatives for compute efficiency, Lua/Julia for advanced interoperability and other languages for concurrency. Thus, since in my job I had to wait hours for my neural networks to train, I decided to create a programming language in the time in between trainings. One language that had all these features, which are of high value for deep learning research. Nowadays, I believe it matured to such a point that it may be extended to other complex problem domains. Julia makes dynamic typying speed reach close to C++ speeds. It also has a mark sweep and channels for parallelism. The idea is very interesting. Let's take a look a in its cuda kernels. function mma kernel Z::CuDeviceMatrix{Float32}, X::CuDeviceMatrix{BFloat16}, Y::CuDeviceMatrix{BFloat16} Grid and block indices bx = blockIdx .x by = blockIdx .y Thread and warp indices tid = threadIdx .x lane = tid - 1 % 32 warp = tid - 1 ÷ 32 STOP Why am I seeing blockIdx .x in my code? Was this supposed to be a high-level scientific language or CUDA in C++28? Besides, it does not expose intrisics like the cp async, which is crucial for high-speed matrix multiplication. They must be explicitly added throgh interop intrisics. And it has the "end" keyword, which in my opinion incurs a lot of code pollution. Mojo has a Python interop, so it did not have to build all libs and frameworks from scratch +1 point. It has Byte-Pair Encoding benchmarks https://github.com/atsentia/mojo-tokenizer +1. It is only the BPE inference, no traning -1 point. It has flash-attention gpu kernels https://www.spheron.network/blog/modular-max-mojo-gpu-cloud-llm-inference/ +1 point. It runs MAX, which allows GPU code portability across different hardware, +2 points. It lacks channels, so I would hardly try to make a parallel dataloader in it. -2 points. It uses Rust ownership +0 points. Now let's look at Mojo kernels for the flash attention. @always inline def fused attention cpu BN: Int, BD: Int Q: LayoutTensor, K: LayoutTensor, V: LayoutTensor, O: LayoutTensor mut=True, ... , : comptime N = K.shape 0 comptime D = K.shape 1 comptime for tile n in range N // BN : var Q tile = Q.tile BN, D tile n, 0 comptime for tile d in range D // BD : var m 1 = LayoutTensor Q tile.dtype, Layout BN, 1 , MutAnyOrigin .stack allocation .fill Scalar Q tile.dtype .MIN var l 1 = LayoutTensor Q tile.dtype, Layout BN, 1 , MutAnyOrigin .stack allocation .fill 0 var O i = LayoutTensor Q tile.dtype, Layout.row major BN, BD , MutAnyOrigin .stack allocation .fill 0 comptime for tile n idx in range N // BN : var K tile = K.tile BN, D tile n idx, 0 var V tile = V.tile BN, BD tile n idx, tile d var S = matmul b transpose Q tile, K tile var m 2 = max m 1, rebind type of m 1 max axis=1 S Quite interesting. It has layouts and tiling, inspired by CuTe and Cutlass. It actually inspired the way Neve layouts and tiles work. Nevertheless, it still has a heavy syntax. Note the keyword comptime appears frequently a sort of metaprogramming . This adds some cognitive overhead. And the layouts are quite verbose. The layout fill and stack allocation can be simplified. Triton has layouts/tiling similar to Mojo, but is dynamically typed and has no comptime headaches. The problem is that Triton does not make Python Dataloaders easier to implement from the systems programming language perspective. We actually need a complete new programming language for this. python @triton.jit def attn fwd inner ... K block ptr = tl.advance K block ptr, 0, lo V block ptr = tl.advance V block ptr, lo, 0 loop over k, v and update accumulator for start kv in range lo, hi, BLOCK SIZE KV : Just let the compiler know that start n is a multiple of BLOCK N, so the compiler can do optimizations start kv = tl.multiple of start kv, BLOCK SIZE KV -- compute qk ---- K block = tl.load K block ptr QK block = tl.dot Q block, K block if STAGE == 2: mask = offs q :, None = start kv + offs kv None, : ... A LAYOUT Q block ptr = tl.make block ptr base=Q + qvk offset, shape= SEQ LEN, HEAD DIM , strides= stride Q seq, stride Q dim , offsets= block index q BLOCK SIZE Q, 0 , block shape= BLOCK SIZE Q, HEAD DIM , order= 1, 0 , ... Algebra -- compute qk ---- K block = tl.load K block ptr QK block = tl.dot Q block, K block Let's see how Neve GPU matrix multiplication looks like. gpu void @ layout<bf16,m,n x, layout<bf16,p,n y, layout<float,m,p,smem z int warp rows = min m+63 //64, 4 int wx = warp%warp rows, wy = warp//warp rows var smem a = layout<bf16,256,2,16,smem var smem b = layout<bf16,128,2,16,smem int warp m = min m//16 16, 64 , warp p = min p//16 16, 64 int m tiles = warp m//16 int p tiles = warp p//8 var c = layout<int,4,8,4 int m cp tiles = min m//16,4 0..m cp tiles i cp async16 smem a{256,8} wx m cp tiles+i, lane , x{256,8} wx m cp tiles+i, lane int p cp tiles = min p//16,4 0..p cp tiles i cp async16 smem b{256,8} wy p cp tiles+i,lane , y{256,8} wy p cp tiles+i,lane cp commit group cp wait group 0 syncthreads var a = layout<int,4,4 var b = layout<int,8,2 0..m tiles i int row = wx m tiles+i 16+lane%16 int col = lane//16 8 ldmatrix x4 a{i,0}, smem a{row, col} 0..p tiles i int row = wy p tiles+i 8+lane%8 int col = lane//8 %2 8 ldmatrix x2 b{i,0}, smem b{row, col} 0..m tiles, 0..p tiles i, j mma 16x8x16 c{i,j,0}, a{i,0}, b{j,0} syncthreads 0..m tiles, 0..p tiles, 0..4 i, j, k int global row = bx 256+ wx m tiles + i 16 + lane // 4 + k // 2 8 int global col = by 128+ wy p tiles + j 8 + lane % 4 2 + k % 2 if global row<m and global col<p z global row p + global col = c i, j, k syncthreads kernel void mma kernel bf16 x, bf16 y, float z, int M, int N, int P var v = layout<bf16, M, N x var u = layout<bf16, P, N y z += v 256,N bx,0 @ u 128,N by,0 Did you like it? It took me one month to replicate the HGEMM https://bruce-lee-ly.medium.com/nvidia-tensor-core-cuda-hgemm-advanced-optimization-5a17eb77dd85 repo to the beta version of frost, which used C++ CUDA kernels instead of Neve. Currently, I still need to do the benchmarks, but I will wait until I finish flash attention. Here you see CUDA instructions like mma 16x8x16, cp async and syncronizations. These are all mandatory when we pretend to implement SOTA matrix multiplications. Now look at the second function, it is almost z += v @ u, but with tiling. The whole previous function is successfully reduced into a new gpu operator. There is also the new multi-loop expression 0..m tiles, 0..p tiles i, j When it comes to writing, Neve syntax is more pythonic than Julia yes to "class" and no to "struct", floating methods and the "end" word and Mojo no comptime headache keywords, no ownership transfer expressions . It is not dynamically typed as Triton, but allows to define dataloaders from scratch. And strong typying is seem as a benefit in systems programming. Neve has high-level kernel algebra, while simultaneously allowing CUDA intrinsics in different hierarchies of complexity. Now let's look at the budget for each programming language. Let's talk about some other features. The optmizer, the backpropagation and parallelism. Optimizers declaration can hardly be more compressed than this Frost implementation. kernel void AdamW k float param, float d, float m, float v, float lr, int n, float beta1, float beta2, float beta1 correction, float beta2 correction, float eps, float wd int i = bx 256 + tx if i =n return float m i = 1-beta1 d i + beta1 m i float v i = 1-beta2 d i d i + beta2 v i m i = m i v i = v i m i/=beta1 correction v i/=beta2 correction param i -= lr m i / sqrt v i + eps + wd param i class AdamW float lr, beta1, beta2, wd, eps array<gpu tensor m, v int ts ctor float lr, float beta1, float beta2 self.lr = lr self.beta1 = beta1 self.beta2 = beta2 self.ts = 1 self.wd = 0.01 self.eps = 0.0000001 def float step int i=0 for param in $optim info.params if i+1 self.v.size self.m.append new gpu tensor param.dims, "zeros" self.v.append new gpu tensor param.dims, "zeros" float beta1 correction = 1.0-pow self.beta1, self.ts float beta2 correction = 1.0-pow self.beta2, self.ts launch param.dims prod/|256 256 AdamW k param.ptr offby 0, param.d offby 0, self.m i .ptr offby 0, self.v i .ptr offby 0, self.lr, param.dims prod, self.beta1, self.beta2, beta1 correction, beta2 correction, self.eps, self.wd param.d=nil i=i+1 self.ts = self.ts+1 $optim info.params.clear tarena reset PyTorch comparison AdamW https://github.com/pytorch/pytorch/blob/main/torch/optim/adamw.py . Torch does not even have the kernel in the same code. All it took were 60 lines Backprop youtube shorts https://www.youtube.com/shorts/3qizuGUWIxY count primes in 10 threads def int is prime int n for i=2, i< n//2+1 if n%i==0 return 0 return 1 def int count prime vec array<int input numbers, channel<int,10 ch int num primes = 0 for i in input numbers num primes = num primes + is prime i print "Found num primes ", num primes, " on thread ", tid ch <- num primes main array<int v = arange int 2,250001 channel<int,10 ch finish split v across 10 threads asyncs 10 count prime vec v, ch var primes = ch.sum print "Total primes: ", primes Here we see two important parallel expressions. The first is the " " operator inside the "asyncs". It splits v across 10 segments and send each one to different threads in count prime vec. The second fathoms the channels. Here, it is used to store the results from the sums of primes for each thread. Then, the main thread uses .sum to aggregate all results. Neve channels are implemented with lock-free structures. They may block threaded operations if their capacity is full capacity 10 in the example . This adds room for a lot of possibilities. def float worker print "Start worker" int yield ptr, bs=self.batch size print "worker ", tid while self.load ch.alive yield ptr = self.increment yield ptr for b=0, b<bs self.getitem w yield ptr+b, b self.load ch <- tid self.x.switch self.y.switch def tuple<gpu tensor,gpu tensor batch int w <- self.load ch var x = self.x.load w var y = self.y.load w x = x.view $cfg.bs, 3, 32, 32 return x, y ... finish asyncs $cfg.num workers ds.worker 0..steps i optimize logic Here channels were used to signal when a worker finished processing its batch. Once the main thread starts processing the batch, it flips its pointer value inside a ping-pong memory buffer, so it may process new data in parallel with data consumption. 3 seeds in a RTX 4090. | Language | Acc | Time | Backend | |---|---|---|---| | Neve | 65.46%+-0.37 | 592s+-6s | Naive Kernel | | Neve | 65.46%+-0.45 | 261s+-9s | Partial cuDNN | | PyTorch | 74.82%+-1.27 | 64s+-6s | cuDNN | Older Neve results surpassed PyTorch old NSK paper https://arxiv.org/pdf/2409.11600 . It won't take too long until the kernels get corrected and optimized. More about Neve Syntax https://dev.to/no saved data/neve-towards-a-unified-programming-model-for-the-complete-deep-learning-stack-1g86