# 原始权重矩阵 W: (4096, 4096) = 1677万参数
# LoRA矩阵 A: (4096, 8) = 3.3万参数
# LoRA矩阵 B: (8, 4096) = 3.3万参数
# 总LoRA参数：6.6万 vs 原始1677万，参数减少99.6%！
# 推理时：W' = W + BA（合并后无额外开销）
#
# 关键：r=8 叫做"秩"，是 LoRA 最重要的超参数
# 秩越大，LoRA 能表达的变化越复杂，但训练参数也越多
# 通常 r=8 到 r=64 之间，r=16 是常用的起点

import torch
import numpy as np

# 假设某层权重矩阵的更新量 ΔW
# 全量微调可能得到任意 (m, n) 矩阵，秩可达 min(m, n)
# 但实验发现，真正有效的更新的秩往往很低

# 示例：用SVD验证实际微调中更新矩阵的秩
# （这里用随机矩阵模拟，实际研究表明真实ΔW的奇异值快速衰减）
torch.manual_seed(42)
delta_W = torch.randn(4096, 4096) * 0.01  # 模拟参数更新量

# 奇异值分解
U, S, Vh = torch.linalg.svd(delta_W)

# 查看奇异值的分布
total_energy = S.sum().item()
top_8_energy = S[:8].sum().item()
top_64_energy = S[:64].sum().item()

print(f"前8个奇异值占总能量: {top_8_energy/total_energy*100:.1f}%")
print(f"前64个奇异值占总能量: {top_64_energy/total_energy*100:.1f}%")
# 实际微调中，有效更新通常集中在前几十个奇异值上
# 这就是为什么 r=8 或 r=16 往往就足够了

import torch
import torch.nn as nn

class LoRALayer(nn.Module):
    def __init__(self, original_layer, rank=8, alpha=16):
        super().__init__()
        self.original_layer = original_layer
        self.rank = rank
        self.alpha = alpha
        # scaling = alpha/rank，控制LoRA的整体贡献强度
        # alpha通常设为2*rank，使得scaling=2，相当于给LoRA更新略微加权
        self.scaling = alpha / rank

        # 冻结原始层：原始预训练权重不参与梯度计算
        for param in original_layer.parameters():
            param.requires_grad = False

        in_features = original_layer.in_features
        out_features = original_layer.out_features

        # LoRA矩阵A：正态分布初始化，用于提取输入的低维特征
        self.lora_A = nn.Parameter(torch.randn(rank, in_features) * 0.01)
        # LoRA矩阵B：零初始化！这是关键设计
        # 确保训练开始时 ΔW = B@A = 0，即LoRA初始不改变模型行为
        # 若A也零初始化，则A更新时梯度全为0，训练无法启动
        self.lora_B = nn.Parameter(torch.zeros(out_features, rank))

    def forward(self, x):
        # 原始层的输出（冻结权重）
        original_output = self.original_layer(x)
        # LoRA增量：先将输入投影到低维（A），再投影回原始维度（B）
        # x: (batch, seq_len, in_features)
        # x @ A.T: (batch, seq_len, rank)        投影到低秩空间
        # ... @ B.T: (batch, seq_len, out_features) 投影回原始空间
        lora_output = (x @ self.lora_A.T @ self.lora_B.T) * self.scaling
        return original_output + lora_output

# 注意力层的作用：控制"关注什么"
# q_proj, k_proj: 决定哪些词之间相互关注（注意力模式）
# v_proj: 决定从被关注的位置提取什么信息
# o_proj: 整合多头注意力的输出

# FFN层的作用：存储"知识"（记忆网络）
# 研究表明FFN层相当于key-value记忆，存储了大量事实性知识

# 微调任务通常是"调整行为风格/任务格式"，而非"灌输新知识"
# 注意力层控制行为模式，更适合微调
# FFN层存储知识，随意修改可能破坏原有知识（灾难性遗忘）

