DEVOURED
Google and Meta Bring Native PyTorch Support to TPUs
Google and Meta launched TorchTPU, an open-source PyTorch backend that brings native eager execution and distributed support to Google's TPU hardware.
What: TorchTPU enables PyTorch developers to run models on TPUs with minimal code changes (e.g., swapping `.to("cuda")` for `.to("tpu")`). It includes a new `DeferAndFuse` eager mode, support for distributed collectives (DDP, FSDP), and integration with tools like vLLM and Hugging Face.
Why it matters: By removing the previous requirement for intrusive source-code modification and JAX-centric workflows, this project significantly lowers the barrier for researchers and engineers to utilize TPU hardware for PyTorch-native development.
Takeaway: If you are running PyTorch models on Google Cloud, test your workload with TorchTPU by installing the library and changing your device target to `tpu` to evaluate potential performance gains.
Deep dive
- Eager vs. Compiled: Supports both standard eager execution and
torch.compilevia TorchDynamo. - DeferAndFuse: A new execution mode that fuses multiple Aten operations to increase throughput and reduce HBM memory pressure.
- Distributed: Provides native Process Group API support, enabling use of FSDP, DTensor, and other PyTorch distributed idioms.
- Kernels: Users can author custom kernels using Pallas or use the Tokamax library of high-performance TPU kernels.
- Compatibility: Validated against 5,000+ Hugging Face models with 98% runtime stability.
- Precision: Maps PyTorch precision settings to TPU-native bfloat16 (BF16) modes for optimal MXU utilization.
Decoder
- TPU (Tensor Processing Unit): Google's custom-developed ASIC designed specifically for accelerating machine learning workloads.
- XLA (Accelerated Linear Algebra): A domain-specific compiler for linear algebra that optimizes PyTorch and TensorFlow graphs.
- StableHLO: A high-level intermediate representation for ML models that serves as a portable interface for hardware targets.
- MXU (Matrix Multiply Unit): A hardware component within a TPU chip specialized for performing massive matrix multiplications efficiently.
- FSDP (Fully Sharded Data Parallel): A training technique that shards model parameters, gradients, and optimizer states across GPUs/TPUs to enable training of massive models.
Original article
Full article content is not available for inline reading.
Explore all 55 articles in today's digest
Fresh dev news, AI-digested daily. Never miss a trend at devoured.fyi