昇腾环境配置(Ver:1.0_alpha001)-CANN8.5.0+torch2.9.0+FlagTree+FlagGems

[winfan1314]

下载，载入并启动容器

wget https://baai-cp-web.ks3-cn-beijing.ksyuncs.com/trans/docker_image_cann-8.2.rc1.alpha003-a3-ubuntu22.04-py3.11-flagtree.tar.gz
docker load -i docker_image_cann-8.2.rc1.alpha003-a3-ubuntu22.04-py3.11-flagtree.tar.gz
docker run -it \
        --name ***起个名字*** \
        --network host \
        --ipc host \
        --privileged \
        --device /dev/davinci0 \
        --device /dev/davinci1 \
        --device /dev/davinci2 \
        --device /dev/davinci3 \
        --device /dev/davinci4 \
        --device /dev/davinci5 \
        --device /dev/davinci6 \
        --device /dev/davinci7 \
        --device /dev/davinci_manager \
        --device /dev/devmm_svm \
        --device /dev/hisi_hdc \
        -v ***你的代码目录***:/workspace \
        -v ***找个可写目录用来映射/root***:/root \
        -v /usr/local/Ascend/driver:/usr/local/Ascend/driver:ro \
        -v /usr/local/bin/npu-smi:/usr/local/bin/npu-smi:ro \
        -v /usr/local/dcmi:/usr/local/dcmi:ro \
        -v /etc/ascend_install.info:/etc/ascend_install.info:ro \
        -w ***设置默认工作目录*** \
        cann:8.2.rc1.alpha003-a3-ubuntu22.04-py3.11-flagtree \
        /bin/bash -lc 'export LD_LIBRARY_PATH=/usr/local/Ascend/driver/lib64/common:/usr/local/Ascend/driver/lib64/driver:$LD_LIBRARY_PATH && bash'

容器内 装CANN, 910B ops

wget https://ascend-repo.obs.cn-east-2.myhuaweicloud.com/CANN/CANN%208.5.0/Ascend-cann-toolkit_8.5.0_linux-aarch64.run
wget https://ascend-repo.obs.cn-east-2.myhuaweicloud.com/CANN/CANN%208.5.0/Ascend-cann-910b-ops_8.5.0_linux-aarch64.run
bash Ascend-cann-toolkit_8.5.0_linux-aarch64.run --install --quiet
bash Ascend-cann-910b-ops_8.5.0_linux-aarch64.run --install --quiet

容器内 装llvm

mkdir -p ~/.flagtree/ascend; cd ~/.flagtree/ascend
wget https://baai-cp-web.ks3-cn-beijing.ksyuncs.com/trans/llvm-a66376b0-ubuntu-aarch64-python311-compat_v0.3.0.tar.gz
tar zxf llvm-a66376b0-ubuntu-aarch64-python311-compat_v0.3.0.tar.gz

容器内 装Flagtree依赖

cd /workspace/container_workspace
git clone https://github.com/flagos-ai/FlagTree.git
cd FlagTree
git checkout -b triton_v3.2.x origin/triton_v3.2.x
wget https://baai-cp-web.ks3-cn-beijing.ksyuncs.com/trans/build-deps-triton_3.2.x-linux-aarch64.tar.gz
sh python/scripts/unpack_triton_build_deps.sh ./build-deps-triton_3.2.x-linux-aarch64.tar.gz

容器内 装FlagTree

cd python
python3 -m pip install -r requirements.txt
export FLAGTREE_BACKEND=ascend
python3 -m pip install . --no-build-isolation -v

容器内 装其他依赖

pip install numpy sqlalchemy scikit-build-core pybind11 ninja cmake -i https://repo.huaweicloud.com/repository/pypi/simple

容器内 装FlagGems

git clone https://github.com/siliconflow/FlagGems.git 
cd FlagGems
pip install --no-build-isolation -e .

容器内 装 torch torch_npu triton-ascend

wget https://download.pytorch.org/whl/cpu/torch-2.9.0%2Bcpu-cp311-cp311-manylinux_2_28_aarch64.whl
wget https://gitcode.com/Ascend/pytorch/releases/download/v7.3.0-pytorch2.9.0/torch_npu-2.9.0-cp311-cp311-manylinux_2_28_aarch64.whl
pip install torch-2.9.0+cpu-cp311-cp311-manylinux_2_28_aarch64.whl torch_npu-2.9.0-cp311-cp311-manylinux_2_28_aarch64.whl -i https://repo.huaweicloud.com/repository/pypi/simple
pip install triton-ascend -i https://repo.huaweicloud.com/repository/pypi/simple

容器内 小测试FlagGems

import torch
import flag_gems
flag_gems.enable()
M, N, K = 1024, 1024, 1024
A = torch.randn((M, K), dtype=torch.float16, device="npu:0")
B = torch.randn((K, N), dtype=torch.float16, device="npu:0")
C = torch.mm(A, B)
print("shape:", C.shape, "dtype:", C.dtype, "device:", C.device)
print("C values:", C[:5, :5].cpu())
s = C.sum()
print("sum:", s.item(), "dtype:", s.dtype, "device:", s.device)
c = input("Press Enter to continue...")

[ShawnXuan]

0417 更新：先安装flagtree后，删除site-packages/triton，然后安装nightly triton-ascend

还是需要删除 site-package下面的triton目录重新安装，并且最后安装 FlagTree 才可以工作，整理了一份 Dockerfile 如下：

