Fix test

Author: balvisio

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

@@ -157,21 +157,39 @@ def _validate_tied_weight_config(self, model: ModelType) -> None:
         ``output_layer`` share the same weight tensor.  Both lists must treat the
         tied pair as a unit: list both layers or neither.  Listing only one side
         raises ``ValueError``.
+
+        The check walks the actual model so that wildcard patterns (e.g.
+        ``"embedding.*"``) are evaluated against real module paths rather than
+        synthetic names.
         """
         if not self._get_is_tied(model):
             return
 
-        target_word_emb = self._matches_lora_target("word_embeddings", "word_embeddings")
-        target_output = self._matches_lora_target("output_layer", "output_layer")
+        targeted_short_names: set[str] = set()
+        skip_frozen_short_names: set[str] = set()
+
+        def _collect(module: nn.Module, name: str | None = None, prefix: str | None = None) -> nn.Module:
+            full_name = f"{prefix}.{name}" if prefix else (name or "")
+            short_name = name or ""
+            if self._matches_lora_target(short_name, full_name):
+                targeted_short_names.add(short_name)
+            if self._matches_skip_freeze(short_name, full_name):
+                skip_frozen_short_names.add(short_name)
+            return module
+
+        self._walk_model(model, _collect)
+
+        target_word_emb = "word_embeddings" in targeted_short_names
+        target_output = "output_layer" in targeted_short_names
         if target_word_emb != target_output:
             raise ValueError(
                 "share_embeddings_and_output_weights is enabled: target_modules must "
                 "include both word_embeddings and output_layer, or neither. "
                 f"word_embeddings matched: {target_word_emb}, output_layer matched: {target_output}."
             )
 
-        skip_word_emb = self._matches_skip_freeze("word_embeddings", "word_embeddings")
-        skip_output = self._matches_skip_freeze("output_layer", "output_layer")
+        skip_word_emb = "word_embeddings" in skip_frozen_short_names
+        skip_output = "output_layer" in skip_frozen_short_names
         if skip_word_emb != skip_output:
             raise ValueError(
                 "share_embeddings_and_output_weights is enabled: skip_freeze_modules must "

bionemo-recipes/recipes/evo2_megatron/tests/bionemo/evo2/models/test_evo2_lora_1.py

@@ -517,19 +517,22 @@ def test_lora_checkpoint_excludes_frozen_embeddings(self, tmp_path: Path, base_c
         assert len(adapter_keys) > 0, "Checkpoint should still contain LoRA adapter keys."
 
     @pytest.mark.parametrize(
-        "skip_module, expected_key_substr, lora_targets",
+        "skip_freeze, expected_key_substr, lora_targets",
         [
-            ("word_embeddings", "word_embeddings", None),
-            ("final_norm", "final_norm", None),
-            ("dense", "mixer.dense.", None),
-            ("linear_fc2", "mlp.linear_fc2.", ["dense_projection", "linear_qkv", "linear_proj", "linear_fc1"]),
+            # word_embeddings and output_layer share a weight tensor when tying is enabled;
+            # both must appear in skip_freeze to satisfy the symmetry contract.
+            (["word_embeddings", "output_layer"], "word_embeddings", None),
+            (["final_norm"], "final_norm", None),
+            (["dense"], "mixer.dense.", None),
+            (["linear_fc2"], "mlp.linear_fc2.", ["dense_projection", "linear_qkv", "linear_proj", "linear_fc1"]),
         ],
+        ids=["word_embeddings", "final_norm", "dense", "linear_fc2"],
     )
     def test_lora_skip_freeze_saves_and_trains_module(
         self,
         tmp_path: Path,
         base_ckpt: Path,
-        skip_module: str,
+        skip_freeze: list[str],
         expected_key_substr: str,
         lora_targets: list[str] | None,
     ):
@@ -538,12 +541,12 @@ def test_lora_skip_freeze_saves_and_trains_module(
 
         from bionemo.evo2.models.evo2_provider import hyena_forward_step
 
-        lora_dir = tmp_path / f"lora_{skip_module}"
+        lora_dir = tmp_path / f"lora_{skip_freeze[0]}"
         cfg = _build_pretrain_config(
             lora_dir,
             train_iters=1,
             lora=True,
-            skip_freeze=[skip_module],
+            skip_freeze=skip_freeze,
             lora_target_modules=lora_targets,
             pretrained_ckpt_dir=str(base_ckpt),
         )