Hợp Nhất & Tách Trọng Số LoRA Cho Zero-Latency Inference

Mục tiêu kiến thức

• Trong giai đoạn Serving/Deployment, việc tính riêng rẽ x W₀^T và (x A^T B^T) · s làm tăng thêm phép nhân ma trận và độ trễ (latency).
• Vì tính chất tuyến tính, ta có thể cộng trực tiếp trọng số LoRA vào trọng số gốc:

W_merged = W₀ + αr (B × A) Khi đó, inference chỉ cần chạy duy nhất x W_merged^T mà không tốn thêm bất kỳ phép tính phụ nào (Zero Inference Latency Overhead).

• Khi tiếp tục fine-tune hoặc muốn phục hồi trọng số gốc, ta tách (unmerge):

W₀ = W_merged - αr (B × A)

Yêu cầu

Xây dựng lớp MergeableLoRALinear:

class MergeableLoRALinear(LoRALinear):
    def __init__(self, in_features: int, out_features: int, r: int = 8, lora_alpha: float = 16.0, bias: bool = True):
        pass

    def merge(self) -> None:
        pass

    def unmerge(self) -> None:
        pass

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        pass

• Quy tắc trạng thái:
• Khởi tạo có cờ self.merged = False.
• merge(): Nếu self.merged == True, raise RuntimeError("Already merged"). Ngược lại, cộng delta_weight vào self.weight.data và gán self.merged = True.
• unmerge(): Nếu self.merged == False, raise RuntimeError("Not merged yet"). Ngược lại, trừ delta_weight khỏi self.weight.data và gán self.merged = False.
• forward(x):
• Nếu đã merge: Chỉ tính F.linear(x, self.weight, self.bias) (không cộng thêm nhánh lora vì đã nằm trong weight).
• Nếu chưa merge: Tính F.linear(x, self.weight, self.bias) + lora_out.

Ràng buộc

• Thời gian chạy tối đa: 3000ms.
• Giới hạn bộ nhớ: 512MB.
• Dữ liệu đầu vào hợp lệ theo đúng kiểu dữ liệu và miền giá trị được mô tả.