Kauldron: Modular, scalable library to train ML models

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 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 ..

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.