Repository: pytorch
Issue: [Inductor] Freezing Add support for Caching Parameter Conversions
Labels: feature,triaged,oncall: pt2,module: inductor

### 🚀 The feature, motivation and pitch

Freezing is a WIP feature within inductor that, similar to JIT freezing, will attempt to inline weights as constants in optimization and run constant folding and other optimizations on them. After freezing, weights can no longer be updated.

We may run transformations on parameters like converting to channels last and padding for mms. Because these transformations would naively duplicate parameter memory use, we have an option to discard the original parameters, [freezing_discards_parameters](https://github.com/pytorch/pytorch/blob/main/torch/_inductor/config.py#L213). 

As a result, currently, there is no way to both avoid duplication of memory of parameters and recompile.

```
import torch
import torchvision

import torch._inductor
import torch._inductor.config as config

config.freezing = True
config.freezing_discard_parameters = True

resnet = torchvision.models.resnet18().eval().cuda()

@torch.compile(mode="reduce-overhead")
def run_model(resnet, inp):
    return resnet(inp)

with torch.no_grad():
    run_model(resnet, torch.rand([8, 3, 255, 255], device="cuda"))

with torch.no_grad():
    run_model(resnet, torch.rand([1, 3, 255, 255], device="cuda"))
# RuntimeError: Trying to Run Pytorch Eager Module After Dynamo Freezing. The original parameters have been discarded
``` 

This isn't a fundamental limitation. The constant-folded conversions we're doing to the parameters are the same regardless of the input shape (at least for vast majority of models). 

If we express all the parameter conversions as transformations in the fx graph, we can cache their final result. The process would look something like :

- On first compilation, record the mapping from real parameters to their [ErasedTensor](https://github.com/pytorch/pytorch/blob/main/torch/_inductor/freezing.py#L186) replacement. When we run constant folding on parameters, record the operators we run, their inputs, and their outputs. Intermediary outputs in a chain of conversions should get mapped to an ErasedTensor just as we did with real parameters. 
- On subsequent recompilation, when we would attempt to run an operator in constant folding, we'll check in the cache if we already have a cached run of that operator with those particular inputs. 

This feature is also important because it would enable us to do real-time parameter updates. For all of the original parameters, we will have a mapping of the series of conversions we need to run in order to replace the final constant that inductor uses in compilation. 

There are a few things that still need investigation, such as getting dynamo to place nicely with the ErasedTensor on recompilation. Additionally, it's possible that some parameters may have different series of transformations from run to run.  This could possibly be resolved with asking users to run with dynamic shapes first, or by having an api to allow users to not constant fold certain parameters. 


 




### Alternatives

_No response_

### Additional context

_No response_

cc @ezyang @msaroufim @wconstab @bdhirsh @anijain2305 @zou3519 @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @peterbell10 @ipiszy @yf225 @chenyang78 @kadeng @muchulee8 @aakhundov @ColinPeppler @Xia-Weiwen @ngimel