Train models with JAX and Flax.

Trained with Dew

176M hybrid DiT · Prompts and settings

green and purple northern lights reflected in a frozen lake, snowy mountains at night, vivid colors, dramatic lighting, highly detailed

a mountain lake reflecting snowy peaks, vivid colors, dramatic lighting, highly detailed

a lighthouse on a rocky coast at sunset, dramatic waves and golden light, vivid colors, dramatic lighting, highly detailed

a snowy mountain village at dusk, warm glowing windows and blue snow, vivid colors, dramatic lighting, highly detailed

a canyon with towering red rock cliffs and a winding river at sunset

a Japanese temple beside a cherry blossom tree on a spring morning, vivid colors, dramatic lighting, highly detailed

a cozy wooden cabin in a snowy forest at night, warm light in the windows, vivid colors, dramatic lighting, highly detailed

a canal in Venice at sunset with gondolas and old buildings

colorful hot air balloons floating above a green valley at sunrise

a red maple tree beside a calm lake in autumn, mountains in the background, vivid colors, dramatic lighting, highly detailed

colorful fireworks over a harbor city at night, reflections in the water

rolling sand dunes in the desert at sunset, deep orange sand and purple sky, vivid colors, dramatic lighting, highly detailed

a turquoise alpine lake surrounded by pine trees and rugged mountains

a dramatic sunset over the ocean, fiery orange and pink clouds

a lush green valley with a winding river and white clouds, Studio Ghibli style landscape, vivid colors, dramatic lighting, highly detailed

a castle on a rocky hill above a misty valley at sunrise, vivid colors, dramatic lighting, highly detailed

a waterfall in a lush green forest, vivid colors, dramatic lighting, highly detailed

a waterfall in a lush green forest, vivid colors, dramatic lighting, highly detailed

Video

Wan 2.1 1.3B's released weights, sampled with Dew's JAX implementation · Prompts and settings

A slow aerial shot gliding over a misty pine forest at sunrise, golden light breaking through the fog between the trees, cinematic, highly detailed

Ocean waves crashing against dark volcanic rocks at sunset, the spray glowing orange in the backlight, slow motion, cinematic

A hummingbird hovering beside bright red flowers in a sunlit garden, its wings a soft blur, shallow depth of field, nature documentary footage

Green northern lights rippling across a starry night sky above a snowy mountain lake, the aurora reflected and shimmering on the still water

Training

A small decoder trains on Tiny Shakespeare on one A100, then writes one sample from each of the prompts ROMEO: and JULIET:.

train.py, editableCopyRun

"""Train a byte decoder on Tiny Shakespeare, evaluate, and generate two samples."""
from dew import Trainer
from dew.config import OptimConfig
from dew.config.data import TextData
from dew.data import ByteTokenizer, Loading, load
from dew.nn.backbones import CausalTransformer
from dew.objectives.lm import LMObjective, Perplexity
from dew.sampling import Sampling, generate
from dew.training.optim import Cosine

steps = 1000
tokens = None


data = (load("hf/winglian/tiny-shakespeare", batch=64, tokenizer="byte", seq_len=256)
        if tokens is None else TextData(source=tokens, seq_len=256).load(
            batch=64, read="memory", validation_records=256, loading=Loading(workers=0)))
model = CausalTransformer(vocab_size=256, emb_features=384, num_layers=6, num_heads=6,
                          mlp_features=1024, max_seq_len=256, dropout_rate=0.2)
optimizer = OptimConfig(schedule=Cosine(peak=1e-3, warmup_steps=10, end=1e-4),
                        b2=0.99, weight_decay=0.1, clip_grads=1.0)
trainer = Trainer(LMObjective(model, seq_len=256), optimizer, key=0)
state = trainer.fit(data, steps=steps, log_every=100, eval_every=500, metrics=(Perplexity(),))

tokenizer = ByteTokenizer()
for key, prompt in enumerate(("ROMEO:", "JULIET:")):
    result = generate(model, state.variables, [tokenizer.encode(prompt)], max_new_tokens=200, key=key,
                      sampling=Sampling(temperature=0.5, top_k=40, eos_token_ids=(46, 33, 63),
                          pad_token_id=32))
    print(tokenizer.decode(result.tokens[0]), end="\n\n")

"""Train a byte decoder on Tiny Shakespeare, evaluate, and generate two samples.""" from dew import Trainer from dew.config import OptimConfig from dew.config.data import TextData from dew.data import ByteTokenizer, Loading, load from dew.nn.backbones import CausalTransformer from dew.objectives.lm import LMObjective, Perplexity from dew.sampling import Sampling, generate from dew.training.optim import Cosine steps = 1000 tokens = None data = (load("hf/winglian/tiny-shakespeare", batch=64, tokenizer="byte", seq_len=256) if tokens is None else TextData(source=tokens, seq_len=256).load( batch=64, read="memory", validation_records=256, loading=Loading(workers=0))) model = CausalTransformer(vocab_size=256, emb_features=384, num_layers=6, num_heads=6, mlp_features=1024, max_seq_len=256, dropout_rate=0.2) optimizer = OptimConfig(schedule=Cosine(peak=1e-3, warmup_steps=10, end=1e-4), b2=0.99, weight_decay=0.1, clip_grads=1.0) trainer = Trainer(LMObjective(model, seq_len=256), optimizer, key=0) state = trainer.fit(data, steps=steps, log_every=100, eval_every=500, metrics=(Perplexity(),)) tokenizer = ByteTokenizer() for key, prompt in enumerate(("ROMEO:", "JULIET:")): result = generate(model, state.variables, [tokenizer.encode(prompt)], max_new_tokens=200, key=key, sampling=Sampling(temperature=0.5, top_k=40, eos_token_ids=(46, 33, 63), pad_token_id=32)) print(tokenizer.decode(result.tokens[0]), end="\n\n")

Run runs the cell on a shared 4-vCPU host; training stops after 20 steps.

$ python train.py

