Train models with JAX and Flax.

A model, an objective, data and a trainer

The model is a plain Flax module. The objective says how to initialize it and what the loss is. The trainer owns everything else, so a new loss is a new objective and the training loop stays as it is.

1. ModelA flax.linen.Module. Built-in ones come from models.build(name, ...).
2. Objectiveinit(key) returns the variables, loss(variables, batch, step) the loss and metrics.
3. DatasetTwo iterator factories, for training batches and one validation pass.
4. TrainerCompiles the step, places the state on the mesh, updates the EMA, writes checkpoints and logs.

regression.py

import itertools

import flax.linen as nn
import jax
import jax.numpy as jnp
import numpy as np
import optax

from dew import Aux, Dataset, Field, InputSpec, Objective, Trainer


class Regression(Objective):
    model = nn.Dense(1)
    inputs = InputSpec(Field("x", (1,)))

    def init(self, key, variables=None):
        return self.model.init(key, jnp.ones((1, 1)))

    def loss(self, variables, batch, step):
        prediction = self.model.apply(variables, batch["x"])
        loss = jnp.mean((prediction - batch["y"]) ** 2)
        return loss, Aux(metrics={"mse": loss})


x = np.linspace(-1, 1, 32, dtype=np.float32).reshape(32, 1)
batch = {"x": x, "y": 2 * x + 1}
data = Dataset(train=lambda partition: itertools.repeat(batch),
               val=None, records=32, batch=32)
objective = Regression()
trainer = Trainer(objective, optax.sgd(0.1), key=jax.random.key(0))
state = trainer.fit(data, steps=100, log_every=50)
print(objective.model.apply(state.params, jnp.array([[0.0], [1.0]])))

Output

Training from step 0 to 100 on {'data': 1, 'expert': 1, 'fsdp': 1, 'tensor': 1, 'sequence': 1, 'stage': 1} (1 process(es))
step 50: loss 0.0002
step 100: loss 0.0000
Goodput: first step after 0.75 s, 20.5% of the wall time in steps
[[0.99999994]
 [2.999413  ]]

What you can train

Each of these is an objective in the library, trained by the same trainer on the same mesh.

Language models

Pretraining on token windows or packed documents, chat fine-tuning that counts only the assistant tokens, and cached generation.

Train a language modelGuide

Image and video diffusion

Diffusion, flow matching and latent diffusion on DiT, U-Net and MMDiT backbones, with 17 solvers and classifier-free guidance.

Train a diffusion modelProcesses and solvers

Diffusion language models

Masked diffusion over tokens (MDLM), LLaDA and Dream continued from their released weights, and block diffusion with Diffusion Gemma.

Guide

Representation learning

I-JEPA and V-JEPA: a context encoder, a moving-average target encoder, a predictor and block masks, with linear and k-NN probes.

Representations with I-JEPAGuide

Post-training

SFT, DPO, GRPO and PPO on the same trainer, asynchronous RL whose rollouts come from Dew, vLLM or SGLang, and Flow-GRPO for images.

Guide

Mixture of experts

Routed and shared experts, bias balancing, routing replay, grouped-matmul kernels, and expert parallelism over an expert mesh axis.

Concepts

Flowers sampled from noise by the 9M-parameter diffusion transformer that tutorial 02 trains in 6,000 steps.

ROMEO:
Tut Tybalt, Gracious fazous to deparal,
Your bring shall be find that is in the sign,
Why both his hardy nundrest to an all ment.

RICHMOND:
Let must be, my lord, if you that day:
But yea, sir, then time and with himself;
Her oft protection to his death to the supper your
mother married to men.

HERMIONE:
And not that his by the duke a heaveness.

Third I change it another canst thou beast;
When

Text from the 14M-parameter byte-level model of tutorial 05, after 1,500 steps on Tiny Shakespeare.

Seven solvers on one checkpoint and one set of starting noises, from tutorial 04.

Pretrained checkpoints, in and out

load_pretrained reads Hugging Face checkpoints of 33 decoder families into Dew's own modules, and Pretrained.save writes trained weights back in the source's layout. This list comes from the code's registries when the site is built.

Dense decoders

• Llama
• Mistral
• Ministral
• Qwen 2
• Qwen 3
• Qwen 3.5, text
• Gemma
• Gemma 2
• Gemma 3, text
• Gemma 3n, text
• Gemma 4, text
• OLMo 3

Mixture of experts

• Mixtral
• Qwen3-MoE
• Qwen 3.5 MoE, text
• gpt-oss
• Llama 4, text
• GLM 4 MoE
• GLM MoE with sparse attention
• DeepSeek V2
• DeepSeek V3
• DeepSeek V3.2
• DeepSeek V4
• Kimi K2
• Kimi K2.5

Hybrid and linear attention

• Qwen3-Next
• GLM 5 Next, text
• Kimi Linear
• Kimi K3, text
• Mamba-2

Diffusion language models

• LLaDA
• Dream
• Diffusion Gemma

Multimodal

• Gemma 3
• Gemma 4
• Gemma 4, unified
• Gemma 3n
• Qwen 3.5
• Llama 4

Diffusion pipelines

• Stable Diffusion
• Stable Diffusion XL
• Stable Diffusion 3
• Flux
• Qwen-Image 2.1
• Stable Diffusion, Flax weights
• Stable Diffusion XL, Flax weights

Each port is checked against its reference implementation in float32, most on fixtures with the release's own configuration and tensor shapes. Qwen3-0.6B, SmolLM2-135M and Mamba-2 130M were also compared at full size. Supported models lists every model_type and how it was checked.

One trainer, from one device to a mesh

MeshSpec names the axes: data, fsdp, expert,tensor, sequence and stage. The model names the logical axes of its parameters, and Layout maps them onto the mesh, so the model code does not change when the mesh does. Attention picks cuDNN on supported GPUs, Pallas splash attention on TPUs and XLA elsewhere.

trainer = Trainer(objective, optax.adamw(3e-4), key=jax.random.key(0),
                  mesh=MeshSpec(fsdp=4, tensor=2))

dew launch --hosts 10.0.0.1 10.0.0.2 -- /opt/dew/.venv/bin/python train.py

MeshSpec(fsdp=4, tensor=2) on eight devices. An MLP weight splits its model width over fsdp and its hidden width over tensor, so each device holds one of eight blocks.

What has been run

• The tutorials on one GPU or a CPU, and the test suite on a CPU and on pools of local processes.
• On one host with four RTX 3090s, a process pool launched by dew launch: every layout of the dense, mixture-of-experts, Mamba-2 hybrid and DiT models, each checked against one device.
• On one TPU v6e chip, the sequence-parallel exchange around the splash attention kernel, forward and backward.

Dew has not been run on two physical nodes yet, and hybrid sharding and the sequence exchanges have no throughput numbers. Train on several nodes keeps the full list.

Tutorials

Notebooks that run top to bottom on Colab, shown here with the outputs of their last complete run.

0101GPUDiffusion models from scratch02GPUTrain a diffusion model with Dew0303GPUText to image with classifier-free guidance04GPUSamplers and schedules05GPUTrain a language model0606GPULearn image representations with I-JEPA0707CPUScale one training run across many devices0808GPULoad a pretrained decoder and keep training it

Built on the JAX stack

Models are Flax modules and their variables are JAX pytrees. Optimizers are Optax transformations, checkpoints are written by Orbax, and data loads through Grain. Dew grew out of FlaxDiff, and Google's TPU Research Cloud supported the larger experiments.

Install DewSource on GitHub

MIT license. Dew is research software before 1.0: the API and checkpoint formats can change.