Training a small language model from scratch

Agents have dramatically changed how code is written, and consequently how a lot of computational research is done. This is particularly true for empirical deep learning research, where the quality of the research largely depends on the practical implementation. Like many other PhD students, I have been ceding more and more of the code to agents, largely because they are much faster and better at writing code and much less prone to the frustration of dealing with awkward errors, scattered compute and silent bugs. However, even though I had access to this incredible system which converts my ideas to working code, the scale and types of ideas I considered remained somewhat constrained. There were several reasons for this hesitation in the past, including the reliability of agents and a preference for controlling specific design decisions.

A few weeks ago I decided to change this and try to increase the scope of problems I considered feasible working alone. I made an empty repository and asked Claude to help me set up the infrastructure to train a ~50M-parameter language model from scratch in JAX while being minimal in the spirit of nanochat . Why a small language model, and is this really that big of a jump in scope? I have worked on different parts of the LLM training stack independently, but I have always felt that there are many interesting questions that are better answered end-to-end rather than looking at any single slice of the training process. Another motivation was updating my understanding and intuitions, as progress in every area has been rapid. A minimal setup that can still replicate interesting phenomena can be useful in quick experimentation. The final motivation was just to see how much new data, recipes and algorithms can push what would have been “large” models a decade ago.

Around 98 hours after my initial prompt, I had a trained ~50M-parameter language model, Femto, which could do simple tasks in the terminal. The final model was pretrained on 10B web tokens followed by a length-extension phase and mid-training for agentic behavior. The model was modestly post-trained for writing basic shell scripts and Python code as well as some general tool calling. The infrastructure developed to train this model includes a clean JAX version of nanochat with a more optimized inference engine, faster RL training, a basic local search index, and a fast sandboxing system. I did not write a single line of code throughout the course of this project. Claude Code (with Opus 5.5) wrote the code for all the experiments and ran them, analyzed the results, suggested next steps, and also made the plots and visualizations in this post.

This post will mostly be a meta-analysis of this short experiment, talking about the process, decisions, and challenges along the way.

Femto in a Terminal

Before talking about the process, I think it is important to set expectations right. Femto is NOT a state-of-the-art model. I would personally say it is barely a good model. There are other models in the sub-1B parameter range which are definitely better in many ways. As an illustration, below are real transcripts of Femto v1 in the CLI, taken from the checks the agent ran on the final checkpoint: the prompts it used for every iteration, sampled with the CLI’s defaults. Search results and the weather response are abbreviated.

Femto v1 (rl_web7), seed 0. Tools run in a fresh sandbox per conversation; search runs over a local Wikipedia index and weather is simulated, with a flag for the live API.

Challenges

While the examples make the model feel, for lack of a better word, underwhelming, there were several constraints that still made it an interesting exercise. First and most important was compute. I am fortunate to have access to decent enough compute for an academic setting, but in the grand scheme of things it still leaves a lot to be desired. For the most part I was constrained to single-node jobs (with 4 A100/H100s) across clusters (a budget of roughly 150 GPU-hours). cluv, a tool made by the incredible Mila IDT, made managing the compute easier but still needed adequate setup. Beyond compute, I also wanted to completely avoid any distillation. Distillation is an easy way to train small models given access to strong larger models. However, in the spirit of treating this as a pipeline of training a model from scratch, I decided to avoid relying on larger models in the pipeline.Unfortunately, in 2026 it is surprisingly hard to ensure that no LLM-generated text enters your training data. Many of the existing curated datasets are built on recent web crawls that contain a lot of AI-generated content that is hard to filter out. The final constraint was limited (Max 5x) access to the Claude agent with no additional API budget.

Learning to Trust the Process

Four days of work. Blue bars: Claude model calls per hour, from the lead session and its subagents. Orange dots: my messages, by kind. Grey bars: workstreams, each usually one agent on its own branch. Shaded bands are midnight to 8 am.

Let’s talk through the process of building this model. Most of my interactions were with a single “lead” session, but most of the work was delegated in pieces to different subagents. Each subagent operated on a different branch or worktree and was allowed to submit experiments to Slurm through cluv. The plans and research notes in the repository acted as shared memory between agents, and across the two session compactions.

Claude prepared a helpful breakdown of my “contributions”:

• Direction (20): the goal, the constraints, and scope changes.
• Decisions (23): choosing between options the agent laid out.
• Status checks (28): checking on the progress of various subagents and runs.
• Logistics (19).

