BakaLLM: Designing new LLM
Tags: llm, machinelearning
Published: 2023-09-25 20:23:03

Hello, fairy dairy diary!

A while ago I ventured on a journey, endless perhaps for a mere mortal without H100 cluster, of making my own LLM, successor to the humanity, with blackjack and fairy maidens! The hopelessness of the endeavor shall not dissuade us from doing it, after all it's all about journey rather than destination.

I might as well write about about it, all under humming of my laptop while training lose of my model approaches perplexity of something you can download from the net! Is it time for chatgpt to panic? Perhaps not!

(no name is solidified yet, so for this blog I will use BakaLLM because ⑨ is the strongest

![2833165892-anime, touhou, 1girl, Cirno in battle pose ### ugly, bad, 2girls, 3girls.](https://dev-to-uploads.s3.amazonaws.com/uploads/articles/4wfv9ftk13i038vb672u.png)
)

(Image generated by NN by [stable horde](https://stablehorde.net/)) So.

The premise is simple: local machines lack memory. Wisdom of the streets says we need to download more RAM! So difficult in the real world, but what if we enhance the memory of the model itself during its design? What if we allow it to dare to reach out of its context window, into lands left so far behind, beyond, no llama can remember it ever happening? Many designs were aimed at this. Some were even published with weights if you know [where to look](https://github.com/facebookresearch/mega). And these designs do exactly that: they give more memory to the model.

Let's stew them together. Who is going to stop us? Sanity? AI police?

In this day of age?

Fairy warning, if you have any idea about ML, you probably need a barf bucket(or two) if you dare to venture in this ongoing story further. You were warned! Though today is simple day of looking on building blocks.

Oh, the style of this blog is intentionally not serious, so easily impressive minds wouldn't take it as a tutorial or educational resource: For teaching-proper requires more knowledge than reading couple of papers and watching youtube videos(with mandatory naps).  It's ramblings of a someone whose GPU is too busy to play video games for a while. But I leave arxiv links.

So let's see what kind of mix of architectures whose new current iteration with `172_062_349` parms looks like this...

![ChatGPT is free from panicking for now](https://dev-to-uploads.s3.amazonaws.com/uploads/articles/sxjmkhan2gkdoy5rdg33.png)

...has under its hood. 

(Spoiler: previous iteration eventually reached loss of ~4.86,  on this wikitext103 dataset(more on that on later date), this iteration will beat it)

The following models were used(misunderstood, misapplied) during this quest to world conquest:

* [Transformer-XL](https://arxiv.org/abs/1901.02860)

![Transformer-XL is pretty and simple](https://dev-to-uploads.s3.amazonaws.com/uploads/articles/wvv252l2l7mof3dq5rfx.png)

The idea of Transformer XL is simple: you reuse KV from the past in the attention phase to create new values. And these values also will be reused later. It's like doubling a context size, but with more whistles from the past.  And less of doubling. Perfect match. More layers = more past.

Implementing it was trickier than I thought:

By default pytorch's `F.scaled_dot_product_attention` aligns causal mask to the left(`torch.ones(L, S, dtype=torch.bool).tril(diagonal=0)`)
Bad torch!
It means if we have more queries than keys, the mask will be used like this:

```txt
1000 0000
1100 0000
1110 0000
1111 0000
```

Not very useful mask. Might-as-well-remove-half-of-the-keys kind of mask. We want "half-causal" kind of mask:

```txt
1111 1000
1111 1100
1111 1110
1111 1111
```

So much prettier! It allows us to speak to the past from any position in the current timeline! Unfortunately doing it is ugly for now.

It will change (https://github.com/pytorch/pytorch/issues/108108), but so far I had to resort to building mask manually, which tanks the performance. Pytorch just can't handle the greatness of this mask.

Alternatives were considered: flash_attention has a lot of quirks which makes it no go: no support for f32, has limits on head_dim(more on that later).  Issue in github states xformers also use right alignment, not sure it's true thought: at least `materialize` method return the same mask. I also considered doing two attention, then just adding them together(lol). It didn't work too well.

* Mega

![If you squint it's like XL](https://dev-to-uploads.s3.amazonaws.com/uploads/articles/gjcddjorymyal2a0zs9u.png)

![If you don't squint looks nothing like XL](https://dev-to-uploads.s3.amazonaws.com/uploads/articles/4fmmnsuorvr1006hls41.png)

From this architecture I borrowed its main mechanism moving average(and feeding it to QK but not V). This adds another source of residual information. It also makes each token to be a function of all of its predecessors on early stages. Neat. Will *remember* this fact for later. 

Now implementing it is even more trickier. Since I'm not that familiar with FFT or linalg, I took [an implementation from lucidrains](https://github.com/lucidrains/Mega-pytorch/tree/main) who seems to implement papers for hobby.

Let's play a game. Can you spot the bug?

```python

def conv1d_fft(x, weights, dim = -2, weight_dim = -1):
    # O(N log(N)) 1d convolution using some fourier trick

    assert weight_dim >= dim

    N = x.shape[dim]
    M = weights.shape[weight_dim]

    fast_len = next_fast_len(N + M - 1)

    f_x = rfft(x, n = fast_len, dim = dim)
    f_weight = rfft(weights, n = fast_len, dim = weight_dim)

    f_v_weight = f_x * append_dims(f_weight.conj(), weight_dim - dim)
    out = irfft(f_v_weight, fast_len, dim = dim)
    out = out.roll(-1, dims = (dim,))

    indices = torch.arange(start = fast_len - N, end = fast_len, dtype = torch.long, device = x.device)
    out = out.index_select(dim, indices)
    return out
```

That's a trick question! The bug is hidden in code deeper than this:  it's an usage of `rfft`, the [Memory leaker](https://github.com/pytorch/pytorch/issues/948930). 

I almost considered switching to JAX. 
I also considered calculating exponents in the loop but it was so slow.
And what if it's cuda issue that exist in JAX too? And I need to learn something new, don't want to become too smart. So I decided to let it leak the memory and just kill training loop every 100 steps and restart it. Then fast-forward to where we left.


![Mem usage: ~1 hour to OoM](https://dev-to-uploads.s3.amazonaws.com/uploads/articles/a00wn52f01l32jigahk8.png)

~1 hour to OoM. Not a pretty picture


* [RMT](https://arxiv.org/pdf/2207.06881.pdf)

![RMT is so simple](https://dev-to-uploads.s3.amazonaws.com/uploads/articles/zbhg1xemabyayw9q0xyu.png)

We need more retention. And RMT provides just that, but from another angle. Essentially it asks the question "what if model autosummarized itself on token level?" by adding memory blocks, model fills them on the end on the go. Then you take memory from the tail and put it in the head of the next chunk. Simple, elegant.

* (Not yet borrowed but will be) [Memorizing Transformers](https://arxiv.org/abs/2203.08913)

They (and [RETRO](https://arxiv.org/abs/2112.04426) and probably many others) use external knowledge DB. Basically if you have context, they look for it, give more context.

I haven't found good implementation, it seems kNN is not so simple. [lucidrains](https://github.com/lucidrains/memorizing-transformers-pytorch) as usual has the implementation. It uses 3rd party library, which has its [quirks](https://github.com/lucidrains/memorizing-transformers-pytorch/blame/272e39bafd2a507d21ac896bd7cf4b593ee9acb7/README.md#L27): `(once it hits this capacity, it will be reset for now, due to limitations in faiss' ability to remove entries)`

Current intention is to use N memory slots and attention mechanism to lookup memory in linear time. This will be left for later as it can be done after training.

Not memory related borrowings which makes BakaLLM not like the other model:

* From [Persimmon](https://huggingface.co/models?other=persimmon)

I took idea of using Layernorm for Q and K. Gave batter result on TinyStories so I kept it.

* From Llama(?): 

Gated MLP. In older GPTs e.g. neox MLP is 

```python
x = up(x) # [b, s, d] -> [b, s, 4d]
x = act(x) #  [b, s, 4d] -> [b, s, 4d]
x = down(x) # [b, s, 4d] -> [b, s, d]
```

Llama uses

```python
x, g = up(x), gate(x) # [b, s, d] -> [b, s, Md]
x = x*act(g) #  [b, s, Md] -> [b, s, Md]
x = down(x) # [b, s, Md] -> [b, s, d]
```
where `M` is a convoluted constant because 四 is a scary number.

I'm not sure if it's a llama novelty, but I haven't seen it before llama(because I didn't look), and since it's easy to implement, idea was borrowed.


* And finally From RetNet I took a claim that this model is a successor. 

In my case - humanity.
Since BakaNet already has Mega, no further borrowing is necessary.  But this one is very important. Some might say it's the most important part. They will be wrong, but we will not judge.


In next part, someday, we'll talk about how it's all glued together and speak of the Greatest Conspiracy Against the Small Values. 

Mandatory nap is mandatory after all!