Kauldron: Modular, scalable library to train ML models

Kauldron is a JAX-oriented library for training machine-learning models, designed to make research workflows faster while keeping components modular. Rather than requiring a fixed end-to-end stack, it connects user-selected datasets, Flax models, and optimizers. Supported data sources include TFDS, Grain, SeqIO, and custom pipelines. Its components are intended to work independently as well as within the framework, so researchers can reuse them elsewhere or replace and customize them to suit an experiment.

The framework aims to provide useful defaults without restricting experimentation. Example configurations can serve as starting points, and users can adapt them for tasks such as sweeping model architectures or changing parameters inside individual layers. A Colab-first workflow supports prototyping, while integrated tools include experiment plots, profiling, post-mortem debugging on Borg, and runtime shape checks. These features target the practical cycle of developing, training, and diagnosing models, not just model definition. The project is distributed under the Apache 2.0 license, and its maintainers clarify that it is not an officially supported Google product.