update docker base to pytorch 2.5

ROCm · Jan 21, 2025 · 4a264fc · 4a264fc
1 parent b618163
commit 4a264fc
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 2 deletions.
diff --git 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
@@ -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