FROM cann:8.2.rc1.alpha003-a3-ubuntu22.04-py3.11-flagtree

ENV DEBIAN_FRONTEND=noninteractive
WORKDIR /workspace

# 基础工具
RUN apt-get update && apt-get install -y \
    wget git curl ca-certificates \
    && rm -rf /var/lib/apt/lists/*

# =========================
# 1. 安装 CANN toolkit + ops
# =========================
RUN wget https://ascend-repo.obs.cn-east-2.myhuaweicloud.com/CANN/CANN%208.5.0/Ascend-cann-toolkit_8.5.0_linux-aarch64.run \
    && wget https://ascend-repo.obs.cn-east-2.myhuaweicloud.com/CANN/CANN%208.5.0/Ascend-cann-910b-ops_8.5.0_linux-aarch64.run \
    && bash Ascend-cann-toolkit_8.5.0_linux-aarch64.run --install --quiet \
    && bash Ascend-cann-910b-ops_8.5.0_linux-aarch64.run --install --quiet \
    && rm -f *.run

# =========================
# 2. 安装 LLVM（flagtree依赖）
# =========================
RUN mkdir -p /root/.flagtree/ascend \
    && cd /root/.flagtree/ascend \
    && wget https://baai-cp-web.ks3-cn-beijing.ksyuncs.com/trans/llvm-a66376b0-ubuntu-aarch64-python311-compat_v0.3.0.tar.gz \
    && tar zxf llvm-a66376b0-ubuntu-aarch64-python311-compat_v0.3.0.tar.gz \
    && rm -f *.tar.gz

# =========================
# 3. 安装 flagtree
# =========================
RUN cd /root \
    && git clone https://gitcode.com/flagos-ai/flagtree.git \
    && cd flagtree \
    && git checkout -b triton_v3.2.x origin/triton_v3.2.x \
    && wget https://baai-cp-web.ks3-cn-beijing.ksyuncs.com/trans/build-deps-triton_3.2.x-linux-aarch64.tar.gz \
    && sed -i '/Press any key to confirm/,+1d' python/scripts/unpack_triton_build_deps.sh \
    && sh python/scripts/unpack_triton_build_deps.sh ./build-deps-triton_3.2.x-linux-aarch64.tar.gz \
    && rm -f *.tar.gz

# Python依赖
RUN pip install --no-cache-dir -r /root/flagtree/python/requirements.txt

ENV FLAGTREE_BACKEND=ascend

# 安装 flagtree（非 editable，避免路径依赖）
# 这里先不安装 RUN pip install --no-build-isolation /root/flagtree/python

# =========================
# 4. 安装基础依赖
# =========================
RUN pip install --no-cache-dir \
    numpy sqlalchemy scikit-build-core pybind11 ninja cmake \
    -i https://repo.huaweicloud.com/repository/pypi/simple

# =========================
# 5. 安装 FlagGems
# =========================
RUN cd /root \
    && git clone https://github.com/siliconflow/FlagGems.git \
    && pip install --no-build-isolation /root/FlagGems
    #&& git clone https://gitcode.com/flagos-ai/FlagGems.git \

# =========================
# 6. 安装 torch + torch_npu + triton-ascend
# =========================
RUN wget https://download.pytorch.org/whl/cpu/torch-2.9.0%2Bcpu-cp311-cp311-manylinux_2_28_aarch64.whl \
    && wget https://gitcode.com/Ascend/pytorch/releases/download/v7.3.0-pytorch2.9.0/torch_npu-2.9.0-cp311-cp311-manylinux_2_28_aarch64.whl \
    && pip install torch-2.9.0+cpu-cp311-cp311-manylinux_2_28_aarch64.whl \
                   torch_npu-2.9.0-cp311-cp311-manylinux_2_28_aarch64.whl \
                   -i https://repo.huaweicloud.com/repository/pypi/simple \
    && pip install --no-build-isolation /root/flagtree/python \
    && rm -rf /usr/local/python3.11.13/lib/python3.11/site-packages/triton \
    && pip install -i https://test.pypi.org/simple/ "triton-ascend<3.2.0rc" --pre --no-cache-dir \
    && rm -f *.whl

# =========================
# 环境变量
# =========================
ENV LD_LIBRARY_PATH=/usr/local/Ascend/driver/lib64/common:/usr/local/Ascend/driver/lib64/driver:$LD_LIBRARY_PATH

CMD ["/bin/bash"]

docker run -it \
        --rm \
        --network host \
        --ipc host \
        --privileged \
        --device /dev/davinci0 \
        --device /dev/davinci1 \
        --device /dev/davinci2 \
        --device /dev/davinci3 \
        --device /dev/davinci4 \
        --device /dev/davinci5 \
        --device /dev/davinci6 \
        --device /dev/davinci7 \
        --device /dev/davinci_manager \
        --device /dev/devmm_svm \
        --device /dev/hisi_hdc \
        -v $PWD:/workspace \
        -v /usr/local/Ascend/driver:/usr/local/Ascend/driver:ro \
        -v /usr/local/bin/npu-smi:/usr/local/bin/npu-smi:ro \
        -v /usr/local/dcmi:/usr/local/dcmi:ro \
        -v /etc/ascend_install.info:/etc/ascend_install.info:ro \
        -w /workspace \
	flag-gems-ascend-cann8.5:v0.1 \
	bash