build: move cudnn-frontend override from Dockerfile.ci to Dockerfile.fw_final

fw_final inherits the Bridge image and re-syncs with fw_pyproject.toml, which would replace a git install made earlier. Install the feat/1.29.0-megatron ref after that sync instead, and drop it from Dockerfile.ci so it lives in one place.

Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>
Signed-off-by: Zhiyu Li <zhiyul@nvidia.com>

Author: ZhiyuLi-Nvidia

docker/Dockerfile.ci

@@ -176,11 +176,6 @@ RUN if [ -n "$BASELINE_MCORE_REF" ]; then \
 ARG MCORE_TRIGGERED_TESTING=false
 ARG UV_CACHE_PRUNE_ARGS
 
-# cuDNN Frontend is replaced after the final sync with this ref of NVIDIA/cudnn-frontend
-# (feat/1.29.0-megatron). Docker-only: outside pyproject.toml and uv.lock. Set empty to keep the
-# locked PyPI wheel.
-ARG CUDNN_FRONTEND_GIT_REF=8924d4b41fd678ba1b98cd7f98bb78b222637573
-
 # TransformerEngine builds from source during uv sync. Target the GPUs CI runs on: Hopper, plus
 # Blackwell, which TE expands to 100a and 103a (B200, GB200, B300). TE's own default adds
 # 75, 80, 89, and 120.
@@ -260,10 +255,6 @@ RUN --mount=type=cache,target=/var/cache/apt,sharing=locked \
         uv sync --link-mode copy --locked --all-extras --all-groups --no-group diffusion --no-install-package mamba-ssm ${BASE_RUNTIME_UV_ARGS}; \
     fi && \
     bash /opt/install_mamba.sh /opt/Megatron-Bridge/uv.lock /opt/mamba.patch && \
-    if [ -n "$CUDNN_FRONTEND_GIT_REF" ] && [ "$MBRIDGE_PRESERVE_BASE_RUNTIME" != "True" ]; then \
-        uv pip install --no-deps --reinstall \
-            "nvidia-cudnn-frontend @ git+https://github.com/NVIDIA/cudnn-frontend.git@${CUDNN_FRONTEND_GIT_REF}"; \
-    fi && \
     python /opt/preserve_base_runtime.py check /opt/base-runtime-packages.json && \
     # uv may retain the baseline editable MCore wheel when only its source tree changes.
     # Rebuild the dispatched checkout so its optional dataset helper is present.

docker/Dockerfile.fw_final

@@ -57,6 +57,11 @@ RUN    pushd /opt/Export-Deploy && \
         echo "Container built with NeMo-Run commit hash: $NEMO_RUN_COMMIT" && \
     popd
 
+# cuDNN Frontend is replaced after the sync below with this ref of NVIDIA/cudnn-frontend
+# (feat/1.29.0-megatron). Docker-only: outside fw_pyproject.toml and uv.lock. Set empty to keep the
+# PyPI wheel.
+ARG CUDNN_FRONTEND_GIT_REF=8924d4b41fd678ba1b98cd7f98bb78b222637573
+
 # Sync environment
 COPY --chmod=644 docker/common/fw_pyproject.toml /opt/NeMo-FW/pyproject.toml
 COPY --chmod=644 docker/common/_version.py docker/common/print_sha.sh /opt/NeMo-FW/
@@ -69,6 +74,10 @@ RUN \
         BASE_RUNTIME_UV_ARGS="--no-install-package nvidia-cudnn-frontend --no-install-package transformer-engine --no-install-package transformer-engine-torch"; \
     fi && \
     uv sync --no-cache-dir --all-groups --inexact ${BASE_RUNTIME_UV_ARGS} && \
+    if [ -n "$CUDNN_FRONTEND_GIT_REF" ] && [ "${MBRIDGE_PRESERVE_BASE_RUNTIME:-False}" != "True" ]; then \
+        uv pip install --no-cache-dir --no-deps --reinstall \
+            "nvidia-cudnn-frontend @ git+https://github.com/NVIDIA/cudnn-frontend.git@${CUDNN_FRONTEND_GIT_REF}"; \
+    fi && \
     python /opt/preserve_base_runtime.py check /opt/base-runtime-packages.json && \
     uv pip list && \
     chmod +x print_sha.sh && \

docker/README.md

@@ -64,7 +64,6 @@ Docker-only dependency, outside `pyproject.toml` and `uv.lock`.
 | `INSTALL_DIFFUSION_DEPS` | When `true`, runs `scripts/install_diffusion_deps.sh` to add WAN codecs to a CI image; defaults to `false` so CVE-carrying codecs stay out of shipped Framework images |
 | `UV_CACHE_PRUNE_ARGS` | Extra arguments forwarded to `uv cache prune` after install |
 | `PRESERVE_BASE_RUNTIME` | Preserve a validated base runtime through dependency installation; defaults to `False`. See [base runtime preservation](#preserving-a-validated-base-runtime). |
-| `CUDNN_FRONTEND_GIT_REF` | Commit of [NVIDIA/cudnn-frontend](https://github.com/NVIDIA/cudnn-frontend) (`feat/1.29.0-megatron`) installed over the locked cuDNN frontend after the final sync; empty keeps the locked PyPI wheel. Docker-only, outside `pyproject.toml` and `uv.lock`; skipped when `PRESERVE_BASE_RUNTIME=True`. |
 
 Use the diffusion dependency opt-in only for CI images that run the diffusion test suite:
 
@@ -251,6 +250,7 @@ docker build \
 | `NEMO_EXPORT_DEPLOY_COMMIT` | NeMo Export-Deploy git commit SHA |
 | `NEMO_EVALUATOR_COMMIT` | NeMo Evaluator git commit SHA |
 | `NEMO_RUN_COMMIT` | NeMo Run git commit SHA |
+| `CUDNN_FRONTEND_GIT_REF` | Commit of [NVIDIA/cudnn-frontend](https://github.com/NVIDIA/cudnn-frontend) (`feat/1.29.0-megatron`) installed over the PyPI cuDNN frontend after the environment sync; empty keeps the PyPI wheel. Docker-only, outside `fw_pyproject.toml`; skipped when `PRESERVE_BASE_RUNTIME=True`. |
 
 ---