fix(peft): reject nonzero warmup initialization for LoRA+

[cuichenx]

Reject nonzero SchedulerConfig.lr_warmup_init when the LoRA+ helper produces different A/B learning-rate bounds. Previously, a requested ratio of 16 became 1, 12.54, 14.64 and 15.52 during warmup. setup_optimizer now raises an actionable ValueError before constructing the optimizer, directing callers to zero warmup initialization or ratio 1.

The guard recognizes the helper's canonical parameter overrides after merging. Omitted bounds inherit the optimizer defaults before comparison. The change is limited to the guard and regression tests.

Validation:

• 84 focused optimizer tests pass, including the guard, supported configurations, unrelated overrides and alternate optimizer factories.
• A real two-layer GPT with 8 A/8 B tensors on one H100 rejects the unsupported configuration. Three supported controls match the baseline exactly over 12 training steps, with finite losses/gradients/weights, unchanged frozen parameters and Adam-reference update error below 3e-8.
• Four partial-override regression cases (omitted maximum or minimum on either group) fail on the previous guard and pass with effective-bound comparison.
• For the three original full-bound controls, adapter, optimizer and scheduler state save/load during warmup matches uninterrupted execution. This covers direct state restoration; full checkpoint-manager integration and explicit scheduler overrides are outside this change.
• All-file pre-commit and git diff --check pass. Targeted type checking introduces no new diagnostics; three existing dependency-typing diagnostics remain.

Additional testing of sparse versus explicit overrides with equal bounds exposed a pre-existing MCore optimizer-state restore collision. The same resumed-loss mismatch reproduces on unchanged source. The guard does not alter that restore behavior; this is not a claim that sparse-override checkpoint resume is fixed.

NVBug 6888439

[github-actions]

Automatic Claude reviews have been retired. To request a pull-request review, post a comment containing:

/review

Add model=claude to use a Claude reviewer (the default is model=codex). mode=light|strict selects the review depth; for example, /review model=claude mode=strict. Comment /review help for all options.

[cuichenx]

/review

[svcnvidia-nemo-ci]

One blocking validation defect: the guard can reject sparse parameter overrides whose effective A/B learning-rate bounds are equal.

Reviewed both changed files and traced override construction, merge order, and scheduler defaults. This was a static review; tests were not executed.

Suggested test cases

• Proposed TestLoraPlusWarmupValidation::test_equal_effective_bounds_with_partial_overrides — omitted bounds inherit the optimizer defaults and must not cause rejection.
• TestLoraPlusWarmupValidation::test_rejects_nonzero_initial_rate_before_optimizer
• TestLoraPlusWarmupValidation::test_rejects_before_alternate_optimizer_construction
• TestLoraPlusWarmupValidation::test_supported_combinations_reach_optimizer
• TestLoraPlusWarmupValidation::test_unrelated_overrides_are_not_rejected
• TestLoraPlusWarmupValidation::test_distinct_minimum_rates_are_rejected
• No perf tests impacted.

Note: formal approval is not available for this request, so this review is posted as a comment only.

Review generated by codex

[svcnvidia-nemo-ci]

Blocking — compare effective bounds rather than optional dictionary entries.

Failure: With optimizer_config.lr=3e-4, min_lr=3e-5, and nonzero warmup initialization, a provider returning the following overrides now raises the LoRA+ ValueError:

{
    ParamKey(name="*.linear_in.weight"):
        ParamGroupOverride(max_lr=3e-4, min_lr=3e-5),
    ParamKey(name="*.linear_out.weight"):
        ParamGroupOverride(max_lr=3e-4),
}

Both groups have the same effective bounds: the omitted B minimum inherits the base minimum. The comparison instead sees 3e-5 != None. The provider contract permits per-group overrides, and _get_scheduler supplies the base maximum/minimum. Equal effective rates should reach optimizer construction, not be classified as an unsupported LoRA+ split. Omitting one maximum while explicitly setting the other to the base maximum has the same problem.

Correction: In setup_optimizer, resolve absent bounds against the scheduler defaults before comparing. Replace the comparison expression at this line with:

and (
    lora_a.get("max_lr", optimizer_config.lr)
    != lora_b.get("max_lr", optimizer_config.lr)
    or lora_a.get("min_lr", optimizer_config.min_lr)
    != lora_b.get("min_lr", optimizer_config.min_lr)
)

The invariant is that omission of a bound must not itself create an A/B difference. This suggested replacement is based on source reasoning, not execution; verify it against the pinned MCore implementation before adopting it.

Preservation: Keep rejection before every optimizer factory for complete helper overrides with ratios 0.5 or 16 and nonzero initialization, and for explicitly different minimum rates. Keep ratio 1, zero initialization, missing canonical keys, and attribute-qualified keys accepted. The helper explicitly emits both fields, including minimum 0.0 when the base minimum is None, so these existing helper cases retain their current comparisons.

Acceptance: Add a proposed regression using the existing mocked setup: pass the mapping above with initial_lr=3e-5; expect no exception and one optimizer/scheduler construction. Also test an omitted maximum against an explicit base maximum. Retain the existing rejection tests' assertions that neither factory nor scheduler runs, the supported-combinations construction assertions, and the distinct-minimum rejection case.

[cuichenx]

/ok to test a13a9f1