Simply

Simply: Minimal Code for End-to-End Frontier LLM Research

*Simply* is a minimal and scalable research codebase in JAX, designed as an environment where both humans and AI agents can rapidly iterate on frontier LLM research.

• Quick to fork and hack for fast iteration. We aim at minimizing the time to implement new ideas (e.g., optimizer, training loss, RL algorithms, etc) by humans and AI agents.
• Minimal abstractions and dependencies for a simple and self-contained codebase. Learn Jax, and you are ready to read and hack the code.
• An environment for automated AI research — An AI agent, which can itself be powered by an LLM served with Simply, can read the code, propose new ideas, run experiments, and iterate autonomously or under the guidance of human researchers. See automated AI research with agents for some simple examples. More on the way.
• That's it, simply get started with hacking now :)

Getting started

Example commands

Local test for debug

EXP=simply_local_test_1; rm -rf /tmp/${EXP}; python -m simply.main --experiment_config lm_test --experiment_dir /tmp/${EXP} --alsologtostderr

Or if you want to debug by printing arrays like normal python code, you can disable `jit` and `use_scan` using the command below.

export JAX_DISABLE_JIT=True; EXP=simply_local_test_1; rm -rf /tmp/${EXP}; python -m simply.main --experiment_config lm_no_scan_test --experiment_dir /tmp/${EXP} --alsologtostderr

Running on Google Cloud TPUs

See the [GCloud Quickstart](gcloud_quickstart.md) to run your first experiment on a Cloud TPU, or the [full GCloud guide](docs/gcloud.md) for multi-host training, preemption handling, and monitoring.

Running on GKE with XPK

Google Kubernetes Engine (GKE) is supported. See [GKE quick start](gcloud_quickstart.md#optional-runn