Attempt to address failing CI

Signed-off-by: John St. John <jstjohn@nvidia.com>

Author: jstjohn

bionemo-recipes/recipes/evo2_megatron/src/bionemo/evo2/models/evo2_provider.py

@@ -18,8 +18,10 @@
 
 
 import math
+import sys
 from dataclasses import dataclass
 from functools import partial
+from pathlib import Path
 from typing import Callable, Iterable, Literal, Optional, Type
 
 import torch
@@ -54,6 +56,30 @@
 from bionemo.evo2.models.megatron.hyena.hyena_utils import hyena_no_weight_decay_cond
 
 
+def _patch_megatron_dataset_helper_compile() -> None:
+    """Skip Megatron's runtime helper build when a wheel already ships the extension."""
+    from megatron.core.datasets import utils as dataset_utils
+
+    original_compile_helpers = dataset_utils.compile_helpers
+    if getattr(original_compile_helpers, "_evo2_prebuilt_helper_guard", False):
+        guarded_compile_helpers = original_compile_helpers
+    else:
+
+        def guarded_compile_helpers() -> None:
+            datasets_dir = Path(dataset_utils.__file__).resolve().parent
+            if not (datasets_dir / "Makefile").exists() and list(datasets_dir.glob("helpers_cpp*.so")):
+                return None
+            return original_compile_helpers()
+
+        guarded_compile_helpers._evo2_prebuilt_helper_guard = True
+        dataset_utils.compile_helpers = guarded_compile_helpers
+
+    bridge_initialize = sys.modules.get("megatron.bridge.training.initialize")
+    if bridge_initialize is not None:
+        bridge_initialize.compile_helpers = guarded_compile_helpers
+
+
+_patch_megatron_dataset_helper_compile()
 register_allowed_target_prefix("bionemo.evo2.")
 
 

bionemo-recipes/recipes/evo2_megatron/src/bionemo/evo2/run/train.py

@@ -408,11 +408,20 @@ def parse_args(args: Optional[List[str]] = None) -> argparse.Namespace:
         help="Use the faster, but maybe less accurate fused form of cross entropy, "
         "which also has bf16 grads internally.",
     )  # DONE
-    parser.add_argument(
-        "--no-fp32-residual-connection",
+    fp32_residual_group = parser.add_mutually_exclusive_group(required=False)
+    fp32_residual_group.add_argument(
+        "--fp32-residual-connection",
+        dest="fp32_residual_connection",
         action="store_true",
-        default=False,
-        help="If set, turn off fp32 residual connections which may be faster but may impact accuracy.",
+        default=None,
+        help="Enable fp32 residual connections. Defaults to the selected model provider setting.",
+    )
+    fp32_residual_group.add_argument(
+        "--no-fp32-residual-connection",
+        dest="fp32_residual_connection",
+        action="store_false",
+        default=None,
+        help="Disable fp32 residual connections. Defaults to the selected model provider setting.",
     )  # DONE
     parser.add_argument(
         "--debug-ddp-parity-freq",
@@ -859,11 +868,11 @@ def train(args: argparse.Namespace) -> None:
         cfg.model.seq_len_interpolation_factor = args.seq_len_interpolation_factor
     cfg.model.calculate_per_token_loss = not args.no_calculate_per_token_loss
     model_type = infer_model_type(args.model_size)
-    if model_type != "hyena" and not args.no_fp32_residual_connection:
+    if args.fp32_residual_connection is not None:
+        cfg.model.fp32_residual_connection = args.fp32_residual_connection
+    if model_type != "hyena" and cfg.model.fp32_residual_connection:
         logger.info("Disabling fp32_residual_connection for non-Hyena model (not compatible with TE layers)")
         cfg.model.fp32_residual_connection = False
-    else:
-        cfg.model.fp32_residual_connection = not args.no_fp32_residual_connection
     cfg.model.cross_entropy_loss_fusion = args.cross_entropy_loss_fusion
     # cfg.model.cuda_graph_impl = "local" # or "transformer_engine"
     # cfg.model.cuda_graph_scope = "full_iteration"

bionemo-recipes/recipes/evo2_megatron/tests/bionemo/evo2/test_model_providers.py

@@ -15,13 +15,16 @@
 
 """Tests for model provider instantiation, naming, and checkpoint converters."""
 
+from pathlib import Path
+
 import pytest
 import torch
 
 from bionemo.evo2.models.evo2_provider import (
     HYENA_MODEL_OPTIONS,
     MODEL_OPTIONS,
     Hyena1bModelProvider,
+    _patch_megatron_dataset_helper_compile,
     infer_model_type,
 )
 from bionemo.evo2.utils.checkpoint.mbridge_to_vortex import _split_fc1, mbridge_to_vortex_state_dict
@@ -63,6 +66,46 @@ def test_infer_model_type_unknown():
         infer_model_type("nonexistent_model")
 
 
+@pytest.mark.parametrize(
+    ("has_makefile", "has_prebuilt_extension", "expected_original_calls"),
+    [
+        (False, True, 0),
+        (True, True, 1),
+        (False, False, 1),
+    ],
+)
+def test_megatron_dataset_helper_compile_guard(
+    monkeypatch: pytest.MonkeyPatch,
+    tmp_path: Path,
+    has_makefile: bool,
+    has_prebuilt_extension: bool,
+    expected_original_calls: int,
+):
+    """Skip Megatron's runtime make step only when a prebuilt helper extension exists."""
+    from megatron.bridge.training import initialize as bridge_initialize
+    from megatron.core.datasets import utils as dataset_utils
+
+    calls = []
+
+    def original_compile_helpers():
+        calls.append("called")
+
+    if has_makefile:
+        (tmp_path / "Makefile").write_text("all:\n")
+    if has_prebuilt_extension:
+        (tmp_path / "helpers_cpp.cpython-312-x86_64-linux-gnu.so").touch()
+
+    monkeypatch.setattr(dataset_utils, "__file__", str(tmp_path / "utils.py"))
+    monkeypatch.setattr(dataset_utils, "compile_helpers", original_compile_helpers)
+    monkeypatch.setattr(bridge_initialize, "compile_helpers", original_compile_helpers)
+
+    _patch_megatron_dataset_helper_compile()
+
+    dataset_utils.compile_helpers()
+    assert bridge_initialize.compile_helpers is dataset_utils.compile_helpers
+    assert len(calls) == expected_original_calls
+
+
 def _make_mock_savanna_sd(pattern: str) -> dict[str, torch.Tensor]:
     """Create a minimal mock savanna state dict for the given pattern.