compile_model

compile_model(model=None, *, mode='default')[source]

A decorator to compile a model using torch-compile.

Parameters:

  • model (torch.nn.Module) – Model to be compiled.

  • mode (str) –

    Mode to be used for torch-compile. Available modes are:

    • disable disables torch-compile

    • default balances performance and overhead

    • reduce-overhead reduces overhead of CUDA graphs (useful for small batches)

    • max-autotune leverages Triton/template based matrix multiplications on GPUs

    • max-autotune-no-cudagraphs similar to “max-autotune” but without CUDA graphs

Returns:

Compiled model.

Return type:

torch.nn.Module