# How to Optimize Large AI Models with PyTorch

Michael Gschwind, Meta Platforms | MLOps Podcast | Episode 274 | 57:44
Hosted by Demetrios Brinkmann

Source: https://www.youtube.com/watch?v=i-Vo0AlaR7g
Channel: MLOps Community, now AAIF Live (https://www.youtube.com/@AAIFLive-x1r). Summarised by MLOps Talks.
Page: https://mlopstalks.com/talks/how-to-optimize-large-ai-models-with-pytorch
Published: 2024-11-26
Tags: benchmarks, edge, inference, open-source

## TL;DR
- TorchChat brings PyTorch's server and on-device optimization work into one example for running large language models.
- Public benchmarks give engineers a shared place to measure performance, find slow paths, and contribute improvements.
- Michael Gschwind recommends getting a system running first, then optimizing only the parts that limit performance, memory, or battery life.

## Summary
Michael Gschwind describes how his work on accelerators led from game consoles and supercomputers to AI systems. He explains that AI fits accelerators well because models perform repeated matrix operations across many neurons and layers. His current project, TorchChat, combines PyTorch work on Better Transformers, Accelerated Transformers, Flash Attention, torch.compile, and ExecuTorch into a path from server inference to on-device AI. TorchChat also gives the community a shared benchmark. Michael argues that measured performance creates a practical feedback loop: engineers see a slow result, investigate it, and submit improvements. He presents collaboration with NVIDIA, ARM, Apple, Qualcomm, and open-source projects such as XNNPACK as necessary for low-level optimization. Server and device deployments share constraints such as memory, bandwidth, and compute, so techniques such as quantization can benefit both. He advises teams to avoid premature specialization and reuse existing work wherever possible.

## Key ideas
### Accelerator programming led Michael Gschwind from gaming systems into AI
[02:10](https://www.youtube.com/watch?v=i-Vo0AlaR7g&t=130s)
Michael Gschwind worked on accelerators used in the PlayStation 3 and Xbox 360, where performance required giving up many conveniences of a general-purpose CPU. Programmers had to manage much more themselves. That experience led him to build software environments, compilers, and systems for accelerators. He later worked on supercomputers including Roadrunner, Blue Gene, and Summit, where GPUs worked closely with CPUs. The connection to AI became clear after ImageNet results showed that convolutional neural networks could use this hardware effectively. Michael says his team trained an AlexNet model in under an hour, before Facebook achieved a better result.

### PyTorch's open-source model lets engineers build on existing work
[10:05](https://www.youtube.com/watch?v=i-Vo0AlaR7g&t=605s)
Michael Gschwind says PyTorch's commitment to open source and community contributions has helped people pursue their own ideas without first building a complete framework. An engineer can work on a better matrix multiplication kernel or quantization method while relying on others for the surrounding infrastructure. He describes research as a team sport because different teams often hold different parts of the answer. Competition can help too, since another group may take an approach and improve it. He gives credit to the people who built the systems and libraries that later work could reuse.

### TorchChat combines several PyTorch optimization efforts into one inference path
[11:47](https://www.youtube.com/watch?v=i-Vo0AlaR7g&t=707s)
Michael describes TorchChat as the result of several development strands in the PyTorch and large language model ecosystem. Better Transformers began the acceleration work, while Accelerated Transformers added Flash Attention, a newer scaled dot product attention operator, and other infrastructure improvements. torch.compile makes it possible to export PyTorch models so they can run outside a Python-hosted environment. ExecuTorch brings optimized inference to on-device and embedded applications. TorchChat combines these pieces, showing how to deploy language models from servers through to on-device applications.

### Shared benchmarks turn performance work into a community feedback loop
[14:08](https://www.youtube.com/watch?v=i-Vo0AlaR7g&t=848s)
TorchChat is also a driver for further optimization because it gives people a place to measure changes. Michael's engineering rule is that what gets measured regularly tends to improve. A visible result can prompt an engineer to ask why it is slow and submit a pull request. Running TorchChat on Mac systems exposed questions about ARM code generation, which led to work with Meta's LLVM and compiler teams, ARM, and Apple. He says the same reference implementation can help people discover ideas that apply beyond language model inference.

### Low-level speedups come from combining hardware, kernels, and software communities
[22:47](https://www.youtube.com/watch?v=i-Vo0AlaR7g&t=1367s)
Michael points to XNNPACK developers and a team engineer who developed a four-bit matrix multiplication for ARM. The implementation used an ARM-specific instruction and combined multiplication, accumulation, quantization, and groupwise scaling in one kernel. He says this made CPU inference for on-device language models much faster. The result depended on several layers working together, including the ARM instruction set, XNNPACK, the kernel, ExecuTorch, and TorchChat. He argues that this kind of result would be difficult to create in isolation.

### Server and on-device inference share many of the same constraints
[29:46](https://www.youtube.com/watch?v=i-Vo0AlaR7g&t=1786s)
Michael rejects treating cloud and on-device inference as completely separate areas. On-device models matter for disconnected use, privacy, and applications such as cars. ExecuTorch and TorchChat allow the same model definition and software infrastructure to move from servers to devices, although the chosen model size may differ. Both settings face limits in memory, compute, and bandwidth. Quantization reduces the number of bits used for weights and activations, which can reduce memory use and the amount of data sent across a bus. The same general method can help a device and a GPU, even when their capacities differ.

### Device-specific libraries can hide much of the hardware detail from developers
[41:16](https://www.youtube.com/watch?v=i-Vo0AlaR7g&t=2476s)
Michael explains that a PyTorch linear operator can use different implementations underneath. GPU workloads may use cuBLAS, AMD systems may use Composable Kernels, and ARM systems may use their own libraries alongside XNNPACK. Triton can generate GPU code when a model is exported. Hardware vendors and software teams contribute these implementations so the system can select an appropriate version for the target device. Developers can then get most of the way to a working solution through common primitives, adding device-specific work only where the target requires it.

### Premature optimization often causes teams to rebuild work that already exists
[45:49](https://www.youtube.com/watch?v=i-Vo0AlaR7g&t=2749s)
Michael's main antipattern is optimizing before a system runs. Teams may assume their problem is unique and build a narrow solution, even though similar constraints have been handled elsewhere in the ecosystem. His advice is to get the system working, identify the actual limits, and optimize only the pieces that affect the target, such as performance or battery life. He gives quantization and model export as examples of components that should usually be reused. A team building an LLM smartwatch could focus on the RISC-V kernels that make the product possible while reusing the rest of the stack.

### Faster experiment turnaround can improve model and system innovation
[54:10](https://www.youtube.com/watch?v=i-Vo0AlaR7g&t=3250s)
Looking ahead, Michael expects more work on Transformer implementations, Mamba, and lower-precision numeric formats. He mentions a progression from FP32 and FP16 toward FP8 and four-bit floating point. Much of this work is experimental, so the number of experiments and the speed of each iteration matter. Faster iteration produces more information and can lead to better system designs. His infrastructure focus is therefore split between improving the systems themselves and reducing the time required to run each experiment.

## Notable quotes
- Michael Gschwind: "Research is a team sport." (09:05)
- Michael Gschwind: "What gets measured gets improved." (14:48)
- Michael Gschwind: "I think the number one magic is collaboration." (25:53)
- Michael Gschwind: "I think the biggest antipattern is premature optimization." (45:49)
- Michael Gschwind: "First get it to run." (49:04)

## Tools & references mentioned
- PyTorch
- TorchChat
- Better Transformers
- Accelerated Transformers
- Flash Attention
- torch.compile
- ExecuTorch
- ImageNet
- AlexNet
- XNNPACK
- ARM
- Apple
- NVIDIA
- Qualcomm
- AMD
- CUDA
- cuBLAS
- Composable Kernels
- Triton
- RISC-V
- Mamba

## Who should watch
- You are choosing between server and on-device inference and want to understand which optimization work can be shared.
- Your team is building custom kernels or hardware-specific code and needs a practical case for collaborating with vendors and open-source projects.
- You are optimizing a model before you have measured the real bottleneck and want a clearer order of operations.

## Related talks

- [PyTorch: Bridging AI Research and Production](https://mlopstalks.com/talks/pytorch-bridging-ai-research-and-production) (Dmytro Dzhulgakov, Facebook, 52:55)
- [Enabling Efficient Trillion Parameter Scale Training for Deep Learning Models](https://mlopstalks.com/talks/enabling-efficient-trillion-parameter-scale-training-for-deep-learning-models) (Tunji Ruwase, Microsoft, 27:36)
- [Efficient Deployment of Models at the Edge](https://mlopstalks.com/talks/efficient-deployment-of-models-at-the-edge) (Krishna Sridhar, Qualcomm, 51:34)
- [Small Models, Big Ideas: The Next Frontier in AI](https://mlopstalks.com/talks/small-models-big-ideas-the-next-frontier-in-ai) (Korri Jones, Chick-fil-A Corporate Support Center & Valdimar Eggertsson, Snjallgögn (Smart Data inc.) & Sophia Skowronski, Breckinridge Capital Advisors & Lihu Chen, Imperial College London & Binoy Perera, MLOps Community, 58:44)
- [Performance Optimization and Software/Hardware Co-design across PyTorch, CUDA, and NVIDIA GPUs](https://mlopstalks.com/talks/performance-optimization-and-software-hardware-co-design-across-pytorch-cuda) (Chris Fregly, AI performance engineer, startup founder, and investor, 1:25:50)