手写LORA

在B站上刷到了好几个写lora的视频 想着自己也写一下

最近状态更新:

上周末去参加了斯坦福的Immerse the bay

Reality hack还会远吗

可能还挺远的… 本来以为是6月份还可以

如果是1月末的话 就算申请通过了 水平也跟不太上吧

回归正题 还是先看完

https://github.com/microsoft/LoRA/tree/main

理解完了 试着自己写出来

class LoRALayer():
    def __init__(
        self, 
        r: int, 
        lora_alpha: int, 
        lora_dropout: float,
        merge_weights: bool,
    ):
        self.r = r
        self.lora_alpha = lora_alpha
        # Optional dropout
        if lora_dropout > 0.:
            self.lora_dropout = nn.Dropout(p=lora_dropout)
        else:
            self.lora_dropout = lambda x: x
        # Mark the weight as unmerged
        self.merged = False
        self.merge_weights = merge_weights

1. Class Embedding

class Embedding(nn.Embedding, LoRALayer):
    # LoRA implemented in a dense layer
    def __init__(
        self,
        num_embeddings: int,
        embedding_dim: int,
        r: int = 0,
        lora_alpha: int = 1,
        merge_weights: bool = True,
        **kwargs
    ):
        nn.Embedding.__init__(self, num_embeddings, embedding_dim, **kwargs)
        LoRALayer.__init__(self, r=r, lora_alpha=lora_alpha, lora_dropout=0,
                           merge_weights=merge_weights)
        # Actual trainable parameters
        if r > 0:
            self.lora_A = nn.Parameter(self.weight.new_zeros((r, num_embeddings)))
            self.lora_B = nn.Parameter(self.weight.new_zeros((embedding_dim, r)))
            self.scaling = self.lora_alpha / self.r
            # Freezing the pre-trained weight matrix
            self.weight.requires_grad = False
        self.reset_parameters()

def train(self, mode: bool = True):
    nn.Embedding.train(self, mode)
    if mode:
        if self.merge_weights and self.merged:
            # Make sure that the weights are not merged
            if self.r > 0:
                self.weight.data -= (self.lora_B @ self.lora_A).transpose(0, 1) * self.scaling
            self.merged = False
    else:
        if self.merge_weights and not self.merged:
            # Merge the weights and mark it
            if self.r > 0:
                self.weight.data += (self.lora_B @ self.lora_A).transpose(0, 1) * self.scaling
            self.merged = True

def forward(self, x: torch.Tensor):
		#同时确保是在training
    if self.r > 0 and not self.merged:
        result = nn.Embedding.forward(self, x)
        after_A = F.embedding(
            x, self.lora_A.transpose(0, 1), self.padding_idx, self.max_norm,
            self.norm_type, self.scale_grad_by_freq, self.sparse
        )
        result += (after_A @ self.lora_B.transpose(0, 1)) * self.scaling
        return result
    else:
      return nn.Embedding.forward(self, x)

还没结束… 这两天更完