I want to torch compile the same model several times, optimized for different shapes, but pointing to the same large model weights. Allowing the same model to be dynamic over the shape ranges resulted in too large a slow down.
compiled_model1 = torch.compile(my_model)
_ = compiled_model1(dummy_input1) # Slow
compiled_model2 = torch.compile(my_model)
_ = compiled_model2(dummy_input2) # Slow
compiled_model3 = torch.compile(my_model)
_ = compiled_model3(dummy_input3) # Fast
compiled_model4 = torch.compile(my_model)
_ = compiled_model4(dummy_input4) # Fast
The ID differs between compiled model variables, but the fact that models 3, 4 and onwards suddenly compile (on first pass) quickly worries me. Does something optimize the graph generation after a few times, or are they all pointing to the same graph, and the reason #2 is slow, is due to it becoming dynamic with the different input size?
Setting dynamic=False in torch.compile() makes every JIT take the full, slow, time.