PyTorch ao

Custom data types and layouts for training and inference

To use experimental optimizers (AdamWFp8, AdamW4bit, AdamW8bit) from Pytorch Ao, please install the package as shown below.

Tip

Some experimental optimizers are already present in regular Pytorch, so please re-check if you actually need this package!

Installation

Stable Release from the PyTorch index

pip install torchao --extra-index-url https://download.pytorch.org/whl/cu121 # full options are cpu/cu118/cu121/cu124

Nightly release

pip install --pre torchao-nightly --index-url https://download.pytorch.org/whl/nightly/cu121 # full options are cpu/cu118/cu121/cu124