# 场景1：轻量微调（指令遵循、格式调整）
# 只微调注意力的Q/V，参数最少，效果通常已足够
lora_config_light = LoraConfig(
    r=8,
    lora_alpha=16,
    target_modules=["q_proj", "v_proj"],  # 最轻量：只有Q和V
    lora_dropout=0.05,
    task_type=TaskType.CAUSAL_LM,
)

# 场景2：标准微调（领域适应、风格迁移）
# 覆盖全部注意力层，平衡效果和参数量
lora_config_standard = LoraConfig(
    r=16,
    lora_alpha=32,
    target_modules=["q_proj", "k_proj", "v_proj", "o_proj"],
    lora_dropout=0.05,
    task_type=TaskType.CAUSAL_LM,
)

# 场景3：深度微调（任务差异大，如代码生成、数学推理）
# 覆盖注意力层 + FFN层，效果更好但参数更多
lora_config_full = LoraConfig(
    r=16,
    lora_alpha=32,
    target_modules=[
        "q_proj", "k_proj", "v_proj", "o_proj",  # 注意力层
        "gate_proj", "up_proj", "down_proj"        # FFN层（LLaMA/Qwen的FFN结构）
    ],
    lora_dropout=0.05,
    task_type=TaskType.CAUSAL_LM,
)

# 如何查看模型有哪些可选的target_modules？
# 打印模型结构：
# for name, module in model.named_modules():
#     if isinstance(module, nn.Linear):
#         print(name)

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, TrainingArguments
from peft import LoraConfig, get_peft_model, TaskType
from trl import SFTTrainer
from datasets import Dataset

# 1. 加载基础模型
model_name = "Qwen/Qwen2.5-7B-Instruct"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype=torch.float16,
    device_map="auto"
)

# 2. 配置LoRA
lora_config = LoraConfig(
    task_type=TaskType.CAUSAL_LM,
    r=16,                    # LoRA秩：影响模型容量，从16开始调试
    lora_alpha=32,           # 缩放因子，通常设为2*r
    lora_dropout=0.1,        # 数据量少时适当增大，防止过拟合
    target_modules=["q_proj", "k_proj", "v_proj", "o_proj",
                    "gate_proj", "up_proj", "down_proj"],
    bias="none"              # 不训练bias项，进一步减少参数
)

model = get_peft_model(model, lora_config)
model.print_trainable_parameters()
# 输出: trainable params: 83,886,080 || all params: 8,030,261,248 || trainable%: 1.04%

# 3. 准备数据集
train_data = [
    {
        "messages": [
            {"role": "system", "content": "你是一个专业的Python编程助手。"},
            {"role": "user", "content": "如何用Python读取CSV文件？"},
            {"role": "assistant", "content": "用pandas读取CSV文件非常简单：\n\n```python\nimport pandas as pd\ndf = pd.read_csv('file.csv')\nprint(df.head())\n```"}
        ]
    }
]

def format_chat(example):
    messages = example["messages"]
    text = tokenizer.apply_chat_template(
        messages, tokenize=False, add_generation_prompt=False
    )
    return {"text": text}

dataset = Dataset.from_list(train_data)
dataset = dataset.map(format_chat)

# 4. 训练配置
training_args = TrainingArguments(
    output_dir="./lora-output",
    num_train_epochs=3,
    per_device_train_batch_size=4,
    gradient_accumulation_steps=4,   # 等效batch_size = 4*4 = 16
    learning_rate=2e-4,              # 比全量微调大10倍：全量用2e-5，LoRA用2e-4
    fp16=True,                       # 混合精度训练（用半精度浮点数做计算，节省约一半显存），节省显存
    logging_steps=10,
    save_steps=100,
    warmup_ratio=0.05,               # 前5%的步数做学习率预热
    lr_scheduler_type="cosine",      # 余弦退火：见第3章梯度下降
    report_to="none"
)

