From 4a264fcae575ebaf8466ff70fc4535fadad8ceee Mon Sep 17 00:00:00 2001 From: Gurpreet Singh Dhami Date: Mon, 20 Jan 2025 14:20:17 -0600 Subject: [PATCH] update docker base to pytorch 2.5 --- Dockerfile_rocm.ci | 4 +++- Dockerfile_rocm.dev | 4 +++- 2 files changed, 6 insertions(+), 2 deletions(-) diff --git a/Dockerfile_rocm.ci b/Dockerfile_rocm.ci index 79cb13e6dd..d397f469c3 100755 --- a/Dockerfile_rocm.ci +++ b/Dockerfile_rocm.ci @@ -1,4 +1,4 @@ -ARG BASE_DOCKER=rocm/pytorch:latest +ARG BASE_DOCKER=rocm/pytorch:rocm6.3_ubuntu22.04_py3.10_pytorch_release_2.5.1_preview FROM $BASE_DOCKER ARG PYTORCH_ROCM_ARCH_OVERRIDE="gfx942" @@ -71,6 +71,8 @@ COPY . Megatron-LM WORKDIR $WORKSPACE_DIR/Megatron-LM RUN pip install -e . +ENV PYTHONPATH=/var/lib/jenkins/triton/python + # record configuration for posterity RUN pip list diff --git a/Dockerfile_rocm.dev b/Dockerfile_rocm.dev index 3999717732..6ab2bbbe4b 100755 --- a/Dockerfile_rocm.dev +++ b/Dockerfile_rocm.dev @@ -1,4 +1,4 @@ -ARG BASE_DOCKER=rocm/pytorch:latest +ARG BASE_DOCKER=rocm/pytorch:rocm6.3_ubuntu22.04_py3.10_pytorch_release_2.5.1_preview FROM $BASE_DOCKER ARG PYTORCH_ROCM_ARCH_OVERRIDE="gfx942" ENV WORKSPACE_DIR=/workspace @@ -73,6 +73,8 @@ RUN git clone https://github.com/ROCm/Megatron-LM.git Megatron-LM &&\ WORKDIR $WORKSPACE_DIR/Megatron-LM +ENV PYTHONPATH=/var/lib/jenkins/triton/python + # record configuration for posterity RUN pip list