Lightning Talk: Building Intermediate Logging for PyTorch
summary
An API captures module arguments, parameters, buffers, outputs and gradients that works across TorchScript, torch.compile, Torch FX, torch.package and plain pickles by abusing forward hooks plus a noop C++ autograd to observe gradients in scripted code. Implementation uses dynamic codegen to emit per-module hooks, tricks to force the compiler to keep buffers/calls, prefers buffers over parameters, and a linecache + importlib loader plus base32-encoded function names to survive pickling. More than 100 tests validate it; cases found gradients off by orders of magnitude and one model is halved in size after fixes.
AI-written summary · about
chapters (13)
- 0:00Intermediate logging in PyTorch
- 1:25PyTorch logging APIs
- 2:27FX transforms and TorchScript issues
- 3:09Forward hooks for gradients
- 4:13Dynamic hook generation
- 5:17Compiler code retention tricks
- 6:21Buffers vs parameters
- 6:43Distributing code with line cache
- 7:45Pure pickling with base32
- 8:47Torch package dependencies
- 9:08Testing transformed models
- 9:29Real-world logging fixes
- 10:09Demo and acknowledgments
related talks
12mLightning Talk: Backpropagation-Free Optimization in PyTorch
11mLightning Talk: But How Do You Autograd Through Your Simulator? Making a...
17mLightning Thunder: Supercharged PyTorch for Modern Hardware
13mLightning Talk: Challenges and Standardization in PyTorch Ecosystem...
08mSponsored Session: Lightning Talk: PyTorch in Production: Boosting LLM Training and...
45mLIGHTNING TALKS
36mDL Compiler Panel Discussion
19mBuilding Trust for Users and Regulators Alike: A Cost-Efficient PyTorch Pat...