[multi-lora] 1/n - 3, fix: launch-time guard for unsafe recompute configs

miles/utils/multi_lora.py

@@ -63,6 +63,24 @@ def targets_expert_leaves(target_modules: Any) -> bool:
     return any(entry.split(".")[-1] in _EXPERT_LEAF_NAMES for entry in entries)
 
 
+def _recompute_source_recognizes_adapters(recompute_module: Any) -> bool:
+    import inspect
+
+    try:
+        source = inspect.getsource(recompute_module.maybe_enable_recompute_inputs_grad)
+    except (AttributeError, OSError, TypeError):
+        return False
+    return ".adapters." in source
+
+
+def _bridge_recompute_patch_recognizes_multi_lora() -> bool:
+    try:
+        from megatron.bridge.peft import recompute
+    except Exception:
+        return False
+    return _recompute_source_recognizes_adapters(recompute)
+
+
 def validate_multi_lora_args(args: Any) -> None:
     """Set ``args.multi_lora``, then validate and default the multi-LoRA arg
     surface. Called from ``miles_validate_args``; a no-op for normal runs."""
@@ -87,6 +105,17 @@ def validate_multi_lora_args(args: Any) -> None:
         "complete adapter to push to the rollout engines, and a pipelined schedule would "
         "recompute activations against a later micro-batch's adapter routing."
     )
