> This page is for AITune, version v0.4.1.
> For other versions, use one of these documentation indexes:
> - Latest (v0.6.0) (default): https://docs.nvidia.com/aitune/latest/llms.txt
> - dev: https://docs.nvidia.com/aitune/dev/llms.txt
> - v0.6.0: https://docs.nvidia.com/aitune/v0.6.0/llms.txt
> - v0.5.0: https://docs.nvidia.com/aitune/v0.5.0/llms.txt
> - v0.4.1: https://docs.nvidia.com/aitune/v0.4.1/llms.txt

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

# aitune.torch.tuning

Main tune function.

## Module Contents

### Functions

| Name                                                                                              | Description                                                    |
| ------------------------------------------------------------------------------------------------- | -------------------------------------------------------------- |
| [`_activate_tuned_modules`](#aitune-torch-tuning-_activate_tuned_modules)                         | Activate all tuned modules after tuning is complete.           |
| [`_clear_cache`](#aitune-torch-tuning-_clear_cache)                                               | Clear the cache.                                               |
| [`_describe_module`](#aitune-torch-tuning-_describe_module)                                       | -                                                              |
| [`_validate_and_normalize_batch_sizes`](#aitune-torch-tuning-_validate_and_normalize_batch_sizes) | Validate and normalize batch sizes for tuning.                 |
| [`load`](#aitune-torch-tuning-load)                                                               | Load the tuned module from a file.                             |
| [`save`](#aitune-torch-tuning-save)                                                               | Save the tuned module to a file.                               |
| [`tune`](#aitune-torch-tuning-tune)                                                               | Tune a callable which runs inference on a pipeline or a model. |

### Data

[`LOG_FORMAT`](#aitune-torch-tuning-LOG_FORMAT)

[`logger`](#aitune-torch-tuning-logger)

### API

```python
aitune.torch.tuning._activate_tuned_modules() -> None
```

Activate all tuned modules after tuning is complete.

```python
aitune.torch.tuning._clear_cache()
```

Clear the cache.

```python
aitune.torch.tuning._describe_module(
    module: aitune.torch.module.wrapper_module.Module
)
```

```python
aitune.torch.tuning._validate_and_normalize_batch_sizes(
    batch_sizes: list[int] | None
) -> list[int]
```

Validate and normalize batch sizes for tuning.

```python
aitune.torch.tuning.load(
    module: torch.nn.Module,
    path: str | pathlib.Path,
    storage: aitune.torch.checkpoint.storage.Storage | None = None,
    device_map: dict[str, torch.device] | None = None,
    disable_external_logging: bool = True
) -> torch.nn.Module | aitune.torch.module.wrapper_module.Module
```

Load the tuned module from a file.

If storage is not provided, a default storage will be used. Check save function for more details.

**Parameters:**

**`module`** `torch.nn.Module`

The module to load.

---

**`path`** `str | Path`

The path to load the module from.

---

**`storage`** `Storage | None` — default: None

The storage to use to load the module. If not provided, a local storage will be used.

---

**`device_map`** `dict[str, torch.device] | None` — default: None

The device map to load modules to. Overrides the device of stored in state dict.

---

**`disable_external_logging`** `bool` — default: True

If True, libraries logging will be suppressed.

---

```python
aitune.torch.tuning.save(
    module: torch.nn.Module,
    path: str | pathlib.Path,
    storage: aitune.torch.checkpoint.storage.Storage | None = None
) -> None
```

Save the tuned module to a file.

If storage is not provided, a default storage will be used.

If you would like to save the module to a different folder, you can use the storage parameter.
For example, if you want to save the module to the `ckpt` folder, you can do the following:

**Parameters:**

**`module`** `torch.nn.Module`

The module to save.

---

**`path`** `str | Path`

The path to save the module to.

---

**`storage`** `Storage | None` — default: None

The storage to use to save the module. If not provided, a local storage will be used.

---

```python
aitune.torch.tuning.tune(
    func: collections.abc.Callable,
    dataset: aitune.torch.dataloader.DatasetLike | aitune.torch.dataloader.DataLoaderFactory | torch.Tensor,
    batch_sizes: list[int] | None = None,
    max_num_batches_per_batch_size: int | None = None,
    device: str | torch.device | None = DEFAULT_DEVICE,
    dry_run: bool = False,
    disable_external_logging: bool = False,
    clear_cache: bool = False,
    ignore_failing_modules: bool = True
) -> None
```

Tune a callable which runs inference on a pipeline or a model.

**Parameters:**

**`func`** `Callable`

The function to tune.

---

**`dataset`** `DatasetLike | DataLoaderFactory | torch.Tensor`

The dataset to tune on. It can be DataLoaderFactory or any dataset/iterable and even torch.Tensor.
Tensor will be treated as a single sample dataset.

---

**`batch_sizes`** `list[int] | None` — default: None

The batch sizes to use for tuning. At least 2 different batch sizes are required for determining
batch axis.

---

**`max_num_batches_per_batch_size`** `int | None` — default: None

The maximum number of batches to use for tuning per batch size.

---

**`device`** `str | torch.device | None` — default: DEFAULT_DEVICE

The device to use for tuning.

---

**`dry_run`** `bool` — default: False

If True, only dry run the tuning.

---

**`disable_external_logging`** `bool` — default: False

If True, libraries logging will be suppressed.

---

**`clear_cache`** `bool` — default: False

If True, the cache will be cleared before tuning.

---

**`ignore_failing_modules`** `bool` — default: True

If True, failing modules will be ignored and tuning will continue.

---

```python
aitune.torch.tuning.LOG_FORMAT = '%(asctime)s - %(levelname)s - %(message)s'
```

```python
aitune.torch.tuning.logger = getLogger(__name__)
```