> For clean Markdown content of this page, append .md to this URL.

# aitune.torch.task.find_max_batch_size

Find max batch size for a model.

## Module Contents

### Functions

| Name                                                                                                        | Description                                                                 |
| ----------------------------------------------------------------------------------------------------------- | --------------------------------------------------------------------------- |
| [`_log_file`](#aitune-torch-task-find_max_batch_size-_log_file)                                             | -                                                                           |
| [`find_max_batch_size`](#aitune-torch-task-find_max_batch_size-find_max_batch_size)                         | Finds max batch size for Torch Compile as a baseline.                       |
| [`find_max_throughput_for_backend`](#aitune-torch-task-find_max_batch_size-find_max_throughput_for_backend) | Profiles a backend to find the batch size that achieves maximum throughput. |
| [`get_throughput_per_batch_size`](#aitune-torch-task-find_max_batch_size-get_throughput_per_batch_size)     | Gets throughput per batch size.                                             |

### Data

[`logger`](#aitune-torch-task-find_max_batch_size-logger)

### API

```python
aitune.torch.task.find_max_batch_size._log_file(
    cache_dir: pathlib.Path,
    filename: str
) -> pathlib.Path
```

```python
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.

---

```python
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.

```python
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.

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