| 22 | |
| 23 | class LoraLinear(nn.Linear): |
| 24 | def __init__( |
| 25 | self, |
| 26 | in_features: int, |
| 27 | out_features: int, |
| 28 | bias: bool = True, |
| 29 | lora_rank = 0 |
| 30 | ): |
| 31 | super().__init__(in_features, out_features, bias) |
| 32 | |
| 33 | self.lora_rank = lora_rank |
| 34 | if self.lora_rank > 0: |
| 35 | self.lora_a = nn.Linear(self.in_features, self.lora_rank, bias=False) |
| 36 | # workaround because trunc_normal_ does not currently support bfloat16 |
| 37 | _ = init.trunc_normal_(self.lora_a.weight.data.to(torch.float32), std=.02) |
| 38 | self.lora_a.weight.data.copy_(_) |
| 39 | self.lora_b = nn.Linear(self.lora_rank, self.out_features, bias=False) |
| 40 | nn.init.zeros_(self.lora_b.weight) |
| 41 | else: |
| 42 | self.lora_a = None |
| 43 | self.lora_b = None |
| 44 | |
| 45 | def forward(self, input_: torch.Tensor) -> torch.Tensor: # type:ignore |
| 46 | # Matrix multiply. |