feat: Dockerfile allows custom mcore install (#2005)

Signed-off-by: Dong Hyuk Chang <donghyukc@nvidia.com>

Author: thomasdhc

.github/workflows/cicd-main.yml

@@ -283,6 +283,7 @@ jobs:
             ${{ env.container-registry }}/megatron-bridge:${{ github.sha }}
           secrets: |
             GH_TOKEN=${{ secrets.PAT }}
+          target: mbridge_final
 
       - name: Verify MCore commit in Docker image
         if: ${{ github.event.inputs.mcore_commit != '' }}

docker/Dockerfile.ci

@@ -13,8 +13,8 @@
 # limitations under the License.
 
 ARG BASE_IMAGE=nvcr.io/nvidia/pytorch:25.11-py3
-FROM ${BASE_IMAGE} AS megatron_bridge
-WORKDIR /opt/Megatron-Bridge
+ARG MCORE_LAYER=submodule_mcore
+FROM ${BASE_IMAGE} AS mbridge_base
 ENV PATH="/root/.local/bin:$PATH"
 ENV UV_PROJECT_ENVIRONMENT=/opt/venv
 ENV PATH="$UV_PROJECT_ENVIRONMENT/bin:$PATH"
@@ -28,12 +28,44 @@ RUN curl -LsSf https://astral.sh/uv/${UV_VERSION}/install.sh | sh && \
 
 COPY pyproject.toml uv.lock /opt/Megatron-Bridge/
 COPY src/megatron/bridge/__init__.py src/megatron/bridge/package_info.py /opt/Megatron-Bridge/src/megatron/bridge/
-COPY 3rdparty/Megatron-LM/pyproject.toml /opt/Megatron-Bridge/3rdparty/Megatron-LM/
-COPY 3rdparty/Megatron-LM/megatron/core/__init__.py 3rdparty/Megatron-LM/megatron/core/package_info.py /opt/Megatron-Bridge/3rdparty/Megatron-LM/megatron/core/
 
+FROM mbridge_base as mbridge_custom_mcore
+
+ARG MCORE_REPO
+ARG MCORE_COMMIT
+ARG MCORE_HOST
+WORKDIR /src/Megatron-LM
+
+# Clone Custom MCore Installation
+RUN --mount=type=secret,id=CI_JOB_TOKEN,target=/run/secrets/CI_JOB_TOKEN \
+    CI_JOB_TOKEN=$(cat /run/secrets/CI_JOB_TOKEN) && \
+    git clone ${MCORE_HOST}:${CI_JOB_TOKEN}@${MCORE_REPO} . && \
+    git pull && \
+    git fetch origin $MCORE_COMMIT && \
+    git checkout FETCH_HEAD && \
+    echo "Container built with Megatron-LM commit hash: $MCORE_COMMIT" && \
+    rm -rf .git && \
+    rm -rf  ~/.netrc && \
+    mkdir -p /tmp/Megatron-LM/megatron/core/ && \
+    cp pyproject.toml /tmp/Megatron-LM/ && \
+    cp megatron/core/__init__.py megatron/core/package_info.py /tmp/Megatron-LM/megatron/core/
+
+FROM mbridge_base as mbridge_submodule_mcore
+
+WORKDIR /src
+COPY 3rdparty/Megatron-LM/pyproject.toml /tmp/Megatron-LM/
+COPY 3rdparty/Megatron-LM/megatron/core/__init__.py 3rdparty/Megatron-LM/megatron/core/package_info.py /tmp/Megatron-LM/megatron/core/
+
+FROM mbridge_${MCORE_LAYER} as mbridge_mcore
+
+FROM mbridge_base as mbridge_final
+
+WORKDIR /opt/Megatron-Bridge
 # Build arg to skip --locked when testing with different MCore versions
 ARG MCORE_TRIGGERED_TESTING=false
 
+COPY --from=mbridge_mcore /tmp/Megatron-LM /opt/Megatron-Bridge/3rdparty/Megatron-LM
+
 RUN --mount=type=cache,target=/var/cache/apt,sharing=locked \
     --mount=type=cache,target=/var/lib/apt,sharing=locked \
     --mount=type=cache,target=/root/.cache/uv \
@@ -48,3 +80,4 @@ RUN --mount=type=cache,target=/var/cache/apt,sharing=locked \
     uv cache prune
 
 COPY . /opt/Megatron-Bridge
+COPY --from=mbridge_mcore /src/Megatron-LM* /opt/Megatron-Bridge/3rdparty