Lightning Talk: Building Intermediate Logging for PyTorch - Kunal Bhalla, Meta

Lightning Talk: Building Intermediate Logging for PyTorch - Kunal Bhalla, Meta
One of the best ways to understand what's going on with your model is to actually look at all the numbers flowing through it. In this talk I'll walk through implementing an API to capture all values as they flow through a PyTorch model: Module arguments, parameters, buffers, return values -- as well as gradients -- in a way that Just Works even if you're using TorchScript, torch.compile, transforming the model with Torch Fx, distributing it with torch.package or just pickling it and passing it around. We'll end up talking about several Python and PyTorch internals along the way. Having all the numbers available opens up a lot of opportunities to understand and debug your model as well, and I'll also talk through some case studies, though I'm far from a domain expert there.