update pytorch base image

pstjohn · pstjohn · commit d63e43b75bc4 · 2026-03-02T06:27:00.000-08:00
Signed-off-by: Peter St. John &lt;pstjohn@nvidia.com&gt;
diff --git a/.devcontainer/recipes/Dockerfile b/.devcontainer/recipes/Dockerfile
@@ -1,6 +1,6 @@
 # Uncomment to use the latest TE from the NGC registry for debugging changes with latest TE.
 # FROM gitlab-master.nvidia.com/dl/transformerengine/transformerengine:main-pytorch-py3-base
-FROM nvcr.io/nvidia/pytorch:26.01-py3
+FROM nvcr.io/nvidia/pytorch:26.02-py3
 
 # FIXME: Fix for "No such file or directory: /workspace/TransformerEngine"
 #  Remove once bug has been addressed in the nvidia/pytorch container.
diff --git a/.github/workflows/unit-tests-recipes.yml b/.github/workflows/unit-tests-recipes.yml
@@ -93,8 +93,8 @@ jobs:
           # Currently, AMPLIFY is the only folder that needs a custom base image, since we have to support both TE and
           # xformers-based models for golden value testing. The rest of the models use the default pytorch image.
 
-          # This uses a squashed version of the pytorch:26.01-py3 image, generated with `docker-squash
-          # nvcr.io/nvidia/pytorch:26.01-py3 -t svcbionemo023/bionemo-framework:pytorch26.01-py3-squashed --output
+          # This uses a squashed version of the pytorch:26.02-py3 image, generated with `docker-squash
+          # nvcr.io/nvidia/pytorch:26.02-py3 -t svcbionemo023/bionemo-framework:pytorch26.02-py3-squashed --output
           # type=registry,compression=zstd,force-compression=true,oci-mediatypes=true,compression-level=15` and pushed
           # to the dockerhub registry. Our github actions are able to cache image pulls from dockerhub but not nvcr, so
           # hopefully this cuts down slightly on CI time at the expense of having a slightly in-directed image location.
@@ -107,8 +107,8 @@ jobs:
                 if . == "bionemo-recipes/models/amplify" then
                   "svcbionemo023/bionemo-framework:amplify-model-devcontainer-082025"
                 else
-                  # "nvcr.io/nvidia/pytorch:26.01-py3"
-                  "svcbionemo023/bionemo-framework:pytorch26.01-py3-squashed"
+                  # "nvcr.io/nvidia/pytorch:26.02-py3"
+                  "svcbionemo023/bionemo-framework:pytorch26.02-py3-squashed"
                 end
               )
             })
diff --git a/bionemo-recipes/models/esm2/Dockerfile b/bionemo-recipes/models/esm2/Dockerfile
@@ -1,4 +1,4 @@
-FROM nvcr.io/nvidia/pytorch:26.01-py3
+FROM nvcr.io/nvidia/pytorch:26.02-py3
 WORKDIR /workspace/bionemo
 COPY . .
 RUN --mount=type=cache,target=/root/.cache/pip \
diff --git a/bionemo-recipes/models/geneformer/Dockerfile b/bionemo-recipes/models/geneformer/Dockerfile
@@ -1,4 +1,4 @@
-FROM nvcr.io/nvidia/pytorch:26.01-py3
+FROM nvcr.io/nvidia/pytorch:26.02-py3
 WORKDIR /workspace/bionemo
 COPY . .
 RUN --mount=type=cache,target=/root/.cache/pip \
