بهینه‌سازی مدل‌های زبانی بزرگ با استفاده از تطبیق‌پذیری رتبه-پایین (LoRA) برای دستگاه‌های لبه

مقدمه

با گسترش روزافزون مدل‌های زبانی بزرگ (LLMs)، یکی از بزرگترین چالش‌های مهندسی، پیاده‌سازی و تنظیم دقیق این مدل‌ها بر روی سخت‌افزارهایی با منابع محدود است. در حالی که مدل‌هایی مانند LLaMA-3 نتایج خیره‌کننده‌ای از خود نشان می‌دهند، اجرای آن‌ها روی دستگاه‌های Edge نیازمند رویکردهای نوین است.

معماری پیشنهادی و فرمولاسیون

به جای بروزرسانی تمام ماتریس‌های وزن در شبکه عصبی، ما تغییرات وزن را به عنوان ضرب دو ماتریس با رتبه پایین (Low-Rank) مدل‌سازی می‌کنیم. اگر ماتریس وزن اولیه را در نظر بگیریم، بروزرسانی به شکل زیر محاسبه می‌شود.

همانطور که در معادلات ریاضی این روش اثبات شده است، فرمول بروزرسانی وزن‌ها به صورت $W = W_0 + \Delta W = W_0 + BA$ در نظر گرفته می‌شود که در آن $B$ و $A$ ماتریس‌هایی با رتبه بسیار پایین‌تر نسبت به $W_0$ هستند.

شکل ۱: مقایسه مصرف حافظه بین Fine-tuning سنتی و روش پیشنهادی LoRA

پیاده‌سازی در PyTorch

برای درک بهتر، در قطعه کد زیر بخشی از پیاده‌سازی لایه خطی جایگزین شده با معماری پیشنهادی را در فریم‌ورک PyTorch مشاهده می‌کنید. برای مشاهده کدهای طولانی‌تر می‌توانید از اسکرول پایین باکس استفاده کنید:

import torch
import torch.nn as nn
import math

class LoRALinearLayer(nn.Module):
    def __init__(self, in_features, out_features, rank=8, alpha=16.0, dropout_rate=0.05):
        super().__init__()
        self.in_features = in_features
        self.out_features = out_features
        
        # Original frozen weights (simulated for this example)
        self.weight = nn.Parameter(torch.Tensor(out_features, in_features), requires_grad=False)
        
        # LoRA specific parameters matrix A and B with extremely long comments to test the horizontal scrollbar behavior properly
        self.lora_A = nn.Parameter(torch.zeros((in_features, rank)))
        self.lora_B = nn.Parameter(torch.zeros((rank, out_features)))
        
        self.scaling = alpha / rank
        self.dropout = nn.Dropout(p=dropout_rate) if dropout_rate > 0 else nn.Identity()
        
        self.reset_parameters()

    def reset_parameters(self):
        nn.init.kaiming_uniform_(self.weight, a=math.sqrt(5))
        nn.init.kaiming_uniform_(self.lora_A, a=math.sqrt(5))
        nn.init.zeros_(self.lora_B) # B is initialized to zero so the initial impact is zero

    def forward(self, x):
        # Base computation + LoRA computation (Applying scaling factor to the low-rank matrices multiplication)
        base_out = nn.functional.linear(x, self.weight)
        lora_out = (self.dropout(x) @ self.lora_A @ self.lora_B) * self.scaling
        return base_out + lora_out

نتیجه‌گیری

آزمایش‌های ما بر روی دستگاه NVIDIA Jetson Orin نشان داد که با ترکیب این روش و تکنیک‌های Quantization 4-bit، می‌توان مدل‌های ۷ میلیارد پارامتری را برای کاربردهای تخصصی صنعتی بهینه‌سازی و مستقر کرد، بدون آنکه نیاز به کلاسترهای ابری گران‌قیمت باشد.