╭─  dew · CausalTransformer  ──────────────────────────────────────────────────────────────────╮
│ 10.7M parameters  ·  on 1 × NVIDIA A100-SXM4-40GB  ·  batch 64  ·  float32                   │
│                                                                                              │
│ ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━  1000/1000  100%  32.9 step/s  0:00:42 │
│                                                                                              │
│ train                                                                                        │
│   loss                         1.225  █▆▆▄▃▃▁▂▁▁      ↓ 43%                                  │
│   ce                           1.225  █▆▆▄▃▃▁▂▁▁      ↓ 43%                                  │
│   perplexity                   3.405  █▅▅▃▃▃▁▂▁▁      ↓ 61%                                  │
│   token_accuracy               62.0%  ▁▄▄▆▆▆█▇██  ↑ 24.5 pt                                  │
│   data_wait_ms               0.02572  ▆▆▃▅▃▃█▆▁▅     ↓ 2.7%                                  │
│   grad_norm                   0.3097  █▄▃▂▂▁▁▁▁▁      ↓ 73%                                  │
│   cache_hits                   11.00  ▁▃▃▃▃█████     ↑ 120%                                  │
│   peak_memory_gib              4.003  ▁▁████████     ↑ 0.1%                                  │
│   compiles 1.000   cache_misses 0                                                            │
│ throughput                                                                                   │
│   step_time_ms                 30.15  ▂▁▁▁▁█▁▁▁▁     ↓ 4.9%                                  │
│   samples_per_sec              2,123  ▇████▁████     ↑ 5.2%                                  │
│   tokens                      16,071  ▇▁▇▄▂█▃▂█▄          →                                  │
│   tokens_per_s               533,040  ▇████▁████     ↑ 5.1%                                  │
│   mfu                          11.6%  ▇████▁████   ↑ 0.6 pt                                  │
│   hlo_flops_util               12.4%  ▇████▁████   ↑ 0.6 pt                                  │
│   hfu                          12.4%  ▇████▁████   ↑ 0.6 pt                                  │
│ optimizer                                                                                    │
│   learning_rate            1.000e-04  ██▇▇▆▅▄▃▂▁      ↓ 90%                                  │
│   accepted 100.0%                                                                            │
│ data                                                                                         │
│   accepted/hf:winglian/tiny-shakespeare@619106eee01474d8eaa5dd400b4b405eb3734ebe:None:train  │
│ hash holdout=[0.0,0.01),held=False,key=None 4,824                                            │
│ attempted/hf:winglian/tiny-shakespeare@619106eee01474d8eaa5dd400b4b405eb3734ebe:None:train   │
│ hash holdout=[0.0,0.01),held=False,key=None 468.0                                            │
│ chunks_per_record/hf:winglian/tiny-shakespeare@619106eee01474d8eaa5dd400b4b405eb3734ebe:None │
│ :train hash holdout=[0.0,0.01),held=False,key=None 0                                         │
│ truncated/hf:winglian/tiny-shakespeare@619106eee01474d8eaa5dd400b4b405eb3734ebe:None:train   │
│ hash holdout=[0.0,0.01),held=False,key=None 0                                                │
│                                                                                              │
│ val · step 1000                                                                              │
│   perplexity       3.741  █▁     ↓ 9.7%                                                      │
╰───────────────────────────────────────────────────────────────────────────────  ✓ finished  ─╯
╭─  ✓ CausalTransformer  ──────────────────────╮
│     trained  1000 steps in 0:00:44           │
│  first step  after 12.65 s                   │
│        then  33.1 step/s                     │
│     goodput  68.4% of the wall time in steps │
│  final loss  1.225                           │
│ val at 1000  perplexity 3.741 ↓ 9.7% █▁      │
╰──────────────────────────────────────────────╯
ROMEO:
And therefore, he shall be the whole of the state,
And who have work and the general of the course,
The execution of the season comes of the childrens
To bear her will the duke of my children:
He dot

JULIET:
Then she is my father's condemn'd the gods
That have I do not with their courts of the world,
Being a wings of some part of men and so welcome to
Her consul.

Recorded on one A100 with JAX 0.11.2.post4 and Dew cb135c1 on 2026-10-14. The run took 60 s. The replay lasts about 10 s. The terminal was 96×200; scroll the panel for earlier lines. The replay needs JavaScript; with reduced motion you see the last frame. Original recording.

Full script

Language modeling

The objective creates the model's variables, computes the loss and runs evaluation. Trainer.fit steps the optimizer, logs progress and saves checkpoints, and every model on this page trains through that one call.

Language modeling, editableCopyRun

from dew import Trainer
from dew.config import OptimConfig
from dew.data import load
from dew.nn.backbones import CausalTransformer
from dew.objectives.lm import LMObjective

data = load("hf/winglian/tiny-shakespeare", batch=8, tokenizer="byte", seq_len=64)
model = CausalTransformer(vocab_size=256, emb_features=32, num_layers=1, num_heads=2,
                          mlp_features=64, max_seq_len=128)
trainer = Trainer(LMObjective(model, seq_len=64), OptimConfig(learning_rate=1e-3), key=0)
state = trainer.fit(data, steps=3)

from dew import Trainer from dew.config import OptimConfig from dew.data import load from dew.nn.backbones import CausalTransformer from dew.objectives.lm import LMObjective data = load("hf/winglian/tiny-shakespeare", batch=8, tokenizer="byte", seq_len=64) model = CausalTransformer(vocab_size=256, emb_features=32, num_layers=1, num_heads=2, mlp_features=64, max_seq_len=128) trainer = Trainer(LMObjective(model, seq_len=64), OptimConfig(learning_rate=1e-3), key=0) state = trainer.fit(data, steps=3)

Run runs the cell on a shared 4-vCPU host; training stops after 20 steps.

$ python lm.py

╭─  dew · CausalTransformer  ──────────────────────────────────────────────────────────────────╮
│ 18,560 parameters  ·  on 1 × cpu  ·  batch 8  ·  float32                                     │
│                                                                                              │
│ ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━  3/3  100%  0:00:02 │
╰───────────────────────────────────────────────────────────────────────────────  ✓ finished  ─╯
╭─  ✓ CausalTransformer  ────────────────────╮
│    trained  3 steps in 0:00:03             │
│ first step  after 2.80 s                   │
│       then  157.0 step/s                   │
│    goodput  0.5% of the wall time in steps │
│ final loss  5.929                          │
╰────────────────────────────────────────────╯

Recorded on a 4-vCPU Cloudflare container with JAX 0.11.2.post4 and Dew 4fbee88 on 2026-10-14. The run took 25 s.

Full script

Representation learning

A JEPA encoder sees part of an image and learns to predict the representations of the patches it cannot see.

JEPA, editableCopyRun

from dew import Field, Trainer
from dew.config import OptimConfig
from dew.config.data import ImageData
from dew.data import Hub, Loading
from dew.objectives.jepa import JepaEncoder, JepaObjective, JepaPredictor, MultiBlockMask

data = ImageData(Hub("uoft-cs/cifar10"), field="img", image_size=32).load(
    batch=64, read="memory", provision=True, loading=Loading(workers=0))
encoder = JepaEncoder(patch_size=4, emb_features=32, num_layers=1, num_heads=2)
predictor = JepaPredictor(grid=(8, 8), emb_features=32, predictor_features=16,
                          num_layers=1, num_heads=2)
objective = JepaObjective(encoder, predictor, mask=MultiBlockMask.for_grid((8, 8)),
                          sample=Field("image", (32, 32, 3)), momentum_steps=3)
state = Trainer(objective, OptimConfig(learning_rate=1e-3), key=0).fit(data, steps=3)

from dew import Field, Trainer from dew.config import OptimConfig from dew.config.data import ImageData from dew.data import Hub, Loading from dew.objectives.jepa import JepaEncoder, JepaObjective, JepaPredictor, MultiBlockMask data = ImageData(Hub("uoft-cs/cifar10"), field="img", image_size=32).load( batch=64, read="memory", provision=True, loading=Loading(workers=0)) encoder = JepaEncoder(patch_size=4, emb_features=32, num_layers=1, num_heads=2) predictor = JepaPredictor(grid=(8, 8), emb_features=32, predictor_features=16, num_layers=1, num_heads=2) objective = JepaObjective(encoder, predictor, mask=MultiBlockMask.for_grid((8, 8)), sample=Field("image", (32, 32, 3)), momentum_steps=3) state = Trainer(objective, OptimConfig(learning_rate=1e-3), key=0).fit(data, steps=3)

