| 21 | |
| 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. |
| 47 | output = F.linear(input_, self.weight, self.bias) |
| 48 | if self.lora_a is not None: |
| 49 | modification = self.lora_b(self.lora_a(input_)) |
| 50 | else: |
| 51 | modification = None |
| 52 | |
| 53 | if modification is not None: |
| 54 | output = output + modification |
| 55 | return output |
| 56 | |
| 57 | |
| 58 | class LoraColumnParallelLinear(ColumnParallelLinear): |