API ReferenceOther APIsAituneTorchTune Strategy

aitune.torch.tune_strategy.multi_backend_strategy

View as Markdown

Share backend selection because strategies can be reused across different module types.

Module Contents

Classes

NameDescription
MultiBackendStrategyKeep candidate lists in strategies and select them once the module is known.

API

class aitune.torch.tune_strategy.multi_backend_strategy.MultiBackendStrategy(
    backends: list[aitune.torch.backend.backend.Backend] | None = None,
    workflow: typing.Literal['aot', 'jit'] = 'aot',
    kwargs = {}
)

Bases: FindMaxBatchSizeMixin

Keep candidate lists in strategies and select them once the module is known.

Rebuild defaults when switching between ordinary and distributed modules, so a reused strategy does not keep incompatible backends. Explicit lists always win, including an empty list. Each strategy must define both workflows explicitly.

_backends
list[Backend]

Return explicit candidates or the defaults for the current module.

_resolved_backends
list[Backend] | None = None
aitune.torch.tune_strategy.multi_backend_strategy.MultiBackendStrategy._configure_for_module(
    module: torch.nn.Module
) -> None

Select distributed defaults without replacing user-supplied candidates.

aitune.torch.tune_strategy.multi_backend_strategy.MultiBackendStrategy._default_aot_backends(
    distributed: bool = False
) -> list[aitune.torch.backend.backend.Backend]
abstract

Return AOT candidates that support the module’s execution requirements.

aitune.torch.tune_strategy.multi_backend_strategy.MultiBackendStrategy._default_jit_backends(
    distributed: bool = False
) -> list[aitune.torch.backend.backend.Backend]
abstract

Return JIT candidates that support the module’s execution requirements.

aitune.torch.tune_strategy.multi_backend_strategy.MultiBackendStrategy._resolve_default_backends() -> list[aitune.torch.backend.backend.Backend]

Build only the candidates for the selected workflow and module type.

classmethod

Create a strategy using AOT candidates and its default batch-size policy.

Match the regular constructor so existing AOT behavior stays unchanged.

Parameters:

kwargs
Defaults to {}

Strategy constructor arguments, including explicit backend overrides.

classmethod

Create a strategy using JIT candidates and no maximum-batch-size discovery.

Use the batches already recorded to limit tuning work during inference.

Parameters:

kwargs
Defaults to {}

Strategy constructor arguments, including explicit backend overrides.