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

The title literally means: "I have parallel datas, 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:

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)

────────────────────────────────────────

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": [
                      "-DPy_LIMITED_API=0x03090000",
                      "-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.

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})

    bx = blockIdx().x
    by = blockIdx().y

    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)!! +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 data 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 Datas 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))

    for start_kv in range(lo, hi, BLOCK_SIZE_KV):
        start_kv = tl.multiple_of(start_kv, BLOCK_SIZE_KV)

        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, :])
[...]
    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),
    )
[...]
        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) 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 datas 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). Torch does not even have the kernel in the same code.

All it took were 60 lines

(Backprop youtube shorts)

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
        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.

Older Neve results surpassed PyTorch ( (old NSK paper) ). It won't take too long until the kernels get corrected and optimized.

More about (Neve Syntax)