Use the PPU pip source in PAI (v1.6.x)

Overview

The PAI-PPU pip source provides frameworks for PPU products, open-source libraries recompiled with the PPU SDK, and customized packages requested by customers. It lets you install packages with the pip install command to quickly set up a production environment. If a required package is not available in the PAI-PPU pip source, pip pulls it from the official PyPI source by default. Python packages in the PAI-PPU pip source are categorized into optimized edition (OE) and community edition (CE). OE versions are released with PAI-PPU image updates, while CE versions are updated weekly.

Pip packages

The list of packages available from the PAI-PPU pip source updates with each PAI-PPU image release. You can view the pip list for a specific release. The list includes two categories: community edition and optimized edition. The community edition contains open-source libraries compiled with the PPU SDK, as well as PPU-compatible versions of the PyTorch and TensorFlow frameworks. The optimized edition contains open-source packages with modified code optimized for PPU performance.

Naming conventions

pip source URL naming convention

For example, consider the URL https://art-pub.eng.t-head.cn/artifactory/api/pypi/ptgai-pypi_ppu_ubuntu_cu118_index/simple/:

pip install command naming convention

Example: pip install flash-attn==2.4.3+ppu1.4.0.ce

• If you omit the .ce or .oe suffix: You can still install the package by running a command like pip install flash-attn==2.4.3+ppu1.4.0. If both a community edition and an optimized edition package are available for your configuration, pip installs the optimized edition by default.

• If you omit the release version: A command like pip install flash-attn==2.4.3 installs the latest version compiled and uploaded to the pip source.

• If you omit the version number: A command like pip install flash-attn installs the latest version of the package by default.

Examples

Set the source per command

Use case: Install the xformers==0.0.25 and flash-attn==2.5.6 customized packages in the ppu:cuda11.8-alios7u2-torch1.13.1-py38 base Docker image.

1. Enter the Docker container and refer to the "customized library" list in the pip package list.

2. Filter for xformers==0.0.25 and flash-attn==2.5.6 to find the installation commands for your environment: AliOS 7u2, CUDA 11.8, PyTorch 1.13.1, and Python 3.8.

pip install flash-attn==2.5.6+ppu1.4.0.oe -i https://art-pub.eng.t-head.cn/artifactory/api/pypi/ptgai-pypi_ppu_alios_cu118_index/simple/
pip install xformers==0.0.25+ppu1.4.0.oe -i https://art-pub.eng.t-head.cn/artifactory/api/pypi/ptgai-pypi_ppu_alios_cu118_index/simple/

Set Docker environment variables

Use case: Set up a production environment that supports PyTorch==2.1.2, torchaudio, and torchvision, based on an official PAI-PPU image.

1. Enter the Docker container, confirm your environment (CUDA 12.2, AliOS 7u2, and Python 3.10), and then set the PAI-PPU pip source URL as an environment variable.

export PIP_INDEX_URL='https://art-pub.eng.t-head.cn/artifactory/api/pypi/ptgai-pypi_ppu_alios_cu122_index/simple/'

2. Check the pip list to ensure the current release includes the packages you want to install.

3. Install the packages from the PAI-PPU pip source. If a package uses PPU but the PAI-PPU pip source does not yet support it, the installation fails. If you do not specify a PPU SDK version, pip installs the latest version of the package from the PAI-PPU pip source by default.

pip install torch==2.1.2+ppu1.4.0
pip install torchaudio==2.1.2+ppu1.4.0
pip install torchvision==0.16.2+ppu1.4.0

To install without specifying a PPU release version:

pip install torch==2.1.2 torchaudio==2.1.2 torchvision==0.16.2

Troubleshooting

If you encounter the following pip source hash error, use the --no-cache-dir option.

pip install --no-cache-dir torch==2.1.2+ppu1.4.0

You can also run pip cache purge to clear the system's entire pip cache.

pip cache purge