W' = W + ΔW = W + BA

# LoRA 的数学表达
# 前向：h = Wx + BAx = (W + BA)x
# 反向：只计算 B 和 A 的梯度，W 的梯度为零

# 初始化策略
# A: 正态分布初始化 N(0, σ²)
# B: 零初始化（确保训练开始时 ΔW = BA = 0）

import torch
import torch.nn as nn

class LoRALayer(nn.Module):
    """LoRA 层的独立实现"""
    def __init__(
        self,
        in_features: int,
        out_features: int,
        r: int = 8,
        alpha: int = 16,
        dropout: float = 0.0,
    ):
        super().__init__()
        self.r = r
        self.scale = alpha / r  # 缩放系数

        # 原始权重（冻结）
        self.base_weight = nn.Parameter(
            torch.randn(out_features, in_features), requires_grad=False
        )

        # LoRA 矩阵
        self.lora_A = nn.Parameter(torch.randn(r, in_features) * 0.01)
        self.lora_B = nn.Parameter(torch.zeros(out_features, r))
        self.dropout = nn.Dropout(dropout)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        base_out = x @ self.base_weight.T
        lora_out = (self.dropout(x) @ self.lora_A.T @ self.lora_B.T) * self.scale
        return base_out + lora_out


# 使用 PEFT 库（推荐）
from peft import LoraConfig, get_peft_model

config = LoraConfig(
    r=16,
    lora_alpha=32,
    lora_dropout=0.05,
    target_modules=["q_proj", "k_proj", "v_proj", "o_proj"],
    bias="none",
    task_type="CAUSAL_LM",
)
model = get_peft_model(base_model, config)

W' = m · (V + ΔV) / ||V + ΔV||

其中：
V = W / ||W||        （方向，冻结）
m = ||W||             （大小，可训练标量）
ΔV = BA / ||BA||     （方向更新，LoRA 低秩分解）

import torch
import torch.nn as nn
import torch.nn.functional as F
from typing import Optional

class DoRALayer(nn.Module):
    """DoRA 层独立实现"""
    def __init__(
        self,
        in_features: int,
        out_features: int,
        r: int = 16,
        alpha: int = 32,
        dropout: float = 0.0,
    ):
        super().__init__()
        self.r = r
        self.scale = alpha / r

        # 原始权重
        self.base_weight = nn.Parameter(
            torch.randn(out_features, in_features), requires_grad=False
        )

        # 计算方向和大小
        with torch.no_grad():
            weight_norm = torch.norm(self.base_weight, dim=1, keepdim=True)
            self.weight_dir = self.base_weight / (weight_norm + 1e-8)
            self.magnitude = nn.Parameter(weight_norm.clone())  # 可训练大小

        # LoRA 分量（方向更新）
        self.lora_A = nn.Parameter(torch.randn(r, in_features) * 0.01)
        self.lora_B = nn.Parameter(torch.zeros(out_features, r))
        self.dropout = nn.Dropout(dropout)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # LoRA 方向更新
        lora_update = (self.dropout(x) @ self.lora_A.T @ self.lora_B.T) * self.scale

        # 合并方向
        updated_dir = self.weight_dir + lora_update  # (out, in)

        # 归一化方向
        updated_dir_norm = F.normalize(updated_dir, dim=1, p=2)

        # 应用大小
        weight = self.magnitude * updated_dir_norm  # (out, in)

        # 前向
        return x @ weight.T


# 使用 PEFT 库
from peft import LoraConfig, get_peft_model

config = LoraConfig(
    r=16,
    lora_alpha=32,
    lora_dropout=0.05,
    target_modules=["q_proj", "k_proj", "v_proj", "o_proj"],
    bias="none",
    task_type="CAUSAL_LM",
    use_dora=True,        # 启用 DoRA（PEFT >= 0.9.0）
)
model = get_peft_model(base_model, config)

# DoRA 的梯度比 LoRA 更健康
# LoRA 梯度：dL/dB = dL/dW · A^T, dL/dA = B^T · dL/dW
# DoRA 梯度：额外有 dL/dm = Σ(dL/dW · V_updated)，使大小独立优化

# 实际效果：DoRA 在低 rank（r=8/16）时优势最明显
# 因为低 rank 下 LoRA 的表达能力受限，DoRA 的解耦提供了额外自由度

