|
22 | 22 | import torch |
23 | 23 | from megatron.core.optimizer import OptimizerConfig, ParamGroupOverride, ParamKey |
24 | 24 |
|
| 25 | +from megatron.bridge.peft.lora import get_lora_plus_config_overrides |
25 | 26 | from megatron.bridge.training.config import SchedulerConfig |
26 | 27 | from megatron.bridge.training.optim import ( |
27 | 28 | memory_efficient_fp32_optimizer_state_loading, |
@@ -530,3 +531,98 @@ def test_skips_none_shard_main_param(self): |
530 | 531 | synced = sync_hybrid_device_optimizer_fp32_master_copies(distrib_opt) |
531 | 532 |
|
532 | 533 | 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