README.md

TorchTitan on TPU

Google's distribution of TorchTitan, adding training on Cloud TPU through TorchTPU. Upstream TorchTitan works as documented; TPU support lives in torchtitan/experiments/tpu and adds TPU recipes, Pallas kernels, and TPU-specific trainer and parallelism paths.

Start here

• Train on TPU: torchtitan/experiments/tpu/README.md covers the environment, building the TorchTPU wheel, a single-host TPU VM run, and multi-host training on GKE.
• Pick a recipe: torchtitan/experiments/tpu/recipes has per-model, per-topology configs, each with its launch block.
• Tune a model: the TorchTPU optimization guide.
• Everything else: upstream's README and docs apply unchanged.

Issues and contributing

File TPU bugs and feature requests in this repository's issue tracker, and send TPU changes here, under torchtitan/experiments/tpu. Anything in shared TorchTitan code goes upstream to pytorch/torchtitan, because we keep this repository in sync with it. See CONTRIBUTING.md for the details.

License

BSD 3-Clause, inherited from upstream TorchTitan. See LICENSE.