Maximizing narrative depth and vocabulary efficiency on 1-hour Colab runs
Hi Arush,
Training a from-scratch causal prototype on a free Colab T4 in a single 1-hour session and bundling ONNX exports is a really neat zero-budget project.
TinyStories is a great sandbox for testing early convergence, but the 1-hour compute ceiling introduces two specific architectural trade-offs:
Narrative continuity versus physical depth:
A 1-hour session on a single T4 caps total training exposure to roughly 10M-20M tokens. Keeping the layer count low helps training speed, but shallow networks struggle with multi-step cause-and-effect reasoning and character persistence across even short stories.
In an open architecture project called Maba (101M reference model: https://huggingface.co/AndrewThompson1233/maba-v1-architecture), we actually evaluated this on TinyStories using deterministic 2-pass block recycling:
Routing activations through physical transformer blocks twice with Split RMSNorm (distinct norm scale vectors for pass 0 and pass 1) doubles effective computational depth to 16 or 24 layers without adding parameter weight. On TinyStories, that extra depth gives the model the representational capacity needed to track characters and resolve plots cleanly.Vocabulary parameter tax on micro-models:
If using a standard 32k or 50k tokenizer, the embedding lookup table alone can consume 30% to 50% of the entire parameter budget on a sub-100M model.
Decoupling the vocabulary via low-rank factorization (projecting token IDs through a rank-64 or rank-128 bottleneck before the hidden dimension) shrinks the embedding footprint down to ~4%, freeing millions of parameters to invest directly into attention heads and SwiGLU width.
What exact layer count, hidden width, and vocabulary size did you settle on for this 1-hour Colab checkpoint?
Best,
Andrew
Sir, I am just a 13-year-old teen, so I am new to this field. My model has 12 layers, around 5M–7M parameters, 16,384 vocab, 12 heads, 192 d_model, and 3× compression using GQA. I am currently thinking of shifting to Mixture of Depth & Exclusive Projection (XSA). My real account is ArushBuilds: https://huggingface.co/ArushBuilds.
Respected sir, can you suggest a good library and tips for maximum GPU speed and utilization? Also, the model uses YaRN for an 8× post-training context increase. Perplexity is 5.9 at the end, and validation accuracy is 58% (low, but good for a 5M baby model).
Hi Arush,
First off, no need to call me sir! We are all just developers and builders here.
Building a custom 12-layer causal model with GQA, YaRN scaling, and ONNX export from scratch at 13 is seriously incredible work. Getting 5.9 validation perplexity on a 5M parameter model within a 1-hour Colab run is genuinely solid convergence.
For maximizing GPU speed and utilization on a free Google Colab T4:
Libraries and framework:
Look into Unsloth (https://github.com/unslothai/unsloth) or LitGPT. They have heavily optimized kernels for consumer cards like the T4 that cut training overhead dramatically.
If you prefer raw PyTorch, wrap your model in torch.compile() with mode="reduce-overhead". On a T4, it fuses operations and speeds up small models significantly.Native PyTorch SDPA:
Make sure your attention implementation calls torch.nn.functional.scaled_dot_product_attention instead of manual matrix multiplications. On the T4 (Turing architecture), SDPA automatically uses memory-efficient attention kernels, saving both VRAM and time.Dataloader throughput:
On free Colab, the CPU often bottlenecks the GPU. Set num_workers=2 and pin_memory=True in your DataLoader, and pre-tokenize your entire TinyStories split into contiguous uint16 binary files using numpy.memmap. That way the T4 spends 100% of its time computing instead of waiting on Python tokenization.Mixed precision:
Train with torch.cuda.amp.autocast(dtype=torch.float16). The T4 Tensor Cores hit peak FLOPS in FP16.
Keep building and experimenting, you are already way ahead of most people starting in this field!
Best,
Andrew
Follow me so that we can build review each other's model and I can learn more from you
Hi Arush,
First off, no need to call me sir! We are all just developers and builders here.
Building a custom 12-layer causal model with GQA, YaRN scaling, and ONNX export from scratch at 13 is seriously incredible work. Getting 5.9 validation perplexity on a 5M parameter model within a 1-hour Colab run is genuinely solid convergence.
For maximizing GPU speed and utilization on a free Google Colab T4:
Libraries and framework:
Look into Unsloth (https://github.com/unslothai/unsloth) or LitGPT. They have heavily optimized kernels for consumer cards like the T4 that cut training overhead dramatically.
If you prefer raw PyTorch, wrap your model in torch.compile() with mode="reduce-overhead". On a T4, it fuses operations and speeds up small models significantly.Native PyTorch SDPA:
Make sure your attention implementation calls torch.nn.functional.scaled_dot_product_attention instead of manual matrix multiplications. On the T4 (Turing architecture), SDPA automatically uses memory-efficient attention kernels, saving both VRAM and time.Dataloader throughput:
On free Colab, the CPU often bottlenecks the GPU. Set num_workers=2 and pin_memory=True in your DataLoader, and pre-tokenize your entire TinyStories split into contiguous uint16 binary files using numpy.memmap. That way the T4 spends 100% of its time computing instead of waiting on Python tokenization.Mixed precision:
Train with torch.cuda.amp.autocast(dtype=torch.float16). The T4 Tensor Cores hit peak FLOPS in FP16.Keep building and experimenting, you are already way ahead of most people starting in this field!
Best,
Andrew
Question, unsloth is for mostly finetuning... I will research , thank you for advice & appreciation...
I don't use SDPA instead we use xformers i benchmarked it , SDPA 2.19x times slower than xformer
Hi Arush,
Gave your profile a follow.
Good catch on xFormers. On Turing GPUs like the T4, xFormers cutlass kernels are indeed faster and more consistent than PyTorch SDPA, which can fall back to slower math paths depending on head dimension and sequence alignment. Sticking with xFormers makes total sense here.
You are also right about Unsloth being focused on fine-tuning. For training small causal models from scratch, LitGPT or a lean custom PyTorch loop with torch.compile and your xFormers setup will give you the cleanest control over training throughput.
Looking forward to seeing your next runs with Mixture of Depths. Feel free to ping anytime you drop a new checkpoint.
Best,
Andrew
Hi Arush,
Checked out the Space, jumping from 5M to 21.5M and maintaining 57% val accuracy on TinyStories is a solid milestone. The generations are noticeably more coherent at this scale compared to the 5M prototype.
Scaling to ~20M parameters gives the network enough width to start forming decent multi-token associations without choking free GPU limits.
Keep pushing forward with the iterations, really fun watching Pragya evolve!
Best,
Andrew
Thank you very much i was amazed by maba ! You should build a massive 124M flagship that proves architecture and is accurate like build your own GPT 2? ... Where do you get compute from ... Also that 21.5M model was trained on Cosmopedia+ Tinystories