Dockerfile_rocm.dev

ARG BASE_DOCKER=rocm/pytorch:latest
FROM $BASE_DOCKER
ARG PYTORCH_ROCM_ARCH_OVERRIDE="gfx942"
ENV WORKSPACE_DIR=/workspace
ENV STAGE_DIR=/workspace/installs
RUN mkdir -p $WORKSPACE_DIR
RUN mkdir -p ${STAGE_DIR}
WORKDIR $WORKSPACE_DIR

RUN pip3 install \
scipy \
transformers \
einops \
flask-restful \
nltk \
pytest \
pytest-cov \
pytest_mock \
pytest-csv \
pytest-random-order \
sentencepiece \
wrapt \
zarr \
wandb \
tensorstore==0.1.45 \
pytest_mock \
pybind11 \
setuptools==69.5.1 \
datasets \
tiktoken \
pynvml

RUN pip3 install "huggingface_hub[cli]"
RUN python3 -m nltk.downloader punkt_tab


# Install Causal-Conv1d and its dependencies
WORKDIR ${STAGE_DIR}
ENV CAUSAL_CONV1D_FORCE_BUILD=TRUE
ENV MAMBA_FORCE_BUILD=TRUE
ENV HIP_ARCHITECTURES=${PYTORCH_ROCM_ARCH_OVERRIDE}
RUN git clone https://github.com/Dao-AILab/causal-conv1d causal-conv1d &&\
    cd causal-conv1d &&\
    git show --oneline -s &&\
    pip install --no-build-isolation .

# Install mamba
WORKDIR ${STAGE_DIR}
RUN git clone https://github.com/state-spaces/mamba mamba &&\
    cd mamba &&\
    git show --oneline -s &&\
    pip install --no-build-isolation .

# Clone TE repo and submodules
WORKDIR ${STAGE_DIR}
ENV NVTE_FRAMEWORK=pytorch
ENV NVTE_ROCM_ARCH=${PYTORCH_ROCM_ARCH_OVERRIDE}
ENV NVTE_USE_HIPBLASLT=1
RUN git clone --recursive https://github.com/ROCm/TransformerEngine &&\
    cd TransformerEngine &&\
    pip install .

RUN git clone https://github.com/caaatch22/grouped_gemm.git &&\
    cd grouped_gemm &&\
    git checkout rocm &&\
    git submodule update --init --recursive &&\
    pip install .

RUN git clone https://github.com/ROCm/flash-attention/ -b v2.7.3-cktile && \
    cd flash-attention && \
    GPU_ARCHS=${PYTORCH_ROCM_ARCH_OVERRIDE} python setup.py install && \
    cd .. &&\
    rm -rf flash-attention

WORKDIR $WORKSPACE_DIR
RUN git clone https://github.com/ROCm/Megatron-LM.git Megatron-LM &&\
    cd Megatron-LM &&\
    git checkout rocm_dev &&\
    pip install -e .

WORKDIR $WORKSPACE_DIR/Megatron-LM

ENV PYTHONPATH=/var/lib/jenkins/triton/python

# record configuration for posterity
RUN pip list