De-mystifying PyTorch for ASICs

Abstract

The PyTorch Conference Europe talk, extended for a community study group. It covers how ASICs and the XLA and Neuron compilers work, PyTorch/XLA on Google TPU and TorchNeuronX on AWS Trainium, the compiler tax, common ASIC errors, and migration notes. It adds where to get TPUs, the Paris experience, and how to speak at Linux Foundation events.

Outline

1. Deep learning computations and how ASICs and XLA work
2. Google TPU with PyTorch/XLA
3. AWS Trainium with TorchNeuronX
4. Performance analysis and the compiler tax
5. Common ASIC errors
6. Migration notes
7. Where to get TPUs and how to apply to the TPU Research Cloud
8. The PyTorch Conference Europe experience and speaking at Linux Foundation events

Key takeaways

• PyTorch stays the standard framework, and ASICs are building on top of it
• Moving to TPU or Trainium means learning new compile semantics and op coverage gaps