---
title: "Integrating TorchTPU into TorchTitan and Optimizing Models for TPU"
description: "TorchTitan TPU is a fork of Meta’s TorchTitan that enables PyTorch model training on Google Cloud TPUs using the TorchTPU native backend with minimal device initialization changes. SimpleFSDP replaces FSDP2 hooks with DTensor parametrization, making all-gathers and reduce-scatters ordinary nodes in the traced graph for compiler scheduling. The Graph Trainer captures the full forward, loss, and backward step as one XLA program via make_fx, achieving 51-52% MFU on Qwen3-8B with functional AdamW. On Llama 3 405B, TorchTitan TPU reaches about 57% MFU on up to 1,024 chips, and a weak-scaling test of Qwen3 235B MoE retained 92% of per-chip throughput from 64 to 8,192 chips. The repository includes an optimization playbook, a coding-agent model-onboarding skill, and supports over 500 Hugging Face models through a Transformers backend."
---

# Integrating TorchTPU into TorchTitan and Optimizing Models for TPU

[Google](https://yomu.fyi/company/google) · Aleksey Vlasenko, Rob Torres, Will Constable · Oct 20, 2026

**Type:** Announcement

## Summary

TorchTitan TPU is a fork of Meta’s TorchTitan that enables PyTorch model training on Google Cloud TPUs using the TorchTPU native backend with minimal device initialization changes. SimpleFSDP replaces FSDP2 hooks with DTensor parametrization, making all-gathers and reduce-scatters ordinary nodes in the traced graph for compiler scheduling. The Graph Trainer captures the full forward, loss, and backward step as one XLA program via make\_fx, achieving 51-52% MFU on Qwen3-8B with functional AdamW. On Llama 3 405B, TorchTitan TPU reaches about 57% MFU on up to 1,024 chips, and a weak-scaling test of Qwen3 235B MoE retained 92% of per-chip throughput from 64 to 8,192 chips. The repository includes an optimization playbook, a coding-agent model-onboarding skill, and supports over 500 Hugging Face models through a Transformers backend.

## Context

Historically, training PyTorch models on Google Cloud TPUs required either a complete codebase rewrite into JAX or use of legacy torch\_xla, but TorchTPU now provides a native backend integrated directly into PyTorch’s core interface.

## Approach / What changed

A fork of TorchTitan with TPU-native execution passes, SimpleFSDP for compiler-friendly sharding, Graph Trainer for full-step capture, and Pallas kernel replacements for attention, loss, and MoE matmuls.

## Takeaways

- TorchTitan TPU achieves 57% MFU on Llama 3 405B across up to 1,024 TPU chips, demonstrating competitive performance with native JAX training.
- SimpleFSDP expresses parameter sharding as DTensor redistributions, enabling the XLA compiler to schedule collectives inside a single traced graph for better overlap.
- The Graph Trainer captures the entire training step as one XLA program, and with scan-over-layers and graph splitting it scales to models like Qwen3 235B MoE.

## FAQ

### What is TorchTitan TPU and how does it differ from Meta's TorchTitan?

TorchTitan TPU is a fork of Meta's TorchTitan adapted for Google Cloud TPUs using the TorchTPU native backend. It adds TPU-specific optimizations like SimpleFSDP, Graph Trainer, and Pallas kernel replacements while retaining compatibility with original model definitions written in pure PyTorch.

### How does SimpleFSDP improve TPU training compared to FSDP2?

SimpleFSDP expresses parameter sharding as DTensor redistributions, making all-gathers and reduce-scatters nodes in the traced graph rather than explicit collective hooks. This allows the XLA compiler to schedule them ahead of matmuls and overlap them with compute, avoiding the breaks in torch.compile regions that occur with FSDP2.

### What performance numbers are available for TorchTitan TPU on large models?

Llama 3 405B achieves about 57% MFU on up to 1,024 TPU chips. Qwen3-8B reaches 51-52% MFU with functional AdamW traced into the Graph Trainer. Qwen3 235B MoE attains 29%+ MFU with scan-over-layers and graph splitting, and a weak-scaling test retained 92% per-chip throughput from 64 to 8,192 chips.

**Tags:** [FSDP](https://yomu.fyi/topic/fsdp), [LLMs](https://yomu.fyi/topic/llm), [Mixture of Experts](https://yomu.fyi/topic/mixture-of-experts), [PyTorch](https://yomu.fyi/topic/pytorch), [TPU](https://yomu.fyi/topic/tpu)

- Source: [Google](https://developers.googleblog.com/integrating-torchtpu-into-torchtitan-and-optimizing-models-for-tpu)
- Source URL: https://developers.googleblog.com/integrating-torchtpu-into-torchtitan-and-optimizing-models-for-tpu
- Ingested by Yomu: 2026-10-20T18:00:19.223Z

[Read original post](https://developers.googleblog.com/integrating-torchtpu-into-torchtitan-and-optimizing-models-for-tpu)