Does warming up the torch.compile necessarily trigger all cases?

torch.compile torch._dynamo
1

my test code is as follows:

@torch.compile(backend=my_compiler)
def my_function(inputs):
    x = inputs["x"]
    y = inputs["y"]
    x = torch.add(x, 0.5)
    if x.mean() > 0.5:
        x = x / 1.1
    return x * y

I debug into the compile process and find that tensor x always save the true value instead of fake tensor, thus different inputs can trigger all code branches.
In theory, x should be a fake tensor, so is it possible that no mattter how many times you warm up models, it will never trigger some corner branches?

2

The direct answer to your question is no: torch.compile is a just-in-time, tracing-based compiler, so in order to fully “warm up” compile (and avoid new compilations / reach a steady state), you need to sure you warm it up with a full set of model inputs that go down all potential branches of your model.

On your specific example: when you branch on “x.mean() > 0.5”, you are attempting to branch on a data-dependent value. Compile will capture a graph up until the branch, compile and execute it, and then use the computed value to “continue” and capture a new graph on the side of the control flow that the value sends it.

If you want to see what the actual graphs look like, you can run with “TORCH_LOGS=graph_code” to see the two distinct graphs that compile produces

1 Like
3

Thanks for your reply. dynamo says it will use faketensors which only include meta data, by in this case real data is captured.

Does this mean that if there is a computational process that is data-dependent, then the values of the real inputs are used, otherwise the faketensor is used?