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

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) 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: Close to Python/SentencePiece in text processing + Byte-Pair Encoding (BPE) training; Competitive with NumPy and OpenBLAS in CPU matrix multiplicaton, but with pure high-level SIMD code. Check exsisting (implementations) Currently working in a better GPU programming interface, towards the implementation of flash-attention. Neve documentation (neve-lang.dev) Neve repo (github.com/NoSavedDATA/Neve). Youtube for updates (youtube.com/@nosaveddata3994). Discord for extensive talks/suggestions (discord.gg/hP5feM7cV) ──────────────────────────────────────── Introduction 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) 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) 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 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) Reinforcement Learning for Atary. A bit before that, I took a glance code of the (Efficient Zero) 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. ──────────────────────────────────────── Thoughts in Other Languages Julia 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 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)!! +1. It is only the BPE inference, no traning -1 point. It has (flash-attention gpu kernels) +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 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. @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) Neve Let's see how Neve GPU matrix multiplication looks like. gpu void @( layout x, layout y, layout z ) int warp_rows = min((m+63)//64, 4) int wx = warp%warp_rows, wy = warp//warp_rows var smem_a = layout() var smem_b = layout() 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 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() var b = layout() [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_rowself.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). Torch does not even have the kernel in the same code. ──────────────────────────────────────── Backpropagation All it took were 60 lines (Backprop youtube shorts) ──────────────────────────────────────── Parallelism # count primes in 10 threads def int is_prime(int n) for i=2, i