Enabling torch.compile in training and inference?

[qchempku2017]

Dear community,
Has anyone performed test before and after applying torch.compile to the pl_module in diffusion/run.py, scripts/finetune.py and scripts/generate.py? Does torch.compile help increasing training and inference speed significantly? Would it worth it to make a pull request for model compilation?

[qchempku2017]

Here I supplement my own test result on an Nvidia 4060Ti 8G. The inductor backend and max-autotune mode was used.

When training chemical_space_enenergy_above_hull model on 500 structures, 4060Ti finished one epoch in around 1 min 15 s with torch compile off, and spent around 1 min 30 s after turning torch compile on. Using torch.compile on pl_module slightly slowed down the training. I'm uncertain with the exact cause.

[StrongerXi]

A few things to try:

1. Upgrade pytorch to 2.7 if you haven't
2. Try mode="default"
3. Run with TORCH_LOGS="graph_breaks" and fix the graph breaks
4. If (3) is too hard, try applying torch.compile to smaller regions rather than the entire model, e.g., apply to each transformer.

Hope this helps!