[WIP] Initial DeepSeek reference implementation

moe_pretraining/nemo/Dockerfile

@@ -0,0 +1,76 @@
+# Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# Permission is hereby granted, free of charge, to any person obtaining a
+# copy of this software and associated documentation files (the "Software"),
+# to deal in the Software without restriction, including without limitation
+# the rights to use, copy, modify, merge, publish, distribute, sublicense,
+# and/or sell copies of the Software, and to permit persons to whom the
+# Software is furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in
+# all copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
+# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
+# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
+# DEALINGS IN THE SOFTWARE.
+
+
+ARG FROM_IMAGE_NAME=nvcr.io/nvidia/pytorch:25.12-py3
+
+FROM ${FROM_IMAGE_NAME}
+
+ENV PIP_CONSTRAINT=""
+
+RUN git config --global user.name "a" && \
+    git config --global user.email "a"
+
+## 0. Pytorch Checkpoint size patch
+WORKDIR /workspace/
+
+
+RUN pip install numcodecs==0.16.3
+RUN pip install nvidia-mathdx==25.1.1
+
+
+## 1. NeMo-Run
+ARG NEMORUN_REVISION=v0.5.0
+ENV CUSTOM_NEMORUN_REVISION ${NEMORUN_REVISION}
+RUN git clone https://github.com/NVIDIA/NeMo-Run.git && \
+    cd NeMo-Run && \
+    git checkout ${NEMORUN_REVISION} && \
+    echo NEMORUN_COMMIT_HASH=$(git rev-parse HEAD) && \
+    pip install -e .
+
+## 2. Megatron-bridge and megatron-core
+ARG MBRIDGE_REVISION=38858ef0ae3b835af236ec1cf8bf1feca4d400fa
+RUN pip uninstall -y megatron-core && \
+    git clone https://github.com/NVIDIA-NeMo/Megatron-Bridge.git /workspace/Megatron-Bridge && \
+    cd /workspace/Megatron-Bridge && \
+    git checkout ${MBRIDGE_REVISION} && \
+    git submodule update --init --recursive && \
+    echo MBRIDGE_COMMIT_HASH=$(git rev-parse HEAD) && \
+    echo $(git rev-parse HEAD) > /MBRIDGE_COMMIT_HASH.env && \
+    # TODO(dfridman): knob for eval batch size. Remove once merged in the upstream
+    git remote add dfridman https://github.com/denys-fridman/Megatron-Bridge.git && \
+    git fetch dfridman && \
+    git cherry-pick 01cc7a6c76f5d8918a981e22f02975ae775c2d97 && \
+    cd /workspace/Megatron-Bridge/3rdparty/Megatron-LM && \
+    echo MCORE_COMMIT_HASH=$(git rev-parse HEAD) && \
+    echo $(git rev-parse HEAD) > /MCORE_COMMIT_HASH.env && \
+    pip install -e .
+
+ENV PYTHONPATH "/workspace/Megatron-Bridge/src:/workspace/Megatron-Bridge/3rdparty/Megatron-LM:${PYTHONPATH}"
+
+## 3. Benchmark dependencies
+RUN pip uninstall transformers -y
+COPY requirements.txt .
+RUN pip install --no-cache-dir -U -r requirements.txt
+
+
+WORKDIR /workspace/code
+
+COPY . .