diff --git a/bionemo-recipes/recipes/README.md b/bionemo-recipes/recipes/README.md
@@ -85,7 +85,7 @@ recipes/{recipe_name}/
 Your `Dockerfile` should create a complete, reproducible training environment:
 
 ```dockerfile
-FROM nvcr.io/nvidia/pytorch:26.01-py3
+FROM nvcr.io/nvidia/pytorch:26.02-py3
 
 # Install dependencies with caching for faster builds
 RUN --mount=type=cache,target=/root/.cache/pip \
diff --git a/bionemo-recipes/recipes/esm2_accelerate_te/Dockerfile b/bionemo-recipes/recipes/esm2_accelerate_te/Dockerfile
@@ -1,4 +1,4 @@
-FROM nvcr.io/nvidia/pytorch:26.01-py3
+FROM nvcr.io/nvidia/pytorch:26.02-py3
 
 RUN --mount=type=cache,target=/root/.cache/pip \
     --mount=type=bind,source=requirements.txt,target=/requirements.txt \
diff --git a/bionemo-recipes/recipes/esm2_native_te/Dockerfile b/bionemo-recipes/recipes/esm2_native_te/Dockerfile
@@ -1,5 +1,5 @@
 # syntax=docker/dockerfile:1.4
-FROM nvcr.io/nvidia/pytorch:26.01-py3
+FROM nvcr.io/nvidia/pytorch:26.02-py3
 
 RUN --mount=type=cache,target=/root/.cache/pip \
     --mount=type=bind,source=requirements.txt,target=/requirements.txt \
diff --git a/bionemo-recipes/recipes/esm2_peft_te/Dockerfile b/bionemo-recipes/recipes/esm2_peft_te/Dockerfile
@@ -1,5 +1,5 @@
 # syntax=docker/dockerfile:1.4
-FROM nvcr.io/nvidia/pytorch:26.01-py3
+FROM nvcr.io/nvidia/pytorch:26.02-py3
 
 RUN --mount=type=cache,target=/root/.cache/pip \
     --mount=type=bind,source=requirements.txt,target=/requirements.txt \
diff --git a/bionemo-recipes/recipes/evo2_megatron/Dockerfile b/bionemo-recipes/recipes/evo2_megatron/Dockerfile
@@ -1,5 +1,5 @@
 # syntax=docker/dockerfile:1.4
-FROM nvcr.io/nvidia/pytorch:26.01-py3
+FROM nvcr.io/nvidia/pytorch:26.02-py3
 
 # 1. Install uv (Method: COPY from official image is cleanest)
 #COPY --from=ghcr.io/astral-sh/uv:latest /uv /uvx /bin/
diff --git a/bionemo-recipes/recipes/geneformer_native_te_mfsdp_fp8/Dockerfile b/bionemo-recipes/recipes/geneformer_native_te_mfsdp_fp8/Dockerfile
@@ -1,5 +1,5 @@
 # syntax=docker/dockerfile:1.4
-FROM nvcr.io/nvidia/pytorch:26.01-py3
+FROM nvcr.io/nvidia/pytorch:26.02-py3
 
 RUN apt-get update && apt-get install -y git
 
diff --git a/bionemo-recipes/recipes/llama3_native_te/Dockerfile b/bionemo-recipes/recipes/llama3_native_te/Dockerfile
@@ -1,5 +1,5 @@
 # syntax=docker/dockerfile:1.4
-FROM nvcr.io/nvidia/pytorch:26.01-py3
+FROM nvcr.io/nvidia/pytorch:26.02-py3
 
 RUN --mount=type=cache,target=/root/.cache/pip \
     --mount=type=bind,source=requirements.txt,target=/requirements.txt \
diff --git a/bionemo-recipes/recipes/vit/Dockerfile b/bionemo-recipes/recipes/vit/Dockerfile
@@ -1,4 +1,4 @@
-FROM nvcr.io/nvidia/pytorch:26.01-py3
+FROM nvcr.io/nvidia/pytorch:26.02-py3
 
 RUN --mount=type=secret,id=netrc,target=/root/.netrc \
     --mount=type=cache,target=/root/.cache/pip \
diff --git a/ci/lepton/model_convergence/configs/base.yaml b/ci/lepton/model_convergence/configs/base.yaml
@@ -11,7 +11,7 @@ template_type: convergence_tests
 # Defines the base Docker image and registry auth needed
 ############################################################
 container:
-  image: nvcr.io/nvidia/pytorch:26.01-py3
+  image: nvcr.io/nvidia/pytorch:26.02-py3
   registry_auth: lepton-nvidia
 
 ############################################################
diff --git a/ci/scripts/recipes_local_test.py b/ci/scripts/recipes_local_test.py
@@ -48,14 +48,14 @@
     "bionemo-recipes/models/amplify": "svcbionemo023/bionemo-framework:amplify-model-devcontainer-082025",
 }
 
-# DEFAULT_CONTAINER = "nvcr.io/nvidia/pytorch:26.01-py3"
+# DEFAULT_CONTAINER = "nvcr.io/nvidia/pytorch:26.02-py3"
 
-# This is a squashed version of the pytorch:26.01-py3 image, generated with
-# docker-squash nvcr.io/nvidia/pytorch:26.01-py3 -t svcbionemo023/bionemo-framework:pytorch26.01-py3-squashed
+# This is a squashed version of the pytorch:26.02-py3 image, generated with
+# docker-squash nvcr.io/nvidia/pytorch:26.02-py3 -t svcbionemo023/bionemo-framework:pytorch26.02-py3-squashed
 # --output type=registry,compression=zstd,force-compression=true,oci-mediatypes=true,compression-level=15
 # and pushed to the dockerhub registry. Our github actions are able to cache image pulls from dockerhub but not nvcr, so
 # hopefully this cuts down slightly on CI time at the expense of having a slightly in-directed image location.
-DEFAULT_CONTAINER = "svcbionemo023/bionemo-framework:pytorch26.01-py3-squashed"
+DEFAULT_CONTAINER = "svcbionemo023/bionemo-framework:pytorch26.02-py3-squashed"
 
 
 def get_git_root() -> str:

Original file line number	Diff line number	Diff line change
`@@ -1,4 +1,4 @@`
`1`		`-FROM nvcr.io/nvidia/pytorch:26.01-py3`
	`1`	`+FROM nvcr.io/nvidia/pytorch:26.02-py3`
`2`	`2`	`WORKDIR /workspace/bionemo`
`3`	`3`	`COPY . .`
`4`	`4`	`RUN --mount=type=cache,target=/root/.cache/pip \`