fix reqs

Dockerfile

@@ -43,7 +43,12 @@ 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
+# We leave it here in case we need to work off of a specific commit in main
+RUN git clone https://github.com/NVIDIA/Megatron-LM.git && \
+  cd Megatron-LM && \
+  git checkout 973330e9c3681604703bf1eb6b5a265d1b9b9b38 && \
+  pip install .
 
 # Distributed Adam support for multiple dtypes
 RUN git clone https://github.com/NVIDIA/apex.git && \
@@ -138,11 +143,3 @@ RUN if [ "${REQUIRE_AIS_CLI}" = true ]; then \
   exit ${INSTALL_CODE}; \
   else echo "AIS CLI installed successfully"; fi \
   else echo "Skipping AIS CLI installation"; fi
-
-
-# Install megatron core, this can be removed once 0.3 pip package is released
-# We leave it here in case we need to work off of a specific commit in main
-RUN git clone https://github.com/NVIDIA/Megatron-LM.git && \
-  cd Megatron-LM && \
-  git checkout 973330e9c3681604703bf1eb6b5a265d1b9b9b38 && \
-  pip install .

requirements/requirements_nlp.txt

@@ -10,7 +10,7 @@ ijson
 jieba
 markdown2
 matplotlib>=3.3.2
-megatron_core==0.3.0
+megatron_core==0.4.0
 nltk>=3.6.5
 opencc<1.1.7
 pangu