diff --git a/Dockerfile b/Dockerfile
index 7d1539ad2db3..06f96a091a22 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -43,6 +43,11 @@ RUN apt-get update && \
   rm -rf /var/lib/apt/lists/*
 
 WORKDIR /workspace/
+# install megatron core, this can be removed once 0.3 pip package is released
+RUN git clone https://github.com/NVIDIA/Megatron-LM.git && \
+  cd Megatron-LM && \
+  git checkout ab0336a5c8eab77aa74ae604ba1e73decbf6d560 && \
+  pip install -e .
 
 WORKDIR /tmp/
 
@@ -52,12 +57,6 @@ RUN git clone https://github.com/NVIDIA/apex.git && \
   git checkout 52e18c894223800cb611682dce27d88050edf1de && \
   pip3 install -v --no-build-isolation --disable-pip-version-check --no-cache-dir --global-option="--cpp_ext" --global-option="--cuda_ext" --global-option="--fast_layer_norm" --global-option="--distributed_adam" --global-option="--deprecated_fused_adam" ./
 
-# install megatron core, this can be removed once 0.3 pip package is released
-RUN git clone https://github.com/NVIDIA/Megatron-LM.git && \
-  cd Megatron-LM && \
-  git checkout ab0336a5c8eab77aa74ae604ba1e73decbf6d560 && \
-  pip install -e .
-
 # uninstall stuff from base container
 RUN pip3 uninstall -y sacrebleu torchtext