De-mystifying PyTorch for ASICs: When (and Why) To Move Your Development To AI Accelerators.

A code-first check on moving PyTorch from GPUs to Google TPUs and AWS Trainium, and when it is worth doing.

The talk

GPU availability and cost are squeezing ML teams, making ASICs like Google TPUs and AWS Trainium attractive alternatives. But does the software stack hold up? The session compares the PyTorch/XLA and TorchNeuron compiler stacks through side-by-side code diffs and benchmarks, fine-tuning Llama 4, Gemma 3 and Qwen 3 and training CNNs and ViTs, and answers three questions.

The code: How much rewriting is actually required?
The performance: Which model architectures thrive on ASICs, and which ones fail due to dynamic shapes?
The debugging: What happens when you hit an out-of-memory error or a compilation hang?

Attendees leave with a Migration Decision Matrix, to judge whether their own workload is ready.

The evidence

The talk's code and assets run from infrastructure and setup through models, training, fine-tuning and inference. Its companion repository fine-tunes three models, SFT then DPO with LoRA, on both GPU and TPU. For a later measurement of serving and training on AWS's chips, read the study below.

63 to 71 ms

per token on Inferentia2 from 1 to 32 streams, while throughput grows 26×

ASICs, July and August 2026

Llama 3.1 8B on AWS Trainium and Inferentia

A fine-tune on Trainium, served on Inferentia2, measured beside an MI300X and an H200.

Where and when

Conference: PyTorch Conference Europe, Paris
Date: Wednesday 8 April 2026, 16:25 to 16:50 CEST
Room: Central Room
Track: Frameworks & Compilers, intermediate level
Speaker: Alpha Romer Coma

Watch and read

The recording The conference listing The talk's code and assets Fine-tuning pipeline code The talk again, at Data Engineering Pilipinas