Train models with JAX and Flax.

Language models, diffusion and JEPA, trained by one trainer on one device or a mesh.

uv pip install "dew-ml @ git+https://github.com/AshishKumar4/dew"

Image and video diffusion

DiT, U-Net and MMDiT backbones, 18 solvers, classifier-free guidance.

Language models

Pretraining, chat fine-tuning and cached generation on one trainer.

One mesh

MeshSpec(fsdp=4, tensor=2) runs the same code on eight devices.

Post-training

SFT, DPO, GRPO and PPO continue the weights you trained.