Example scripts

The examples/ folder holds complete programs, each a single file you can copy and change. The descriptions below are the scripts’ own docstrings, so they match the code in the repository. Every script has a command-line interface; run it with --help to see the options.

The first group trains one kind of model on one dataset and is the quickest way to see a whole workflow. The second group runs a full job, from data to scored or exported weights. Each of those takes a --smoke flag that swaps the real settings for the repository’s tiny fixtures, a few steps and one CPU device, which is how the test suite runs them. End-to-end runs gives both command lines for each.

Train one model

readme_demo.py

Train a toy language model, resume it, learn preferences, and generate images.

Run from an installed Dew checkout, without downloads: JAX_PLATFORMS=cpu python examples/readme_demo.py —out runs/readme-demo

Use a new output directory for each invocation. The generated data demonstrates mechanics, not language ability, preference quality, or useful image generation.

Source on GitHub

train_lm.py

Train a byte-level language model on a directory of token files, then generate.

curl -o data/shakespeare.txt —create-dirs \ https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt python tools/tokenize_text.py —input data/shakespeare.txt —out data/shakespeare —tokenizer byte python examples/train_lm.py —tokens data/shakespeare —epochs 4 python examples/train_lm.py —tokens data/shakespeare —steps 20 —sequence-length 32 # smoke run

Source on GitHub

train_flowers.py

Train a pixel-space diffusion model on prepared Oxford Flowers ArrayRecords.

CUDA_VISIBLE_DEVICES=0 JAX_PLATFORMS=cuda python examples/train_flowers.py \
--data ~/.cache/dew/datasets/oxford_flowers102/2.1.1 --steps 1000

Source on GitHub

train_diffusion.py

Train a text-to-image diffusion model on Oxford Flowers, sample from it, export the weights.

python examples/train_diffusion.py --data-path /data/oxford_flowers102/2.1.1 --epochs 200
python examples/train_diffusion.py --data-path /data/oxford_flowers102/2.1.1 --steps 20 --image-size 32

Source on GitHub

train_jepa.py

Train an I-JEPA encoder on Oxford Flowers, probe it, save the encoder.

python examples/train_jepa.py --data-path /data/oxford_flowers102/2.1.1 --epochs 300
python examples/train_jepa.py --data-path /data/oxford_flowers102/2.1.1 --steps 20 --image-size 32 --patch-size 4

Source on GitHub

Run a whole job

train_flowers_tpu.py

Text-to-image DiT from scratch on Oxford Flowers, across a TPU slice.

Prepare the records once, at the resolution the run trains at, so a training read is a memcpy rather than a JPEG decode:

python -c "import tensorflow_datasets as tfds; tfds.builder('oxford_flowers102', \
    data_dir='~/.cache/dew/datasets').download_and_prepare(\
    file_format='array_record')"
python tools/prepare_images.py --dataset oxford_flowers102 \
    --data-path ~/.cache/dew/datasets/oxford_flowers102/2.1.1 \
    --split all --image-size 256 --out prepared/flowers-256

Then launch the same file on every worker of the slice:

python examples/train_flowers_tpu.py --data prepared/flowers-256 --steps 200000

--data reads whichever of the two layouts it is given: the TFDS version directory loads as oxford_flowers102, the prepare_images.py output as array_record_images. The smoke run writes a handful of synthetic records in that second layout and trains on them on one CPU device:

JAX_PLATFORMS=cpu python examples/train_flowers_tpu.py --smoke --out /tmp/flowers-smoke

Source on GitHub

sft_gemma4.py

Full-weight SFT of a Gemma 4 text decoder on a Hub chat dataset.

python examples/sft_gemma4.py --model google/gemma-4-E2B \
    --dataset allenai/tulu-3-sft-mixture --steps 4000 --out runs/gemma4-sft

The run packs whole conversations into windows, counts the loss on assistant targets alone, shards the weights over the visible devices and accumulates micro-batches into one update. It ends by exporting the trained weights to the Hugging Face layout, so transformers and load_pretrained both read them, and the run directory itself scores through the harness:

python -m dew.eval --model dew --model_args run=runs/gemma4-sft/gemma4-sft \
    --tasks hellaswag --limit 64

dew.data.ChatMessages reads the Hub dataset itself, so the id on the command line is what the run trains on: it renders every conversation with the checkpoint’s own chat template and packs them into windows.

JAX_PLATFORMS=cpu python examples/sft_gemma4.py --smoke --out /tmp/gemma4-smoke

Source on GitHub

sft_diffusion_gemma.py

LoRA SFT of DiffusionGemma on chat data, with the base weights host-streamed.

The adapter is the only thing the optimizer moves, and Layout(host=("params",)) keeps the whole train state on the host between steps, so a 26B-A4B base fits beside its factors on one accelerator:

python examples/sft_diffusion_gemma.py \
    --model google/diffusiongemma-26B-A4B-it \
    --chat data/tulu-3-sft.parquet --steps 2000

--chat is a Hub dataset id, a .jsonl file or a parquet file of conversations, which is what dew.data.ChatMessages reads and renders with the checkpoint’s own chat template. The run writes two things: the PEFT adapter directory LoRA.save produces, which transformers loads, and the merged checkpoint in the source’s own layout, which dew.pipeline generates from.

JAX_PLATFORMS=cpu python examples/sft_diffusion_gemma.py --smoke --out /tmp/dg-smoke

