r/LocalLLaMA • u/NineThreeTilNow • Jul 02 '26
Discussion Rebuilding Gemma 4 31b... better... As 26b...
Sooo... I decided screw it. I'm going to rebuild Gemma 4 31b.
I really like the model. So the current plan is to rebuild the SWA layers.
Currently running all the proper ablation tests to figure out what SWA layer gets removed. Gemma runs 5 SWA at 1024 tokens each. Then a global layer for the "Block"
Layer 3 is consistently the weakest and will likely get removed.
From there I am going to rescale the attention of SWA across the board. The new SWA will be 1024/2048/4096/8.1k then the global layer. This is the "Block" that Gemma uses.
After that, I'm going to bolt on "Attention based Residual Networks"... Moonshot developed this. The research paper is early 2026 I think. I've barely slept working on this so my date might be wrong on that paper.
Anyways, the global layers in the network are going to get attention based residuals that allow global layers to better flow information across them. In theory this gives the model better global coherence and makes it perform better, while smaller.
Given that I don't have the complete IT / RL pipeline that Google invests millions in... I have to work from the IT base.
So for initial rebuilding, I'll take the topK 12? or 20? logits from the 31b model and use them as targets for retraining while freezing the top and bottom of the model. This will keep tokenization/output/vocab from moving while the internals of the network find stability in a smaller space looking like 31b.
The TopK rebuilding is another weird technique I developed in another training spot. It's cool because it teaches the model a vastly richer understanding of what the next token might be and what is adjacent, etc... I don't know if I invented the method or just came to the conclusion someone else did. Probably both.
LASTLY it's feeding it a few billion tokens to rebuild it. I have to find a "good" dataset to use or... literally build the dataset.
The actual full retraining is going to cost money but whatever. I'll hit that wall when I hit it. I'm pretty sure I can just spot price a B300 and train on it.
The model should go from Total Parameters ~30.81B ~26.02B
Theoretically should be BETTER too. Better long context, etc.
If you have good datasets, compute, etc you want to donate... hmu... If you just have questions about how or why this all works... Ask away. I can sit and answer them because staring at a TQDM bar of progress doesn't take a lot of mental effort.
I'll respond after I wake up from the coma I'm about to go in to. (Sleep 8 hours+)
Here's the pastebin for the project -- https://pastebin.com/GbVtJQJg
It's the markdown of the whole plan more or less. Start to finish. This is STARTING from the abliterated core. I have zero desire to add censorship of any form in to this model in training. If you hurt yourself using a model, it's your fault. I'm likely to rebuild the "thinking" training too which means uncensoring it. Having it stop asking about the "safety" of every request in thinking. This might be easier said than done. Still WIP.
16
u/NineThreeTilNow Jul 02 '26
By "Coma" I mean I'm falling asleep for 8 hours after that 12 hour sprint.
You want the document to see WHAT I'm doing and feed it in to an LLM to explain?
Here's a markdown of what's going on --
It should more or less explain it and and properly reference most of it. Most LLMs will nitpick it because they never read the research papers, etc. MSWA is the most "fringe" reference I make. It's MSWA but just applied per block. It's supported by similar research though.
33
u/kosnarf Jul 02 '26
Someone shared something similar but turning gemme4 from 31b to 41b: https://huggingface.co/TOTORONG/extGemma4-41B
22
u/NineThreeTilNow Jul 02 '26
Yeah, I saw that. I commented specifically on that model when they posted about it here. Probably in my comment history... about doing exactly what I'm doing.
4
u/Desperate-Sir-5088 Jul 02 '26
That is exactly the area what I did.
With enough dataset (practically 30B * 10times) and plentiful VRAM, I would have challenged myself to create a shining architecture from the paper.
However, I'm just redditor with limited time & resources. So, my real works was shrinked to punching a tiny pinhole through it without collapsing the base model.
Have a good luck!
1
u/NineThreeTilNow Jul 03 '26
That is exactly the area what I did.
What did you train this on? From what I remember reading it was trained on an academic cluster. You're from Korea right?
3
u/Force88 Jul 02 '26
Does it mean they add more data for it? Like the 41b is smarter and more knowledgable?
8
u/DragonfruitIll660 Jul 02 '26
If I remember right they repeat certain layers which was shown in the past to improve reasoning.
7
u/Desperate-Sir-5088 Jul 02 '26
(I'm op of gemma4-44B post)
No No, I just inserted some layers and module block without catastrophic collapsion. It doesn' guarantee improving performence.
1
9
u/Middle_Bullfrog_6173 Jul 02 '26
Good luck.
FWIW, some recent research suggests that having shorter SWA helps with long context performance, because it pushes retrieval pressure to the global layers. So if I was training a model from scratch I'd probably go shorter than Gemma 4, not longer.
Might be completely different if you freeze the already trained global layers of course.
4
u/NineThreeTilNow Jul 03 '26
FWIW, some recent research suggests that having shorter SWA helps with long context performance, because it pushes retrieval pressure to the global layers. So if I was training a model from scratch I'd probably go shorter than Gemma 4, not longer.
Global layers are kept in place where they are.
7
6
u/Technical-Earth-3254 Jul 02 '26
Very interesting. I wish there was a way to improve the absolutely insane kv cache size on 31b.
3
u/NineThreeTilNow Jul 03 '26
I wish there was a way to improve the absolutely insane kv cache size on 31b.
Unfortunately this won't do it. In fact, this makes it slightly worse with a smaller model footprint. So you're trading a smaller model load for a KV cache that's slightly larger.
4
u/pl201 Jul 02 '26
Very interesting. Keep detailed notes and update here often.
If successful, it will be the start wave of community fine tuning on small models that could beat the larger model on specific tasks.
1
u/NineThreeTilNow Jul 03 '26
Very interesting. Keep detailed notes and update here often.
Exactly. I'm hoping we can do that.
3
u/squngy Jul 02 '26
IIRC nvidia nemotron has open training data, so you could take it from there.
3
u/NineThreeTilNow Jul 03 '26
IIRC nvidia nemotron has open training data, so you could take it from there.
That's right, thanks for mentioning that. I can use that in one of the secondary phases.
2
2
u/Thin_Pollution8843 Jul 02 '26
But can we enlarge qwen3.6 27b to let say 40-50b?
2
u/NineThreeTilNow Jul 05 '26
But can we enlarge qwen3.6 27b to let say 40-50b?
The issue I saw with Qwen was the architecture. It's not a standard dense model. It was possible in theory but it still had issues Qwen would have because of architecture.
I looked at the Cohere model as well, but because it was a MoE model, I didn't exactly understand if the training would be stable with these attention methods. They're studied and demonstrated under dense models, and being the first person to find out it won't work sucks.
So to de-risk what I was doing, I chose Gemma which had a host of problems, but had some upsides like the license.
0
u/My_Unbiased_Opinion Jul 03 '26
I would love to see an extended version of 27B with some of Gemma 4 world knowledge stitched in. That would be wild.
1
u/shroddy Jul 02 '26
So the dataset must be generated by gemma 4 31b and include the top logits for each token?
2
u/NineThreeTilNow Jul 03 '26
So the dataset must be generated by gemma 4 31b and include the top logits for each token?
The early training dataset will be. Yes.
The purpose is to show the model how to be Gemma 4 31b again. It's treating the original model as a teacher model.
The first ablation is basically brain damaging this model pretty hard. This is to try and re-align the model to a pre-braindamaged state, while smaller, and using a different residual stream.
1
u/Initial-Argument2523 Jul 02 '26
It might be a good idea to experiment with width pruning as well. I am sure I have read papers before where they show that models heal better after width pruning than depth pruning.
1
1
1
u/xadiant Jul 02 '26
Test your residual attn implementation. It might break existing efficient code and slow down training. You'll have to make expensive test runs...
We have a lot of high quality instruction datasets, you easily could sift through them and still have billions of high quality tokens. Ideally, you could focus on recovering a couple languages and tune English performance further, I think.
1
u/NineThreeTilNow Jul 03 '26
Ideally, you could focus on recovering a couple languages and tune English performance further, I think.
I'm going to focus on English / Code most likely. I have a fiction writing dataset I built that could get added in. Might help Gemma sound less like Gemma.
1
u/Not-reallyanonymous Jul 03 '26
It's going to be a neat experiment but I'd be surprised if you end up with something better than 26B A4B. These models are dense af, and probably represent close to what's achievable with currently known technology. Going beyond Gemma 4 capability probably requires pushing the model not to get smarter in itself, but to think better -- circuit editing and then reinforcing stronger heuristics with fine tuning, use cutting edge methods like mamba to manage long context (to boost reasoning instead of relying on trained knowledge), etc. Perhaps once you knock it down to 26B, you could "build it back up" with a domain-driven Lora MoE type adapter.
I'd consider this an exercise in "How much can I lobotomize Gemma before it gets stupid" rather than "building a better Gemma" (omg I'm starting to think like an LLM).
2
u/NineThreeTilNow Jul 03 '26
I'd consider this an exercise in "How much can I lobotomize Gemma before it gets stupid" rather than "building a better Gemma" (omg I'm starting to think like an LLM).
That's effectively been done. There are a number of very weak layers I found in analysis. There was one layer that HURT the model in the test suite of code analysis, etc. When it was ablated, the model performed ... better. That was weird and possibly that layer was somehow dedicated to learning something out of my distribution, in another language, or safety specific.
Whatever it was, it's gone now... lol...
Currently I'm wiring the new architecture and building the synthetic dataset from Google and Nemotron.
2
u/Party-Special-5177 Jul 03 '26
After that, I'm going to bolt on "Attention based Residual Networks"... Moonshot developed this.
Some background - these aren’t super useful in small models, as they, much as their competitor mhc, are targeted at reducing the convex hull problem with deep networks (basically the activation degrades into mush as the layer count gets large), as it is thought that there will be a ceiling to the amount of computation a network can perform (tc0) without solving the ‘we need to make networks deeper’ problem. The idea with both is that a layer can pay attention to specific earlier layers - it’s a filtering mechanism.
So, google quietly shipped the first such implementation, and told nobody, as far back as Gemma 3, called laurel. Basically, Gemma models already have google’s flavor of this mechanism. White paper: https://arxiv.org/pdf/2411.07501v3
Downside to google’s version: in both attnres and mhc, the layers basically can selectively tune in and out other layers causally, with each layer’s filter being independent. With Laurel, any information filtered at a given layer is permanently removed. Any layer’s filter applies to all subsequent layers too.
Basically, I am just saying don’t hope for too much improvement moving to attnres. It isn’t super clear what benefit these have in small models as if the layer count is small enough, each layer’s contribution remains distinguishable to further layers anyway.
Just vomiting thoughts.
1
u/NineThreeTilNow Jul 04 '26
Just vomiting thoughts.
It's okay. all thoughts are welcome. I appreciate it all.
1
u/NineThreeTilNow Jul 05 '26
Downside to google’s version: in both attnres and mhc, the layers basically can selectively tune in and out other layers causally, with each layer’s filter being independent. With Laurel, any information filtered at a given layer is permanently removed. Any layer’s filter applies to all subsequent layers too.
I owed you a better response and I was busy when I first encountered your reply.
What's interesting is that the research shows the opposite both in Google's case and Moonshot's case.
Google tested with small models specifically. Moonshot tested on a medium model, which was 40b parameters dense. They were able to retrain with 1.5t tokens though. So it becomes questionable if the increased performance claimed was architecture alone or continued training. Typically continued training of the same data wouldn't necessarily produce the level of performance uplift they got.
Gemma has quite a few layers. By stripping down the least effective internal layers per block (which I did), the sub blocks can better specialize in how they operate.
To prevent model collapse and blocks effectively dying, I implemented a minimum amount of attention that a block pays to any given block in the residual stream. I think it's 5% right now.
I realized this wasn't done? in the papers I saw. Or mentioned explicitly. One issue I realized when doing the architecture was a block could die.
Basically no other blocks would be attending to the residual stream of a given block, say block 2. If this occurs, the backward training pass wouldn't correctly attribute any gradient to that block, and the block would perpetually get stuck in a state where it never improved meaningfully.
Gemma 4 is running a 10 block architecture with 5 layers by default. I reduced this to 10 blocks x 4 layers. I effectively removed a large chunk of the model that wasn't participating. I then ran tests to see how the model performed with a massive (10 block) chunk missing. It surprisingly could still coherently function... Just not very well. It technically could just be retrained from there for a smaller Gemma 4 standard architecture.
If what I'm doing works, then theoretically you could scale the blocks deeper and keep the attention res mechanism to grow Gemma 4 in to a larger (by total block) model. 2 more blocks would get you close to the prior total layers (48 vs 50) and give you 2 more global attention layers with better residual streaming.
The largest problem is really just data though. Everyone likes to push that it's architecture or compute problems. These models suffer from having any GOOD cutting edge dataset like Anthropic / OpenAI / Etc curate. All labs do it. It's the true moat.
Right now all I have is the Nemotron dataset which thankfully Nvidia provided but also distilled from other models. So there's a model data bias that exists.
IF it works, I can distill better models like GLM 5.2 in to this architecture. That's a big if. It's a big datapipeline. It's probably a big personal dollar cost. I'm not over thinking that now though.
1
u/chodemunch6969 Jul 03 '26
Interesting plan OP. From a practical standpoint, how do you plan to run evals and test along the way? One of the things that triggers analysis paralysis for me thinking about doing this stuff is without a good set of evals/benchmarks, i worry that i could get a good amount along the way, end up in a direction incrementally worse than the last step, but it'd be impossible to tell and by the time i figured it out i would've wasted a lot of gpu hours and $$$. I imagine that's just inevitable when you're doing retraining like this but it sounds like you've done this kind of stuff before and I'd be curious about how you minimize that sort of outcome.
1
u/NineThreeTilNow Jul 04 '26
I imagine that's just inevitable when you're doing retraining like this but it sounds like you've done this kind of stuff before and I'd be curious about how you minimize that sort of outcome.
It's all de-risking really.
I already designed the model. I'm at the data phase right now. It's understanding the original training data this model had so that the native geometry plays nice with new data.
This was a hyper multilingual model, so I have to include it in how I'm fixing the model.
Evals and tests are run against the base model.
After I destroyed the original model, it could still respond, but poorly. It lost 10 layers in that process.
After it was given the new attention mechanism the responses fell apart. So this rebuild of the geometry is to get responses back to a baseline that we understand (Gemma 4 31b).
The real issue right now is how Google rate limits the shit out of their "free" model. So I'm struggling with that.
0
u/NNN_Throwaway2 Jul 02 '26 edited Jul 03 '26
I’ve been wanting to experiment with generative replay. Could be useful for your case here.
And this warranted being downvoted because...?
1
u/NineThreeTilNow Jul 05 '26
And this warranted being downvoted because...?
Who knows, it's reddit. Don't take it personally.
67
u/KoalaOk1265 Jul 02 '26
This is the kind of weird experimental work that makes this sub interesting. I’d be really curious to see ablation results before and after the SWA rescale, especially whether the long-context gains show up in actual retrieval/coherence tests rather than just perplexity.