PyTorch Compilation Tutorial
This tutorial is designed for ECE6980 Spring 2023 @ Cornell. It covers the basic compilation techniques in the PyTorch ecosystem.
- Basic PyTorch
- Hierarchical module construction
- TorchScript
- Tracing mode
- Scripting mode
- torch.fx
- Symbolic trace
- Intermediate representation (IR)
- Graph traversal and manipulation
- limitation
- slapo
- Progressive optimization (manual + compiler optimization)
- Subgraph matching
- Operator fusion
- Module replacement
- Quantization (optional)
- PyTorch 2.0 (optional)
Set up environment
Please use a clean environment and install the following package.
pip3 install torch numpy slapo tabulate transformers
Reading materials
Paper
- Adam Paszke et al., PyTorch: An Imperative Style, High-Performance Deep Learning Library, NeurIPS, 2019.
- James Reed et al., Torch.fx: Practical Program Capture and Transformation for Deep Learning in Python, MLSys, 2022.
- Hongzheng Chen et al., Decoupled Model Schedule for Deep Learning Training, arXiv:2302.08005, 2023.
Blogs and Documents
- Yuxin Wu, TorchScript: Tracing vs. Scripting
- James Reed, Introduction to TorchScript
- PyTorch, JIT Technical Overview