Merge branch 'main' into savitha/og2-readme-metrics-update

Author: savitha-eng

bionemo-recipes/recipes/esm2_accelerate_te/tests/test_accelerate_amplify.py

@@ -20,6 +20,8 @@
 recipe tests.
 """
 
+import pytest
+
 # Local helper function import, resolved in conftest.py
 from launch import launch_accelerate, requires_multi_gpu
 
@@ -39,6 +41,10 @@ def test_te_with_fp8_config(tmp_path):
     assert train_loss < 3.0, f"Final train_loss {train_loss} should be less than 3.0"
 
 
+@pytest.mark.xfail(
+    strict=True,
+    reason="BIO-466: AMPLIFY HF model does not implement get_input_embeddings, required by accelerate FSDP2 (transformers>=5.6).",
+)
 def test_te_with_fsdp2_config(tmp_path):
     train_loss = launch_accelerate("fsdp2_te.yaml", tmp_path, 1, "L0_sanity_amplify")
     assert train_loss < 3.0, f"Final train_loss {train_loss} should be less than 3.0"
@@ -56,6 +62,10 @@ def test_te_with_fp8_config_two_gpus(tmp_path):
     assert train_loss < 3.0, f"Final train_loss {train_loss} should be less than 3.0"
 
 
+@pytest.mark.xfail(
+    strict=True,
+    reason="BIO-466: AMPLIFY HF model does not implement get_input_embeddings, required by accelerate FSDP2 (transformers>=5.6).",
+)
 @requires_multi_gpu
 def test_te_with_fsdp2_config_two_gpus(tmp_path):
     train_loss = launch_accelerate("fsdp2_te.yaml", tmp_path, 2, "L0_sanity_amplify")

bionemo-recipes/recipes/esm2_accelerate_te/tests/test_accelerate_esm2.py

@@ -13,8 +13,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import pytest
-
 # Local helper function import, resolved in conftest.py
 from launch import launch_accelerate, requires_multi_gpu
 
@@ -49,7 +47,6 @@ def test_hf_with_default_config(tmp_path):
     assert train_loss < 3.0, f"Final train_loss {train_loss} should be less than 3.0"
 
 
-@pytest.mark.xfail(reason="BIONEMO-3331: FSDP2 and HF model failing with 25.11+ torch container.")
 def test_hf_with_fsdp2_config(tmp_path):
     train_loss = launch_accelerate("fsdp2_hf.yaml", tmp_path, 1, "L0_sanity", "model_tag=facebook/esm2_t6_8M_UR50D")
     assert train_loss < 3.0, f"Final train_loss {train_loss} should be less than 3.0"
@@ -85,7 +82,6 @@ def test_hf_with_fsdp1_config_two_gpus(tmp_path):
     assert train_loss < 3.0, f"Final train_loss {train_loss} should be less than 3.0"
 
 
-@pytest.mark.xfail(reason="BIONEMO-3331: FSDP2 and HF model failing with 25.11+ torch container.")
 @requires_multi_gpu
 def test_hf_with_fsdp2_config_two_gpus(tmp_path):
     train_loss = launch_accelerate("fsdp2_hf.yaml", tmp_path, 2, "L0_sanity", "model_tag=facebook/esm2_t6_8M_UR50D")