Run runs the cell on a shared 4-vCPU host; training stops after 20 steps.

$ python jepa.py

validation: 256 records held out of train
╭─  dew · JepaEncoder  ────────────────────────────────────────────────────────────────────────╮
│ 18,736 parameters  ·  on 1 × cpu  ·  batch 64  ·  float32                                    │
│                                                                                              │
│ ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━  3/3  100%  0:00:07 │
╰───────────────────────────────────────────────────────────────────────────────  ✓ finished  ─╯
╭─  ✓ JepaEncoder  ──────────────────────────╮
│    trained  3 steps in 0:00:09             │
│ first step  after 9.36 s                   │
│       then  32.1 step/s                    │
│    goodput  0.7% of the wall time in steps │
│ final loss  1.499                          │
╰────────────────────────────────────────────╯

Recorded on a 4-vCPU Cloudflare container with JAX 0.11.2.post4 and Dew 10d1092 on 2026-10-14. The run took 47 s.

Fine-tuning

Load a Hugging Face checkpoint as a Flax model, fine-tune it with the same trainer, and save it back in the Hugging Face layout with its original tensor names. The first cell fine-tunes SmolLM2-135M on Tiny Shakespeare and compares its text before and after.

Fine-tune SmolLM2-135M, editableCopyRun

from dataclasses import replace

from dew import Trainer
from dew.config import OptimConfig
from dew.data import load
from dew.interop import PretrainedDecoder
from dew.objectives.lm import LMObjective
from dew.sampling import Sampling

name = "HuggingFaceTB/SmolLM2-135M-Instruct"
model = PretrainedDecoder.load(name, dtype="float32", max_seq_len=256)
data = load("hf/winglian/tiny-shakespeare", batch=1, tokenizer=name, seq_len=64)
state = Trainer(LMObjective(model, seq_len=64), OptimConfig(learning_rate=1e-4), key=0).fit(
    data, steps=20, log_every=5)
before = model.text_generation(sampling=Sampling(temperature=0))
after = replace(model, variables=state.variables).text_generation(sampling=Sampling(temperature=0))
print("Before:", before("ROMEO:", 32, key=0).text[0])
print("After: ", after("ROMEO:", 32, key=0).text[0])

from dataclasses import replace from dew import Trainer from dew.config import OptimConfig from dew.data import load from dew.interop import PretrainedDecoder from dew.objectives.lm import LMObjective from dew.sampling import Sampling name = "HuggingFaceTB/SmolLM2-135M-Instruct" model = PretrainedDecoder.load(name, dtype="float32", max_seq_len=256) data = load("hf/winglian/tiny-shakespeare", batch=1, tokenizer=name, seq_len=64) state = Trainer(LMObjective(model, seq_len=64), OptimConfig(learning_rate=1e-4), key=0).fit( data, steps=20, log_every=5) before = model.text_generation(sampling=Sampling(temperature=0)) after = replace(model, variables=state.variables).text_generation(sampling=Sampling(temperature=0)) print("Before:", before("ROMEO:", 32, key=0).text[0]) print("After: ", after("ROMEO:", 32, key=0).text[0])

Run runs the cell on a shared 4-vCPU host; training stops after 20 steps.

$ python finetune.py

