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.
- Train on TPU:
torchtitan/experiments/tpu/README.mdcovers the environment, building the TorchTPU wheel, a single-host TPU VM run, and multi-host training on GKE. - Pick a recipe:
torchtitan/experiments/tpu/recipeshas 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.
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.
BSD 3-Clause, inherited from upstream TorchTitan. See LICENSE.