flaxchat

flaxchat

A minimal, end-to-end LLM training harness for Google Cloud TPU pods, built on JAX/Flax NNX.

Port of nanochat to the JAX ecosystem with full feature parity plus speculative decoding.

Part of the 2026 Q1 TPU Research Sprint, supported by the Google AI Developer Programs team.

Quick Start

pixi install
pixi run test          # 148 tests
pixi run -- python -m scripts.run_tinystories --depth=8 --steps=5000

Core Features

GPT Architecture (all nanochat features)

• Rotary Embeddings (RoPE), Group-Query Attention (GQA), QK Normalization
• ReLU^2 MLP, Value Embeddings (ResFormer), Sliding Window (SSSL pattern)
• Per-layer Residual Scaling, Smear, Backout, Logit Soft-capping
• Gradient Checkpointing (nnx.remat, dots_saveable)

Optimizer

• Mixed Muon + AdamW via optax.multi_transform
• Muon: Polar Express orthogonalization + NorMuon variance reduction
• Per-group LR, betas, weight decay. Warmup -> constant -> warmdown schedule

Inference Engine (4 modes)

Tool Use (streaming)

engine = Engine(model, tokenizer)
for tokens, masks in engine.generate(prompt_ids, num_samples=3):
    # Automatic: calculator for math, sandboxed Python for code
    pass

Sandboxed Code Execution

from flaxchat.execution import execute_code
result = execute_code("print(sum(range(10)))", timeout=5.0)
# ExecutionResult(success=True, stdout="45\n")

Parallelism (default, not optional)

mesh = compute_init()  # auto mesh over ALL devices
# Data parallel, FSDP, multi-host — all automatic via JAX SPMD

Depth-Based Config

config = FlaxChatConfig.from_depth(depth=12)
# -> 12 layers, 768 dims, 6 heads, ~79M params

Pipeline

Evaluation Tasks

Test Suite

148 tests across 10 files covering model, engine (all 4 gen modes + speculative + tool use), optimizer, config, evaluation, execution sandbox, checkpoints, tokenizer, dataloader, and distributed utilities.

pixi run test  # all tests pass on CPU, GPU, and TPU

Remote Execution

Kaggle GPU (via kgz)

from flaxchat.remote import KaggleRunner
runner = KaggleRunner("https://...")
runner.run_pipeline(depth=8, steps=5000)

GCP TPU (via tpuz)

from tpuz import TPU
tpu = TPU("my-tpu", accelerator="v6e-8")
tpu.up()
tpu.setup(extra_pip="flaxchat")
tpu.run("python -m scripts.pretrain --depth=12", sync=".")

Verified Results

Full Pipeline: Pretrain -> SFT -> RL (Kaggle TPU v5e-8)

Model: 12L/768d/6h (GQA: 3kv) = 203.7M params

Chinchilla Scaling Law (TRC TPU v6e-8)

Nanochat architecture trained at Chinchilla-optimal budgets (20x params) on C4:

Baselines

W&B: irf-sic/flaxchat

Acknowledgments

This project is part of the 2026 Q1 TPU Sprint, supported by the Google AI Developer Programs team.

• Google AI Developer Programs for issuing GCP credits
• TPU Research Cloud (TRC) for providing free access to Cloud TPU v4, v5e, and v6e
• Kaggle for free TPU v5e access for prototyping

Built on nanochat, JAX, Flax, Optax, Orbax, tpuz, kgz.

Blog Posts

• Training a 19M Parameter Story Generator from Scratch on 2xT4 GPUs — April 01, 2026 We trained a GPT language model from scratch on the TinyStories dataset using flaxchat, our JAX/Flax NNX port of nanochat. The entire pipeline ran end-to-end on 2x NVIDIA T4 GPUs…

Posts

• Apr 1, 2026

  Training a 19M Parameter Story Generator from Scratch on 2xT4 GPUs

subscribe via RSS