PyTorch/XLA documentation

torch_xla is a Python package that implements XLA as a backend for PyTorch.

Familiar APIs Create and train PyTorch models on TPUs, with only minimal changes required.

High Performance Scale training jobs across thousands of TPU cores while maintaining high MFU.

Cost Efficient TPU hardware and the XLA compiler are optimized for cost-efficient training and inference.

Getting Started

Install with pip.

pip install torch torch_xla[tpu]

Verify the installation:

python -c "import torch_xla; print(torch_xla.__version__)"
python -c "import torch; import torch_xla; print(torch.tensor(1.0, device='xla').device)"

Tutorials

Learn the Basics

• PyTorch on XLA Devices
• Pytorch/XLA Overview

Distributed Training on TPU

• Learn about TPUs
• SPMD User Guide
• SPMD advanced topics
• Distributed Checkpointing
• Support for Torch Distributed
• Distributed Data Parallel (DDP)
• Fully Sharded Data Parallel (FSDP) with One Process Per Accelerator
• Fully Sharded Data Parallel (FSDP) using SPMD

Advanced Techniques

• Custom Kernels via Pallas
• Torch Export to StableHLO
• Automatic Mixed Precision
• Dynamic Shapes
• TorchDynamo Integration
• Quantized Operations
• Optimizing Repeated Layers with scan and scan_layers
• Optimize Memory Utilization with while_loop
• Speed Up Tracing with @assume_pure

Troubleshooting

• Troubleshooting Basics
• Eager Mode + Compile API
• Source of recompilations in torch_xla
• Troubleshooting recompilations

Training on GPU

• Learn about GPUs
• Custom GPU Kernels via Triton
• Running SPMD on GPU

Contributing

• Building with Bazel
• Configure A Development Environment
• C++ Debugging
• Op Lowering Guide
• Codegen Migration Guide
• Custom Hardware Plugins

API Reference

• PyTorch/XLA API
  • torch_xla
  • runtime
  • xla_model
  • distributed
  • spmd
  • experimental
  • debug