【大模型后训练专题】 DoRA微调原理及实战项目

import torch
import torch.nn as nn
import torch.nn.functional as F

class DoRALinear(nn.Module):
    def __init__(
        self,
        in_features: int,
        out_features: int,
        r: int = 8,
        lora_alpha: float = 8.0,
        lora_dropout: float = 0.0,
        bias: bool = True,
        eps: float = 1e-6,
    ):
        super().__init__()
        self.in_features = in_features
        self.out_features = out_features
        self.r = r
        self.lora_alpha = lora_alpha
        self.scaling = lora_alpha / r if r > 0 else 1.0
        self.eps = eps

        # 冻结的预训练权重（W0）
        self.weight = nn.Parameter(torch.empty(out_features, in_features))
        if bias:
            self.bias = nn.Parameter(torch.zeros(out_features))
        else:
            self.bias = None

        # LoRA 分支（作用在方向上）
        if r > 0:
            self.lora_A = nn.Parameter(torch.zeros(r, in_features))
            self.lora_B = nn.Parameter(torch.zeros(out_features, r))
            self.lora_dropout = nn.Dropout(p=lora_dropout)
        else:
            self.register_parameter("lora_A", None)
            self.register_parameter("lora_B", None)
            self.lora_dropout = nn.Identity()

        self.magnitude = nn.Parameter(torch.ones(1, in_features))

        self.reset_parameters()

        self.weight.requires_grad = False

    def reset_parameters(self):
        nn.init.kaiming_uniform_(self.weight, a=5**0.5)
        if self.bias is not None:
            nn.init.zeros_(self.bias)

        if self.r > 0:
            nn.init.kaiming_uniform_(self.lora_A, a=5**0.5)
            nn.init.zeros_(self.lora_B)

        with torch.no_grad():
            # [out, in] -> 按列求范数 -> [1, in]
            col_norm = torch.norm(self.weight, dim=0, keepdim=True)  # ||W0||_c
            col_norm = torch.clamp(col_norm, min=self.eps)
            self.magnitude.copy_(col_norm)

    def _effective_weight(self):
        W0 = self.weight  # [out, in]

        if self.r > 0:
            # ΔV = scaling * B @ A   (注意这里是“权重更新”，与前向的 x 无关)
            delta_V = self.scaling * (self.lora_B @ self.lora_A)  # [out, in]
            V = W0 + delta_V
        else:
            V = W0

        # 按列归一化：|| · ||_c
        col_norm = torch.norm(V, dim=0, keepdim=True)  # [1, in]
        col_norm = torch.clamp(col_norm, min=self.eps)
        direction = V / col_norm 

        W_eff = direction * self.magnitude  # broadcast: [out, in] * [1, in]
        return W_eff

    def forward(self, x):
        """
        x: [B, in_features]
        """
        W_eff = self._effective_weight()
        return F.linear(x, W_eff, self.bias)

pip install "transformers>=4.44.0" \
            "peft>=0.11.0" \
            "accelerate" \
            "bitsandbytes>=0.45.0"

from transformers import AutoModelForCausalLM, AutoTokenizer

model_name = "meta-llama/Llama-3-8b"

tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    device_map="auto",
    torch_dtype="bfloat16",  # 或 "auto"
)

from peft import LoraConfig, get_peft_model, TaskType

lora_config = LoraConfig(
    task_type=TaskType.CAUSAL_LM,
    r=8,
    lora_alpha=16,
    lora_dropout=0.05,
    bias="none",
    target_modules=[
        "q_proj", "k_proj", "v_proj", "o_proj",
    ],
    use_dora=True,
)

model = get_peft_model(model, lora_config)
model.print_trainable_parameters()

from transformers import TrainingArguments, Trainer

training_args = TrainingArguments(
    output_dir="./dora-llama-sft",
    per_device_train_batch_size=4,
    gradient_accumulation_steps=4,
    learning_rate=1e-4,
    num_train_epochs=3,
    logging_steps=10,
    save_steps=500,
    fp16=True,
)

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=train_dataset, 
)

trainer.train()

model.save_pretrained("./dora-llama-sft")

from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
from peft import LoraConfig, get_peft_model, TaskType
import torch

model_name = "meta-llama/Llama-3-8b"

tokenizer = AutoTokenizer.from_pretrained(model_name)

bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_compute_dtype=torch.bfloat16,  
    bnb_4bit_use_double_quant=True,       
    bnb_4bit_quant_type="nf4",          
)

base_model = AutoModelForCausalLM.from_pretrained(
    model_name,
    quantization_config=bnb_config,
    device_map="auto",
)

lora_config = LoraConfig(
    task_type=TaskType.CAUSAL_LM,
    r=8,
    lora_alpha=16,
    lora_dropout=0.05,
    bias="none",
    target_modules=["q_proj", "k_proj", "v_proj", "o_proj"],
    use_dora=True,  
)

model = get_peft_model(base_model, lora_config)
model.print_trainable_parameters()