+    recompute_modules = list(getattr(args, "recompute_modules", None) or [])
+    risky_full = getattr(args, "recompute_granularity", None) == "full"
+    risky_moe = "moe" in recompute_modules and targets_expert_leaves(args.target_modules)
+    if risky_full or risky_moe:
+        bridge_fixed = _bridge_recompute_patch_recognizes_multi_lora()
+        assert (
+            not risky_full or bridge_fixed
+        ), "Full recompute requires Megatron-Bridge#27 ('.adapters.' aware); upgrade or use selective recompute"
+        assert (
+            not risky_moe or bridge_fixed
+        ), "Expert targets with MoE recompute require Megatron-Bridge#27; upgrade or recompute core_attn and moe_act"
     # Per-slot token spans assume sequence-major contiguous sample packing, which only 'thd' provides.
     assert getattr(args, "qkv_format", "thd") == "thd", (
         "Multi-LoRA requires --qkv-format thd: per-adapter token spans assume the "

tests/fast/utils/test_multi_lora_recompute_guard.py

@@ -0,0 +1,157 @@
+"""Launch-time multi-LoRA recompute guards: unsupported recompute shapes fail at launch unless the Bridge patch is present."""
+
+import importlib.util
+import sys
+from types import SimpleNamespace
+
+import pytest
+
+import miles.utils.multi_lora as multi_lora_module
+from miles.utils.multi_lora import (
+    _bridge_recompute_patch_recognizes_multi_lora,
+    _recompute_source_recognizes_adapters,
+    validate_multi_lora_args,
+)
+
+
+def _args(**overrides) -> SimpleNamespace:
+    base = dict(
+        multi_lora_n_adapters=2,
+        lora_rank=8,
+        target_modules=["linear_qkv"],
+        train_backend="megatron",
+        pipeline_model_parallel_size=1,
+        qkv_format="thd",
+        experts_shared_outer_loras=False,
+        optimizer="adam",
+        colocate=False,
+        indep_dp=False,
+        ft_components=[],
+        offload_train=False,
+        enable_witness=False,
+        sglang_tokenizer_worker_num=1,
+        calculate_per_token_loss=False,
+        disable_rollout_trim_samples=False,
+        use_dynamic_global_batch_size=False,
+        megatron_to_hf_mode="bridge",
+        rollout_global_dataset=False,
+        rollout_function_path=None,
+        data_source_path="miles.rollout.data_source.RolloutDataSourceWithBuffer",
+        multi_lora_max_coalesce_wait_s=0.5,
+        multi_lora_max_adapter_global_batch_size=None,
+        recompute_granularity=None,
+        recompute_modules=None,
+    )
+    base.update(overrides)
+    return SimpleNamespace(**base)
+
+
+EXPERT_TARGETS = ["gate_proj", "up_proj", "down_proj"]
+
+PROBE_NAME = "_bridge_recompute_patch_recognizes_multi_lora"
+
+
+@pytest.fixture
+def unfixed_bridge(monkeypatch):
+    monkeypatch.setattr(multi_lora_module, PROBE_NAME, lambda: False)
+
+
+@pytest.fixture
+def fixed_bridge(monkeypatch):
+    monkeypatch.setattr(multi_lora_module, PROBE_NAME, lambda: True)
+
+
+@pytest.fixture
+def probe_must_not_run(monkeypatch):
+    def _boom():
+        raise AssertionError("bridge probe ran for a recompute shape that never needs it")
+
+    monkeypatch.setattr(multi_lora_module, PROBE_NAME, _boom)
+
+
+class TestUnfixedBridgeRefusals:
+    def test_full_recompute_is_refused_for_any_targets(self, unfixed_bridge):
+        validate_multi_lora_args(_args())
+        with pytest.raises(AssertionError, match=r"Megatron-Bridge#27.*selective"):
+            validate_multi_lora_args(_args(recompute_granularity="full"))
+
+    def test_moe_module_with_expert_targets_is_refused(self, unfixed_bridge):
+        with pytest.raises(AssertionError, match=r"Megatron-Bridge#27.*moe_act"):
+            validate_multi_lora_args(
+                _args(
+                    recompute_granularity="selective",
+                    recompute_modules=["core_attn", "moe"],
+                    target_modules=EXPERT_TARGETS,
+                )
+            )
+
+
+class TestFixedBridgePassThrough:
+    def test_full_recompute_is_allowed(self, fixed_bridge):
+        validate_multi_lora_args(_args(recompute_granularity="full"))
+
+    def test_moe_module_with_expert_targets_is_allowed(self, fixed_bridge):
+        validate_multi_lora_args(
+            _args(
+                recompute_granularity="selective",
+                recompute_modules=["core_attn", "moe"],
+                target_modules=EXPERT_TARGETS,
+            )
+        )
+
+    def test_pass_through_still_runs_the_rest_of_validation(self, fixed_bridge):
+        with pytest.raises(AssertionError, match="qkv-format thd"):
+            validate_multi_lora_args(_args(recompute_granularity="full", qkv_format="bshd"))
+
+
+class TestShapesThatNeverProbeTheBridge:
+    def test_no_recompute_is_allowed(self, probe_must_not_run):
+        validate_multi_lora_args(_args(target_modules=EXPERT_TARGETS))
+
+    def test_moe_module_without_expert_targets_is_allowed(self, probe_must_not_run):
+        validate_multi_lora_args(
+            _args(
+                recompute_granularity="selective",
+                recompute_modules=["core_attn", "moe"],
+                target_modules=["linear_qkv"],
+            )
+        )
+
+
+def _load_module_file(tmp_path, name: str, body: str):
+    path = tmp_path / f"{name}.py"
+    path.write_text(body)
+    spec = importlib.util.spec_from_file_location(name, path)
+    module = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(module)
+    return module
+
+
+class TestSourceProbe:
+    FIXED_BODY = (
+        "def maybe_enable_recompute_inputs_grad(model):\n"
+        '    names = ["x.adapter.w", "x.adapters.0.w"]\n'
+        '    return any(".adapter." in n or ".adapters." in n for n in names)\n'
+    )
+    UNFIXED_BODY = (
+        "def maybe_enable_recompute_inputs_grad(model):\n"
+        '    names = ["x.adapter.w"]\n'
+        '    return any(".adapter." in n for n in names)\n'
+    )
+
+    def test_fixed_source_is_recognized(self, tmp_path):
+        module = _load_module_file(tmp_path, "probe_fixed_bridge_recompute", self.FIXED_BODY)
+        assert _recompute_source_recognizes_adapters(module) is True
+
+    def test_unfixed_source_is_not_recognized(self, tmp_path):
+        module = _load_module_file(tmp_path, "probe_unfixed_bridge_recompute", self.UNFIXED_BODY)
+        assert _recompute_source_recognizes_adapters(module) is False
+
+    def test_module_without_the_patch_function_fails_closed(self, tmp_path):
+        module = _load_module_file(tmp_path, "probe_empty_bridge_recompute", "X = 1\n")
+        assert _recompute_source_recognizes_adapters(module) is False
+
+    def test_unimportable_bridge_fails_closed(self, monkeypatch):
+        monkeypatch.setitem(sys.modules, "megatron.bridge.peft", None)
+        monkeypatch.delitem(sys.modules, "megatron.bridge.peft.recompute", raising=False)
+        assert _bridge_recompute_patch_recognizes_multi_lora() is False