Beyond high-level goals and design choices, I left much of the implementation to the agent. In my experience, it is much better to let the agent figure out the details rather than micromanaging them, as that is likely “out-of-distribution” relative to its largely autonomous training. The agent was quite a bit more rigorous in how it wrote code than I would be for a quick idea, but at some point it reached the “going overboard” stage where I had to ask it to do a complete cleanup. The final codebase ended up being around 9k LOC, with a lot of it being the eval, data, sandbox infra and inference engine.

Model design choices

To recap, there is nothing particularly special about the model; the goal of the exercise was mainly to see how much a relatively unsupervised agent can implement. Nevertheless, below is a summary of the model design and reasoning for certain choices made. (This section is mainly written by the agent and lightly edited by me.)

Design details

Base architecture

The starting point was Andrej Karpathy’s nanochat , ported to plain JAX with no Flax or optax, with Muon and AdamW implemented directly. Femto is a 10-layer, 512-wide decoder with an untied 16,384-token vocabulary, 48,234,516 parameters in total; three sliding-window layers alternate with one global layer.

Architecture of the 48.2M-parameter model. Left: the 10-layer SSSL stack alternates three 512-token sliding-window RoPE layers (S, blue) with one global causal NoPE layer (L, orange), with the normalized input embedding x₀ mixed back into every block. Right: internal structure of each 3.15M-parameter block. Bottom: parameter allocation across the untied embeddings, attention projections, and ReLU² MLPs.

The first baseline run was on the cluster within an hour of starting the project. With a 2.5B-token baseline pretraining corpus for faster iteration, it scored 0.116 on CORE (GPT-2 124M gets 0.114 for reference). Training on 2.5B tokens took 15–60 minutes (depending on the GPU) and we maintained this setup for many of the ablations downstream. In addition to CORE, the agent also built a TOOL benchmark to measure agentic performance.

No positional encoding on the global layers

One of the first architectural questions that I wanted to investigate was the choice of positional encodings on the three global layers (while the seven sliding-window layers keep RoPE to handle local token order). The agent designed a small task to ablate the potential choices: RoPE , NoPE and partial RoPE .

Three positional-encoding arms trained at 2,048 tokens and evaluated out to 32k with no long-context training. Left: passkey retrieval by context length and by how deep in the context the key sits (1.0 = right before the question). Right: loss by position in long documents, with eval-time attention scaling on the global layers. At 2k the three arms are indistinguishable (val bpb within 0.001).

At the training length the arms tie, but beyond it RoPE degrades while NoPE does not: it retrieves the key at every depth out to 32k. We kept NoPE, which is also why the later 16k stage needed no RoPE rescaling . The heatmap shows a second thing the agent caught. Partial RoPE fails only at depth 1.0, where the key is the last thing before the question; the model there continues the prompt (“…is the pass key. What is the pass key? The pass key is” → ” the pass key.”). Passkey is a template test as much as a memory test, so long-context evaluation moved to numbers hidden in real documents.

Code, and then agentic data

The next important question was how to give the base model some world knowledge while also making it a good base for tool-use training. The agent trained a code-aware tokenizer (12–29% better compression on shell and Python, 0.6% worse on web text, no change in loss) and did a sweep of the code fraction in the main pretraining mix.

The code-share sweep: ClimbMix web text plus shell, Python and sysadmin Q&A at 0–20% of documents. Drag the slider.

The tool metrics keep improving with the code share in the mix, so the final run used 20%. In addition to the code mix in pretraining, the tool-calling abilities also came from the anneal (the last 15% of training) where the learning rate decays and the mix switches to higher-quality data :

Last 15% of training	CORE	TOOL	emits a tool call	well-formed
(a) keep the main mix	0.106	0.157	0%	–
(b) high-quality text	0.111	0.162	0%	–
(c) (b) plus 9% agentic data	0.111	0.251	93%	72%

The agentic data consists of terminal-agent trajectories (Nemotron-Terminal-Corpus, OpenThoughts-Agent) compacted into Femto’s chat format, natural-language-to-bash pairs , and small shell and Python tasks that were actually executed so the model sees real outputs. Arms (a) and (b), given a chat-formatted task, answer in prose; their loss on the special tool-call tokens is about 25 nats/token (near the logit-softcap ceiling).

Final pretraining run

