Clone/Copy compiled model

torch.compile
1

I have two models with the same architecture that I torch.compile() separately. And then during the first run of the model when the actual data is run through the model, the graphs are built. This process takes some time.

model_1 = SomeModel()
model_1_compiled = torch.compile(model_1)

model_2 = SomeModel()
model_2_compiled = torch.compile(model_2)


pred_1 = model_1_compiled(input_1)
pred_2 = model_2_compiled(input_2)

I was wondering if it is possible to do save some time by compiling only one model, and then copy the built graph for the second model?

Something like this:

model_1 = Model()
model_1_compiled = torch.compile(model_1)

pred_1 = model_1_compiled(input_1)

model_2_compiled = COPY_COMPILED_MODEL(model_1_compiled)

pred_2 = model_1_compiled(input_2)
2

We’re looking at having something sort of like with saving the inductor compilation cache, there’s a POC here PT 2.0 - Are compiled models savable - #7 by marksaroufim and I’ll be working on cleaning it up to ship in time for next release

1 Like