Skip to content
 
 

Repository files navigation

autoresearch

teaser

Autoresearch lets an AI agent run autonomous LLM training experiments overnight, even if you don't have a local GPU.

By leveraging Modal's cloud infrastructure, you can effortlessly spin up an on-demand H100 Sandbox, upload this codebase, and let your AI agent run actual neural network experiments directly in the cloud.

The agent modifies the code, pushes it to the remote Sandbox, trains for a fixed 5-minute budget, checks if the validation loss (val_bpb) improved, keeps or discards the changes, and repeats. You wake up in the morning to a log of experiments and (hopefully) a better model.

The core idea is that you're not touching Python files. Instead, you program the program.md file that provides instructions to the AI agent. This training code is a single-GPU implementation of nanochat.

How it works

The repo is deliberately kept small and only really has a three files that matter:

  • prepare.py — fixed constants, one-time data prep (downloads training data, trains a BPE tokenizer), and runtime utilities (dataloader, evaluation). Not modified.
  • train.py — the single file the agent edits. Contains the full GPT model, optimizer (Muon + AdamW), and training loop. Everything is fair game: architecture, hyperparameters, optimizer, batch size, etc. This file is edited and iterated on by the agent.
  • run_sandbox.py — custom Modal script added to spin up an interactive, cloud-based H100 Sandbox environment. This is the magic that allows you to run high-end GPU experiments from any laptop.
  • program.md — baseline instructions for one agent. Point your agent here and let it go. This file is edited and iterated on by the human.

By design, training runs for a fixed 5-minute time budget (wall clock, excluding startup/compilation), regardless of the details of your compute. The metric is val_bpb (validation bits per byte) — lower is better, and vocab-size-independent so architectural changes are fairly compared.

If you are new to neural networks, this "Dummy's Guide" looks pretty good for a lot more context.

Quick start

Requirements: You must have a Modal account and Python 3.10+ with uv installed.

# 1. Install dependencies (this includes modal)
uv sync

# 2. Authenticate with Modal, you should register a Modal account first and then authorize it to the terminal app.
uv run modal setup

# 3. Spin up the sandbox (it will run for 24 hours allowing you to exec commands), don't close the terminal.
uv run modal run run_sandbox.py

# 4. In a new terminal, connect to the sandbox using the Sandbox ID printed in the previous step
# e.g. modal exec sb-xxxx bash, and then run `python prepare.py` followed by `python train.py`

If the above commands all work ok, your setup is working and you can go into autonomous research mode. Note that the run_sandbox.py script specifically provisions an H100 with a mounted volume at /root/.cache/autoresearch to ensure the dataset cache persists across sandbox lifecycles. It also copies your local repo into the Sandbox at creation time.

Running the agent

Simply spin up your Claude/Codex or whatever you want in this repo (and disable all permissions), then you can prompt something like:

Hi have a look at program.md and let's kick off a new experiment! let's do the setup first.

The program.md file is essentially a super lightweight "skill".

Project structure

prepare.py      — constants, data prep + runtime utilities (do not modify)
train.py        — model, optimizer, training loop (agent modifies this)
run_sandbox.py  — script to start a Modal Sandbox for remote H100 execution
program.md      — agent instructions
pyproject.toml  — dependencies

Design choices

  • Single file to modify. The agent only touches train.py. This keeps the scope manageable and diffs reviewable.
  • Fixed time budget. Training always runs for exactly 5 minutes, regardless of your specific platform. This means you can expect approx 12 experiments/hour and approx 100 experiments while you sleep. There are two upsides of this design decision. First, this makes experiments directly comparable regardless of what the agent changes (model size, batch size, architecture, etc). Second, this means that autoresearch will find the most optimal model for your platform in that time budget. The downside is that your runs (and results) become not comparable to other people running on other compute platforms.
  • Self-contained. No external dependencies beyond PyTorch and a few small packages. No distributed training, no complex configs. One GPU, one file, one metric.

Platform support

This code currently requires that you have a single NVIDIA GPU. In principle it is quite possible to support CPU, MPS and other platforms but this would also bloat the code. I'm not 100% sure that I want to take this on personally right now. People can reference (or have their agents reference) the full/parent nanochat repository that has wider platform support and shows the various solutions (e.g. a Flash Attention 3 kernels fallback implementation, generic device support, autodetection, etc.), feel free to create forks or discussions for other platforms and I'm happy to link to them here in the README in some new notable forks section or etc.

Seeing as there seems to be a lot of interest in tinkering with autoresearch on much smaller compute platforms than an H100, a few extra words. If you're going to try running autoresearch on smaller computers (Macbooks etc.), I'd recommend one of the forks below. On top of this, here are some recommendations for how to tune the defaults for much smaller models for aspiring forks:

  1. To get half-decent results I'd use a dataset with a lot less entropy, e.g. this TinyStories dataset. These are GPT-4 generated short stories. Because the data is a lot narrower in scope, you will see reasonable results with a lot smaller models (if you try to sample from them after training).
  2. You might experiment with decreasing vocab_size, e.g. from 8192 down to 4096, 2048, 1024, or even - simply byte-level tokenizer with 256 possibly bytes after utf-8 encoding.
  3. In prepare.py, you'll want to lower MAX_SEQ_LEN a lot, depending on the computer even down to 256 etc. As you lower MAX_SEQ_LEN, you may want to experiment with increasing DEVICE_BATCH_SIZE in train.py slightly to compensate. The number of tokens per fwd/bwd pass is the product of these two.
  4. Also in prepare.py, you'll want to decrease EVAL_TOKENS so that your validation loss is evaluated on a lot less data.
  5. In train.py, the primary single knob that controls model complexity is the DEPTH (default 8, here). A lot of variables are just functions of this, so e.g. lower it down to e.g. 4.
  6. You'll want to most likely use WINDOW_PATTERN of just "L", because "SSSL" uses alternating banded attention pattern that may be very inefficient for you. Try it.
  7. You'll want to lower TOTAL_BATCH_SIZE a lot, but keep it powers of 2, e.g. down to 2**14 (~16K) or so even, hard to tell.

I think these would be the reasonable hyperparameters to play with. Ask your favorite coding agent for help and copy paste them this guide, as well as the full source code.

Notable forks

License

MIT

About

AI agents running research on Modal sandbox nanochat training automatically for no-local-GPU users

Resources

Stars

Watchers

Forks

Releases

Packages

Contributors

Languages