Skip to content

Commit 6ff7723

Browse files
committed
Merge branch 'main' into yuya/qwen3-vl-packed-attention-tail-padding
Signed-off-by: yaoyu-33 <yaoyu.094@gmail.com>
2 parents b15898e + 96db955 commit 6ff7723

4 files changed

Lines changed: 169 additions & 2 deletions

File tree

‎docs/training/peft.md‎

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -127,6 +127,26 @@ lora_config = LoRA(
127127
)
128128
```
129129

130+
#### LoRA+
131+
132+
[LoRA+](https://arxiv.org/abs/2402.12354) trains the B matrix (`linear_out`) at a higher learning rate than A (`linear_in`) by a fixed ratio `lr_B = lora_plus_ratio * lr_A` (paper default 16). The ratio is set on the optimizer, not the `LoRA` adapter: when it differs from `1.0`, `get_lora_plus_config_overrides` (`megatron.bridge.peft.lora`) builds Megatron-Core per-group overrides giving the `*.linear_out.weight` group `max_lr`/`min_lr` scaled by the ratio. Pass `1.0` (or skip the call) to keep a single learning rate.
133+
134+
```python
135+
from megatron.bridge.peft.lora import get_lora_plus_config_overrides
136+
from megatron.core.optimizer import OptimizerConfig, get_megatron_optimizer
137+
138+
opt_config = OptimizerConfig(optimizer="adam", lr=1e-4, bf16=True)
139+
140+
# LoRA+ LR split: lr_B = lora_plus_ratio * lr_A; pass 1.0 (or skip) to disable.
141+
lora_plus_ratio = 16
142+
config_overrides = get_lora_plus_config_overrides(opt_config, lora_plus_ratio)
143+
optimizer = get_megatron_optimizer(
144+
config=opt_config,
145+
model_chunks=model_chunks,
146+
config_overrides=config_overrides,
147+
)
148+
```
149+
130150
### Canonical LoRA: Performant vs Canonical Variants
131151

132152
There are two variants of LoRA implemented in Megatron Bridge: "performant LoRA" (`LoRA`) and "canonical LoRA" (`CanonicalLoRA`).

‎scripts/conversion/gpu_backend.py‎

Lines changed: 10 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -79,7 +79,16 @@ def _prepare_distributed_output(path: str, *, overwrite: bool, source_paths: Ite
7979
validate_output_path(path, source_paths=source_paths)
8080
if torch.distributed.get_rank() == 0:
8181
prepare_output_directory(path, overwrite=overwrite, source_paths=source_paths)
82-
torch.distributed.barrier()
82+
# Without device_ids, torch guesses the barrier's device as
83+
# rank % local_gpu_count. Slurm assigns SLURM_PROCID cyclically across
84+
# nodes (RANK = num_nodes * local_rank + node_id, not block-wise), so
85+
# that guess collides across ranks on the same node once more than one
86+
# node is used, and NCCL aborts with "Multiple ranks detected using the
87+
# same GPU on this node." Passing the real device (already set via
88+
# torch.cuda.set_device in _ensure_distributed_initialized) avoids the
89+
# guess, matching the pattern used by training/initialize.py's own
90+
# first post-init barrier.
91+
torch.distributed.barrier(device_ids=[torch.cuda.current_device()])
8392

8493

8594
def _maybe_generate_pipeline_layout(bridge: AutoBridge, model_provider: GPTModelProvider, pp: int) -> bool:

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

Lines changed: 43 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -14,11 +14,13 @@
1414

1515
import logging
1616
from dataclasses import dataclass, field
17-
from typing import List, Literal, Optional
17+
from typing import Dict, List, Literal, Optional
1818

1919
import torch
2020
import torch.nn as nn
2121
from megatron.core import parallel_state
22+
from megatron.core.optimizer import OptimizerConfig, ParamKey, get_standard_config_overrides
23+
from megatron.core.optimizer_param_scheduler import ParamGroupOverride
2224
from megatron.core.transformer.moe.router import TopKRouter
2325
from megatron.core.utils import unwrap_model
2426

@@ -46,6 +48,46 @@
4648
logger = logging.getLogger(__name__)
4749

4850

51+
def get_lora_plus_config_overrides(
52+
config: OptimizerConfig,
53+
lora_plus_ratio: float,
54+
) -> Dict[ParamKey, ParamGroupOverride]:
55+
"""Build Megatron-Core optimizer ``config_overrides`` for LoRA+.
56+
57+
LoRA+ (https://arxiv.org/abs/2402.12354) trains ``B`` (``linear_out``) at
58+
``lora_plus_ratio`` times the LR of ``A`` (``linear_in``); the ratio is fixed
59+
and only the base LR is tuned. Standard overrides (bias/1-D ``wd_mult=0.0``,
60+
decoupled LR) are preserved as the starting point.
61+
62+
Args:
63+
config: base optimizer config; ``lr`` / ``min_lr`` are A's learning rates.
64+
lora_plus_ratio: ``lr_B / lr_A`` (e.g. 16.0). Must be > 0.
65+
66+
Returns:
67+
Per-group overrides ready for ``get_megatron_optimizer(config_overrides=...)``.
68+
"""
69+
if lora_plus_ratio <= 0:
70+
raise ValueError(f"lora_plus_ratio must be > 0, got {lora_plus_ratio}")
71+
72+
config_overrides = get_standard_config_overrides(config)
73+
74+
base_lr = config.lr
75+
base_min_lr = config.min_lr if config.min_lr is not None else 0.0
76+
77+
config_overrides[ParamKey(name="*.linear_in.weight")] = ParamGroupOverride(max_lr=base_lr, min_lr=base_min_lr)
78+
config_overrides[ParamKey(name="*.linear_out.weight")] = ParamGroupOverride(
79+
max_lr=base_lr * lora_plus_ratio, min_lr=base_min_lr * lora_plus_ratio
80+
)
81+
82+
logger.info(
83+
"LoRA+ enabled: lora_plus_ratio=%s (linear_in/A max_lr=%s, linear_out/B max_lr=%s)",
84+
lora_plus_ratio,
85+
base_lr,
86+
base_lr * lora_plus_ratio,
87+
)
88+
return config_overrides
89+
90+
4991
@dataclass
5092
class LoRA(PEFT, ModuleMatcher):
5193
"""

‎tests/unit_tests/training/test_optim.py‎

Lines changed: 96 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@
2222
import torch
2323
from megatron.core.optimizer import OptimizerConfig, ParamGroupOverride, ParamKey
2424

25+
from megatron.bridge.peft.lora import get_lora_plus_config_overrides
2526
from megatron.bridge.training.config import SchedulerConfig
2627
from megatron.bridge.training.optim import (
2728
memory_efficient_fp32_optimizer_state_loading,
@@ -530,3 +531,98 @@ def test_skips_none_shard_main_param(self):
530531
synced = sync_hybrid_device_optimizer_fp32_master_copies(distrib_opt)
531532

532533
assert synced is True
534+
535+
536+
class TestGetLoraPlusConfigOverrides:
537+
"""Tests for ``get_lora_plus_config_overrides`` (LoRA+ A/B LR split)."""
538+
539+
def _make_config(self, lr=3e-5, min_lr=0.0, **kwargs):
540+
return OptimizerConfig(optimizer="adam", lr=lr, min_lr=min_lr, bf16=True, **kwargs)
541+
542+
def _get_override(self, overrides, name):
543+
"""Find a ParamKey by its ``name`` glob and return its ParamGroupOverride."""
544+
for key, val in overrides.items():
545+
if (
546+
getattr(key, "name", None) == name
547+
and not getattr(key, "predicate", None)
548+
and not getattr(key, "with_name_predicate", None)
549+
and not getattr(key, "attr", None)
550+
):
551+
return val
552+
return None
553+
554+
@pytest.mark.parametrize("ratio", [0, -1, -0.5])
555+
def test_non_positive_ratio_raises(self, ratio):
556+
with pytest.raises(ValueError, match="lora_plus_ratio must be > 0"):
557+
get_lora_plus_config_overrides(self._make_config(), ratio)
558+
559+
def test_a_keeps_base_lr(self):
560+
overrides = get_lora_plus_config_overrides(self._make_config(lr=3e-5, min_lr=1e-6), 16.0)
561+
a = self._get_override(overrides, "*.linear_in.weight")
562+
assert a is not None
563+
assert a["max_lr"] == 3e-5
564+
assert a["min_lr"] == 1e-6
565+
566+
def test_b_is_ratio_times_base_lr(self):
567+
overrides = get_lora_plus_config_overrides(self._make_config(lr=3e-5, min_lr=1e-6), 16.0)
568+
b = self._get_override(overrides, "*.linear_out.weight")
569+
assert b is not None
570+
assert b["max_lr"] == pytest.approx(3e-5 * 16)
571+
assert b["min_lr"] == pytest.approx(1e-6 * 16)
572+
573+
def test_none_min_lr_defaults_to_zero(self):
574+
cfg = OptimizerConfig(optimizer="adam", lr=3e-5, min_lr=None, bf16=True)
575+
overrides = get_lora_plus_config_overrides(cfg, 16.0)
576+
a = self._get_override(overrides, "*.linear_in.weight")
577+
b = self._get_override(overrides, "*.linear_out.weight")
578+
assert a["min_lr"] == 0.0
579+
assert b["min_lr"] == 0.0
580+
assert b["max_lr"] == pytest.approx(3e-5 * 16)
581+
582+
def test_b_schedule_is_ratio_scaled_copy_of_a(self):
583+
"""B's max_lr/min_lr are both exactly ratio x A's — the whole schedule scales."""
584+
ratio = 8.0
585+
overrides = get_lora_plus_config_overrides(self._make_config(lr=5e-5, min_lr=2e-6), ratio)
586+
a = self._get_override(overrides, "*.linear_in.weight")
587+
b = self._get_override(overrides, "*.linear_out.weight")
588+
assert b["max_lr"] / a["max_lr"] == pytest.approx(ratio)
589+
assert b["min_lr"] / a["min_lr"] == pytest.approx(ratio)
590+
591+
def test_ratio_one_keeps_equal_lr(self):
592+
"""ratio == 1.0 still produces valid overrides with A == B (no-op split)."""
593+
overrides = get_lora_plus_config_overrides(self._make_config(lr=3e-5, min_lr=0.0), 1.0)
594+
a = self._get_override(overrides, "*.linear_in.weight")
595+
b = self._get_override(overrides, "*.linear_out.weight")
596+
assert a["max_lr"] == b["max_lr"] == 3e-5
597+
assert a["min_lr"] == b["min_lr"] == 0.0
598+
599+
def test_preserves_standard_bias_wd_override(self):
600+
"""The standard bias/1-D wd_mult=0.0 override is preserved alongside LoRA+."""
601+
overrides = get_lora_plus_config_overrides(self._make_config(), 16.0)
602+
wd_skips = [v for v in overrides.values() if v.get("wd_mult") == 0.0]
603+
assert wd_skips, "standard wd_mult=0.0 override should be preserved"
604+
605+
def test_preserves_decoupled_lr_override(self):
606+
"""A configured decoupled_lr is preserved by get_standard_config_overrides."""
607+
cfg = self._make_config(decoupled_lr=1e-4, decoupled_min_lr=1e-5)
608+
overrides = get_lora_plus_config_overrides(cfg, 16.0)
609+
# The decoupled entry is keyed by attr, not name.
610+
decoupled = [v for k, v in overrides.items() if getattr(k, "attr", None)]
611+
assert decoupled, "decoupled_lr override should be preserved"
612+
assert decoupled[0]["max_lr"] == 1e-4
613+
614+
def test_returns_param_key_keyed_mapping(self):
615+
"""Result is keyed by ParamKey; values are dict-like (ParamGroupOverride is a TypedDict)."""
616+
overrides = get_lora_plus_config_overrides(self._make_config(), 16.0)
617+
assert all(isinstance(k, ParamKey) for k in overrides)
618+
assert all(isinstance(v, dict) for v in overrides.values())
619+
# Only the LoRA A/B entries carry the LR fields; standard entries may not.
620+
assert self._get_override(overrides, "*.linear_in.weight")["max_lr"] is not None
621+
assert self._get_override(overrides, "*.linear_out.weight")["max_lr"] is not None
622+
623+
def test_only_two_lora_name_globs_added(self):
624+
"""Exactly two name-glob entries are added: linear_in and linear_out."""
625+
overrides = get_lora_plus_config_overrides(self._make_config(), 16.0)
626+
name_keys = [k for k in overrides if getattr(k, "name", None) and not getattr(k, "attr", None)]
627+
names = {k.name for k in name_keys}
628+
assert {"*.linear_in.weight", "*.linear_out.weight"}.issubset(names)

0 commit comments

Comments
 (0)