Merge branch 'main' into pbinder/remove_core

Author: polinabinder1

.devcontainer/recipes/requirements.txt

@@ -7,6 +7,6 @@ megatron-fsdp==0.1.0rc0
 torchmetrics
 tqdm
 transformer_engine
-transformers @ git+https://github.com/huggingface/transformers.git
+transformers
 typer
 wandb

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

@@ -89,7 +89,30 @@ jobs:
             fi
           fi
 
-          echo "dirs=$DIRS" >> $GITHUB_OUTPUT
+          # 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.
+
+          # This uses 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.
+
+          DIRS_WITH_IMAGES=$(echo "$DIRS" | jq -c '
+            map({
+              dir: .,
+              image: (
+                if . == "models/amplify" then
+                  "svcbionemo023/bionemo-framework:amplify-model-devcontainer-082025"
+                else
+                  "svcbionemo023/bionemo-framework:pytorch25.06-py3-squashed-zstd"
+                end
+              )
+            })
+          ')
+          echo "dirs=$DIRS_WITH_IMAGES" >> $GITHUB_OUTPUT
+
       - name: Show output
         run: |
           echo "=== Changed Files Analysis ==="
@@ -106,31 +129,43 @@ 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: Show GPU info
+        run: nvidia-smi
       - 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:
-          python-version: "3.12"
+          sparse-checkout: "${{ matrix.recipe.dir }}"
+          sparse-checkout-cone-mode: false
 
-      - name: Install ci script dependencies
+      - name: Install dependencies
+        working-directory: ${{ matrix.recipe.dir }}
         run: |
-          python -m pip install --upgrade pip
-          pip install platformdirs
+          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
 
       - name: Run tests
-        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
+        working-directory: ${{ matrix.recipe.dir }}
+        run: pytest -v .
 
   verify-recipe-tests:
     # This job checks the status of the unit-tests matrix and fails if any matrix job failed or was cancelled.

ci/benchmarks/partial-conv/evo2_pretrain.yaml

@@ -7,15 +7,15 @@ key_segments:
   lr: False
   min_lr: False
   wu_steps: False
-  artefacts_url: False
+  pckg_url: False
   file_name_wheel: False
 script_args:
   # All arguments referenced in the script string must be specified here.
   # Arguments not referenced in the script string must have the 'arg' field specified.
   # See jet/core/configs.py for the specification of the configuration class
   workspace: /workspace/bionemo2
   data_path: /data/evo2
-  artefacts_url: https://__token__:${{JET_GITLAB_TOKEN}}@gitlab-master.nvidia.com/api/v4/projects/180496/packages/pypi/simple
+  pckg_url: gitlab-master.nvidia.com/api/v4/projects/180496/packages/pypi/simple/
   file_name_wheel: subquadratic-ops
   model: evo2
   variant: train
@@ -40,7 +40,7 @@ script_args:
 script: |-
   INSTALL_FLAG="/tmp/install_done_${{SLURMD_NODENAME}}";
   if [ "$SLURM_LOCALID" = "0" ]; then
-      pip install ${file_name_wheel} --index-url ${artefacts_url}
+      pip install ${file_name_wheel} --index-url https://oauth2:$JET_GITLAB_TOKEN@${pckg_url} --extra-index-url https://pypi.org/simple/
       touch $INSTALL_FLAG
   fi
   # All ranks wait until install flag file appears

models/amplify/pyproject.toml

@@ -19,8 +19,8 @@ dependencies = [
     "nvidia_resiliency_ext",
     "omegaconf",
     "pytest",
-    "torch",
-    # "transformer_engine[pytorch]",
+    "torch==2.6.0a0+ecf3bae40a.nv25.01",
+    "transformer_engine[pytorch]",
     "transformers",
     "xformers",
 ]

recipes/esm2_accelerate/hydra_config/L0_sanity.yaml

@@ -1,14 +1,17 @@
 defaults:
   - defaults
+  - _self_
 
 model_tag: "nvidia/esm2_t6_8M_UR50D"
-stop_after_n_steps: 4
+stop_after_n_steps: 250
+
 trainer:
   run_name: "esm2_t6_8M_UR50D_sanity"
   per_device_train_batch_size: 2
   per_device_eval_batch_size: 2
-  save_steps: 2
-  eval_steps: 2
-  logging_steps: 1
+  save_steps: 1000
+  eval_steps: 1000
+  logging_steps: 10
   report_to: "none"
   dataloader_num_workers: 0
+  warmup_steps: 0

recipes/esm2_accelerate/requirements.txt

@@ -3,5 +3,5 @@ datasets
 deepspeed
 hydra-core
 torchmetrics
-transformers @ git+https://github.com/huggingface/transformers.git
+transformers
 wandb

recipes/esm2_accelerate/test_train.py

@@ -14,6 +14,7 @@
 # limitations under the License.
 
 import os
+import random
 import re
 import shutil
 import subprocess
@@ -40,6 +41,38 @@
 )
 
 
