diff --git a/relax/utils/megatron_peft_utils.py b/relax/utils/megatron_peft_utils.py index 7fd78a9ab..f4c7bc25d 100644 --- a/relax/utils/megatron_peft_utils.py +++ b/relax/utils/megatron_peft_utils.py @@ -70,6 +70,19 @@ def count_adapter_parameters(model) -> Tuple[int, int, float]: } +def _leaf_module_name(pattern: str) -> str: + """Reduce a Bridge target-module pattern to the module name it selects. + + Bridge's matcher accepts a path pattern so LoRA can be anchored to part of the + tree — ``scripts/training/sft/run-qwen3.5-35B-A3B-pokemon-lora-mtp-8xgpu.sh`` + uses ``*decoder.layers.*.linear_qkv`` to keep the MTP layers frozen. HF target + module names are bare and position-independent, so only the trailing segment + survives the translation; the anchoring is carried by which weights the exported + adapter actually contains. + """ + return pattern.rsplit(".", 1)[-1] + + def convert_megatron_to_hf_target_modules(megatron_modules: list[str]) -> list[str]: """Expand Megatron-style LoRA target module names to HF-style names. @@ -80,16 +93,32 @@ def convert_megatron_to_hf_target_modules(megatron_modules: list[str]) -> list[s so translation happens at export time via this one-to-many expansion. Args: - megatron_modules: List of Megatron-style module names - (e.g. ``["linear_qkv", "linear_proj"]``). + megatron_modules: List of Megatron-style module names, optionally written as + Bridge path patterns (e.g. ``["linear_qkv", "*decoder.layers.*.linear_proj"]``). Returns: - List of HF-style module names with duplicates removed. Unknown names are + List of HF-style module names with duplicates removed. A pattern contributes + its trailing segment, since HF names carry no position. Unknown names are passed through unchanged (already HF-style or custom). + + Raises: + ValueError: If a pattern's trailing segment is itself a wildcard, which has no + HF equivalent. """ hf_target_modules = [] for module in megatron_modules: - hf_target_modules.extend(MEGATRON_TO_HF_MODULES.get(module, [module])) + name = _leaf_module_name(module) + if "*" in name: + # A wildcard inside the trailing segment (e.g. "linear_*") has no HF + # equivalent. Passing it through would put a glob in adapter_config.json, + # where neither PEFT nor SGLang matches anything, so fail loudly instead of + # emitting a config that silently selects nothing. + raise ValueError( + f"LoRA target module {module!r} cannot be translated to an HF module name: " + f"the trailing segment {name!r} is itself a wildcard. Anchor the pattern on a " + f"concrete module name instead, e.g. '*decoder.layers.*.linear_qkv'." + ) + hf_target_modules.extend(MEGATRON_TO_HF_MODULES.get(name, [name])) # Remove duplicates while preserving order return list(dict.fromkeys(hf_target_modules)) diff --git a/tests/utils/test_megatron_peft_utils.py b/tests/utils/test_megatron_peft_utils.py index f16219e21..f0aed6a34 100644 --- a/tests/utils/test_megatron_peft_utils.py +++ b/tests/utils/test_megatron_peft_utils.py @@ -51,6 +51,34 @@ def test_convert_megatron_to_hf_passes_unknown_through(self): """Unknown / already-HF names are passed through unchanged.""" assert convert_megatron_to_hf_target_modules(["q_proj", "custom_mod"]) == ["q_proj", "custom_mod"] + def test_convert_megatron_to_hf_expands_path_patterns(self): + """A Bridge path pattern contributes its trailing module name. + + ``scripts/training/sft/run-qwen3.5-35B-A3B-pokemon-lora-mtp-8xgpu.sh`` scopes + LoRA this way to keep the MTP layers frozen; without the reduction the glob + reached adapter_config.json verbatim. + """ + result = convert_megatron_to_hf_target_modules( + ["*decoder.layers.*.linear_qkv", "*decoder.layers.*.linear_proj"] + ) + assert result == ["q_proj", "k_proj", "v_proj", "o_proj"] + + def test_convert_megatron_to_hf_dedups_pattern_and_bare_name(self): + """A pattern and the bare name it ends with collapse to one entry.""" + result = convert_megatron_to_hf_target_modules(["*decoder.layers.*.linear_proj", "linear_proj"]) + assert result == ["o_proj"] + + def test_convert_megatron_to_hf_accepts_exact_paths(self): + """A fully qualified module path works like a pattern.""" + result = convert_megatron_to_hf_target_modules(["decoder.layers.0.self_attention.linear_qkv"]) + assert result == ["q_proj", "k_proj", "v_proj"] + + def test_convert_megatron_to_hf_rejects_trailing_wildcard(self): + """A glob in the trailing segment has no HF equivalent and must not be + written into adapter_config.json.""" + with pytest.raises(ValueError, match="trailing segment"): + convert_megatron_to_hf_target_modules(["*decoder.layers.*.linear_*"]) + def test_convert_megatron_to_hf_empty(self): assert convert_megatron_to_hf_target_modules([]) == []