Model-Changing Transforms With Torch.compile
summary
Proposes using torch.compile’s three outputs — graph, guards, and wrapper — as a transformable contract so quantization, distribution and other changes act on model, inputs and guards without rewriting models. Builds a prototype that uses Dynamo + Inductor, composes transforms for sharding/materialization (FSDP-style), state_dict load/save and quantization, and runs end-to-end by swapping the trace via convert_frame monkeypatch and a custom backend; the hacky ~500 lines of glue proves the idea works though it needs a safer API.
AI-written summary · about
chapters (13)
- 0:00Introduction
- 1:45Removing friction in model optimization
- 3:56Source code to torch.compile
- 5:24Composable optimizations in PyTorch
- 7:10Challenges in transforming code
- 8:36torch.compile toolkit
- 10:22FSDP and materialization transforms
- 12:53Composable transforms for optimizations
- 13:35Internal mechanics of guards and wrappers
- 15:25Hacking frame conversion for model switching
- 17:14Implementation and demo results
- 19:49Benefits of model-changing transforms
- 20:52Future directions
related talks
36mDL Compiler Panel Discussion
26mThunder: Distribute and Optimize Your PyTorch Models With Zero...
20mPyTorch x Transformers Journey
08mPyTorch 2.0: TorchDynamo
21mThe New Vision of Compiler-Accelerated
24mAdvancements in High-Performance AI/ML through PyTorch's Python Compiler
09mTogether Goes Brrr: Threading Research & Production with Torch Compile
25mTransformers: Standardizing Model Definitions Across the PyTorch Ecosystem