# PEFT — 参数高效微调概述
from peft import LoraConfig, get_peft_model, TaskType

# 1. 定义 LoRA 配置
config = LoraConfig(
    task_type=TaskType.CAUSAL_LM,
    r=8,                       # 秩
    lora_alpha=16,              # 缩放因子
    lora_dropout=0.05,
    target_modules=["q_proj", "v_proj"],  # 目标层
)

# 2. 包装基础模型 — 冻结原始权重，仅训练适配器
model = get_peft_model(base_model, config)
model.print_trainable_parameters()
# trainable params: 1,572,864 || all params: 675,235,840 || trainable%: 0.23%

import torch
import torch.nn as nn

class LoRALinear(nn.Module):
    def __init__(self, in_features, out_features, r=8, alpha=16):
        super().__init__()
        self.W0 = nn.Linear(in_features, out_features, bias=False)
        self.W0.weight.requires_grad = False          # 冻结原始权重

        self.B = nn.Parameter(torch.zeros(out_features, r))    # d_out × r
        self.A = nn.Parameter(torch.randn(r, in_features) * 0.01)  # r × d_in
        self.scaling = alpha / r

    def forward(self, x):
        # W = W₀ + (α/r) · B·A
        return self.W0(x) + self.scaling * (x @ self.A.T) @ self.B.T

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

# 1. 4-bit NF4 量化配置
bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",               # NormalFloat4
    bnb_4bit_compute_dtype=torch.float16,
    bnb_4bit_use_double_quant=True,           # 双重量化
)

# 2. 加载 4-bit 量化模型
model = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Llama-2-7B",
    quantization_config=bnb_config,
    device_map="auto",
)

# 3. 准备 k-bit 训练 + 应用 LoRA
model = prepare_model_for_kbit_training(model)
model = get_peft_model(model, LoraConfig(r=64, lora_alpha=16,
    target_modules=["q_proj","k_proj","v_proj","o_proj"]))

import torch.nn as nn

class Adapter(nn.Module):
    def __init__(self, d_model=768, bottleneck=64):
        super().__init__()
        self.down = nn.Linear(d_model, bottleneck)        # 降维 d → r
        self.act  = nn.GELU()                            # 非线性激活
        self.up   = nn.Linear(bottleneck, d_model)        # 升维 r → d
        self.scale = nn.Parameter(torch.ones(1))         # 残差缩放

    def forward(self, h):
        # 残差连接：h + scale · up(act(down(h)))
        return h + self.scale * self.up(self.act(self.down(h)))

class PrefixTuning(nn.Module):
    def __init__(self, num_layers=12, num_heads=12,
                 prefix_len=20, d_head=64):
        super().__init__()
        self.prefix_len = prefix_len
        d_kv = num_heads * d_head

        # 可学习前缀参数（重参数化）
        self.prefix_keys = nn.Parameter(torch.randn(num_layers, prefix_len, d_kv))
        self.prefix_vals = nn.Parameter(torch.randn(num_layers, prefix_len, d_kv))

    def forward(self, layer_idx, K, V):
        # 在第 layer_idx 层的 K/V 前拼接前缀
        pk = self.prefix_keys[layer_idx]    # [prefix_len, d_kv]
        pv = self.prefix_vals[layer_idx]
        K = torch.cat([pk.unsqueeze(0), K], dim=1)  # 拼接
        V = torch.cat([pv.unsqueeze(0), V], dim=1)
        return K, V

def sft_loss(logits, labels, attention_mask):
    # 损失掩码：仅在回复 token 上计算损失
    # labels 中 instruction 部分设为 -100 (忽略)
    shift_logits = logits[..., :-1, :].contiguous()
    shift_labels = labels[..., 1:].contiguous()

    loss_fct = nn.CrossEntropyLoss(ignore_index=-100)
    loss = loss_fct(
        shift_logits.view(-1, shift_logits.size(-1)),
        shift_labels.view(-1),
    )
    return loss

# 构造带掩码的标签
def build_labels(input_ids, response_start):
    labels = input_ids.clone()
    labels[:, :response_start] = -100   # 掩码指令部分
    return labels

import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP

def setup_ddp():
    dist.init_process_group(backend="nccl")
    torch.cuda.set_device(dist.get_rank())

def train_step(model, optimizer, batch):
    # model 是 DDP 包装后的模型
    inputs, labels = batch
    outputs = model(inputs)                    # 前向（含自动同步准备）
    loss = loss_fn(outputs, labels)

    optimizer.zero_grad()
    loss.backward()                          # 反向 → AllReduce 同步梯度
    optimizer.step()                           # 各 GPU 独立更新（梯度相同）
    return loss.item()

# FSDP 等价 — 只需替换包装器
# from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
# model = FSDP(model, sharding_strategy=ShardingStrategy.FULL_SHARD)

class EWC:
    def __init__(self, model, fisher=None, params_star=None):
        self.fisher = fisher              # Fisher 信息矩阵
        self.params_star = params_star    # 旧任务最优参数

    def penalty(self, model):
        # L_ewc = Σ F_i · (θ_i - θ*_i)²
        loss = 0
        for n, p in model.named_parameters():
            if n in self.fisher:
                loss += (self.fisher[n] * (p - self.params_star[n]) ** 2).sum()
        return loss

# 训练时：总损失 = 任务损失 + λ · EWC惩罚
# total_loss = task_loss + lambda_ewc * ewc.penalty(model)