ARG OS=openeuler24.03
ARG BASE_IMAGE_VERSION=9.1.0-beta.3
ARG NPU_TYPE=910b
ARG PYTHON_VERSION=3.12
ARG BASE_IMAGE=""
FROM ${BASE_IMAGE:-swr.cn-south-1.myhuaweicloud.com/ascendhub/cann:${BASE_IMAGE_VERSION}-${NPU_TYPE}-${OS}-py${PYTHON_VERSION}} AS base
USER root
ARG OS_FAMILY=openeuler
SHELL ["/bin/bash", "-c"]
RUN ARCH=$(uname -m) && \
echo "Detected CPU architecture: ${ARCH}" && \
if [ "$ARCH" != "x86_64" ] && [ "$ARCH" != "aarch64" ]; then \
echo "ERROR: Unsupported architecture: $ARCH"; \
exit 1; \
fi
RUN if [ -f /etc/resolv.conf ]; then cat /etc/resolv.conf; fi && \
echo "nameserver 114.114.114.114" > /etc/resolv.conf && \
echo "nameserver 8.8.8.8" >> /etc/resolv.conf && \
echo "nameserver 223.5.5.5" >> /etc/resolv.conf
COPY configure_repo.sh /tmp/configure_repo.sh
RUN chmod +x /tmp/configure_repo.sh && \
bash /tmp/configure_repo.sh && \
rm /tmp/configure_repo.sh
RUN echo "Installing system dependencies..." && \
if [ "$OS_FAMILY" = "openeuler" ]; then \
yum install -y git iproute wget curl gcc gcc-c++ make cmake patch && \
yum clean all && rm -rf /var/cache/yum; \
elif [ "$OS_FAMILY" = "ubuntu" ]; then \
apt-get install -y git iproute2 wget curl gcc g++ make cmake patch && \
apt-get clean && rm -rf /var/lib/apt/lists/*; \
else \
echo "ERROR: Unsupported OS: $OS_FAMILY"; \
exit 1; \
fi
FROM base AS builder
ARG TORCH_VERSION=2.9.0
ARG TORCH_NPU_VERSION=2.9.0
ARG TORCH_WHL_URL=""
ARG TORCH_NPU_WHL_URL="https://gitcode.com/Ascend/pytorch/releases/download/v26.1.0-beta.2-pytorch2.9.0/torch_npu-2.9.0.post5-cp312-cp312-manylinux_2_28_aarch64.whl"
ARG TRITON_ASCEND_VERSION=3.2.1
WORKDIR /tmp
RUN echo "Using Python from the CANN base image..." && \
if command -v python3 >/dev/null 2>&1; then \
PYTHON_BIN="$(command -v python3)"; \
elif command -v python >/dev/null 2>&1; then \
PYTHON_BIN="$(command -v python)"; \
else \
echo "ERROR: Python is not found in the base image"; \
exit 1; \
fi && \
ln -sf "${PYTHON_BIN}" /usr/local/bin/python && \
python --version && \
if ! python -m pip --version >/dev/null 2>&1; then \
python -m ensurepip --upgrade; \
fi && \
python -m pip install --no-cache-dir --upgrade pip setuptools wheel
RUN python -m pip config set global.index-url https://repo.huaweicloud.com/repository/pypi/simple && \
python -m pip config set global.trusted-host "repo.huaweicloud.com"
RUN echo "Installing PyTorch ${TORCH_VERSION} and TorchNPU ${TORCH_NPU_VERSION}..." && \
if [ -n "${TORCH_WHL_URL}" ]; then \
python -m pip install --no-cache-dir "${TORCH_WHL_URL}"; \
else \
ARCH=$(uname -m) && \
if [ "$ARCH" = "x86_64" ]; then \
python -m pip install --no-cache-dir torch==${TORCH_VERSION} torchvision==0.24.0 torchaudio --index-url https://download.pytorch.org/whl/cpu; \
else \
python -m pip install --no-cache-dir torch==${TORCH_VERSION} torchvision==0.24.0 torchaudio; \
fi; \
fi && \
if [ -n "${TORCH_NPU_WHL_URL}" ]; then \
python -m pip install --no-cache-dir "${TORCH_NPU_WHL_URL}"; \
else \
python -m pip install --no-cache-dir torch-npu==${TORCH_NPU_VERSION}; \
fi
RUN python -m pip install --no-cache-dir triton-ascend==${TRITON_ASCEND_VERSION} \
--extra-index-url=https://triton-ascend.osinfra.cn/pypi/simple \
--trusted-host triton-ascend.osinfra.cn
FROM builder AS final
ARG MINDSPEED_BRIDGE_BRANCH=master
ARG MINDSPEED_BRANCH=core_r0.16.0
ARG MINDSPEED_OPS_BRANCH=master
ARG MEGATRON_BRANCH=core_v0.16.1
ARG MEGATRON_BRIDGE_BRANCH=v0.3.1
ARG FLASH_LINEAR_ATTENTION_NPU_BRANCH=v26.1.0
ARG FLA_NPU_SOC=ascend910b
ARG FLA_NPU_OPS="causal_conv1d,chunk_bwd_dv_local,chunk_bwd_dqkwg,chunk_gated_delta_rule_bwd_dhu,prepare_wy_repr_bwd_da,prepare_wy_repr_bwd_full,chunk_fwd_o,chunk_gated_delta_rule_fwd_h,recurrent_gated_delta_rule,recompute_wu_fwd"
WORKDIR /workspace
RUN git config --global http.sslVerify false
RUN git clone https://github.com/NVIDIA/Megatron-LM.git && \
cd Megatron-LM && \
git checkout ${MEGATRON_BRANCH} && \
python -m pip install -e .
RUN git clone https://github.com/NVIDIA-NeMo/Megatron-Bridge.git && \
cd Megatron-Bridge && \
git checkout ${MEGATRON_BRIDGE_BRANCH}
RUN git clone https://github.com/flashserve/flash-linear-attention-npu.git && \
cd flash-linear-attention-npu && \
git checkout ${FLASH_LINEAR_ATTENTION_NPU_BRANCH}
RUN cd flash-linear-attention-npu && \
if [ -f /usr/local/Ascend/ascend-toolkit/set_env.sh ]; then \
source /usr/local/Ascend/ascend-toolkit/set_env.sh; \
elif [ -f /usr/local/Ascend/cann/set_env.sh ]; then \
source /usr/local/Ascend/cann/set_env.sh; \
else \
echo "ERROR: CANN set_env.sh not found"; \
exit 1; \
fi && \
bash build.sh --soc=${FLA_NPU_SOC} --pkg --ops=${FLA_NPU_OPS} && \
./build_out/cann-*.run && \
cd torch_custom/fla_npu && \
bash build.sh
RUN git clone https://gitcode.com/ascend/MindSpeed.git && \
cd MindSpeed && \
git checkout ${MINDSPEED_BRANCH} && \
python -m pip install -r requirements.txt && \
python -m pip install -e .
RUN git clone https://gitcode.com/ascend/MindSpeed-Ops.git && \
cd MindSpeed-Ops && \
git checkout ${MINDSPEED_OPS_BRANCH} && \
python -m pip install -e . --no-build-isolation --no-deps
RUN git clone https://gitcode.com/ascend/MindSpeed-Bridge.git && \
cd MindSpeed-Bridge && \
git checkout ${MINDSPEED_BRIDGE_BRANCH} && \
python -m pip install -r requirements.txt && \
python -m pip install -e . --no-deps
WORKDIR /workspace/Megatron-Bridge
RUN echo 'cd /workspace/Megatron-Bridge' >> /root/.bashrc && \
echo 'export PYTHONPATH=/workspace/Megatron-Bridge:/workspace/MindSpeed-Bridge:${PYTHONPATH}' >> /root/.bashrc
ENV ASCEND_TOOLKIT_HOME=/usr/local/Ascend/ascend-toolkit
ENV PYTHONPATH=/workspace/Megatron-Bridge:/workspace/MindSpeed-Bridge
