google-deepmind/simply

google-deepmind/simply

Description: Minimal and scalable research codebase in JAX, designed for rapid iteration on frontier research in LLM and other autoregressive models.

Language: Python

License: Apache-2.0

Stars: 530

Forks: 63

Open issues: 16

Created: 2025-03-13T05:18:20Z

Pushed: 2026-06-03T23:21:41Z

Default branch: main

Fork: no

Archived: no

README:

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](#getting-started)* 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](#automated-ai-research-with-agents) for some simple examples. More on the way.
• That's it, *simply* [get started](#getting-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-running-on-gke-with-xpk) to run your first experiment on GKE, and [GKE section in the full GCloud guide](docs/gcloud.md#11-running-on-gke-with-xpk) for details.

Automated AI research with agents

You can use agents like Google Antigravity, Claude Code, or Gemini CLI to run automated research experiments. For example, paste the following prompt into your agent from the repo root to have it design and benchmark new optimizers on a toy setting:

You are an AI research agent.
Design and benchmark new optimizers for training a small transformer.
Read simply/utils/optimizers.py to understand the interface.
First, run the Adam baseline using the lm_test config and record the final loss.
In multiple research iterations, propose novel optimizers, tune their hyperparameters, and run experiments to compare with the baseline.
In each iteration, you can propose and run 3 experiments, wait for results and then start the next iteration based on the results of finished experiments. Keep running new iterations until you have finished 15 experiments or proposed 10 new optimizers.
Write a report to /tmp/optimizer_report.md.

See the [full guide](docs/automated_ai_research_example.md) for more interesting examples including RL algorithm search for post-training. Have fun playing around with different prompts and interacting with the agent for longer research.

Simply also includes a built-in agent harness (simply/agent/) with Bash tool and context management for autonomous long-running research tasks. Quick test:

pip install ".[agent]"
python -m simply.agent.main \
--task_file=simply/agent/example_tasks/code_stats.md \
--env="Local:." \
--llm="LiteLLM:vertex_ai/gemini-2.5-pro"

See the [agent README](simply/agent/README.md) for setup and configuration details.

Dependencies

The main dependencies are: Jax for model and training. Orbax for checkpoint management. Grain for data pipeline.

Install dependencies:

# JAX installation is environment-specific. See https://docs.jax.dev/en/latest/installation.html
# CPU:
pip install -U jax
# GPU:
pip install -U "jax[cuda13]"
# TPU:
pip install -U "jax[tpu]"

# Install simply and its dependencies:
pip install .
# With optional dependencies:
pip install ".[tfds]" # for TensorFlow Datasets
pip install ".[math-eval]" # for simply/utils/math_eval.py
pip install ".[dev]" # for testing (pytest)

Setup Model Checkpoints and Datasets

Download datasets and model checkpoints in format supported by Simply from HuggingFace:

# Install huggingface_hub
pip install huggingface_hub

# Download both models and datasets
python setup/setup_assets.py

# Or download only models/datasets
python setup/setup_assets.py --models-only
python setup/setup_assets.py --datasets-only

This will download models to ~/.cache/simply/models/ and datasets to ~/.cache/simply/datasets/. You can customize locations with --models-dir and --datasets-dir flags, or set environment variables SIMPLY_MODELS and SIMPLY_DATASETS. (Currently we only included a few datasets and models for testing, and will add more soon.)

Citation

If you find *Simply* helpful, please cite the following BibTeX:

@misc{Liang2025Simply,
author = {Chen Liang and Da Huang and Chengrun Yang and Xiaomeng Yang and Andrew Li and Xinchen Yan and {Simply Contributors}},
title = {{Simply: an experiment to accelerate and automate AI research}},
year = {2025},
howpublished = {GitHub repository},
url = {https://github.com/google-deepmind/simply}
}

Contributors list: Chiyuan Zhang, Robert Dyro, Alex Zhai, Xingjian Zhang, Jiaxi Tang, Lizhang Chen, Ran Tian

License

Copyright 2025 Google LLC

All software is licensed under the…