The 67M untied vocabulary footprint and context scaling on 6GB RTX 3050 budgets
Hi AobanZ,
Training a 105M causal model from scratch on a local RTX 3050 6GB in 1.5 hours and driving perplexity down from 139 to 21 on Booksum is a really cool indie engineering achievement.
Looking at your architectural accounting (12 layers, 512 hidden size, ~65,830 vocab) and the 230-token sequence ceiling:
Untied vocabulary parameter tax (64.2% of your entire model):
With a 65,830 vocabulary at hidden dimension 512 and untied weights, your input lookup table and output LM head consume ~33.7M parameters each, totaling ~67.4M parameters.3.15M per layer).
Out of your 105M budget, that leaves only ~37.6M parameters for the entire 12-layer transformer backbone (
In other words, static embedding lookup tables consume almost two-thirds of your entire parameter count. That static memory costs more than 21 full transformer layers of active compute.
Using low-rank input factorization (projecting 65,830 -> 128 -> 512 = ~8.5M params) and tying or factorizing the head reclaims over 58M parameters. Reallocating those weights directly into the layer stack would allow you to double or triple physical depth (from 12 to 24+ layers) within the exact same 105M ceiling.
FP32 precision and the 230-token context limit:
The RTX 3050 (Ampere architecture) has dedicated Tensor Cores optimized for half-precision. Training in full FP32 doubles VRAM consumption and leaves substantial Tensor Core throughput on the table.
Wrapping the forward pass in torch.cuda.amp.autocast(dtype=torch.float16) cuts activation and weight memory in half.
For Booksum, where book chapters and summaries span several thousand words, a 230-token window limits the model to roughly 160-180 words, making multi-paragraph plot continuity hard to track. Freeing VRAM with FP16 easily opens up headroom to scale sequence length to 1,024 or 2,048 tokens on the exact same 6GB card.
Macro-topology and O(1) sequence scaling:
In an open architecture project called Maba v2 (101M reference release: https://huggingface.co/AndrewThompson1233/maba-v2-architecture), we target this exact 100M envelope on consumer GPUs.
We keep the vocabulary tax under 5% using factorized projections and use a 3:1 hybrid stack (75% DGDA linear recurrence / 25% latent sparse attention).
DGDA updates an associative recurrent state in strictly bounded O(1) memory, allowing models to ingest thousands of tokens on 6GB consumer cards without running into quadratic KV-cache bottlenecks.
Did memory pressure from FP32 training enforce the 230-token context boundary during the run?
Best,
Andrew
Hello. Yes, it used 90% of my memory. My RTX was at 70c lol.
And to clarify, i kinda forgot to change the copy paste. The real training time was 6 hours. This took a lot of my memory, so i couldnt do anything else on my PC.
The 230-token seq length, about that i wanted to increase from 210 without overloading my machine, not necessarily a hard VRAM limit. I was training in FP32, so memory usage was definitely a constraint, but I haven’t tested the same architecture with FP16 and a substantially longer context yet.
The vocabulary overhead is something I agree is worth looking at. With the vocabulary taking such a large fraction of the parameter budget, factorized embeddings or weight tying could free a lot of parameters for the transformer itself.
For the next version, I’m considering increasing the model size to around 200M while experimenting with FP16, although i’ve trained larger models before(400M-2B) on my consumer GPU, they had a smaller dataset.
Best,
AobanZ
Hi AobanZ,
Haha, 6 hours of full GPU lockup at 70C while the whole PC is unusable is the true rite of passage for local pretraining! :)
Switching to FP16 with torch.cuda.amp.autocast() in your next run will be a game changer on that 3050. The Ampere Tensor Cores will slice memory in half and practically double your step speed, so you'll easily be able to run 1,024+ context without your system locking up.
Tying the embedding weights or adding a low-rank bottleneck for that upcoming 200M run is definitely the way to go - reclaiming those 50M+ parameters will give your transformer blocks way more depth and expressive capacity.
Best of luck with the 200M run, looking forward to testing the next checkpoint!
Best,
Andrew