From c561fb0f36599422890a0bb3d859d8dfe7a6bbc7 Mon Sep 17 00:00:00 2001 From: Yusheng Su Date: Sun, 23 Aug 2026 16:17:56 -0700 Subject: [PATCH] fix: refuse unsafe multi-LoRA recompute configs at launch Full recompute, and selective MoE recompute with expert LoRA targets, silently corrupt gradients unless the Megatron-Bridge PEFT input-gradient patch recognizes .adapters. parameters. Probe the installed Bridge source at validation time and fail the launch for unsupported shapes. --- miles/utils/multi_lora.py | 29 ++++ .../utils/test_multi_lora_recompute_guard.py | 157 ++++++++++++++++++ 2 files changed, 186 insertions(+) create mode 100644 tests/fast/utils/test_multi_lora_recompute_guard.py diff --git a/miles/utils/multi_lora.py b/miles/utils/multi_lora.py index 8a0a2d42f57..7e227110402 100644 --- a/miles/utils/multi_lora.py +++ b/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 " diff --git a/tests/fast/utils/test_multi_lora_recompute_guard.py b/tests/fast/utils/test_multi_lora_recompute_guard.py new file mode 100644 index 00000000000..4c94502dc08 --- /dev/null +++ b/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