╭─  dew · CausalTransformer  ──────────────────────────────────────────────────────────────────╮
│ 134.5M parameters  ·  on 1 × cpu  ·  batch 1  ·  float32                                     │
│                                                                                              │
│ ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━  20/20  100%  1.6 step/s  0:00:32 │
│                                                                                              │
│ train                                                                                        │
│   loss                         3.941  ▃▆█▁     ↓ 1.7%                                        │
│   ce                           3.941  ▃▆█▁     ↓ 1.7%                                        │
│   perplexity                   51.45  ▃▆█▁     ↓ 6.7%                                        │
│   token_accuracy               35.9%  ▁▂▇█   ↑ 9.4 pt                                        │
│   data_wait_ms               0.02623  █▁▂▁      ↓ 34×                                        │
│   grad_norm                    8.476  █▂▆▁      ↓ 23%                                        │
│   compiles 1.000   cache_hits 0   cache_misses 0                                             │
│ throughput                                                                                   │
│   step_time_ms                 680.8  █▁▁▁      ↓ 31%                                        │
│   samples_per_sec              1.469  ▁███      ↑ 45%                                        │
│   tokens                       64.00  ██▁█          →                                        │
│   tokens_per_s                 94.00  ▁█▇█      ↑ 45%                                        │
│ optimizer                                                                                    │
│   accepted 100.0%   learning_rate 1.000e-04                                                  │
│ data                                                                                         │
│   attempted/files:parqu…       23.00  ▁▄▆█     ↑ 188%                                        │
│   accepted/files:parque…       23.00  ▁▄▆█     ↑ 188%                                        │
│   truncated/files:parquet:('hf://datasets/winglian/tiny-shakespeare@30eb0a2816080ba55bd65d85 │
│ ac6c9a185a5a9c80/default/train/0000.parquet',) hash holdout=[0.0,0.01),held=False,key=None 0 │
│ chunks_per_record/files:parquet:('hf://datasets/winglian/tiny-shakespeare@30eb0a2816080ba55b │
│ d65d85ac6c9a185a5a9c80/default/train/0000.parquet',) hash                                    │
│ holdout=[0.0,0.01),held=False,key=None 0                                                     │
╰───────────────────────────────────────────────────────────────────────────────  ✓ finished  ─╯
╭─  ✓ CausalTransformer  ─────────────────────╮
│    trained  20 steps in 0:00:34             │
│ first step  after 20.69 s                   │
│       then  1.5 step/s                      │
│    goodput  38.5% of the wall time in steps │
│ final loss  3.941                           │
╰─────────────────────────────────────────────╯
Before:

I'm glad you're interested in learning more about the history of the Eiffel Tower. The tower was
 built in 1889 as a
After:

First Citizen:
First Citizen:
First Citizen:
First Citizen:
First Citizen:
First Citizen:
First Citizen:
First Citizen

Recorded on a 4-vCPU Cloudflare container with JAX 0.11.2.post4 and Dew 10d1092 on 2026-10-14. The run took 100 s.

A larger model: Qwen3-0.6B, one step, then saved in its original layout.

Fine-tune Qwen3-0.6B, editableCopyRun in Colab

import jax.numpy as jnp

from dew import Trainer
from dew.config import OptimConfig
from dew.data import load
from dew.interop import PretrainedDecoder
from dew.objectives.lm import LMObjective
from dew.sampling import Sampling

bundle = PretrainedDecoder.load("Qwen/Qwen3-0.6B", dtype=jnp.bfloat16, param_dtype=jnp.bfloat16,
                                max_seq_len=128)
task = bundle.text_generation(sampling=Sampling(temperature=0))
print(task("The capital of France is", 12, key=0).text[0])

data = load("hf/winglian/tiny-shakespeare", batch=1, tokenizer="Qwen/Qwen3-0.6B", seq_len=64)
trainer = Trainer(LMObjective(bundle, seq_len=64), OptimConfig(learning_rate=1e-5), key=0)
state = trainer.fit(data, steps=1)
bundle.save("export", variables=state.variables)

import jax.numpy as jnp from dew import Trainer from dew.config import OptimConfig from dew.data import load from dew.interop import PretrainedDecoder from dew.objectives.lm import LMObjective from dew.sampling import Sampling bundle = PretrainedDecoder.load("Qwen/Qwen3-0.6B", dtype=jnp.bfloat16, param_dtype=jnp.bfloat16, max_seq_len=128) task = bundle.text_generation(sampling=Sampling(temperature=0)) print(task("The capital of France is", 12, key=0).text[0]) data = load("hf/winglian/tiny-shakespeare", batch=1, tokenizer="Qwen/Qwen3-0.6B", seq_len=64) trainer = Trainer(LMObjective(bundle, seq_len=64), OptimConfig(learning_rate=1e-5), key=0) state = trainer.fit(data, steps=1) bundle.save("export", variables=state.variables)

This cell needs more memory than a shared host has, so Run in Colab opens it on a GPU.

$ python pretrained.py

 Paris. The capital of France is also the capital of the
╭─  dew · CausalTransformer  ──────────────────────────────────────────────────────────────────╮
│ 596.0M parameters  ·  on 1 × cpu  ·  batch 1  ·  bfloat16                                    │
│                                                                                              │
│ ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━  1/1  100%  0:00:47 │
╰───────────────────────────────────────────────────────────────────────────────  ✓ finished  ─╯
╭─  ✓ CausalTransformer  ────────────────────╮
│    trained  1 steps in 0:00:49             │
│ first step  after 48.97 s                  │
│    goodput  0.0% of the wall time in steps │
│ final loss  3.826                          │
╰────────────────────────────────────────────╯

Recorded on a 4-vCPU Cloudflare container with JAX 0.11.2.post4 and Dew 10d1092 on 2026-10-14. The run took 138 s.

Supported families: Qwen, Llama, Gemma, Mistral, DeepSeek, Mamba and multimodal decoders.

Inference

A TextGeneration task holds the model, its weights and the sampling settings. Run the cell to generate on a shared CPU host.

generate.pyRun

from dew.interop import PretrainedDecoder
from dew.sampling import Sampling

model = PretrainedDecoder.load("HuggingFaceTB/SmolLM2-135M-Instruct",
                               dtype="float32", max_seq_len=256)
task = model.text_generation(sampling=Sampling(temperature=0))
print(task("The capital of France is", 24, key=0).text[0])

from dew.interop import PretrainedDecoder from dew.sampling import Sampling model = PretrainedDecoder.load("HuggingFaceTB/SmolLM2-135M-Instruct", dtype="float32", max_seq_len=256) task = model.text_generation(sampling=Sampling(temperature=0)) print(task("The capital of France is", 24, key=0).text[0])

 Paris. Paris is the largest city in France and the capital of the French department of the Espace de Paris.

Recorded on a 4-vCPU Cloudflare container with Dew 898a4ec on 2026-10-08, in 40 s. Run the cell on a live CPU kernel.

Server runs many requests through one task in shared batches, and each request keeps its own random key and KV cache:

Batched decoding, editableCopyRun in Colab

import jax.numpy as jnp

from dew.inference.serving import Server
from dew.interop import PretrainedDecoder
from dew.sampling import Sampling

bundle = PretrainedDecoder.load("Qwen/Qwen3-0.6B", dtype=jnp.bfloat16, param_dtype=jnp.bfloat16,
                                max_seq_len=128)
task = bundle.text_generation(sampling=Sampling(temperature=0))
server = Server.from_task(task, slots=4, capacity=128)
results = server(["The capital of France is", "The capital of Japan is"], 24, key=0)
print([result.text[0] for result in results])

import jax.numpy as jnp from dew.inference.serving import Server from dew.interop import PretrainedDecoder from dew.sampling import Sampling bundle = PretrainedDecoder.load("Qwen/Qwen3-0.6B", dtype=jnp.bfloat16, param_dtype=jnp.bfloat16, max_seq_len=128) task = bundle.text_generation(sampling=Sampling(temperature=0)) server = Server.from_task(task, slots=4, capacity=128) results = server(["The capital of France is", "The capital of Japan is"], 24, key=0) print([result.text[0] for result in results])

This cell needs more memory than a shared host has, so Run in Colab opens it on a GPU.

$ python serving.py

[' Paris. The capital of France is also the capital of the Republic of France. The capital of Fr
ance is also the capital', ' Tokyo. The capital of the United States is Washington, D.C. The cap
ital of France is Paris. The capital of']

Recorded on a 4-vCPU Cloudflare container with JAX 0.11.2.post3 and Dew 898a4ec on 2026-10-08. The run took 64 s.

Store weights in int8 or fp8:

Weight-only int8 and fp8, editableCopyRun in Colab

import jax.numpy as jnp

from dew.interop import PretrainedDecoder
from dew.sampling import Sampling
from dew.quantization import Quantization

bundle = PretrainedDecoder.load("Qwen/Qwen3-0.6B", dtype=jnp.bfloat16, param_dtype=jnp.bfloat16,
                                max_seq_len=128)
task = bundle.text_generation(sampling=Sampling(temperature=0))
for dtype in ("int8", "fp8"):
    quantized = task.quantized(Quantization(dtype=dtype, weight_only=True))
    print(dtype, quantized("The capital of France is", 12, key=0).text[0])

import jax.numpy as jnp from dew.interop import PretrainedDecoder from dew.sampling import Sampling from dew.quantization import Quantization bundle = PretrainedDecoder.load("Qwen/Qwen3-0.6B", dtype=jnp.bfloat16, param_dtype=jnp.bfloat16, max_seq_len=128) task = bundle.text_generation(sampling=Sampling(temperature=0)) for dtype in ("int8", "fp8"): quantized = task.quantized(Quantization(dtype=dtype, weight_only=True)) print(dtype, quantized("The capital of France is", 12, key=0).text[0])

This cell needs more memory than a shared host has, so Run in Colab opens it on a GPU.

$ python formats.py

int8  Paris. The capital of Italy is Rome. The capital of
fp8  Paris. The capital of France is also the capital of the

Recorded on a 4-vCPU Cloudflare container with JAX 0.11.2.post4 and Dew 6a553a6 on 2026-10-19. The run took 79 s.

The serving guide covers the paged KV cache, prefix sharing, constrained decoding, and the error each quantized format adds.

Decision models

A decision model answers typed questions about a state in one forward pass: each question names its options, and the answer is a probability distribution over exactly those options. Laya, a ModernBERT-large decision model, loads as a Dew model, and systemone takes and returns Jev's wire format.

Route a support ticket, editableCopyRun

import json

from dew.decision import Decide

decide = Decide.from_pretrained("convaiinnovations/laya")
response = decide.systemone({
    "state": "Hi, we were billed twice for March and I want it reversed today.",
    "questions": {
        "department": {"type": "choice", "instructions": "Which department should handle this?",
                       "criteria": {"billing": "invoices, refunds", "technical": "bugs, outages"}},
        "urgency": {"type": "score", "instructions": "How urgent is this?",
                    "criteria": ["not urgent", "soon", "blocking"]},
        "churn_risk": {"type": "noul", "instructions": "Does the user threaten to leave?"},
    },
})
print(json.dumps(response["answers"], indent=2))

import json from dew.decision import Decide decide = Decide.from_pretrained("convaiinnovations/laya") response = decide.systemone({ "state": "Hi, we were billed twice for March and I want it reversed today.", "questions": { "department": {"type": "choice", "instructions": "Which department should handle this?", "criteria": {"billing": "invoices, refunds", "technical": "bugs, outages"}}, "urgency": {"type": "score", "instructions": "How urgent is this?", "criteria": ["not urgent", "soon", "blocking"]}, "churn_risk": {"type": "noul", "instructions": "Does the user threaten to leave?"}, }, }) print(json.dumps(response["answers"], indent=2))

Run runs the cell on a shared 4-vCPU host; training stops after 20 steps.

$ python decide.py

{
  "department": {
    "type": "choice",
    "choice": "billing",
    "probabilities": {
      "billing": 0.9897,
      "technical": 0.0103
    },
    "confidence": 0.9794
  },
  "urgency": {
    "type": "score",
    "score": 1.7091,
    "legend": {
      "0": "not urgent",
      "1": "soon",
      "2": "blocking"
    },
    "probabilities": {
      "0": 0.0483,
      "1": 0.1942,
      "2": 0.7574
    },
    "confidence": 0.5636
  },
  "churn_risk": {
    "type": "noul",
    "noul": 0.0818
  }
}

Recorded on a 4-vCPU Cloudflare container with JAX 0.11.2.post3 and Dew 898a4ec on 2026-10-08. The run took 62 s.

The decision models guide covers calibration, abstention and fine-tuning a decision head on any backbone.

Diffusion

The process sets how training adds noise and what the model learns to predict. At generation time a solver runs the process in reverse, and you can switch solvers without retraining.

sample.py, editableRun

from dew.sampling import CFG, DPMSolverMultistep, TextToImage

pipe = TextToImage.from_pretrained("dewml/hybrid-dit-176m", revision="4b322022189373fd9086313241269aa5cea92b86")
prompt = "the northern lights over a frozen lake at night, vivid colors, dramatic lighting, highly detailed"
negative = ("letterbox, white border, black border, frame, text, watermark, collage, blurry, lowres, "
            "low quality, dull colors, washed out, low contrast, grainy")
inputs = pipe.prepare([prompt], key=3, steps=15, unconditional=negative)
result = pipe(
    inputs,
    key=3,
    steps=15,
    solver=DPMSolverMultistep(),
    guidance=CFG(6.0, interval=(0.15, 0.9)),
)
result.pil()[0]

from dew.sampling import CFG, DPMSolverMultistep, TextToImage pipe = TextToImage.from_pretrained("dewml/hybrid-dit-176m", revision="4b322022189373fd9086313241269aa5cea92b86") prompt = "the northern lights over a frozen lake at night, vivid colors, dramatic lighting, highly detailed" negative = ("letterbox, white border, black border, frame, text, watermark, collage, blurry, lowres, " "low quality, dull colors, washed out, low contrast, grainy") inputs = pipe.prepare([prompt], key=3, steps=15, unconditional=negative) result = pipe( inputs, key=3, steps=15, solver=DPMSolverMultistep(), guidance=CFG(6.0, interval=(0.15, 0.9)), ) result.pil()[0]

Run sends the cell's Dew calls to a process on the host that already holds the model. The same cell runs as is with Dew installed.

Run opens an isolated Python context backed by the shared Dew model on a 4-vCPU CPU host.

Recorded on a Colab A100-SXM4-40GB with Dew ff6c4dc on 2026-10-16, in 102 s.

Or open it in Colab to run it on a GPU and change anything.

Define and train a DiT, editableCopyRun

from dew import Field, InputSpec, Trainer
from dew.config import OptimConfig
from dew.config.data import ImageData
from dew.data import Hub, Loading
from dew.diffusion.presets import Flow
from dew.nn.backbones import SimpleDiT
from dew.objectives.diffusion import DiffusionObjective
from dew.sampling import Euler

data = ImageData(Hub("uoft-cs/cifar10"), field="img", image_size=32).load(
    batch=64, read="memory", provision=True, loading=Loading(workers=0))
model = SimpleDiT(patch_size=4, emb_features=32, num_layers=1, num_heads=2, mlp_ratio=2)
objective = DiffusionObjective(model, Flow(), InputSpec(Field("image", (32, 32, 3))),
                               solver=Euler(), steps=4)
state = Trainer(objective, OptimConfig(learning_rate=1e-3), key=0).fit(data, steps=3)

from dew import Field, InputSpec, Trainer from dew.config import OptimConfig from dew.config.data import ImageData from dew.data import Hub, Loading from dew.diffusion.presets import Flow from dew.nn.backbones import SimpleDiT from dew.objectives.diffusion import DiffusionObjective from dew.sampling import Euler data = ImageData(Hub("uoft-cs/cifar10"), field="img", image_size=32).load( batch=64, read="memory", provision=True, loading=Loading(workers=0)) model = SimpleDiT(patch_size=4, emb_features=32, num_layers=1, num_heads=2, mlp_ratio=2) objective = DiffusionObjective(model, Flow(), InputSpec(Field("image", (32, 32, 3))), solver=Euler(), steps=4) state = Trainer(objective, OptimConfig(learning_rate=1e-3), key=0).fit(data, steps=3)

Run runs the cell on a shared 4-vCPU host; training stops after 20 steps.

$ python diffusion.py

validation: 256 records held out of train
╭─  dew · SimpleDiT  ──────────────────────────────────────────────────────────────────────────╮
│ 26,320 parameters  ·  on 1 × cpu  ·  batch 64  ·  float32                                    │
│                                                                                              │
│ ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━  3/3  100%  0:00:06 │
╰───────────────────────────────────────────────────────────────────────────────  ✓ finished  ─╯
╭─  ✓ SimpleDiT  ────────────────────────────╮
│    trained  3 steps in 0:00:08             │
│ first step  after 8.19 s                   │
│       then  45.0 step/s                    │
│    goodput  0.5% of the wall time in steps │
│ final loss  0.5997                         │
╰────────────────────────────────────────────╯

Recorded on a 4-vCPU Cloudflare container with JAX 0.11.2.post4 and Dew 10d1092 on 2026-10-14. The run took 50 s.

A 176M-parameter text-to-image model at 256×171, trained with Dew. Weights: dewml/hybrid-dit-176m.

“the northern lights over a frozen lake at night”

“a lighthouse on a rocky coast at sunset, dramatic waves and golden light”

“a snowy mountain village at dusk, warm glowing windows and blue snow, vivid colors, dramatic lighting, highly detailed”

A style LoRA

Stable Diffusion 1.5 with a rank-16 LoRA on its attention projections, trained with Dew for 160 steps on 20 public-domain Hiroshige prints from the Library of Congress, on one A100. Each pair shares its prompt and seed; drag the slider to compare. Prompts, seeds and settings.

BeforeAfter

Compare

a picture in dewprint style of a mountain lake with a small sailboat, snowy peaks, dawn
Seed 1042

BeforeAfter

Compare

a picture in dewprint style of a winding river through a village in autumn, red maple trees
Seed 2073

BeforeAfter

Compare

a picture in dewprint style of a wooden bridge in the rain, people crossing with umbrellas
Seed 3104

BeforeAfter

Compare

a picture in dewprint style of a coastal fishing village, a red sunset over the sea
Seed 4135

BeforeAfter

Compare

a picture in dewprint style of a snow-covered pine tree beside a quiet temple
Seed 5166

BeforeAfter

Compare

a picture in dewprint style of cherry blossom trees beside a river, Mount Fuji in the distance
Seed 6197

Distributed

To use more devices, give Trainer a mesh; the objective and data stay as they are. FSDP shards parameters and optimizer state, tensor parallelism splits each layer's weight matrices, and an expert axis puts an MoE's experts on different devices. The devices left over form the data-parallel axis.

Four simulated CPU devices, editableCopyRun

import jax

jax.config.update("jax_num_cpu_devices", 4)

from dew import MeshSpec, Trainer
from dew.config import OptimConfig
from dew.data import load
from dew.nn.backbones import CausalTransformer
from dew.objectives.lm import LMObjective

data = load("hf/winglian/tiny-shakespeare", batch=8, tokenizer="byte", seq_len=64)
model = CausalTransformer(vocab_size=256, emb_features=32, num_layers=1, num_heads=2,
                          mlp_features=64, max_seq_len=128)
trainer = Trainer(LMObjective(model, seq_len=64), OptimConfig(learning_rate=1e-3), key=0,
                  mesh=MeshSpec(fsdp=2, tensor=2))
state = trainer.fit(data, steps=3)

import jax jax.config.update("jax_num_cpu_devices", 4) from dew import MeshSpec, Trainer from dew.config import OptimConfig from dew.data import load from dew.nn.backbones import CausalTransformer from dew.objectives.lm import LMObjective data = load("hf/winglian/tiny-shakespeare", batch=8, tokenizer="byte", seq_len=64) model = CausalTransformer(vocab_size=256, emb_features=32, num_layers=1, num_heads=2, mlp_features=64, max_seq_len=128) trainer = Trainer(LMObjective(model, seq_len=64), OptimConfig(learning_rate=1e-3), key=0, mesh=MeshSpec(fsdp=2, tensor=2)) state = trainer.fit(data, steps=3)

Run runs the cell on a shared 4-vCPU host; training stops after 20 steps.

$ python mesh.py

╭─  dew · CausalTransformer  ──────────────────────────────────────────────────────────────────╮
│ 18,560 parameters  ·  on 4 × cpu  ·  mesh fsdp 2 × tensor 2, 0% of the parameters' bytes spl │
│                                                                                              │
│ ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━  3/3  100%  0:00:02 │
╰───────────────────────────────────────────────────────────────────────────────  ✓ finished  ─╯
╭─  ✓ CausalTransformer  ────────────────────╮
│    trained  3 steps in 0:00:03             │
│ first step  after 3.36 s                   │
│       then  145.1 step/s                   │
│    goodput  0.4% of the wall time in steps │
│ final loss  5.929                          │
╰────────────────────────────────────────────╯

Recorded on a 4-vCPU Cloudflare container with JAX 0.11.2.post4 and Dew 10d1092 on 2026-10-14. The run took 26 s.

Placement and sharding · Several hosts · Cloud TPU

Expert routing

With the experts split over devices, each MoE layer's router sends tokens to experts on other devices through all-to-all exchanges. This is the routing examples/moe_mesh.py recorded over 300 training steps.

PlayStep

training step 1 of 300

Eight simulated CPU devices: data 2 × expert 2 × FSDP 2. Arrows show the recorded assignments sent between devices; each circle shows the assignments kept locally. Full routing record.

Post-training

GRPO samples a group of completions for each prompt and scores them with a reward function. The reward here is the fraction of adjacent tokens that count up by one, modulo 13.

Rollout and reward, editableCopyRun

import itertools

import numpy as np

from dew import LocalTracker, Trainer
from dew.config import OptimConfig
from dew.config.data import PromptData
from dew.data import Loading
from dew.nn.backbones import CausalTransformer
from dew.objectives.rl import GRPOObjective, SampledRollout
from dew.sampling import Sampling

def reward(completion, record):
    tokens = [int(token) for token in completion.split()]
    pairs = list(itertools.pairwise(tokens))
    return sum(b == (a + 1) % 13 for a, b in pairs) / max(len(pairs), 1)

rng = np.random.default_rng(0)
records = tuple({"prompt": rng.integers(0, 13, 4).tolist()} for _ in range(512))
data = PromptData(records, tokenizer="byte", max_prompt_len=4).load(
    batch=8, read="memory", loading=Loading(workers=0))
model = CausalTransformer(vocab_size=13, emb_features=64, num_layers=2, num_heads=4,
                          head_dim=16, mlp_features=128, max_seq_len=16)
objective = GRPOObjective(model, seq_len=11, beta=0.02)
rollout = SampledRollout(objective, reward=reward, groups=4, max_new_tokens=8,
                         sampling=Sampling(temperature=1.0))
tracker = LocalTracker("tracking")
trainer = Trainer(objective, OptimConfig(learning_rate=1e-3), key=0, rollout=rollout, tracker=tracker)
state = trainer.fit(data, steps=150, log_every=1)
tracker.close()

import itertools import numpy as np from dew import LocalTracker, Trainer from dew.config import OptimConfig from dew.config.data import PromptData from dew.data import Loading from dew.nn.backbones import CausalTransformer from dew.objectives.rl import GRPOObjective, SampledRollout from dew.sampling import Sampling def reward(completion, record): tokens = [int(token) for token in completion.split()] pairs = list(itertools.pairwise(tokens)) return sum(b == (a + 1) % 13 for a, b in pairs) / max(len(pairs), 1) rng = np.random.default_rng(0) records = tuple({"prompt": rng.integers(0, 13, 4).tolist()} for _ in range(512)) data = PromptData(records, tokenizer="byte", max_prompt_len=4).load( batch=8, read="memory", loading=Loading(workers=0)) model = CausalTransformer(vocab_size=13, emb_features=64, num_layers=2, num_heads=4, head_dim=16, mlp_features=128, max_seq_len=16) objective = GRPOObjective(model, seq_len=11, beta=0.02) rollout = SampledRollout(objective, reward=reward, groups=4, max_new_tokens=8, sampling=Sampling(temperature=1.0)) tracker = LocalTracker("tracking") trainer = Trainer(objective, OptimConfig(learning_rate=1e-3), key=0, rollout=rollout, tracker=tracker) state = trainer.fit(data, steps=150, log_every=1) tracker.close()

Run runs the cell on a shared 4-vCPU host; training stops after 20 steps.

$ python grpo.py

╭─  dew · CausalTransformer  ──────────────────────────────────────────────────────────────────╮
│ 83,136 parameters  ·  on 1 × cpu  ·  batch 8  ·  float32                                     │
│                                                                                              │
│ ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━  150/150  100%  19.4 step/s  0:00:21 │
│                                                                                              │
│ train                                                                                        │
│   loss                       0.04121  █▅▆▃▄▇▅▃▅▅▅█▄▂▄▄█▃▂▃▂▁▃▄▃▃▅▅▆▅▅▅▂▅▅▁▅▆▆▃█▂▆     ↓ 7.0% │
│   actor/ppo_kl             2.049e-08  ▆▅▆▆▄▆▆▅▅▅▆█▄▆▄▅▆▃▆▆▅▆▅▅▆▅▅▇▅▅▅▆▆▆▇▄▁▅▅▇▅▂▆      ↓ 27% │
│   kl                           2.061  █▅▆▃▄▇▅▃▅▅▅█▄▂▄▄█▃▂▃▂▁▃▄▃▃▅▅▆▅▅▅▂▅▅▁▅▆▆▃█▂▆     ↓ 7.0% │
│   pg                       1.304e-08  █▇▅▅▅█▁▆▅▅▅▇▆▅▅▂▆▅▇▅█▅▆▅▅▂▄▅▅▂▇▅▅▅▅▅▄▅▅▅▅▅▆      ↓ 78% │
│   data_wait_ms               0.03979  ▃▅▄▁▇▅▂▄▃▃▃▄▃▂▅▄▂▃▃▂▆▄█▃▃▃▃▂▅▆▁▄▁▃▃▃▄▄▁▂▂▄▁      ↓ 32% │
│   grad_norm                   0.3026  ▇▆▂▂▂▆▆▆▆▂▇█▇▂▁▇▇▁▆▁▆▃█▁▂▆▇▁▁▆▆▁▁▁▁▁▇▂▁▂▇▃▇      ↓ 45% │
│   actor/pg_clipfrac 0   actor/pg_clipfrac_lower 0   mismatch/ess 1.000   mismatch/k3_kl 0    │
│ mismatch/kl 0   compiles 1.000   cache_hits 0   cache_misses 0                               │
│ throughput                                                                                   │
│   step_time_ms                 15.19  ▇▇▇▇▇▇▇▇█▇▇▇▆▆▆▇▇▆▇▇▇▇▇▇█▇▇█▇▇▇▇█▁█▇▇▆▇▇▆▇▅      ↓ 66% │
│   samples_per_sec              2,107  ▂▂▂▂▂▂▂▂▁▂▂▂▃▃▃▂▂▃▂▂▂▂▂▂▁▂▂▁▂▂▂▂▁█▁▂▂▃▂▂▃▂▄     ↑ 196% │
│   rollout_seconds            0.03458  ▇▃▃▆▄▃▆▃▃▇▂▄▅▁▂▂▅▃▂▆▂▃▆▃▂█▃▃█▃▃▆▃▅▄▃▆▃▂▇▂▂▂      ↓ 49% │
│ optimizer                                                                                    │
│   accepted 100.0%   learning_rate 0.001000                                                   │
│ rollout                                                                                      │
│   reward/mean                 0.8527  ▄▆▃▃▇▇▃▆▅▅▅▄▅▅▅▃▃▅▆▇▁▅▃▇▃▃▅▃██▄▅▇▅▅▅▄█▃▂▇▅█     ↑ 5.5% │
│   length/mean 8.000   status/truncated 100.0%                                                │
│ data                                                                                         │
│   accepted/memory:ba5333df301c9ae60de68de5c585a95827443e7975a787385fc977ac66c8ba8c 512.0     │
│ attempted/memory:ba5333df301c9ae60de68de5c585a95827443e7975a787385fc977ac66c8ba8c 512.0      │
│ chunks_per_record/memory:ba5333df301c9ae60de68de5c585a95827443e7975a787385fc977ac66c8ba8c 0  │
│ truncated/memory:ba5333df301c9ae60de68de5c585a95827443e7975a787385fc977ac66c8ba8c 0          │
╰───────────────────────────────────────────────────────────────────────────────  ✓ finished  ─╯
╭─  ✓ CausalTransformer  ─────────────────────╮
│    trained  150 steps in 0:00:24            │
│ first step  after 15.83 s                   │
│       then  19.2 step/s                     │
│    goodput  32.8% of the wall time in steps │
│ final loss  0.04121                         │
╰─────────────────────────────────────────────╯

Recorded on a 4-vCPU Cloudflare container with JAX 0.11.2.post4 and Dew d8eaf5c on 2026-10-19. The run took 59 s.

Mean reward of the sampled completions at each of the recorded run's 150 updates (raw log). SFT, DPO, PPO and Flow-GRPO train with the same Trainer.

Reliability

A checkpoint stores the training state and the position in the data stream. When the job receives SIGTERM, fit stops at the next step every process agrees on, writes a checkpoint and exits with status 143. Run the same command again to resume.

Save and resume, editableCopyRun

from dew import Checkpoints, Trainer
from dew.config import OptimConfig
from dew.data import load
from dew.nn.backbones import CausalTransformer
from dew.objectives.lm import LMObjective
from dew.training import prepare_process

prepare_process(multi_host=False, xla_flags="--xla_gpu_deterministic_ops=true")
data = load("hf/winglian/tiny-shakespeare", batch=8, tokenizer="byte", seq_len=64)
model = CausalTransformer(vocab_size=256, emb_features=32, num_layers=1, num_heads=2,
                          mlp_features=64, max_seq_len=128)
objective = LMObjective(model, seq_len=64)
state = Trainer(objective, OptimConfig(learning_rate=1e-3), key=0,
                checkpoints=Checkpoints("checkpoints")).fit(data, steps=2, checkpoint_every=1)
# A new trainer on the same directory resumes from step 2.
resumed = Trainer(objective, OptimConfig(learning_rate=1e-3), key=0,
                  checkpoints=Checkpoints("checkpoints")).fit(data, steps=3, checkpoint_every=1)

from dew import Checkpoints, Trainer from dew.config import OptimConfig from dew.data import load from dew.nn.backbones import CausalTransformer from dew.objectives.lm import LMObjective from dew.training import prepare_process prepare_process(multi_host=False, xla_flags="--xla_gpu_deterministic_ops=true") data = load("hf/winglian/tiny-shakespeare", batch=8, tokenizer="byte", seq_len=64) model = CausalTransformer(vocab_size=256, emb_features=32, num_layers=1, num_heads=2, mlp_features=64, max_seq_len=128) objective = LMObjective(model, seq_len=64) state = Trainer(objective, OptimConfig(learning_rate=1e-3), key=0, checkpoints=Checkpoints("checkpoints")).fit(data, steps=2, checkpoint_every=1) # A new trainer on the same directory resumes from step 2. resumed = Trainer(objective, OptimConfig(learning_rate=1e-3), key=0, checkpoints=Checkpoints("checkpoints")).fit(data, steps=3, checkpoint_every=1)

Run runs the cell on a shared 4-vCPU host; training stops after 20 steps.

$ python reliability.py

╭─  dew · CausalTransformer  ──────────────────────────────────────────────────────────────────╮
│ 18,560 parameters  ·  on 1 × cpu  ·  batch 8  ·  float32                                     │
│                                                                                              │
│ ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━  2/2  100%  0:00:02 │
╰───────────────────────────────────────────────────────────────────────────────  ✓ finished  ─╯
╭─  ✓ CausalTransformer  ────────────────────╮
│    trained  2 steps in 0:00:03             │
│ first step  after 2.88 s                   │
│       then  10.4 step/s                    │
│    goodput  3.0% of the wall time in steps │
│ final loss  6.069                          │
╰────────────────────────────────────────────╯
Resumed from step 2 in /tmp/cell-reliability-7sl8xvg2/checkpoints
╭─  dew · CausalTransformer  ──────────────────────────────────────────────────────────────────╮
│ 18,560 parameters  ·  on 1 × cpu  ·  batch 8  ·  float32                                     │
│                                                                                              │
│ ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━  3/3  100%  0:00:02 │
╰───────────────────────────────────────────────────────────────────────────────  ✓ finished  ─╯
╭─  ✓ CausalTransformer  ────────────────────╮
│    trained  1 steps in 0:00:03             │
│ first step  after 2.51 s                   │
│    goodput  7.0% of the wall time in steps │
│ final loss  5.929                          │
╰────────────────────────────────────────────╯

Recorded on a 4-vCPU Cloudflare container with JAX 0.11.2.post4 and Dew 10d1092 on 2026-10-14. The run took 29 s.

JAX/XProf capture, editableCopyRun

from dew import ProfileWindow, Trainer
from dew.config import OptimConfig
from dew.data import load
from dew.nn.backbones import CausalTransformer
from dew.objectives.lm import LMObjective

data = load("hf/winglian/tiny-shakespeare", batch=8, tokenizer="byte", seq_len=64)
model = CausalTransformer(vocab_size=256, emb_features=32, num_layers=1, num_heads=2,
                          mlp_features=64, max_seq_len=128)
trainer = Trainer(LMObjective(model, seq_len=64), OptimConfig(learning_rate=1e-3), key=0,
                  profile=ProfileWindow("profile", steps=2))
trainer.fit(data, steps=5)

from dew import ProfileWindow, Trainer from dew.config import OptimConfig from dew.data import load from dew.nn.backbones import CausalTransformer from dew.objectives.lm import LMObjective data = load("hf/winglian/tiny-shakespeare", batch=8, tokenizer="byte", seq_len=64) model = CausalTransformer(vocab_size=256, emb_features=32, num_layers=1, num_heads=2, mlp_features=64, max_seq_len=128) trainer = Trainer(LMObjective(model, seq_len=64), OptimConfig(learning_rate=1e-3), key=0, profile=ProfileWindow("profile", steps=2)) trainer.fit(data, steps=5)

Run runs the cell on a shared 4-vCPU host; training stops after 20 steps.

$ python profile.py

Wrote profile for 2 steps to profile
╭─  dew · CausalTransformer  ──────────────────────────────────────────────────────────────────╮
│ 18,560 parameters  ·  on 1 × cpu  ·  batch 8  ·  float32                                     │
│                                                                                              │
│ ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━  5/5  100%  0:00:02 │
╰───────────────────────────────────────────────────────────────────────────────  ✓ finished  ─╯
╭─  ✓ CausalTransformer  ────────────────────╮
│    trained  5 steps in 0:00:03             │
│ first step  after 2.89 s                   │
│       then  21.6 step/s                    │
│    goodput  6.0% of the wall time in steps │
│ final loss  5.767                          │
╰────────────────────────────────────────────╯

Recorded on a 4-vCPU Cloudflare container with JAX 0.11.2.post4 and Dew 10d1092 on 2026-10-14. The run took 28 s.

Checkpoint guarantees · Profiling

Try the 176M model

• Live samplerEdit the cell and run it on a CPU kernel.
• Colab notebookLoad the model and sample it on a GPU.
• dewml/hybrid-dit-176mWeights and model card on the Hub.

Tutorials

1. 01Diffusion from scratchGPU
2. 02Train a diffusion modelGPU
3. 03Text to image with guidanceGPU
4. 04Samplers and schedulesGPU
5. 05Train a language modelGPU
6. 06Representations with I-JEPAGPU
7. 07Scale across devicesCPU
8. 08Continue a pretrained decoderGPU
9. 09SFT for LLMsGPU
10. 10RL with verifiable rewardsGPU
11. 11LLMs in an agent harnessGPU

How each example above was recorded

• lm: a 4-vCPU Cloudflare container, JAX 0.11.2.post4, Dew 4fbee88, 2026-10-14, 25.4 s.
• jepa: a 4-vCPU Cloudflare container, JAX 0.11.2.post4, Dew 10d1092, 2026-10-14, 47 s.
• finetune: a 4-vCPU Cloudflare container, JAX 0.11.2.post4, Dew 10d1092, 2026-10-14, 100.4 s.
• decide: a 4-vCPU Cloudflare container, JAX 0.11.2.post3, Dew 898a4ec, 2026-10-08, 61.8 s.
• diffusion: a 4-vCPU Cloudflare container, JAX 0.11.2.post4, Dew 10d1092, 2026-10-14, 50.4 s.
• mesh: a 4-vCPU Cloudflare container, JAX 0.11.2.post4, Dew 10d1092, 2026-10-14, 26.4 s.
• grpo: a 4-vCPU Cloudflare container, JAX 0.11.2.post4, Dew d8eaf5c, 2026-10-19, 58.7 s.
• reliability: a 4-vCPU Cloudflare container, JAX 0.11.2.post4, Dew 10d1092, 2026-10-14, 29.1 s.
• profile: a 4-vCPU Cloudflare container, JAX 0.11.2.post4, Dew 10d1092, 2026-10-14, 27.7 s.
• pretrained: a 4-vCPU Cloudflare container, JAX 0.11.2.post4, Dew 10d1092, 2026-10-14, 137.9 s.
• serving: a 4-vCPU Cloudflare container, JAX 0.11.2.post3, Dew 898a4ec, 2026-10-08, 64.4 s.
• formats: a 4-vCPU Cloudflare container, JAX 0.11.2.post4, Dew 6a553a6, 2026-10-19, 79.1 s.

MIT license. Research software before 1.0.All examples as one scriptGitHub