Is there a way to export a warmed up graph so that it is ready to go from cold boot?

[amazingvince]

I have not tried straight pickle but did not see any documentation about exporting model.

[pommedeterresautee]

PyTorch 2.0 has a way to save on disk and reload compiled Triton kernels.
Using their API should save lots of time during the warmup.
We just merged a merge to a recent version of PyTorch 2.0 so I guess it's very doable to just call their API to launch Triton kernels instead of the original Triton one.

However, a CUDA graph can't be exported and reused.
The main reason is that kernels are re-executed with their parameters.
Those parameters include gpu memory addresses.
And it's very likely that your tensor memory address will change from launch to launch (if you quit Python session and the CUDA pool is freed).

It's certainly possible in theory to update the graph parameters, but it seems very hard, so probably better to just rerun the warmup (without recompiling triton kernel ofc).

[pommedeterresautee]

@amazingvince FYI we have found a way to speedup warmup significantly in #234.

[github-actions]

This issue is marked as stale because it has been open for 30 days with no activity.

[github-actions]

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.