基座模型 (FP16, 14GB) → NF4 量化 (4bit, 3.5GB)
→ 冻结量化权重
→ 添加 LoRA 适配器 (FP16)
→ 用分页 AdamW 训练
→ 总显存: ~6GB (7B 模型)

import torch
from transformers import BitsAndBytesConfig, AutoModelForCausalLM
from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training

# 4bit 量化配置
bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.bfloat16,
    bnb_4bit_use_double_quant=True,
)

# 加载量化模型
model = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Meta-Llama-3-8B",
    quantization_config=bnb_config,
    device_map="auto",
)

# 准备 kbit 训练
model = prepare_model_for_kbit_training(model)

# LoRA 配置
config = LoraConfig(
    r=64,
    lora_alpha=128,
    target_modules=["q_proj", "k_proj", "v_proj", "o_proj",
                    "gate_proj", "up_proj", "down_proj"],
    lora_dropout=0.05,
    bias="none",
    task_type="CAUSAL_LM",
)
model = get_peft_model(model, config)

LoRA:  W' = W + (α/r) · BA
DoRA:  W' = m · normalize(W/||W|| + (α/r) · BA)
QLoRA: W' = dequant(NF4(W)) + (α/r) · BA

# QDoRA: 4bit 量化基座 + DoRA 适配器
config = LoraConfig(
    r=64,
    lora_alpha=128,
    target_modules=["q_proj", "k_proj", "v_proj", "o_proj"],
    lora_dropout=0.05,
    bias="none",
    task_type="CAUSAL_LM",
    use_dora=True,        # DoRA
)

# 模型加载时使用 4bit 量化
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    quantization_config=bnb_config,  # NF4 量化
    device_map="auto",
)
model = prepare_model_for_kbit_training(model)
model = get_peft_model(model, config)  # 添加 DoRA 适配器

显存够不够加载 FP16 模型？
├── 不够（<20GB）→ QLoRA 或 QDoRA
│   ├── 追求效果 → QDoRA
│   └── 追求速度 → QLoRA
└── 够 → LoRA 或 DoRA
    ├── r <= 16 → DoRA（优势明显）
    ├── r >= 64 → LoRA（差距缩小）或 DoRA
    └── 追求全量微调效果 → DoRA r=64

# LoRA/DoRA 可以合并到基座权重，推理时零额外开销
merged_model = model.merge_and_unload()
merged_model.save_pretrained("./merged_model")

# 注意：QLoRA 不能直接合并（精度不匹配）
# 需要先反量化再合并，或保持 adapter 形式推理

# 同一基座模型加载多个 LoRA adapter，按需切换
from peft import PeftModel

base_model = AutoModelForCausalLM.from_pretrained(...)
model = PeftModel.from_pretrained(base_model, "./lora_chat")
model.load_adapter("./lora_code", adapter_name="code")
model.load_adapter("./lora_math", adapter_name="math")

# 切换 adapter
model.set_adapter("code")   # 切到代码 adapter
model.set_adapter("math")   # 切到数学 adapter
model.set_adapter("default") # 切回默认

# 经验值：
# - 通用对话/指令跟随：r=16~64
# - 代码生成：r=32~128（代码需要更高表达力）
# - 领域适应（如医疗→通用）：r=64~256
# - 风格迁移：r=8~16（变化幅度小）
#
# 建议：从 r=16 开始，效果不够翻倍增加

# alpha 控制 LoRA 更新的缩放
# 常见设置：
# - alpha = r（保守，更新幅度小）
# - alpha = 2*r（推荐，大多数场景最优）
# - alpha = 4*r（激进，适合大 rank）
#
# 实际有效学习率 = lr * alpha / r
# 所以 alpha=2*r, lr=2e-4 等价于 alpha=r, lr=4e-4

# 不一定。以下情况 DoRA 优势不明显：
# 1. r >= 128：高 rank 下 LoRA 表达力已足够
# 2. 微调数据与预训练数据分布接近
# 3. 简单任务（如文本分类）
#
# DoRA 优势最大的场景：
# 1. r <= 16 的低 rank 场景
# 2. 领域跨度大的迁移学习
# 3. 需要 close-the-gap 到全量微调