Merge remote-tracking branch 'origin/main' into dorotat/fix-jet-evo2-pretrain

Author: dorotat-nv

.github/workflows/unit-tests-framework.yml

@@ -30,10 +30,18 @@ jobs:
           fetch-depth: 0
           submodules: "recursive"
 
+      - name: Get merge-base commit
+        id: merge-base
+        run: |
+          # Get the merge-base between current branch and main
+          MERGE_BASE=$(git merge-base HEAD origin/main)
+          echo "merge-base=$MERGE_BASE" >> $GITHUB_OUTPUT
+          echo "Merge-base commit: $MERGE_BASE"
+
       - uses: step-security/changed-files@v46
         id: changed-files
         with:
-          base_sha: main
+          base_sha: ${{ steps.merge-base.outputs.merge-base }}
           files: |
             **
             !models/**
@@ -42,6 +50,7 @@ jobs:
             !.github/**
             !.gitignore
             !.devcontainer/**
+            !ci/scripts/recipes_local_test.py
             .github/workflows/unit-tests-framework.yml
 
       - name: Show output
@@ -239,8 +248,6 @@ jobs:
       - name: Run notebook tests
         env:
           BIONEMO_DATA_SOURCE: ngc
-          # this variable should be used in the notebooks to run a subset of the model layers or a smaller model/dataset
-          FAST_CI_MODE: true
         run: |
           chmod +x ./ci/scripts/run_pytest_notebooks.sh
           ./ci/scripts/run_pytest_notebooks.sh

.github/workflows/unit-tests-recipes.yml

@@ -35,13 +35,21 @@ jobs:
         with:
           fetch-depth: 0
 
+      - name: Get merge-base commit
+        id: merge-base
+        run: |
+          # Get the merge-base between current branch and main
+          MERGE_BASE=$(git merge-base HEAD origin/main)
+          echo "merge-base=$MERGE_BASE" >> $GITHUB_OUTPUT
+          echo "Merge-base commit: $MERGE_BASE"
+
       - name: Get changed files
         id: changed-files
         uses: step-security/changed-files@v46
         with:
           json: true
           matrix: true
-          base_sha: main
+          base_sha: ${{ steps.merge-base.outputs.merge-base }}
           dir_names: true
           dir_names_max_depth: 2
           files: |
@@ -81,24 +89,15 @@ jobs:
             fi
           fi
 
-          # Assign Docker images to the selected directories
-          # 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.
-          DIRS_WITH_IMAGES=$(echo "$DIRS" | jq -c '
-            map({
-              dir: .,
-              image: (
-                if . == "models/amplify" then
-                  "svcbionemo023/bionemo-framework:amplify-model-devcontainer-082025"
-                else
-                  "nvcr.io/nvidia/pytorch:25.06-py3"
-                end
-              )
-            })
-          ')
-          echo "dirs=$DIRS_WITH_IMAGES" >> $GITHUB_OUTPUT
+          echo "dirs=$DIRS" >> $GITHUB_OUTPUT
       - name: Show output
         run: |
+          echo "=== Changed Files Analysis ==="
+          echo "Current branch: ${{ github.ref_name }}"
+          echo "Merge-base commit: ${{ steps.merge-base.outputs.merge-base }}"
+          echo "Changed files compared to merge-base:"
+          echo '${{ steps.changed-files.outputs.all_changed_files }}' | jq -r '.[]' | sed 's/^/  - /'
+          echo "Total changed files: $(echo '${{ steps.changed-files.outputs.all_changed_files }}' | jq '. | length')"
           echo '${{ toJSON(steps.changed-files.outputs) }}'
           echo '${{ toJSON(steps.set-dirs.outputs) }}'
         shell: bash
@@ -107,40 +106,31 @@ jobs:
     needs: changed-dirs
     runs-on: linux-amd64-gpu-l4-latest-1
     if: ${{ needs.changed-dirs.outputs.dirs != '[]' }}
-    container:
-      image: ${{ matrix.recipe.image }}
-    strategy:
-      matrix:
-        recipe: ${{ fromJson(needs.changed-dirs.outputs.dirs) }}
-      fail-fast: false
 
     steps:
       - name: Setup proxy cache
         uses: nv-gha-runners/setup-proxy-cache@main
+
       - name: Checkout repository
         uses: actions/checkout@v4
+
+      - name: Setup python
+        uses: actions/setup-python@v5
         with:
-          sparse-checkout: "${{ matrix.recipe.dir }}"
-          sparse-checkout-cone-mode: false
+          python-version: "3.12"
 
-      - name: Install dependencies
-        working-directory: ${{ matrix.recipe.dir }}
-        #
+      - name: Install ci script dependencies
         run: |
-          if [ -f pyproject.toml ] || [ -f setup.py ]; then
-            PIP_CONSTRAINT= pip install -e .
-            echo "Installed ${{ matrix.recipe.dir }} as editable package"
-          elif [ -f requirements.txt ]; then
-            PIP_CONSTRAINT= pip install -r requirements.txt
-            echo "Installed ${{ matrix.recipe.dir }} from requirements.txt"
-          else
-            echo "No pyproject.toml, setup.py, or requirements.txt found in ${{ matrix.recipe.dir }}"
-            exit 1
-          fi
+          python -m pip install --upgrade pip
+          pip install platformdirs
 
       - name: Run tests
-        working-directory: ${{ matrix.recipe.dir }}
-        run: pytest -v .
+        env:
+          DIRS_JSON: ${{ needs.changed-dirs.outputs.dirs }}
+        run: |
+          # Convert JSON array to space-separated arguments
+          DIRS_ARGS=$(echo "$DIRS_JSON" | jq -r '.[]' | tr '\n' ' ')
+          ./ci/scripts/recipes_local_test.py $DIRS_ARGS
 
   verify-recipe-tests:
     # This job checks the status of the unit-tests matrix and fails if any matrix job failed or was cancelled.

3rdparty/NeMo

@@ -49,6 +49,7 @@ apt-get install -qyy \
   curl \
   pre-commit \
   sudo \
+  emacs-nox \
   gnupg \
   unzip \
   libsqlite3-dev

Dockerfile

@@ -23,26 +23,39 @@
 from pathlib import Path
 from typing import List, Optional
 
+from platformdirs import user_cache_dir
+
+
+PIP_CACHE_DIR = user_cache_dir(appname="bionemo-pip-cache", appauthor="nvidia")
 
 logger = logging.getLogger(__name__)
 logger.setLevel(logging.DEBUG)
 
 DOCKER_RUN_ARGS = [
     "--rm",
-    "-it",
     "--gpus",
     "all",
     "--ipc=host",
     "--ulimit",
     "memlock=-1",
     "--ulimit",
     "stack=67108864",
+    "-v",
+    f"{PIP_CACHE_DIR}:/workspace/.cache/pip",
 ]
 
 CUSTOM_CONTAINERS = {
     "models/amplify": "svcbionemo023/bionemo-framework:amplify-model-devcontainer-082025",
 }
-DEFAULT_CONTAINER = "nvcr.io/nvidia/pytorch:25.06-py3"
+
+# DEFAULT_CONTAINER = "nvcr.io/nvidia/pytorch:25.06-py3"
+
+# This is a squashed version of the pytorch:25.06-py3 image, generated with
+# docker-squash nvcr.io/nvidia/pytorch:25.06-py3 -t svcbionemo023/bionemo-framework:pytorch25.06-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:pytorch25.06-py3-squashed"
 
 
 def get_git_root() -> str:
@@ -89,24 +102,26 @@ def run_tests_in_docker(work_dir: str) -> bool:
     install_and_test_script = textwrap.dedent("""
         set -e  # Exit on any error
 
-        echo "Checking for dependency files..."
+        # Ensure image-embedded constraints do not leak into local recipe installs
+        unset PIP_CONSTRAINT || true
 
+        echo "Checking for dependency files..."
         # Install dependencies based on available files
         if [ -f pyproject.toml ] || [ -f setup.py ]; then
             echo "Installing package in editable mode..."
-            PIP_CONSTRAINT= pip install -e .
+            PIP_CACHE_DIR=/workspace/.cache/pip pip install -e .
             echo "Installed package as editable package"
         elif [ -f requirements.txt ]; then
             echo "Installing from requirements.txt..."
-            PIP_CONSTRAINT= pip install -r requirements.txt
+            PIP_CACHE_DIR=/workspace/.cache/pip pip install -r requirements.txt
             echo "Installed from requirements.txt"
         else
             echo "No pyproject.toml, setup.py, or requirements.txt found"
             exit 1
         fi
 
         echo "Running tests..."
-        pytest -v .
+        python -m pytest -v .
         """)
 
     relative_path = Path(work_dir).relative_to(git_root).as_posix()
@@ -166,6 +181,8 @@ def main():
     if args.debug:
         logging.getLogger().setLevel(logging.DEBUG)
 
+    logger.info(f"Caching pip installations to: {PIP_CACHE_DIR}")
+
     # Get directories to test
     test_dirs = get_test_directories(args.directories)
 

ci/recipes_local_test.py ci/scripts/recipes_local_test.pyci/recipes_local_test.py renamed to ci/scripts/recipes_local_test.py

@@ -19,4 +19,4 @@
 # Enable strict mode with better error handling
 set -euox pipefail
 
-pytest -v --nbval-lax -x -p no:python docs/ sub-packages/
+FAST_CI_MODE=true pytest -v --nbval-lax -x -p no:python docs/ sub-packages/

ci/scripts/run_pytest_notebooks.sh

@@ -14,6 +14,3 @@ wandb_init_args:
 # Learning rate scheduler config
 lr_scheduler_kwargs:
   num_warmup_steps: 0
-
-adamw_kwargs:
-  lr: 1e-2

recipes/esm2_native_te_mfsdp/hydra_config/L0_sanity.yaml

@@ -2,8 +2,8 @@ defaults:
   - defaults
 
 # Training config
-model_name: esm2_t33_650M_UR50D
-micro_batch_size: 32
+model_name: nvidia/esm2_t36_3B_UR50D
+micro_batch_size: 16
 num_train_steps: 10_000
 
 # WandB config

…ive_te_mfsdp/hydra_config/L1_3B_ddp.yaml …_native_te_mfsdp/hydra_config/L1_3B.yamlrecipes/esm2_native_te_mfsdp/hydra_config/L1_3B_ddp.yaml renamed to recipes/esm2_native_te_mfsdp/hydra_config/L1_3B.yaml recipes/esm2_native_te_mfsdp/hydra_config/L1_3B_ddp.yaml renamed to recipes/esm2_native_te_mfsdp/hydra_config/L1_3B.yaml

@@ -3,12 +3,12 @@ defaults:
   - _self_
 
 # Training config
-model_name: esm2_t33_650M_UR50D
-micro_batch_size: 16
+model_name: nvidia/esm2_t33_650M_UR50D
+micro_batch_size: 4
 num_train_steps: 200
 
 # WandB config
 wandb_init_args:
-  name: "esm2_t33_650M_UR50D_nvfsdp"
+  name: "esm2_t33_650M_UR50D_mfsdp"
   project: "bionemo-recipes-pstjohn"
   mode: "offline"

recipes/esm2_native_te_mfsdp/hydra_config/L1_650M.yaml

@@ -5,15 +5,19 @@ max_seq_length: 1024
 data_path: .
 num_train_steps: ???
 
+# TODO: Once BIONEMO-2583 and BIONEMO-2719 are fixed, enable this by default and simplify training scripts to remove the
+# meta-device conditional.
+use_meta_device: false
+
 # WandB config
 wandb_init_args:
   name: ???
 
-# nvFSDP config
+# mFSDP config
 fully_shard_kwargs:
   zero_dp_strategy: "optim_grads_params"
   calculate_per_token_loss: false
-  init_model_with_meta_device: false
+  init_model_with_meta_device: ${use_meta_device}
   check_for_nan_in_grad: true
   grad_reduce_in_fp32: false
   preserve_fp32_weights: true