User manual for aipu - flagos-ai/FlagTree GitHub Wiki

💫 ARM China（安谋科技）aipu3.3

  • Based on Triton 3.3, x64/arm64

1. Quick start

1.1 Use the preinstalled image (for the x64 CPU simulation environment)

Tips: You can use apt install aria2; aria2c -c ... instead of wget ... to download large files with resume support.

If you use this preinstalled image, you do not need to perform the later installation steps.

# Plan A: docker pull (36.9GB)
IMAGE=harbor.baai.ac.cn/flagtree/flagtree-aipu-py310-torch2.6.0-clang16-lld16-ubuntu22.04:202603
docker pull ${IMAGE}
# Plan B: docker load (17GB)
IMAGE=flagtree-aipu-py310-torch2.6.0-clang16-lld16-ubuntu22.04:202603
wget https://baai-cp-web.ks3-cn-beijing.ksyuncs.com/trans/flagtree-aipu-py310-torch2.6.0-clang16-lld16-ubuntu22.04.202603.tar.gz
docker load -i flagtree-aipu-py310-torch2.6.0-clang16-lld16-ubuntu22.04.202603.tar.gz
CONTAINER=flagtree-dev-xxx
docker run -dit \
    --net=host --uts=host --ipc=host --privileged \
    --shm-size 100gb --ulimit memlock=-1 \
    --security-opt seccomp=unconfined --security-opt apparmor=unconfined \
    -v /etc/localtime:/etc/localtime:ro \
    -v /data:/data -v /home:/home \
    -w /root --name ${CONTAINER} ${IMAGE} bash
docker exec -it ${CONTAINER} /bin/bash

1.2 Source-free Installation

# Note: First install PyTorch, then execute the following commands
python3 -m pip uninstall -y triton  # Repeat the cmd until fully uninstalled
RES="--index-url=https://resource.flagos.net/repository/flagos-pypi-hosted/simple"
python3.10 -m pip install flagtree===0.5.0+aipu3.3 $RES

flagtree is already installed in the preinstalled image.

2. Build from Source

2.1 Manually download the FlagTree dependencies

If your network connection is available, you do not need to download the dependencies which will be fetched automatically during the build.

mkdir -p ~/.flagtree/aipu; cd ~/.flagtree/aipu
wget https://baai-cp-web.ks3-cn-beijing.ksyuncs.com/trans/llvm-a66376b0-ubuntu-x64-clang16-lld16_v0.4.0.tar.gz
tar zxvf llvm-a66376b0-ubuntu-x64-clang16-lld16_v0.4.0.tar.gz

2.2 Manually download the Triton dependencies

The Triton dependencies are already downloaded and installed in the preinstalled image. If your network connection is available, you do not need to download the dependencies which will be fetched automatically during the build.

cd ${YOUR_CODE_DIR}/FlagTree
# For Triton 3.3 (x64)
wget https://baai-cp-web.ks3-cn-beijing.ksyuncs.com/trans/build-deps-triton_3.3.x-linux-x64.tar.gz
sh python/scripts/unpack_triton_build_deps.sh ./build-deps-triton_3.3.x-linux-x64.tar.gz

After executing the above script, the original ~/.triton directory will be renamed, and a new ~/.triton directory will be created to store the pre-downloaded packages. Note that the script will prompt for manual confirmation during execution.

2.3 Build Commands

Before building, you need to execute source ~/env_setup.sh. The content of this script is as follows:

SDK_PATH=~/AI610-SDK-dev-4.0.7
export PATH=${PATH}:${SDK_PATH}/AI610-SDK-1002/simulator/bin:${SDK_PATH}/AI100-SDK-0006/opencl-tool-chain/compiler/aipu_opencl_toolchain/bin
export LD_LIBRARY_PATH=${SDK_PATH}/AI610-SDK-1002/simulator/lib:${SDK_PATH}/AI610-SDK-1012/Linux-driver/bin/sim/release:${LD_LIBRARY_PATH}
export ZHOUYI_LINUX_DRIVER_HOME=${SDK_PATH}/AI610-SDK-1012/Linux-driver
export PYTHONPATH=~/.flagtree/aipu/llvm-a66376b0-ubuntu-x64-clang16-lld16/python_packages/mlir_core:${PYTHONPATH}
cd ${YOUR_CODE_DIR}/FlagTree/python
git checkout -b triton_v3.3.x origin/triton_v3.3.x
export FLAGTREE_BACKEND=aipu
MAX_JOBS=32 python3 -m pip install . --no-build-isolation -v

3. Testing and validation

Before testing, you need to execute source ~/env_setup.sh. The content of this script is shown above.

Refer to Tests of aipu backend

⚠️ **GitHub.com Fallback** ⚠️