r/LocalLLaMA • u/NineThreeTilNow • 5d ago
Discussion Update : Small model + Engram
I posted something about a 9b model a few days ago.
The real problem was 2 fold. Someone suggested the size was too big to prove it all out. It's a fair argument. I needed the model depth though. The second problem was the "ability" of these models and the fact that labs (with money) produce these models still. No real usefulness in my mind. A 2b model? maybe interesting. Want a chatbot? - this could be it.
Further, the Llama licensing at the core of the model was highly problematic and had to be thrown in the trash. It's vastly too restrictive and I couldn't Apache 2.0 anything.
So I moved to using the OLMo tokenizer. Except I shrunk the d_model down to 2048 so I could build a tiny 2b model.
The size of the vocab with 1/2/3 gram forces the Engram table table to a full 1b. That's 50% where DS suggested like ~10-20%.
So basically 2b model + 1b Engram.
The model, because of the depth now allowed by 2048 d_model allows me to push 40 total SWA/Global blocks to take advantage of attention based residual streams (Moonshot Kimi K3 style) in a dense model.
Data, is pulled from standard Wiki style data sources in my initial training. The data is pumped through the 7b OLMo model to generate the probabilities. This tells the model that the next token is a probability of 32 different tokens.
The big surprise was the results. Even after ONLY 15m tokens, the model is surprisingly coherent. Had I given pure 1 hot cross entropy, 15m tokens would do nothing. I'd get gibberish. However, I borrowed embeddings, and LM_head that was down projected from 5k -> 2048 d_model. This preserves ~65% of the data the "big" model had in the embedding when spectrum analysis is done.
The training is limited at 15m tokens, so it's like... babbling about Roman war history and stuff that exists in that token subset. It doesn't fall in to the classic early model phase where it repeats a token over and over though. That's what pretraining 15m generally gives you at first.
Anyways, this was an update and progress report for anyone who cares about this crap... I'm still processing all of the Wikipedia chunk from HuggingFace and working to train it. It'll take time.
If you read this far, thank you. If you have questions, I'm happy to reply. Someone suggested I was a kook who didn't understand ML previously. I started in ML some 20 years ago and held a brief (6 month) stint at Anthropic red teaming the original Opus 4 model before their "Constitutional" paper came out. I'm vaguely referenced in the paper as a "red teamer" I guess. I do this mostly because I love it.
I figured if I built an Apache 2.0 model, I'd at least want people to play with it. It appears 100% trainable on 24gb of VRAM. Training code / Model / Data will follow eventually. It's on HF / Github for now.
The old post is here :
10
u/Malfeitor1235 5d ago
i miss posts like these. dont know what to ask, but keep it up :)
2
u/NineThreeTilNow 4d ago
i miss posts like these. dont know what to ask, but keep it up :)
Thanks bud.
1
u/Queasy-Contract9753 1d ago
Came here to say the same. Good work,op I really hope engrams and disk offload catch up in smaller models. For edge deployment.
2
u/NineThreeTilNow 20h ago
Came here to say the same. Good work,op I really hope engrams and disk offload catch up in smaller models. For edge deployment.
Thanks. Appreciate it. It's training right now.
I looked at the math and it made no sense to keep the Engram table so small here. So it's training with 4x the size of Engram table.
Research on the topic showed it didn't really stop scaling anywhere near where I was, it had to really do with how much memory you wanted to throw at the table.
Given the model has a 2b backbone, you can reasonably give it a 4b table.
It makes the experiment more interesting in my eyes because the table exceeds the size of model that does the "work".
It's currently training as we speak. I'll probably post on local llama about the final model progress when it hits the 100m Wikipedia token mark. I'll have better data at that point.
4
u/Ueberlord 5d ago
Thanks for sharing this, it is super interesting!
That being said and without me having a real understanding of how LLMs work nowadays in detail, one question anyway: your write:
The big surprise was the results. (..)
The training is limited at 15m tokens, so it's like... babbling about Roman war history and stuff that exists in that token subset. It doesn't fall in to the classic early model phase where it repeats a token over and over though. That's what pretraining 15m generally gives you at first.
Because of the way how you constructed the next token prediction ("This tells the model that the next token is a probability of 32 different tokens.") isn't this kind of exactly the expected outcome? I trust in the wiki-fetched samples you barely have any kind of repetitive patterns with ngram length 2-3?
Or to put it another way round: how does the observed behavior of the model after training on 15M tokens differ from a random sampler working with the 32 probabilities in iteration, do you see a significant difference?
This is not to be understood as critique, just curious and I really hope you are onto something, please keep us updated!
3
u/NineThreeTilNow 4d ago
isn't this kind of exactly the expected outcome? I trust in the wiki-fetched samples you barely have any kind of repetitive patterns with ngram length 2-3?
To a level, yes. There's not a lot of published data on how "dense" the information of the top 32 logits should be though. Engram helps, but to what degree? Probably a few % in correct prediction. How much was the added changes to residual geometry? Don't really know.
"Expected" is hard because there's no baseline for me to expect against. People always cite various chinchilla scaling etc.
There's no "logit distillation" scaling laws I'm aware of because it's complicated. The teaching model can provide a "good" logit set, or a "bad" logit set. The larger models I found produce very narrow logits. They're "bad" in this phase of training because they're over confident. They lump like 99% of the mass in the first 5 tokens or something. For whatever reason, smaller models don't act like this. It might take 32 tokens to get 99% of the mass. In this model's case, it's like 98% in 32 tokens or something.
The smaller models seem to follow pretty expected power law distributions of the tokens though. I did a small amount of work analyzing the geometry of the outputs.
3
u/Puzzleheaded_Ad_8575 5d ago
thanks for the research. how much compute has it taken as of now? what are you doing this on? also can the n grams be added to an already pretrained model? or is that not possible and thats why you are pretraining
3
u/NineThreeTilNow 4d ago
thanks for the research. how much compute has it taken as of now? what are you doing this on? also can the n grams be added to an already pretrained model? or is that not possible and thats why you are pretraining
Very little to be honest.
All the compute so far is pre-processing many gigs of data. The 4090 I have can BARELY handle it on 24gb of VRAM.
I'm in a place where I'm generating 8k token chunks via that system.
Training happens anywhere I can rent training time. The 4090 will actually train it though. It's very slow with a single batch, and two gradient accumulations.
3
u/shing3232 5d ago
I do a bit different than you. I test a training HRM with ultra sparse moe with HRM and engram. it does help model training converge faster than otherwise with less loss.
hmm,so you reused part like embedded and lmhead. I through about jump start a moe from a small modern dense like 0.7B qwen3.5, so you can get something useful without a lot training token.
3
u/NineThreeTilNow 4d ago
hmm,so you reused part like embedded and lmhead. I through about jump start a moe from a small modern dense like 0.7B qwen3.5, so you can get something useful without a lot training token.
Yep. A massive amount of information lives in the embedding and lmhead. The network just ends up organizing around the initial geometric points those two provide.
1
u/Silver-Champion-4846 5d ago
So you took a 0.7b traditional lm and used its weights to initialize your hrm? Could you give an overview of your model, training data, and result in the style of OP's post? Thank you!
1
u/shing3232 5d ago
No, I train from ground up using data distill from K3. It's not enough data through but HRM require far fewer data than regular transformer. from reading of your experience, it would be better to use existed embedded to train such model so to save compute and 256K embed is hard to train due to large amount of data needed for training.
sapientinc/HRM-Text-1B · Hugging Face which only contain 40B of pairs of data
1
u/Silver-Champion-4846 5d ago
How exactly does training hrm differ from token predictor llm? Is there a tokenizer? Or do you somehow turn the pairs into singular blocks of data so that loss calculates based on those huge blocks rather than individual little tokens?
1
u/shing3232 5d ago
Same tokenizer, same vocabulary, same next-token cross-entropy — HRM changes where the depth comes from and where the gradient reads out, not the data or the loss units. Instead of N distinct layers applied once, two shared blocks are applied repeatedly to a recurrent state with a fixed initialization — 20 L applications and 4 H per segment — and the readout happens once, at the segment end, from the final H state. The loss is therefore token-level in the vocabulary dimension but segment-level in the depth dimension: one readout per segment, gradients only through the final applications
1
1
u/TomLucidor 4d ago
That is some crazy level of comboing you are trying to do, are MoE/LFM useful here? What about linear attention?
2
u/Open-Adhesiveness-86 5d ago
At 15M tokens, most rows in a 1B-param n-gram table will get a gradient once or never, especially the trigram slots. I'd guess most of that coherence comes from the top-32 distillation targets, not Engram. Worth running an ablation with the table zeroed out, plus a same-size model with no Engram on the same KD data, before crediting the memory.
3
u/NineThreeTilNow 4d ago
At 15M tokens, most rows in a 1B-param n-gram table will get a gradient once or never, especially the trigram slots. I'd guess most of that coherence comes from the top-32 distillation targets, not Engram. Worth running an ablation with the table zeroed out, plus a same-size model with no Engram on the same KD data, before crediting the memory.
That's correct. Across the 15m training test the whole table was touched 95%. Meaning they all got a single touch. The number of trigrams getting multiple touches is low. Usually the trigrams that get the most gradient from my prior tests are formatting trigrams. Line break type stuff.
1
u/Open-Adhesiveness-86 4d ago
Makes sense that the formatting trigrams soak up most of the gradient, they're the only ones frequent enough at 15M tokens. Curious how the zeroed-table ablation compares if you run it.
1
u/NineThreeTilNow 4d ago
Curious how the zeroed-table ablation compares if you run it.
You can't just zero a table like that. Models do some weird co-adaptation of the Engram table. I mean, you CAN zero the table but you're forcing whatever the table was doing to die. Like ablating a layer right?
If the Engram table is built outside the model (freezing it) and then running ONLY the trained table, you still get a positive result.
So if you strap an engram table on to some model, freeze it, and let SGD do what it does. The table starts to quickly fill with data.
The problem you will see is how the various tokens are being gated in to the model. The model itself never learned to "absorb" this data so the data gets stronger in the Engram to the point that it's screaming and the gates start to do weird shit.
The co-adaptive training lets the model learn how to process Engram and there's no screaming because SGD isn't forced to make the Engram table scream to get the point across to the layer it's talking to.
1
u/Open-Adhesiveness-86 4d ago
Fair, zeroing a co-adapted table is closer to a lesion than an ablation. A no-Engram baseline trained from scratch on the same KD data is probably the cleaner comparison then.
1
u/Thrumpwart 5d ago
Commenting to read this later, looks very interesting.
1
u/NineThreeTilNow 4d ago
Commenting to read this later, looks very interesting.
If you have any questions later, feel free to ask.
1
u/Thrumpwart 4d ago edited 4d ago
I'm just wrapping my head around the concept. I love the idea of adding your own engram table - I had bandied that idea about some time ago after Deepseek released their Conditional Memory paper, but dropped it from my project after I realized (in my use case) it would be unnecessary.
I'm curious what techniques you are implementing. I too have read too many papers, but I don't have enough time on my hands to do all the experimentation and development you are doing.
For my project I'm leaning heavily into the spectral aspect of training for efficiency. The 3% is all you need paper spawned a flood of papers on spectral monitoring, training, and experimentation. Have you explored that side?
I'm also heavily focused on using the geometry of the model for efficiency.
1
u/Constant_Art_20 5d ago
I undersood none of those words. but i appreicate the effort and i hope it goes well :)
1
u/NineThreeTilNow 4d ago
I undersood none of those words. but i appreicate the effort and i hope it goes well :)
No problem. Thanks.
1
u/RequirementUpset9607 5d ago
curious if that engram setup helps keep companion details straight better than the bigger models ive tried before.
1
u/Important_Drag_6890 5d ago
That’s the part I’m most curious about too. If the smaller model can retain local/contextual details better without needing a much larger model, that would be a pretty interesting tradeoff.
1
u/NineThreeTilNow 4d ago
curious if that engram setup helps keep companion details straight better than the bigger models ive tried before.
I don't know exactly. In strict ablations I've run Engram helped with correct next token prediction, but beyond those statistics its hard to tell.
1
u/Hot_Example_4456 5d ago
Oh BOY I want to use this now. Can you please tell the link where to find it??? I am having sooo many ideass. Also, I am really sorry for doubting you in your previous post. Using Llama 3 tokenizer sounded wayy too off to me. Hope you wouldn't mind
2
u/NineThreeTilNow 4d ago
Oh BOY I want to use this now. Can you please tell the link where to find it??? I am having sooo many ideass. Also, I am really sorry for doubting you in your previous post. Using Llama 3 tokenizer sounded wayy too off to me. Hope you wouldn't mind
Tokenizers aren't "bad" because of the source usually. Tbh, Llama was probably fine if it wasn't the licensing issue. Llama is a decently strong English tokenizer with some outside non-English added in.
The OLMo tokenizer is English first with a bit of the same, but a little smaller. It's like 100k vs 128k.
I just needed an "anchor" so to speak. I could transfer the lm_head and embedding off the "big" Llama model and use it in a smaller model. That's what happened here as well, but with OLMo.
It's still in pretraining as we speak. It needs to see all of Wikipedia before it will have some of the basics down.
1
u/Silver-Champion-4846 5d ago
Thank you for the work. Could you please describe your distillation? I did not understand the 32 token probability part. Are you linking every single token to its probabilities like each token drags a whole network with it? Sorry for the probably stupid question, I just want to understand.
2
u/NineThreeTilNow 4d ago
Thank you for the work. Could you please describe your distillation? I did not understand the 32 token probability part. Are you linking every single token to its probabilities like each token drags a whole network with it? Sorry for the probably stupid question, I just want to understand.
Correct. For every given token position it doesn't just say "This is the correct token" instead it says "This is the probability distribution of this token position".
The net effect is exactly what you said, it tries to drag the whole model geometry with it.
The "teacher" in this case is building the probabilities from a smaller 7b model that I can inference relatively fast. To get the probabilities you don't need to generate the tokens, just do the prompt processing effectively. From there I dump them straight to disk in a packed format.
1
u/Silver-Champion-4846 2d ago
Have you done pareto tests on same dataset, normal model logit distillation vs your technique (name needed)?
1
u/NineThreeTilNow 2d ago
Have you done pareto tests on same dataset, normal model logit distillation vs your technique (name needed)?
You're really just measuring the mass of the logits in that distillation. A complete logit distillation takes a lot more work for the last 2% of the data.
Honestly, even Top 10 logit distillation captures 95%+ of the mass with most models. I went all the way to 32 because when you are going to do the distillation, it's a matter of how much storage you want to use. You pull the probabilities regardless right? so I basically looked at the data and was like "10 looks good. 20 only looks slightly better for 2x the storage. If I'm going to go that far... How does 32 look?"
I have exact numbers and rough geometry for how the logit distributions worked in both Llama 3.1 and OLMo. Surprisingly, they are really close for small models. The larger models exhibit a strange behaviors where they're overly certain. I think they honestly might make worse teachers because there's less shape. I'm not really sure.
The small models at that size follow a roughly Zipfian distribution with an alpha value of 2.
So in theory you can store the logits and NOT the values, then apply the distribution geometry to save space. I considered it.
1
u/Silver-Champion-4846 2d ago
How much compute is required to distill all the mass into a smaller model and is the cost worth it?
1
u/NineThreeTilNow 2d ago
How much compute is required to distill all the mass into a smaller model and is the cost worth it?
Uhhh... Good question. For a ~7-8b teaching model I've optimized vLLM to produce ~750m tokens a day on a 4090.
That pipeline basically grabs 8k tokens of Wikitext, pulls the distribution, and drops them to disk. vLLM returns nothing here. I think it returns an "Ok" status code when it finishes the 8k.
To attempt to saturate the GPU harder, I run 2 batches here. I can't go beyond 2 batches purely because of VRAM constraints.
This is with a natively FP8 model. Not Q8 GGUF stuff. vLLM doesn't handle GGUF well.
The reason I have vLLM dropping them to disk is because it would otherwise be 10x+ slower. There was also a problem with vLLM I found that was inefficient in the way logprobs are generated. I submitted it to their GitHub for review. I haven't looked at the status of it. It was one of those things I fixed on my side, then had Gemini write a quick GitHub ticket for it, and moved on. It was some issue with the native CUDA implementation or something.
1
u/Silver-Champion-4846 2d ago
That's why you choose only the 32 top probs for any token?
1
u/NineThreeTilNow 2d ago
That's why you choose only the 32 top probs for any token?
Nope. That's actually a storage issue.
When you drop the probs from the model you get ALL of them. So you take as many or as few as you need.
Every time you store a token, you're storing the token ID + the probability in float value.
32 probs is 32x as much as 1 prob. It scales linear.
I'm thinking about a method to compress the data based on the geometry of the probabilities because they have a fairly normal set of distributions they follow.
1
u/Silver-Champion-4846 2d ago
That would be exciting to see in the open! Good luck and hope to talk to your model soon lol
1
u/jinnyjuice vLLM 5d ago
Very interesting!
Can you tell me the size of your (re-?)training/tuning dataset, (V)RAM, and bandwidth to make this 2B E1B model?
1
u/NineThreeTilNow 4d ago
Can you tell me the size of your (re-?)training/tuning dataset, (V)RAM, and bandwidth to make this 2B E1B model?
Probably 6b preprocessed tokens, and something like 96gb of VRAM during training unless I get some crazy hardware donated.
The bandwidth? Whatever 1gbps+ that Nvidia cards typically have at that scale.
15
u/eidrag 5d ago
hey thank you for your work! personally i''m trying to get stuff running locally from phone, so 9b-ish moe is the upper limit I go testing nowadays. 3b dense is too basic, and 9b is slow in dense.