226 lines
6.6 KiB
Markdown
226 lines
6.6 KiB
Markdown
# Lesson 18 笔记:LoRA 原理与实现
|
||
|
||
## 目录
|
||
|
||
- [setattr(module, "lora", lora) 的含义](#setattrmodule-lora-lora-的含义)
|
||
- [LoRA 生命周期工具链](#lora-生命周期工具链)
|
||
- [总体流程图](#总体流程图)
|
||
- [apply_lora(初始化时)](#apply_lora初始化时)
|
||
- [save_lora(训练中/后)](#save_lora训练中后)
|
||
- [load_lora(恢复/推理前)](#load_lora恢复推理前)
|
||
- [merge_lora(部署前)](#merge_lora部署前)
|
||
- [四个函数角色分工](#四个函数角色分工)
|
||
- [三种典型使用模式](#三种典型使用模式)
|
||
- [一句话理解整个系统](#一句话理解整个系统)
|
||
|
||
---
|
||
|
||
## setattr(module, "lora", lora) 的含义
|
||
|
||
给 `module` 动态新增属性 `lora`,等价于 `module.lora = lora`。
|
||
|
||
**为什么用 setattr 而不是直接赋值?**
|
||
- 属性名可以是变量,更灵活
|
||
- 在循环里更通用
|
||
|
||
**LoRA 代码里的作用:**
|
||
|
||
```python
|
||
setattr(module, "lora", lora)
|
||
```
|
||
|
||
在每个被改造的 Linear 层上额外挂一个 LoRA 子模块。
|
||
|
||
**结构变成:**
|
||
```
|
||
Linear layer:
|
||
├── weight (原始参数,冻结)
|
||
├── bias
|
||
└── lora ← 新增的模块
|
||
```
|
||
|
||
**挂上去有什么用?**
|
||
|
||
1. **方便保存 checkpoint**:`model.state_dict()` 会自动包含 `layer.lora.A.weight`、`layer.lora.B.weight`
|
||
2. **方便 freeze/unfreeze**:可以统一对 `module.lora.requires_grad_(True)`
|
||
3. **方便 debug/inspect**:可以直接 `print(module.lora)`
|
||
4. **方便 LoRA merge**:推理时 `W = W + BA`,需要访问 lora 参数
|
||
|
||
**为什么不 setattr 会怎样?**
|
||
|
||
仍然可以 `module.forward = ...`,但:
|
||
- state_dict 不一定能找到
|
||
- 不能统一 freeze
|
||
- 不能 merge
|
||
- 不好保存结构
|
||
|
||
**一句话总结:**
|
||
|
||
`setattr(module, "lora", lora)` 本质是在原始 PyTorch layer 上动态挂一个 LoRA 子模块,让它成为模型结构的一部分,便于训练、保存和管理。
|
||
|
||
**为什么 LoRA 要"挂在 layer 上"?**
|
||
|
||
因为 LoRA 本质是每个 Linear layer 的"附加模块",不是独立模型:
|
||
|
||
```
|
||
Transformer Block
|
||
├── Linear
|
||
│ ├── W (冻结)
|
||
│ └── LoRA (训练)
|
||
```
|
||
|
||
---
|
||
|
||
## LoRA 生命周期工具链
|
||
|
||
**理解为一整套"LoRA 插件系统":安装 → 训练 → 保存 → 加载 → 合并部署**
|
||
|
||
### 总体流程图
|
||
|
||
```
|
||
┌──────────────┐
|
||
│ apply_lora │ ← 初始化模型(插插件)
|
||
└──────┬───────┘
|
||
↓
|
||
训练(只更新 LoRA)
|
||
↓
|
||
┌──────────────┐
|
||
│ save_lora │ ← 保存 LoRA 参数
|
||
└──────┬───────┘
|
||
↓
|
||
训练结束 / 断点
|
||
↓
|
||
┌──────────────┐
|
||
│ load_lora │ ← 加载 LoRA 继续训练/推理
|
||
└──────┬───────┘
|
||
↓
|
||
┌──────────────┐
|
||
│ merge_lora │ ← 合并进 W(部署用)
|
||
└──────────────┘
|
||
```
|
||
|
||
### apply_lora(初始化时)
|
||
|
||
```python
|
||
def apply_lora(model, rank=16):
|
||
for name, module in model.named_modules():
|
||
if isinstance(module, nn.Linear) and module.in_features == module.out_features:
|
||
lora = LoRA(module.in_features, module.out_features, rank=rank).to(model.device)
|
||
setattr(module, "lora", lora)
|
||
original_forward = module.forward
|
||
|
||
def forward_with_lora(x, layer1=original_forward, layer2=lora):
|
||
return layer1(x) + layer2(x)
|
||
|
||
module.forward = forward_with_lora
|
||
```
|
||
|
||
**作用:** 给模型每个 Linear 层外挂 LoRA 分支,改写 forward。
|
||
|
||
**用在:** 训练开始前
|
||
|
||
---
|
||
|
||
### save_lora(训练中/后)
|
||
|
||
```python
|
||
def save_lora(model, path):
|
||
raw_model = getattr(model, '_orig_mod', model)
|
||
state_dict = {}
|
||
for name, module in raw_model.named_modules():
|
||
if hasattr(module, 'lora'):
|
||
clean_name = name[7:] if name.startswith("module.") else name
|
||
lora_state = {f'{clean_name}.lora.{k}': v.cpu().half() for k, v in module.lora.state_dict().items()}
|
||
state_dict.update(lora_state)
|
||
torch.save(state_dict, path)
|
||
```
|
||
|
||
**作用:** 只保存 LoRA 参数(A/B矩阵),不保存原模型
|
||
|
||
**用在:** checkpoint 或训练结束保存 adapter
|
||
|
||
---
|
||
|
||
### load_lora(恢复/推理前)
|
||
|
||
```python
|
||
def load_lora(model, path):
|
||
state_dict = torch.load(path, map_location=model.device)
|
||
state_dict = {(k[7:] if k.startswith('module.') else k): v for k, v in state_dict.items()}
|
||
|
||
for name, module in model.named_modules():
|
||
if hasattr(module, 'lora'):
|
||
lora_state = {k.replace(f'{name}.lora.', ''): v for k, v in state_dict.items() if f'{name}.lora.' in k}
|
||
module.lora.load_state_dict(lora_state)
|
||
```
|
||
|
||
**作用:** 把 LoRA 权重重新塞回 model.lora
|
||
|
||
**用在:** 继续训练、推理前加载 adapter
|
||
|
||
**前提:** 必须先 `apply_lora(model)`
|
||
|
||
---
|
||
|
||
### merge_lora(部署前)
|
||
|
||
```python
|
||
def merge_lora(model, lora_path, save_path):
|
||
load_lora(model, lora_path)
|
||
raw_model = getattr(model, '_orig_mod', model)
|
||
state_dict = {k: v.cpu().half() for k, v in raw_model.state_dict().items() if '.lora.' not in k}
|
||
for name, module in raw_model.named_modules():
|
||
if isinstance(module, nn.Linear) and '.lora.' not in name:
|
||
state_dict[f'{name}.weight'] = module.weight.data.clone().cpu().half()
|
||
if hasattr(module, 'lora'):
|
||
state_dict[f'{name}.weight'] += (module.lora.B.weight.data @ module.lora.A.weight.data).cpu().half()
|
||
torch.save(state_dict, save_path)
|
||
```
|
||
|
||
**作用:** 把 LoRA 融合进权重 W,生成纯模型
|
||
|
||
**核心步骤:**
|
||
1. `load_lora()` 加载 LoRA
|
||
2. `W' = W + B @ A` 合并权重
|
||
3. 保存纯模型
|
||
|
||
**用在:** 推理部署前
|
||
|
||
---
|
||
|
||
### 四个函数角色分工
|
||
|
||
| 函数 | 作用 | 什么时候用 |
|
||
|------|------|------------|
|
||
| apply_lora | 给模型插 LoRA | 训练前 |
|
||
| save_lora | 存 LoRA 参数 | 训练中/结束 |
|
||
| load_lora | 恢复 LoRA | 继续训练/推理 |
|
||
| merge_lora | 合并成完整模型 | 推理部署前 |
|
||
|
||
---
|
||
|
||
### 三种典型使用模式
|
||
|
||
**模式1:训练 LoRA**
|
||
```
|
||
apply_lora → train → save_lora
|
||
```
|
||
|
||
**模式2:继续训练**
|
||
```
|
||
apply_lora → load_lora → train
|
||
```
|
||
|
||
**模式3:推理部署**
|
||
```
|
||
apply_lora (可选) → load_lora → merge_lora → inference
|
||
```
|
||
|
||
---
|
||
|
||
### 一句话理解整个系统
|
||
|
||
- `apply_lora` 是"装插件"
|
||
- `save/load_lora` 是"存插件"
|
||
- `merge_lora` 是"拆插件并写死进模型"
|