Hugging Face on PyTorch / XLA TPUs

Summary

PyTorch / XLA integrates Cloud TPU training into Hugging Face transformers while preserving the standard Trainer interface. The implementation adds an xla device type and uses lazy execution to trace intermediate representation graphs before lowering them to XLA Higher Level Operations for compilation. Parallel data loading through MpDeviceLoader overlaps CPU tracing with TPU execution, while xm.optimizer_step coordinates gradient consolidation across data-parallel replicas. Checkpoints use xm.save to transfer tensors via CPU storage and restrict disk writes to the master process. Benchmarks on a Cloud TPU v3-8 system training bert-large-uncased on WikiText-103 demonstrate completion times of 178.4 minutes in FP32 precision and 106.4 minutes in BF16 precision.

Context

PyTorch users required a way to run and scale transformer models on Google Cloud TPUs without abandoning the familiar Hugging Face Trainer interface.

Approach / What changed

Hugging Face and the PyTorch / XLA team integrated the xla device type into the Trainer module, handled gradient synchronization via xm.optimizer_step, implemented overlapped tracing and execution using MpDeviceLoader, and routed checkpoint saving through CPU tensors via xm.save.

Takeaways

• PyTorch / XLA tensors execute lazily, accumulating operations into an intermediate representation graph until explicitly evaluated or stepped.
• The MpDeviceLoader pipelines data loading to trace step n+1 on the host CPU while step n executes on TPU hardware.
• Training bert-large-uncased on WikiText-103 using a v3-8 Cloud TPU took 178.4 minutes in FP32 and 106.4 minutes in BF16.

Related reading