Skip to content

fix(lora): preserve alpha warmup across EMA lifecycle - #177

Merged
isLinXu merged 1 commit into
Tencent:mainfrom
SidKC:fix/lora-alpha-warmup-ema-lifecycle
Jul 27, 2026
Merged

fix(lora): preserve alpha warmup across EMA lifecycle#177
isLinXu merged 1 commit into
Tencent:mainfrom
SidKC:fix/lora-alpha-warmup-ema-lifecycle

Conversation

@SidKC

@SidKC SidKC commented Jul 27, 2026

Copy link
Copy Markdown
Contributor

背景

这个改动接着 #170 往下做。#170 处理了 fallback RS-LoRA 的 scaling 和保存加载;alpha warmup 还涉及 EMA、验证和断点恢复,所以当时单独拆了出来。

目前 fallback LoRA wrapper 的 lora_A 是 bare nn.Parameter,现有 warmup 发现逻辑不会识别它。配置 lora_alpha_warmup 后,fallback adapter 实际上仍会直接使用完整 scaling。

同时,fallback wrapper 的 scaling 是普通 Python 属性,不在 state_dict 中。online model 更新 scaling 后,EMA model 不会自动同步,验证、保存 checkpoint 和恢复训练时可能使用不同的 scaling。

修改内容

  • 让 alpha warmup 能识别 bare-Parameter fallback wrapper;
  • 每个 epoch 更新 warmup 后,同步 online model 和 EMA model 的 scaling;
  • 在验证、保存 checkpoint 和恢复训练时再次同步;
  • 同步时检查 wrapper 类型、use_rslora、rank 和 alpha 是否一致;
  • 在运行参数中记录 requested/effective RS-LoRA mode。

这里只同步 fallback wrapper 的 scaling,EMA 原有的参数和 buffer 更新方式保持不变。

测试

新增 9 个测试,覆盖:

  • ManualLoRAConvFewShotLoRAConv
  • 普通 LoRA 与 RS-LoRA;
  • 各个 epoch 的 warmup scaling;
  • 验证和 checkpoint 中的 EMA scaling;
  • 从 warmup 中途或结束后恢复训练;
  • 非 fallback backend。

同一组测试在未修改的 main 上为 8 failed, 1 passed,应用修复后为 9 passed

相关 LoRA、adapter、EMA、recovery 和 planner 测试为 159 passed, 1 deselected。Ruff、py_compilegit diff --check 也已通过。

说明

这次没有改 fallback adapter 的 merge 逻辑,也没有调整 EMA 对参数和 buffer 的正常更新方式。

@SidKC
SidKC marked this pull request as ready for review July 27, 2026 04:17
Copilot AI review requested due to automatic review settings July 27, 2026 04:17

Copilot AI left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.

@isLinXu

isLinXu commented Jul 27, 2026

Copy link
Copy Markdown
Collaborator

LGTM

@isLinXu
isLinXu merged commit c2523ca into Tencent:main Jul 27, 2026
12 of 18 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants