Make torch.compile the default

[sophie-xhonneux]

Is your feature request related to a problem? Please describe.

There is a flag in the config to compile the model with torch.compile() I believe by default it is set to False. I think we should make the default True, to ensure we don't change the model in ways that break that functionality even if at the moment the performance gain is not the greatest.

Describe the solution you'd like

Set cf.compile_model to True by default

Describe alternatives you've considered

No response

Additional context

No response

Organisation

No response

[sophie-xhonneux]

@clessig @tjhunter any objections?

[clessig]

Is the current code compiling for you? I get:

  File "/lus/h2resw01/hpcperm/nacl/WeatherGenerator/.venv/lib/python3.12/site-packages/torch/_dynamo/variables/constant.py", line 55, in __init__
    assert ConstantVariable.is_base_literal(
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
AssertionError: 
Cannot construct `ConstantVariable` for value of type <class 'torch._dynamo.variables.tensor.TensorVariable'>.

This failure likely due to PyTorch-internal use of `ConstantVariable` on
non-literal python values, please try using `VariableTracker.build` instead. If
you believe it's a necessary and legitimate use case (the value is immutable and
can't easily be represented with another `VariableTracker` class), please add
its type to `common_constant_types`.


from user code:
   File "/lus/h2resw01/hpcperm/nacl/WeatherGenerator/src/weathergen/model/model.py", line 388, in torch_dynamo_resume_in_forward_at_382
    for fstep in range(forecast_offset, forecast_offset + forecast_steps):

Set TORCH_LOGS="+dynamo" and TORCHDYNAMO_VERBOSE=1 for more information


You can suppress this exception and fall back to eager by setting:
    import torch._dynamo
    torch._dynamo.config.suppress_errors = True

> /lus/h2resw01/hpcperm/nacl/WeatherGenerator/.venv/lib/python3.12/site-packages/torch/_dynamo/variables/constant.py(55)__init__()

flash attention also didn't used to be torch.compile compatible, although this might have changed.