Source on GitHub

train_rlvr.py

RLVR: GRPO on programs that must pass their tests, rolled out asynchronously.

python examples/train_rlvr.py --backend native --steps 40 --out runs/rlvr-native
python examples/train_rlvr.py --backend vllm --steps 40 --out runs/rlvr-vllm
python examples/train_rlvr.py --backend sglang --steps 40 --out runs/rlvr-sglang

Every prompt asks for a Python program that reads two integers from stdin and prints a stated function of them. A completion’s reward is the fraction of three hidden test cases its program passes, run in a SandboxFleet. --runner container runs each program in a network-less python:3.12-slim container. --runner process (the default, and the only choice where Docker is absent, as on Colab) runs it as a process with a wall clock, CPU time and memory cap; that process has your user’s filesystem and network.

Rollouts run on a rollout server while the trainer updates. --backend native serves them from Dew’s own continuous-batching Server in this process, with weights pushed in place. --backend vllm and --backend sglang start that engine’s OpenAI-compatible server on an export of the same checkpoint, sample from it by token ids, and push weights by writing safetensors and asking the engine to reload them (vLLM’s development endpoints, VLLM_SERVER_DEV_MODE=1; SGLang’s /update_weights_from_disk). --vllm and --sglang name the executables, which may live in their own environments. A RolloutScheduler over a PromptSource draws one batch ahead of the update, so each batch is at most one update stale; the GRPO objective’s importance cap corrects for it. A completion that runs out of --new-tokens is still scored and trained on (truncation="score"): a closed code block followed by cut-off prose can pass every test.

--turns N gives each task up to N attempts through an EnvironmentSource, Dew’s in-process multi-turn session source. An attempt that fails a test gets back how many of the three it passed and tries again; the reward is the last attempt’s. A harness would render the report as a user turn with the chat template instead, and tools/audit_template.py tells whether that stays append-only.

The run prints one line per update and writes rewards.json to --out with the per-update reward, policy version and lag, the mean reward of the first and last --window updates, and the seconds each weight push took.

JAX_PLATFORMS=cpu python examples/train_rlvr.py --smoke --out /tmp/rlvr-smoke
JAX_PLATFORMS=cpu python examples/train_rlvr.py --smoke --turns 2 --out /tmp/rlvr-smoke-turns

Source on GitHub

train_harbor.py

Agentic GRPO: a harness runs in Harbor sandboxes, a gateway records its model calls, and Dew trains on them.

python examples/train_harbor.py --tasks path/to/harbor/task ... --gateway http://127.0.0.1:9090 \
    --sandbox-gateway http://proxy:9091 --engines http://127.0.0.1:8011 --served runs/harbor/served
JAX_PLATFORMS=cpu python examples/train_harbor.py --smoke --out /tmp/harbor-smoke

Each update draws --prompts Harbor tasks and runs --groups trials of each through HarborSource: Harbor starts the task’s sandbox and the harness (mini-swe-agent by default), the harness’s model calls go through rllm-model-gateway to the engines, and the gateway’s traces become the session’s calls. RolloutScheduler admits complete groups at most --max-lag updates stale and packs them; the GRPO objective trains on the sampled ids only. Publication pushes each version to every engine with SafetensorsReload and then stamps the gateway, so every call carries the version it was sampled under.

The example writes the policy to --served and then waits, up to --ready-timeout seconds, for the gateway to route to an engine: launch the engines on --served (vLLM with VLLM_SERVER_DEV_MODE=1) and the gateway in front of them while it waits, as docs/concepts/post_training.md describes (warm each engine, size the gateway’s health interval, keep its admin routes away from the sandboxes).

--smoke needs neither Harbor nor an engine. A stand-in harbor script runs a two-turn harness against an in-process gateway whose engine is Dew’s own server on the committed tiny Qwen2, so the ids and likelihoods it records are the policy’s own draws. The reward is the share of vowels in the replies. Two updates on CPU.

Source on GitHub

evaluate_and_serve.py

Score a finished run four ways, then compare it against a served model.

Perplexity over held-out tokens, an lm-evaluation-harness suite, greedy continuations, and, for a diffusion run, FID and CLIPScore of a sampled grid against a reference set. Everything is read through dew.pipeline, so a run directory, a published checkpoint and a Hub repository all work.

python examples/evaluate_and_serve.py --run runs/shakespeare/lm-shakespeare \
    --tokens data/shakespeare --tasks hellaswag arc_easy --harness-limit 200
python examples/evaluate_and_serve.py --run runs/shakespeare/lm-shakespeare \
    --image-run runs/flowers-tpu/checkpoints/flowers-256 \
    --reference-images data/flowers-heldout

The same suites run from the command line, which is the harness’s own entry point with Dew’s model registered:

python -m dew.eval --model dew --model_args run=runs/shakespeare/lm-shakespeare \
    --tasks hellaswag --limit 200

--openai-base-url adds a served comparison: the same prompts through the OpenAI SDK (a vLLM endpoint speaks it too), or --ollama-host through ollama’s. Both are optional extras; without them, or with the endpoint unreachable, the report says so and the rest of the run is unaffected. The OpenAI key is OPENAI_API_KEY when set; a vLLM server started without --api-key takes any key, so an unset one sends vLLM’s placeholder.

JAX_PLATFORMS=cpu python examples/evaluate_and_serve.py --smoke --out /tmp/eval-smoke

Source on GitHub