API ReferenceOther APIsAituneTorchTask

aitune.torch.task.find_max_batch_size

View as Markdown

Find max batch size for a model.

Module Contents

Functions

NameDescription
_log_file-
find_max_batch_sizeFinds max batch size for Torch Compile as a baseline.
find_max_throughput_for_backendProfiles a backend to find the batch size that achieves maximum throughput.
get_throughput_per_batch_sizeGets throughput per batch size.

Data

logger

API

aitune.torch.task.find_max_batch_size._log_file(
    cache_dir: pathlib.Path,
    filename: str
) -> pathlib.Path
 
aitune.torch.task.find_max_batch_size.find_max_batch_size(
    module: torch.nn.Module,
    name: str,
    graph_spec: aitune.torch.module.graph_spec.GraphSpec,
    data: list[aitune.torch.module.recording_module.Sample],
    profiling_config: aitune.torch.task.profiling.config.ProfilingConfig,
    device: torch.device,
    cache_dir: pathlib.Path
) -> tuple[int, float, aitune.torch.task.profiling.profiling.ProfilingResults]

Finds max batch size for Torch Compile as a baseline.

Uses profiling with max throughput strategy to find max batch size.

Note: This function expects user to set profiling_config.max_batch_size to the highest batch size they want to profile.

Parameters:

module
nn.Module

Model to find max batch size for.

name
str

Name of the model.

graph_spec
GraphSpec

Graph spec of the model.

data
list[Sample]

Data to profile.

profiling_config
ProfilingConfig

Profiling configuration.

torch_backend

Backend to use for the find max batch size. If not provided, Torch Eager backend will be used.

device
torch.device

Device to use for the calculation.

cache_dir
Path

Cache directory to store the backend artifacts.

aitune.torch.task.find_max_batch_size.find_max_throughput_for_backend(
    backend: aitune.torch.backend.backend.Backend,
    name: str,
    graph_spec: aitune.torch.module.graph_spec.GraphSpec,
    data: list[aitune.torch.module.recording_module.Sample],
    profiling_config: aitune.torch.task.profiling.config.ProfilingConfig
) -> tuple[int, float, aitune.torch.task.profiling.profiling.ProfilingResults]

Profiles a backend to find the batch size that achieves maximum throughput.

Parameters:

module

Model to calculate maximum throughput for.

name
str

Name of the model.

graph_spec
GraphSpec

Graph spec of the model.

data
list[Sample]

Data to profile.

profiling_config
ProfilingConfig

Profiling configuration.

backend
Backend

Backend to use for the calculation.

device

Device to use for the calculation.

Returns: tuple[int, float, ProfilingResults]

Tuple containing:

  • Batch size with maximum throughput.
  • Throughput for the batch size.
  • Backend used for the calculation.
  • Profiling results.
aitune.torch.task.find_max_batch_size.get_throughput_per_batch_size(
    profiling_results: list[aitune.torch.task.profiling.events.ProfilingResultEvent],
    measuring_stop_strategy: aitune.torch.task.profiling.measuring_stop_strategy.MeasuringStopStrategy
) -> list[tuple[int, float]]

Gets throughput per batch size.

aitune.torch.task.find_max_batch_size.logger = logging.getLogger(__name__)