+def extract_final_train_loss(output_text: str) -> float:
+    """
+    Parse the training output to extract the final train_loss value.
+
+    Args:
+        output_text: Combined stdout and stderr from training process
+
+    Returns:
+        Final train_loss value as float
+
+    Raises:
+        ValueError: If no train_loss found or parsing fails
+    """
+    # Look for dictionary-like patterns containing train_loss
+    # Pattern matches: {'key': value, 'train_loss': value, ...}
+    pattern = r'\{[^{}]*[\'"]train_loss[\'"]:\s*([0-9.]+)[^{}]*\}'
+
+    matches = re.findall(pattern, output_text)
+
+    if not matches:
+        # Fallback: try to find train_loss in any context
+        simple_pattern = r'[\'"]train_loss[\'"]:\s*([0-9.]+)'
+        matches = re.findall(simple_pattern, output_text)
+
+    if not matches:
+        raise ValueError("No train_loss found in training output")
+
+    # Return the last (final) train_loss value found
+    final_train_loss = float(matches[-1])
+    return final_train_loss
+
+
 def test_train_can_resume_from_checkpoint(monkeypatch, tmp_path: Path):
     """Test that train.py runs successfully with sanity config and creates expected outputs."""
 
@@ -51,11 +84,20 @@ def test_train_can_resume_from_checkpoint(monkeypatch, tmp_path: Path):
     monkeypatch.setenv("RANK", "0")
     monkeypatch.setenv("WORLD_SIZE", "1")
     monkeypatch.setenv("MASTER_ADDR", "localhost")
-    monkeypatch.setenv("MASTER_PORT", "29500")
+    monkeypatch.setenv("MASTER_PORT", f"{random.randint(20000, 40000)}")
     monkeypatch.setenv("WANDB_MODE", "disabled")
 
     with initialize_config_dir(config_dir=str(recipe_dir / "hydra_config"), version_base="1.2"):
-        sanity_config = compose(config_name="L0_sanity", overrides=[f"trainer.output_dir={tmp_path}"])
+        sanity_config = compose(
+            config_name="L0_sanity",
+            overrides=[
+                f"trainer.output_dir={tmp_path}",
+                "stop_after_n_steps=4",
+                "trainer.do_eval=False",
+                "trainer.save_steps=2",
+                f"hydra.run.dir={tmp_path}/outputs",
+            ],
+        )
 
     main(sanity_config)
 
@@ -155,11 +197,15 @@ def test_accelerate_launch(accelerate_config, model_tag, tmp_path):
         str(accelerate_config_path),
         "--num_processes",
         "1",
+        "--main_process_port",
+        f"{random.randint(20000, 40000)}",
         str(train_py),
         "--config-name",
         "L0_sanity.yaml",
         f"model_tag={model_tag}",
         f"trainer.output_dir={tmp_path}",
+        f"hydra.run.dir={tmp_path}/outputs",
+        "trainer.do_eval=False",
     ]
 
     result = subprocess.run(
@@ -176,6 +222,17 @@ def test_accelerate_launch(accelerate_config, model_tag, tmp_path):
         print(f"STDERR:\n{result.stderr}")
         pytest.fail(f"Command:\n{' '.join(cmd)}\nfailed with exit code {result.returncode}")
 
+    # Parse the training output to check final train_loss
+    combined_output = result.stdout + result.stderr
+    try:
+        final_train_loss = extract_final_train_loss(combined_output)
+        print(f"Final train_loss: {final_train_loss}")
+        assert final_train_loss < 3.0, f"Final train_loss {final_train_loss} should be less than 3.0"
+    except ValueError as e:
+        print(f"STDOUT:\n{result.stdout}")
+        print(f"STDERR:\n{result.stderr}")
+        pytest.fail(f"Failed to extract train_loss from output: {e}")
+
 
 @requires_multi_gpu
 @pytest.mark.parametrize(
@@ -186,9 +243,11 @@ def test_accelerate_launch(accelerate_config, model_tag, tmp_path):
         # modeling_esm_te import seems to fix it.
         # ("fsdp1_te.yaml", "nvidia/esm2_t6_8M_UR50D"),
         ("fsdp2_te.yaml", "nvidia/esm2_t6_8M_UR50D"),
-        ("default.yaml", "facebook/esm2_t6_8M_UR50D"),
-        ("fsdp1_hf.yaml", "facebook/esm2_t6_8M_UR50D"),
-        ("fsdp2_hf.yaml", "facebook/esm2_t6_8M_UR50D"),
+        # TODO: (BIONEMO-2761). These tests were broken by https://github.com/huggingface/transformers/pull/40370, but
+        # oddly the single-GPU tests still seem to pass. Changing the attention_backend doesn't seem to help.
+        # ("default.yaml", "facebook/esm2_t6_8M_UR50D"),
+        # ("fsdp1_hf.yaml", "facebook/esm2_t6_8M_UR50D"),
+        # ("fsdp2_hf.yaml", "facebook/esm2_t6_8M_UR50D"),
     ],
 )
 def test_accelerate_launch_multi_gpu(accelerate_config, model_tag, tmp_path):