# 5. 开始训练
trainer = SFTTrainer(
    model=model,
    args=training_args,
    train_dataset=dataset,
    tokenizer=tokenizer,
    dataset_text_field="text",
    max_seq_length=2048,
)
trainer.train()

# 6. 保存LoRA权重（只保存LoRA参数，几十MB，非常轻量）
model.save_pretrained("./lora-weights")
tokenizer.save_pretrained("./lora-weights")

from peft import PeftModel

base_model = AutoModelForCausalLM.from_pretrained(
    model_name, torch_dtype=torch.float16, device_map="auto"
)
# 加载LoRA权重并叠加到基础模型上
model = PeftModel.from_pretrained(base_model, "./lora-weights")

# merge_and_unload()：将LoRA权重合并进基础模型，然后移除LoRA结构
model = model.merge_and_unload()

方式1：保持LoRA结构（不合并）
  推理时计算：output = W(x) + BA(x) × scaling
  优点：可以随时切换不同的LoRA权重（服务多个客户的个性化模型）
  缺点：每次推理需要两次矩阵乘法，速度略慢；内存中同时存着W和B/A

方式2：merge_and_unload()（合并后推理）
  合并步骤：W_merged = W + B @ A × scaling
  推理时计算：output = W_merged(x)
  优点：推理速度与原始模型完全相同，无任何额外开销
  缺点：合并后无法再切换LoRA权重，需要重新加载基础模型

import torch
import torch.nn as nn

# merge_and_unload() 的核心逻辑（简化示意）
def merge_lora_weights(base_weight, lora_A, lora_B, scaling):
    """
    将LoRA权重合并进基础模型权重

    Args:
        base_weight: 原始模型权重 W，形状 (out_features, in_features)
        lora_A: LoRA矩阵A，形状 (rank, in_features)
        lora_B: LoRA矩阵B，形状 (out_features, rank)
        scaling: alpha / rank

    Returns:
        merged_weight: W + B@A*scaling，形状与base_weight相同
    """
    # B @ A 得到与 W 同形状的矩阵，再乘以缩放因子
    delta_W = (lora_B @ lora_A) * scaling  # (out_features, rank) @ (rank, in_features)
    merged_weight = base_weight + delta_W   # 直接相加，形状不变
    return merged_weight

# 合并后，lora_A 和 lora_B 就可以丢弃了（unload）
# 推理时的计算量与使用原始模型完全相同

# 验证合并前后的输出一致性
def verify_merge():
    torch.manual_seed(0)
    in_f, out_f, rank = 512, 512, 8

    W = torch.randn(out_f, in_f)
    A = torch.randn(rank, in_f) * 0.01
    B = torch.zeros(out_f, rank)
    scaling = 2.0
    x = torch.randn(4, 10, in_f)  # batch=4, seq_len=10

    # LoRA推理（未合并）
    output_lora = x @ W.T + (x @ A.T @ B.T) * scaling

    # 合并后推理
    W_merged = merge_lora_weights(W, A, B, scaling)
    output_merged = x @ W_merged.T

    # 两者结果应完全相同
    max_diff = (output_lora - output_merged).abs().max().item()
    print(f"最大差异: {max_diff:.2e}")  # 应该非常接近0（浮点误差级别）

verify_merge()

def chat(prompt, system="你是一个专业的Python编程助手。"):
    messages = [
        {"role": "system", "content": system},
        {"role": "user", "content": prompt}
    ]
    text = tokenizer.apply_chat_template(
        messages, tokenize=False, add_generation_prompt=True
    )
    inputs = tokenizer(text, return_tensors="pt").to(model.device)

    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=512,
            temperature=0.7,
            do_sample=True  # temperature > 0 时需要 do_sample=True
        )

    # 只解码新生成的token，跳过输入部分
    return tokenizer.decode(
        outputs[0][inputs['input_ids'].shape[1]:],
        skip_special_tokens=True
    )