【実践】LoRA/PEFTで効率的にLLMをファインチューニングする

import torch
import torch.nn as nn
import math


class LoRALayer(nn.Module):
    """LoRA層"""
    def __init__(self, in_features, out_features, rank=8, alpha=16):
        super().__init__()
        self.rank = rank
        self.alpha = alpha
        self.scaling = alpha / rank

        # 低ランク行列
        self.lora_A = nn.Parameter(torch.zeros(rank, in_features))
        self.lora_B = nn.Parameter(torch.zeros(out_features, rank))

        # 初期化
        nn.init.kaiming_uniform_(self.lora_A, a=math.sqrt(5))
        nn.init.zeros_(self.lora_B)

    def forward(self, x):
        # x: (batch, seq_len, in_features)
        # LoRAの寄与を計算
        return (x @ self.lora_A.T @ self.lora_B.T) * self.scaling


class LinearWithLoRA(nn.Module):
    """LoRA付き線形層"""
    def __init__(self, linear_layer, rank=8, alpha=16):
        super().__init__()
        self.linear = linear_layer
        self.lora = LoRALayer(
            linear_layer.in_features,
            linear_layer.out_features,
            rank=rank,
            alpha=alpha
        )
        # 元の重みは凍結
        for param in self.linear.parameters():
            param.requires_grad = False

    def forward(self, x):
        return self.linear(x) + self.lora(x)

def apply_lora_to_model(model, rank=8, alpha=16, target_modules=None):
    """
    モデルにLoRAを適用

    Args:
        model: 対象のTransformerモデル
        rank: LoRAのランク
        alpha: スケーリング係数
        target_modules: LoRAを適用するモジュール名のリスト
    """
    if target_modules is None:
        target_modules = ['query', 'value']  # デフォルト

    for name, module in model.named_modules():
        if any(target in name for target in target_modules):
            if isinstance(module, nn.Linear):
                # 親モジュールを取得
                parent_name = '.'.join(name.split('.')[:-1])
                child_name = name.split('.')[-1]
                parent = model.get_submodule(parent_name) if parent_name else model

                # LoRA付き層に置換
                lora_layer = LinearWithLoRA(module, rank=rank, alpha=alpha)
                setattr(parent, child_name, lora_layer)

    # 学習可能パラメータの確認
    trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
    total = sum(p.numel() for p in model.parameters())
    print(f"Trainable: {trainable:,} / {total:,} ({100*trainable/total:.2f}%)")

    return model

def train_with_lora(model, train_loader, optimizer, num_epochs=3):
    """
    LoRAでのファインチューニング
    """
    model.train()
    for epoch in range(num_epochs):
        total_loss = 0
        for batch in train_loader:
            input_ids = batch['input_ids']
            labels = batch['labels']

            optimizer.zero_grad()
            outputs = model(input_ids, labels=labels)
            loss = outputs.loss
            loss.backward()
            optimizer.step()

            total_loss += loss.item()

        avg_loss = total_loss / len(train_loader)
        print(f"Epoch {epoch + 1}: Loss = {avg_loss:.4f}")


# 使用例
from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained("gpt2")
model = apply_lora_to_model(model, rank=8, alpha=16)

# LoRAパラメータのみを最適化
optimizer = torch.optim.AdamW(
    [p for p in model.parameters() if p.requires_grad],
    lr=1e-4
)

def merge_lora_weights(model):
    """
    LoRAの重みを元の重みにマージ
    """
    for name, module in model.named_modules():
        if isinstance(module, LinearWithLoRA):
            # LoRAの寄与を計算
            delta_w = (module.lora.lora_B @ module.lora.lora_A) * module.lora.scaling

            # 元の重みに加算
            module.linear.weight.data += delta_w

            # LoRAをゼロにリセット（または削除）
            module.lora.lora_A.data.zero_()
            module.lora.lora_B.data.zero_()

    print("LoRA weights merged successfully")
    return model

def save_lora_weights(model, path):
    """LoRAの重みのみを保存"""
    lora_state_dict = {}
    for name, param in model.named_parameters():
        if 'lora' in name:
            lora_state_dict[name] = param.data
    torch.save(lora_state_dict, path)
    print(f"Saved LoRA weights to {path}")


def load_lora_weights(model, path):
    """LoRAの重みを読み込み"""
    lora_state_dict = torch.load(path)
    model_state_dict = model.state_dict()

    for name, param in lora_state_dict.items():
        if name in model_state_dict:
            model_state_dict[name] = param

    model.load_state_dict(model_state_dict)
    print(f"Loaded LoRA weights from {path}")

元の層: x → Layer → y
Adapter: x → Layer → Adapter(y) + y → output

class Adapter(nn.Module):
    """Adapterモジュール"""
    def __init__(self, hidden_size, bottleneck_size):
        super().__init__()
        self.down_proj = nn.Linear(hidden_size, bottleneck_size)
        self.up_proj = nn.Linear(bottleneck_size, hidden_size)
        self.activation = nn.GELU()

    def forward(self, x):
        down = self.down_proj(x)
        activated = self.activation(down)
        up = self.up_proj(activated)
        return x + up  # 残差接続

class PrefixTuning(nn.Module):
    """Prefix Tuning"""
    def __init__(self, num_layers, num_heads, head_dim, prefix_length):
        super().__init__()
        self.prefix_length = prefix_length
        # 各層・各ヘッドに対するプレフィックス
        self.prefix_k = nn.Parameter(
            torch.randn(num_layers, prefix_length, num_heads * head_dim)
        )
        self.prefix_v = nn.Parameter(
            torch.randn(num_layers, prefix_length, num_heads * head_dim)
        )

    def get_prefix(self, layer_idx, batch_size):
        prefix_k = self.prefix_k[layer_idx].unsqueeze(0).expand(batch_size, -1, -1)
        prefix_v = self.prefix_v[layer_idx].unsqueeze(0).expand(batch_size, -1, -1)
        return prefix_k, prefix_v

class PromptTuning(nn.Module):
    """Prompt Tuning"""
    def __init__(self, num_tokens, embed_dim):
        super().__init__()
        self.prompt_embeddings = nn.Parameter(
            torch.randn(num_tokens, embed_dim)
        )

    def forward(self, input_embeddings):
        # input_embeddings: (batch, seq_len, embed_dim)
        batch_size = input_embeddings.size(0)
        prompt = self.prompt_embeddings.unsqueeze(0).expand(batch_size, -1, -1)
        return torch.cat([prompt, input_embeddings], dim=1)

from transformers import BitsAndBytesConfig

# 4ビット量子化設定
quantization_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_compute_dtype=torch.float16,
    bnb_4bit_use_double_quant=True,
    bnb_4bit_quant_type="nf4"
)

# 量子化モデルの読み込み
model = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Llama-2-7b-hf",
    quantization_config=quantization_config
)

# LoRAを適用
model = apply_lora_to_model(model, rank=16, alpha=32)