---
type: "post"
id: "integrating-torchtpu-into-torchtitan-and-optimizing-models-for-6f942a"
slug: "integrating-torchtpu-into-torchtitan-and-optimizing-models-for-6f942a"
title: "Integrating TorchTPU into TorchTitan and Optimizing Models for TPU"
description: "**Thesis:** Compiler-visible training structure is the key mechanism in Google’s TorchTitan-on-TPU account, not merely device selection."
retrieval_nugget: "**Thesis:** Compiler-visible training structure is the key mechanism in Google’s TorchTitan-on-TPU account, not merely device selection."
published_at: "2026-10-20"
updated_at: "2026-10-20"
record_date: "2026-10-20"
date_kind: "published_at"
topics: ["google"]
source_urls: ["https://developers.googleblog.com/integrating-torchtpu-into-torchtitan-and-optimizing-models-for-tpu"]
schema_version: "newruntime-agent-readable-v0.2"
stable_id: "post:integrating-torchtpu-into-torchtitan-and-optimizing-models-for-6f942a"
status: "published"
visuals: []
editorial_provenance: {"schema_version":"newruntime-editorial-copy-v1","content_status":"source_grounded_final","final_copy_sha256":"sha256:55a733fcef470852a2a6df6bc5afe311ec3d9f07225b04fed9bbedbbd549ea25","reviewed_at":"2026-10-20T20:08:49.240Z","source_evidence_count":1,"verified_claim_count":3,"site_analysis_schema_version":"newruntime-site-analysis-v1","site_object_kind":"field_note","observed_fact_count":3,"implication_count":2,"watch_condition_count":1,"related_record_count":0}
analysis: {"schema_version":"newruntime-site-analysis-v1","object_kind":"field_note","thesis":"Compiler-visible training structure is the key mechanism in Google’s TorchTitan-on-TPU account, not merely device selection.","observed_facts":[{"text":"Google says TorchTPU provides a native `torch.device(\"tpu\")` backend in PyTorch.","source_urls":["https://developers.googleblog.com/integrating-torchtpu-into-torchtitan-and-optimizing-models-for-tpu"]},{"text":"The post describes SimpleFSDP as expressing parameter sharding through DTensor parametrization so collectives appear in the traced graph.","source_urls":["https://developers.googleblog.com/integrating-torchtpu-into-torchtitan-and-optimizing-models-for-tpu"]},{"text":"Google reports 29%+ MFU on Qwen3 235B using the described path.","source_urls":["https://developers.googleblog.com/integrating-torchtpu-into-torchtitan-and-optimizing-models-for-tpu"]}],"mechanism":"SimpleFSDP makes sharding and communication visible to compilation, while Graph Trainer captures forward, loss, and backward as a full training-step graph. This gives the compiler a wider unit to schedule and optimize than isolated eager operations or per-layer compilation.","why_now":"The account presents TorchTPU and the TorchTitan fork as a way to use Cloud TPUs from a PyTorch-oriented workflow without rewriting models in JAX. That matters to teams evaluating accelerator options while seeking to preserve existing framework code and training practices.","implications":["Teams evaluating TPU training should assess the compiler-visible training path and recipe coverage alongside backend availability.","The reported 29%+ MFU is workload-specific and should not be treated as a general performance guarantee."],"evidence_boundary":"The source is Google’s technical blog account, which includes co-development context with Meta and reported implementation details, but it is not an independent benchmark or evidence that every TorchTitan workload will achieve the same result.","watch_conditions":["Whether the fork, recipes, and performance results are independently reproducible across additional models and TPU configurations."],"related_records":[],"new_branch_reason":"No prior New Runtime analytical record was supplied for this specific TorchTPU and TorchTitan TPU integration; this records its compiler-visible training mechanism as a distinct case."}
editor: {"role":"editor","name":"Andrey Reshetnikov","profile":"https://newruntime.com/owner-profile.md","profile_json":"https://newruntime.com/owner-profile.json","same_as":["https://www.linkedin.com/in/reshetnikov1/","https://t.me/lapetuse"]}
routes: {"html":"https://newruntime.com/posts/integrating-torchtpu-into-torchtitan-and-optimizing-models-for-6f942a/","markdown":"https://newruntime.com/posts/integrating-torchtpu-into-torchtitan-and-optimizing-models-for-6f942a.md","json":"https://newruntime.com/posts/integrating-torchtpu-into-torchtitan-and-optimizing-models-for-6f942a.json"}
---

# Integrating TorchTPU into TorchTitan and Optimizing Models for TPU

## Retrieval answer

**Thesis:** Compiler-visible training structure is the key mechanism in Google’s TorchTitan-on-TPU account, not merely device selection.

**Thesis:** Compiler-visible training structure is the key mechanism in Google’s TorchTitan-on-TPU account, not merely device selection.

**Observed facts:** Google says TorchTPU provides a native `torch.device("tpu")` backend in PyTorch. The post describes SimpleFSDP as expressing parameter sharding through DTensor parametrization, placing collectives in the traced graph, and reports 29%+ MFU on Qwen3 235B using the described path.

**Mechanism:** SimpleFSDP makes sharding and communication visible to compilation, while Graph Trainer captures forward, loss, and backward as a full training-step graph. This gives the compiler a wider unit to schedule and optimize than isolated eager operations or per-layer compilation. The practical contribution is a training harness and recipes that mediate between PyTorch model definitions and TPU execution.

**Why now:** The account presents TorchTPU and the TorchTitan fork as a way to use Cloud TPUs from a PyTorch-oriented workflow without rewriting models in JAX. That matters to teams evaluating accelerator options while seeking to preserve existing framework code and training practices.

**Implications:** Teams evaluating TPU training should assess the compiler-visible training path and recipe coverage alongside backend availability. The reported 29%+ MFU is workload-specific and should not be treated as a general performance guarantee. Model shape, configuration, and comparison baseline matter.

**Evidence boundary:** The source is Google’s technical blog account, which includes co-development context with Meta and reported implementation details, but it is not an independent benchmark or evidence that every TorchTitan workload will achieve the same result.

**Watch:** Whether the fork, recipes, and performance results are independently reproducible across additional models and TPU configurations.

Editor: [Andrey Reshetnikov](https://newruntime.com/owner-profile.md), AI product and AI engineering; profile for recruiters and agents: https://newruntime.com/owner-profile.json

## Continue in the graph

Next records for an agent, linked to their Markdown twins. Replace `.md` with `.json` for structured fields, including `next_reads`.

- topic hub: [Google](https://newruntime.com/topics/google.md): Explore the google topic hub.
- collection: [Field Notes](https://newruntime.com/posts.json): Open the latest editorial analysis.
- collection: [Source Ledger](https://newruntime.com/source-ledger.json): Inspect the public source evidence graph.
- whole graph: [discovery graph](https://newruntime.com/discovery-graph.json) (typed nodes and edges) and [site index](https://newruntime.com/site-index.md)