The final pretraining run used NoPE global layers, the v2 tokenizer, 20% code and the agentic anneal, with a total of 10B tokens, and ran on two A100s in four hours. It reached CORE 0.127. A 315M-token stage at 16k context, chosen from eight arms (learning rate, document masking, training-time attention scaling, twice the tokens, 32k), made the loss keep falling with position out to 32k and lifted retrieval of four numbers hidden in real 32k-token documents from 0.18 to 0.76.

The 10B base model before and after the 16k stage (learning rate ×0.3 and ×0.1). Without it the loss rises past 8k; with it, the loss keeps falling to 32k. The shaded band is beyond the 16k training length.

Because 20% of the pretraining mix went to code rather than web text, the base model’s headline CORE gain over the web-only 2.5B baseline came mostly from BoolQ, whose centered score swings with the model’s yes/no bias; without it, scaling to 10B mainly recovered the web-knowledge drop from adding code.

Tool use

The base model emitted well-formed tool calls in every terminal episode and solved 7% of the tasks. The agent planned the post-training pipeline from a literature review it ran while pretraining finished.

Mid-training (0.9B tokens of replay, agentic trajectories, chat, math written as Python , multiple choice) more than doubled terminal success and took GSM8K-with-Python from 0.2% to 9%, at the cost of about 0.013 on CORE. A short SFT stage after it made everything worse. Greedy decoding looped on open-ended prompts, which was fixed with sampling (temperature 0.7, top-p 0.9, a light repetition penalty in the CLI).

RL used vanilla GRPO in nanochat’s simplified form, with real tool execution (every rollout runs its bash and Python in a fresh sandbox) and reference-solution injection when an entire rollout group fails so the gradient does not vanish. One early decision kept it stable: the KV-cached inference engine reuses the trainer’s forward function in fp32, where log-prob mismatch between rollout decoding and training was mostly non-existent. Throughout post-training, the agent fixed various data contamination issues and bugs in the initial code over several test runs.

Breaking the zero-pass plateau. (a) Two seeds of plain GRPO stop at 0.56 held-out terminal success; with reference solutions injected into groups that have no pass, every trained task type is solved by step ~60. (b) Almost all of the ~230 injected references come in the first 30 steps.

Beyond simple terminal use, I also wanted the model to be good at some general tool calling, including things like search and basic API calls, and four subagents built the required pieces in parallel: a BM25 index over 20.8M Wikipedia passages, a weather simulator with Open-Meteo’s response format, function-calling data from existing datasets , and evaluations (BFCL-style function calling , open-domain QA, weather). RL then trained the model on a mixture of all the tasks. There were several reward-hacking behaviors that the agent uncovered over the course of training and fixed. Most of the fixes came from fixing the reward and related infra. Overall, mid-training and RL combined accounted for only ~10% of the FLOPs of the final run.

RL results for the final recipe: the tool mid-trained model, after 600 RL steps (rl_web5) and after 300 more (Femto, rl_web7), on greedy, execution-checked evaluations. Hover a row for what it measures. RL is about 1% of the training FLOPs but took about half as many GPU-hours as pretraining: it runs in fp32, decodes token by token and executes every tool call in a sandbox. Every post-training run is a single seed.

What did I learn?

While the final model and recipe is nothing special, being able to get to a working codebase written from scratch in such a short time is quite powerful. I have been able to push new ideas beyond what is discussed here and the model (and the infra around it) has only gotten better in the time that it took me to finally get to writing this post (some of these will probably appear here in future posts). The setup has also been useful to reproduce findings in the community at a smaller scale – validating its utility as a testbed.

Some of my takeaways from this experiment are below.

• Agents are really really good at writing code. Something at this scale would have certainly taken me much longer to write from scratch by myself so they sped things up quite a lot. I am certainly not the only one to have noticed this, with results in even larger settings.

• The current generation of agents lacks good experiment-design taste (and tends to set up many unnecessary experiments) and is very susceptible to getting stuck in dead ends. They latch onto things mentioned in passing and never revisit them unless explicitly prompted. The agent would take paper findings without any skepticism or uncertainty or if a cluster was down temporarily it would never consider checking and scheduling jobs unless prompted.

• I did learn quite a bit about parts of the training stack that I was less familiar with, but I did end up with the feeling of having mostly been the audience in the agent’s show. If I had been working on the project more hands-on I would have certainly learned more about the details, but maybe would have missed part of the bigger picture. The trade-off is hard to quantify but it is hard not to feel that something was lost in this process.

• Collaborating with humans is still a lot more fun than working alone with agents.