Kauldron: Modular, scalable library to train ML models

Google Research released Kauldron 1.3.0, an open-source machine learning training library built for research velocity and modularity, authored by Klaus Greff, Etienne Pot and Mehdi S. M. Sajjadi. Kauldron lets researchers plug in any dataset (TFDS, Grain, SeqIO or custom pipelines), any Flax model and any optimizer, with every component usable independently and fully overwritable, and ships example configs, a Colab-first workflow, integrated XM plots, a profiler, post-mortem debugging on Borg and runtime shape checking. The repository notes it is not an officially supported Google product.

Kauldron is a library for training machine learning models, optimized for research velocity and modularity . Modularity : - All parts of Kauldron are self-contained, so can be used independently outside Kauldron. - Use any dataset TFDS, Grain, SeqIO, your custom pipeline , any flax model, any optimizer,... Kauldron provides the glue that link everything together. - Everything can be customized and overwritten e.g. sweep over models architecture, overwrite any inner layer parameter,... Research velocity : - Everything should work out-of the box. The example configs https://github.com/google-research/kauldron/tree/main/examples/mnist autoencoder.py can be used and customized as a starting point. - Colab-first workflow for easy prototyping and fast iteration - Polished user experience integrated XM plots, profiler, post-mortem debugging on borg, runtime shape checking, and many others... . Open an issue https://github.com/google-research/kauldron/issues .. If Kauldron was helpful for a publication, please cite this repository: @software{kauldron2025github, author = {Klaus Greff and Etienne Pot and Mehdi S. M. Sajjadi}, title = {{Kauldron}: A neural network training framework, optimized for research velocity and modularity.}, url = {https://github.com/google-research/kauldron}, version = {1.3.0}, year = {2025}, } This is not an officially supported Google product.