@@ -211,11 +270,15 @@ def test_accelerate_launch_multi_gpu(accelerate_config, model_tag, tmp_path):
         str(accelerate_config_path),
         "--num_processes",
         "2",
+        "--main_process_port",
+        f"{random.randint(20000, 40000)}",
         str(train_py),
         "--config-name",
         "L0_sanity.yaml",
         f"model_tag={model_tag}",
         f"trainer.output_dir={tmp_path}",
+        f"hydra.run.dir={tmp_path}/outputs",
+        "trainer.do_eval=False",
     ]
 
     result = subprocess.run(
@@ -231,3 +294,14 @@ def test_accelerate_launch_multi_gpu(accelerate_config, model_tag, tmp_path):
         print(f"STDOUT:\n{result.stdout}")
         print(f"STDERR:\n{result.stderr}")
         pytest.fail(f"Command:\n{' '.join(cmd)}\nfailed with exit code {result.returncode}")
+
+    # Parse the training output to check final train_loss
+    combined_output = result.stdout + result.stderr
+    try:
+        final_train_loss = extract_final_train_loss(combined_output)
+        print(f"Final train_loss: {final_train_loss}")
+        assert final_train_loss < 3.0, f"Final train_loss {final_train_loss} should be less than 3.0"
+    except ValueError as e:
+        print(f"STDOUT:\n{result.stdout}")
+        print(f"STDERR:\n{result.stderr}")
+        pytest.fail(f"Failed to extract train_loss from output: {e}")

recipes/esm2_native_te_nvfsdp_thd/hydra_config/L1_3B_ddp.yaml

@@ -2,7 +2,7 @@ defaults:
   - defaults
 
 # Training config
-model_name: esm2_t33_650M_UR50D
+model_name: esm2_t36_3B_UR50D
 micro_batch_size: 32
 num_train_steps: 10_000
 

recipes/esm2_native_te_nvfsdp_thd/test_thd_format.py

@@ -226,10 +226,10 @@ def test_mlm_data_collator_integration():
         if mlm_prob == 0.0:
             # No masking - all labels should be -100
             assert (sample["labels"] == -100).all(), "With mlm_probability=0.0, all labels should be -100"
-        else:
-            # Some masking should occur
-            masked_count = (sample["labels"] != -100).sum()
-            assert masked_count > 0, f"With mlm_probability={mlm_prob}, some tokens should be masked"
+        # TODO: This is a very flaky test with such a small input batch, we should make it larger if we want to ensure a
+        # token is masked
+        # else: # Some masking should occur masked_count = (sample["labels"] != -100).sum() assert
+        #     masked_count > 0, f"With mlm_probability={mlm_prob}, some tokens should be masked"
 
 
 if __name__ == "__main__":

recipes/esm2_native_te_nvfsdp_thd/test_train.py

@@ -16,12 +16,16 @@
 from pathlib import Path
 
 import pytest
+import torch
 from hydra import compose, initialize_config_dir
 
 from train import main
 
 
-@pytest.mark.xfail(reason="CUDNN padded packed sequences not supported on all hardware currently.")
+@pytest.mark.xfail(
+    torch.cuda.get_device_capability() == (12, 0),
+    reason="CUDNN padded packed sequences not supported on all hardware currently (nvbugs/5458694).",
+)
 def test_main_invocation(monkeypatch, tmp_path):
     """Test that the main function can be invoked with the correct arguments."""
 
@@ -43,7 +47,10 @@ def test_main_invocation(monkeypatch, tmp_path):
     main(sanity_config)
 
 
-@pytest.mark.xfail(reason="CUDNN padded packed sequences not supported on all hardware currently.")
+@pytest.mark.xfail(
+    torch.cuda.get_device_capability() == (12, 0),
+    reason="CUDNN padded packed sequences not supported on all hardware currently (nvbugs/5458694).",
+)
 def test_main_invocation_ddp(monkeypatch, tmp_path):
     """Test that the main function can be invoked wrapping the model in DDP."""