NeuroBLAST v46 · 150M (pretraining in public)

This repo holds a 151.8M-parameter language model as it trains from scratch on a single Colab A100 GPU. The training curves and eval table below update themselves at every eval point, so you can watch a small model learn language step by step.

Work in progress. The files checkpoint-* are intermediate training states (model, optimizer, scheduler), not a finished model. Expect a base model that continues text, not a chat assistant.

Live training progress

Status (2026-11-10 18:57 UTC): step 114,791 of 114,791 (100.0%) · ~7.55B tokens seen · train loss 2.897

Updated automatically at every eval point. Tokens seen are summed from the trainer's num_tokens log (non-pad tokens), so they are approximate across resumes.

The model at a glance. Tokens pass through three stages. Each stage mixes information over a longer range than the one before:

• Sensory layers see a few neighbouring tokens: short convolutions plus attention.
• Associative layers carry a compressed recurrent memory of the whole text (GLA), with attention every third layer.
• Motor layers attend over everything to choose the next token.

Small gated bottlenecks connect the stages. Two cross-layer shortcuts (MGVR and GCLA) let later attention layers reuse what earlier layers computed. Details are in Architecture.

How to read the chart

• Loss is the mean cross-entropy per token, in nats. A model that guesses uniformly over the 65,538-token vocabulary starts at ln(65,538) ≈ 11.1. Perplexity is exp(loss), roughly how many tokens the model is still "choosing between" on average.
• Train loss (blue) is measured on the batches being trained on. The thin line is raw and the thick line is smoothed. Eval loss (orange) is measured on held-out text the model never trains on. Because the data streams in a single pass, train and eval stay close; a widening gap would mean overfitting.
• Learning rate follows a warmup-stable-decay schedule: a 2,000-step warmup, then a long flat phase, then a linear decay over the last 20% of steps. Most of the loss drop during the decay comes "for free" as the noise from a high learning rate settles. Watch for the curve bending down after step ~91.8k.
• The dotted line marks where training is right now.

The plan

Why two context stages? Short sequences are cheaper per token: attention cost grows with length, and short batches let the GPU process 192 sequences at once. The model learns most of grammar and facts at 512 tokens, and the final 25% at 2,048 tokens teaches it to use longer context. Every step processes the same 98,304 token positions in both stages, so the learning-rate schedule carries on unchanged across the switch. Expect the loss to jump at step 70,424. With more context, later tokens become easier to predict, so the loss usually steps down.

Why split long documents instead of truncating? In FineWeb-Edu, 37% of all tokens sit in documents longer than 2,048 tokens. Cutting every document at 512 tokens would throw away about 60% of the text, so long documents are split into consecutive pieces.

Architecture

NeuroBLAST is an experimental hybrid decoder. Its 16 layers form three stages connected by bottleneck gates:

• Sensory: short causal convolutions (kernel 7) alternating with attention.
• Associative: gated linear attention (GLA) blocks mixed with latent (MLA-style) attention.
• Motor: attention.

To save memory, the training loss is computed by a fused, chunked LM-head + cross-entropy kernel, so the full logits matrix of shape (tokens × 65,538) never exists in memory.

Things we learned speeding this up

At the start, the A100 ran at about 5% of its peak FLOPs. These fixes made training much faster without changing the math:

1. Pad the vocabulary to a multiple of 64 inside the LM-head kernel. With 65,538 rows (not divisible by 8), cuBLAS fell back to slow kernels that use 2-element memory alignment. The output head took about 90 of every 120 ms. Padding the rows with masked columns fixed it, and the result is identical in fp32.
2. torch.compile each decoder layer separately rather than the whole model. Layers of different types compile independently, and the recompile limit has to be raised from its default of 8.
3. Bigger batches of shorter sequences (stage 1 above) keep the GPU busy for small models.

Usage

The src/ folder holds a snapshot that loads on its own: weights, tokenizer, config, and the custom modeling code (model_type: neuroblast_v43, so trust_remote_code=True is required). It is refreshed from time to time during training, so it lags the live run. Expect rough text until the run is further along.

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

repo = "mkurman/neuroblast-v46-150m-v41tok"
tok = AutoTokenizer.from_pretrained(repo, subfolder="src")
model = AutoModelForCausalLM.from_pretrained(repo, subfolder="src", trust_remote_code=True,
                                             torch_dtype=torch.bfloat16).eval()

inputs = tok("The water cycle begins when", return_tensors="pt")
out = model.generate(**inputs, max_new_tokens=60, do_sample=True, temperature=0.8, top_p=0.95)
print(tok.decode(out[0], skip_special_tokens=True))

Generation is cached (NeuroBLASTv43Cache, on by default). It keeps the attention keys/values, the recurrent GLA state and the short-conv tails, so each new token costs one position, not a pass over the whole prefix. On CUDA, single-token steps run as one captured CUDA graph. Two limits: beam search falls back to uncached generation, and batched prompts must be the same length (no left padding). The checkpoint-* folders are full training states for resuming the run and do not include the modeling code.

Continuing pretraining

Only the newest checkpoints stay in the main branch. Two milestones are kept as tags, each with the full training state (weights, optimizer, scheduler, RNG):

The second tag is the one to continue from if more compute becomes available. With a warmup-stable-decay schedule, the learning rate is still at its peak there. You can resume training on more tokens at that rate and then run your own decay, without paying for the first decay. The tags appear in the repo's branch/tag list once training passes each milestone.

from huggingface_hub import snapshot_download
snapshot_download("mkurman/neuroblast-v46-150m-v41tok", revision="pre-decay-step-90000",
                  allow_patterns=["checkpoint-90000/*"], local_dir="neuroblast-v46-150m")

Limitations

This is a small research model trained on educational web text. It knows little, makes facts up, can reproduce biases in its training data, and has had no instruction tuning or safety tuning. Use it for research and learning, not for decisions.