Skip to content

Commit 4edf36d

Browse files
committed
Support LoRA for targeted parameter-backed projections
Allow LoRA to promote explicitly targeted bare 2-D parameters using the existing target_modules matcher. This supports replicated DSV4 grouped output projections without adding a model-specific parameter name check or a new configuration field. The original parameter object is preserved as the wrapped linear weight, while unrelated parameters and non-LoRA checkpoint layouts remain unchanged. Also add mapping fallbacks for promoted `.to_wrap.weight` names and regression coverage for parameter-backed LoRA targets. Signed-off-by: Hollow Man <hollowman@opensuse.org>
1 parent e8202eb commit 4edf36d

5 files changed

Lines changed: 114 additions & 1 deletion

File tree

‎src/megatron/bridge/models/conversion/model_bridge.py‎

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2376,7 +2376,14 @@ def _validate_conversion_mappings(
23762376
hf_key_set = set(hf_keys) if hf_keys is not None else None
23772377

23782378
for global_name in global_param_names:
2379-
mapping = mapping_registry.megatron_to_hf_lookup(self._get_lora_unwrapped_name(global_name))
2379+
unwrapped_name = self._get_lora_unwrapped_name(global_name)
2380+
mapping = mapping_registry.megatron_to_hf_lookup(unwrapped_name)
2381+
if mapping is None and global_name.endswith(".to_wrap.weight"):
2382+
# A parameter-backed LoRA target (for example DeepSeek V4's
2383+
# replicated ``linear_o_group_proj``) is registered without a
2384+
# ``.weight`` leaf. The wrapper exposes it as
2385+
# ``.to_wrap.weight`` after PEFT promotion.
2386+
mapping = mapping_registry.megatron_to_hf_lookup(unwrapped_name[: -len(".weight")])
23802387
if mapping is None:
23812388
missing_mappings.append(global_name)
23822389
continue

‎src/megatron/bridge/models/conversion/peft_bridge.py‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -196,6 +196,10 @@ def _resolve_hf_adapter_param_name(
196196
)
197197

198198
base_mapping = mapping_registry.megatron_to_hf_lookup(f"{global_base_prefix}{base_suffix}")
199+
if base_mapping is None and base_suffix == ".weight":
200+
# DeepSeek V4's grouped output projection is a bare replicated
201+
# parameter, so its base mapping is registered without `.weight`.
202+
base_mapping = mapping_registry.megatron_to_hf_lookup(global_base_prefix)
199203
assert base_mapping is not None, (
200204
f"Expected mapping for adapter base '{global_base_prefix}{base_suffix}' but none found"
201205
)
@@ -224,6 +228,11 @@ def _get_base_hf_param_names_for_adapter(
224228
"""Return all HF base parameter names associated with this adapter."""
225229

226230
base_mapping = mapping_registry.megatron_to_hf_lookup(f"{global_base_prefix}{base_suffix}")
231+
if base_mapping is None and base_suffix == ".weight":
232+
# See `_resolve_hf_adapter_param_name`: the grouped output
233+
# projection's base is a plain parameter without a `.weight`
234+
# suffix in the mapping registry.
235+
base_mapping = mapping_registry.megatron_to_hf_lookup(global_base_prefix)
227236
if base_mapping is None:
228237
return []
229238

‎src/megatron/bridge/peft/lora.py‎

Lines changed: 58 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -43,6 +43,7 @@
4343
is_expert_linear,
4444
is_grouped_expert_linear,
4545
is_modelopt_linear,
46+
wildcard_match,
4647
)
4748

4849

@@ -159,6 +160,63 @@ class LoRA(PEFT, ModuleMatcher):
159160
experts_shared_outer_loras: bool = False
160161
use_transformer_engine_op_fuser: bool = False
161162

163+
def __call__(self, model, training: bool = True):
164+
"""Apply LoRA, promoting targeted parameter-backed projections first.
165+
166+
Parameter-backed projections use the same explicit ``target_modules``
167+
patterns as ordinary modules. Turning such a parameter into the regular
168+
LoRA module only when it is targeted lets the existing LoRA wrapper and
169+
adapter export path handle it without changing the base model's
170+
checkpoint layout for non-LoRA runs.
171+
"""
172+
173+
self._promote_parameter_targets(model)
174+
return super().__call__(model, training=training)
175+
176+
def _promote_parameter_targets(self, model) -> None:
177+
"""Expose supported bare parameters as linear modules before walking.
178+
179+
The promoted module keeps the original ``nn.Parameter`` object as its
180+
weight, so initialization, freezing, and optimizer ownership remain
181+
unchanged. Parameters that do not match ``target_modules`` retain the
182+
existing PEFT behavior.
183+
"""
184+
185+
target_patterns = self.target_modules
186+
if isinstance(target_patterns, str):
187+
target_patterns = (target_patterns,)
188+
else:
189+
target_patterns = tuple(target_patterns or ())
190+
if not target_patterns:
191+
return
192+
193+
models = model if isinstance(model, list) else [model]
194+
for root in models:
195+
if isinstance(root, torch.nn.parallel.DistributedDataParallel):
196+
root = root.module
197+
for prefix, parent in list(root.named_modules()):
198+
for parameter_name, parameter in list(parent.named_parameters(recurse=False)):
199+
full_name = f"{prefix}.{parameter_name}" if prefix else parameter_name
200+
if not any(
201+
parameter_name == pattern or wildcard_match(pattern, full_name) for pattern in target_patterns
202+
):
203+
continue
204+
if parameter.ndim != 2:
205+
raise ValueError(
206+
f"{full_name} must be a 2-D matrix for LoRA, got shape {tuple(parameter.shape)}"
207+
)
208+
promoted = nn.Linear(
209+
parameter.shape[1],
210+
parameter.shape[0],
211+
bias=False,
212+
device=parameter.device,
213+
dtype=parameter.dtype,
214+
)
215+
promoted.weight = parameter
216+
delattr(parent, parameter_name)
217+
setattr(parent, parameter_name, promoted)
218+
logger.info("Promoted parameter-backed LoRA target: %s", full_name)
219+
162220
def transform(self, module: nn.Module, name: Optional[str] = None, prefix: Optional[str] = None) -> nn.Module:
163221
"""
164222
Applies LoRA to a specific module within the model architecture.

‎tests/unit_tests/models/test_model_bridge_lora.py‎

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,7 @@
3131
from megatron.bridge.models.conversion.param_mapping import (
3232
AutoMapping,
3333
ColumnParallelMapping,
34+
ReplicatedMapping,
3435
RowParallelMapping,
3536
_fuse_gdn_separate_to_grouped,
3637
merge_gdn_linear_weights,
@@ -754,6 +755,20 @@ def named_parameters(self):
754755
assert names == ["decoder.layers.0.mlp.linear_fc1.to_wrap.weight"]
755756

756757

758+
def test_conversion_mapping_falls_back_for_parameter_backed_lora_target():
759+
bridge = DummyBridge()
760+
base_name = "decoder.layers.0.self_attention.linear_o_group_proj"
761+
mapping = ReplicatedMapping(base_name, "model.layers.0.attn.wo_a.weight")
762+
763+
class Registry:
764+
def megatron_to_hf_lookup(self, key):
765+
return mapping if key == base_name else None
766+
767+
names = [f"{base_name}.to_wrap.weight"]
768+
resolved = bridge._validate_conversion_mappings(Registry(), names)
769+
assert resolved[names[0]] is mapping
770+
771+
757772
def test_megatron_global_adapters_info_all_pp_ranks(monkeypatch):
758773
bridge = DummyBridge()
759774

‎tests/unit_tests/peft/test_lora.py‎

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -63,6 +63,15 @@ def __init__(self):
6363
self.layernorm = nn.LayerNorm(512)
6464

6565

66+
class ParameterBackedProjectionModel(nn.Module):
67+
"""Model fixture with an explicitly targetable bare projection weight."""
68+
69+
def __init__(self):
70+
super().__init__()
71+
self.projection_weight = nn.Parameter(torch.randn(8, 4))
72+
self.unmarked_matrix = nn.Parameter(torch.randn(8, 4))
73+
74+
6675
class NestedModel(nn.Module):
6776
"""Model with nested structure for testing pattern matching."""
6877

@@ -303,6 +312,21 @@ def test_lora_transform_simple_model(self):
303312
assert isinstance(transformed_model.embedding, nn.Embedding)
304313
assert isinstance(transformed_model.layernorm, nn.LayerNorm)
305314

315+
def test_lora_promotes_targeted_bare_parameter(self):
316+
"""An explicitly targeted bare projection receives the regular LoRA wrapper."""
317+
model = ParameterBackedProjectionModel()
318+
base_weight = model.projection_weight
319+
lora = LoRA(target_modules=["projection_weight"], dim=2, alpha=4)
320+
321+
transformed_model = lora(model, training=True)
322+
323+
assert isinstance(transformed_model.projection_weight, LoRALinear)
324+
assert transformed_model.projection_weight.to_wrap.weight is base_weight
325+
assert not base_weight.requires_grad
326+
assert transformed_model.projection_weight.adapter.linear_in.weight.requires_grad
327+
assert transformed_model.projection_weight.adapter.linear_out.weight.requires_grad
328+
assert isinstance(transformed_model.unmarked_matrix, nn.Parameter)
329+
306330
def test_lora_transform_with_exclude_modules(self):
307331
"""Test LoRA transformation with exclude_modules parameter."""
308332
model = SimpleModel()

0 commit comments

Comments
 (0)