aitune.torch.backend.torch_inductor_jit_backend

Torch Inductor JIT backend.

Module Contents

Classes

Name	Description
`TorchInductorJitBackend`	Backend that does torch compilation with Inductor.
`TorchInductorJitBackendConfig`	Configuration for torch.compile with inductor backend.

Data

logger

API

class aitune.torch.backend.torch_inductor_jit_backend.TorchInductorJitBackend(
    config: aitune.torch.backend.torch_inductor_jit_backend.TorchInductorJitBackendConfig | None = None
)

Bases: Backend

Backend that does torch compilation with Inductor.

STATE_COMPILE_DYNAMIC

= 'compile_dynamic'

STATE_CONFIG

= 'config'

STATE_DATA

= 'data'

STATE_DEVICE

= 'device'

STATE_ORIG_MODULE

= 'orig_module'

STATE_OUTPUT_DTYPE

= 'output_dtype'

STATE_TYPE

= 'type'

_compile_dynamic

= self._config.dynamic

_config

= config or TorchInductorJitBackendConfig()

aitune.torch.backend.torch_inductor_jit_backend.TorchInductorJitBackend._activate()

Activates backend.

aitune.torch.backend.torch_inductor_jit_backend.TorchInductorJitBackend._build(
    module: torch.nn.Module,
    graph_spec: aitune.torch.module.graph_spec.GraphSpec,
    data: list[aitune.torch.module.recording_module.Sample],
    cache_dir: pathlib.Path
) -> aitune.torch.backend.backend.Backend

Builds the model with torch.compile.

aitune.torch.backend.torch_inductor_jit_backend.TorchInductorJitBackend._compile()

aitune.torch.backend.torch_inductor_jit_backend.TorchInductorJitBackend._deactivate()

Deactivates backend.

aitune.torch.backend.torch_inductor_jit_backend.TorchInductorJitBackend._deploy()

Deploys the backend.

aitune.torch.backend.torch_inductor_jit_backend.TorchInductorJitBackend._get_required_casting_dtype(
    module: torch.nn.Module,
    data: list[aitune.torch.module.recording_module.Sample]
) -> torch.dtype | None

Get the required casting dtype of the module by running a sample inference with and without autocast.

If the dtype of the output is different with and without autocast, return the dtype of the output without autocast. Otherwise, return None.

Parameters:

module

nn.Module

The module to get the dtype from.

data

list[Sample]

List of sample inputs to run through the module.

Returns: torch.dtype | None

torch.dtype: The required casting dtype. Returns None if no casting is required.

aitune.torch.backend.torch_inductor_jit_backend.TorchInductorJitBackend._infer(
    args: typing.Any = (),
    kwargs: typing.Any = {}
) -> typing.Any

Runs inference with the given arguments. Does not use autocast.

It can be replaced at runtime by _infer_with_autocast.

Parameters:

*args

AnyDefaults to ()

inference arguments

**kwargs

AnyDefaults to {}

inference keyword arguments

Returns: Any

The result of the inference.

aitune.torch.backend.torch_inductor_jit_backend.TorchInductorJitBackend._infer_with_autocast(
    args: typing.Any = (),
    kwargs: typing.Any = {}
) -> typing.Any

Runs inference with the given arguments.

Parameters:

*args

AnyDefaults to ()

inference arguments

**kwargs

AnyDefaults to {}

inference keyword arguments

Returns: Any

The result of the inference.

aitune.torch.backend.torch_inductor_jit_backend.TorchInductorJitBackend._save_config(
    cache_dir: pathlib.Path
)

Store the backend configuration to a file.

aitune.torch.backend.torch_inductor_jit_backend.TorchInductorJitBackend.describe() -> str

Returns the description of the backend.

aitune.torch.backend.torch_inductor_jit_backend.TorchInductorJitBackend.from_dict(
    module: torch.nn.Module | None,
    state_dict: dict
)

classmethod

Creates a backend from a state_dict.

aitune.torch.backend.torch_inductor_jit_backend.TorchInductorJitBackend.key() -> str

Returns the key of the backend.

aitune.torch.backend.torch_inductor_jit_backend.TorchInductorJitBackend.to_dict()

Returns the state_dict of the backend.

class aitune.torch.backend.torch_inductor_jit_backend.TorchInductorJitBackendConfig(
    fullgraph: bool = False,
    dynamic: bool | None = None,
    mode: aitune.torch.libs.torch_compile.TorchCompileMode | None = None,
    options: dict[str, str | int | bool] | None = None,
    autocast_enabled: bool = False,
    autocast_dtype: torch.dtype | None = None
)

Dataclass

Bases: BackendConfig

Configuration for torch.compile with inductor backend.

Parameters:

fullgraph

boolDefaults to False

If False (default), torch.compile attempts to discover compileable regions in the function it will tune. If True, then we require the entire function to be captured into a single graph. If this is not possible (that is, if there are graph breaks), then this will raise an error.

dynamic

bool or NoneDefaults to None

Use dynamic shape tracing. When this is True, we will up-front attempt to generate a kernel that is as dynamic as possible to avoid recompilations when sizes change. This may not always work as some operations/optimizations will force specialization; use TORCH_LOGS=dynamic to debug overspecialization. When this is False, we will NEVER generate dynamic kernels, we will always specialize. By default (None), we automatically detect if dynamism has occurred and compile a more dynamic kernel upon recompile.

mode

TorchCompileMode or NoneDefaults to None

Can be either “default”, “reduce-overhead”, “max-autotune” or “max-autotune-no-cudagraphs”.

“default” is the default mode, which is a good balance between performance and overhead
“reduce-overhead” is a mode that reduces the overhead of python with CUDA graphs, useful for small batches. Reduction of overhead can come at the cost of more memory usage, as we will cache the workspace memory required for the invocation so that we do not have to reallocate it on subsequent runs. Reduction of overhead is not guaranteed to work; today, we only reduce overhead for CUDA only graphs which do not mutate inputs. There are other circumstances where CUDA graphs are not applicable; use TORCH_LOG=perf_hints to debug.
“max-autotune” is a mode that leverages Triton or template based matrix multiplications on supported devices and Triton based convolutions on GPU. It enables CUDA graphs by default on GPU.
“max-autotune-no-cudagraphs” is a mode similar to “max-autotune” but without CUDA graphs
To see the exact configs that each mode sets you can call torch._inductor.list_mode_options()

options

dictDefaults to None

A dictionary of options to pass to the backend.

To see the full list of configs that it supports by calling torch._inductor.list_options()

autocast_enabled

boolDefaults to False

If True, enable autocast.

autocast_dtype

torch.dtypeDefaults to None

The dtype to use for autocast.

Note

inference is done with torch.no_grad() context. The torch.inference_mode() context must not be used

autocast_dtype

dtype | None = None

autocast_enabled

bool = False

dynamic

bool | None = None

fullgraph

bool = False

mode

TorchCompileMode | None = None

options

dict[str, str | int | bool] | None = None

aitune.torch.backend.torch_inductor_jit_backend.TorchInductorJitBackendConfig.__post_init__()

Post init.

aitune.torch.backend.torch_inductor_jit_backend.logger = getLogger(__name__)