De-mystifying PyTorch for ASICs: When (and Why) To Move Your Development To AI A... Alpha Romer Coma

This talk explores the transition of machine learning development from traditional GPUs to AI accelerators like ASICs, specifically focusing on Google TPUs and AWS Trainium. The speaker provides a practical, code-first examination of migrating PyTorch workloads to these ASICs, delving into the underlying compiler stacks like PyTorch/XLA and TorchNeuron. The session addresses the reality of the 'Compiler Tax,' evaluates the necessary code modifications, and identifies which model architectures perform well on ASICs and which struggle due to dynamic shapes. By the end, attendees will gain a comprehensive ‘Migration Decision Matrix’ to assess the readiness of their workloads for this transition.