| 497 | |
| 498 | |
| 499 | class LoRALinearLayer(nn.Module): |
| 500 | def __init__(self, in_features, out_features, rank=4): |
| 501 | super().__init__() |
| 502 | |
| 503 | if rank > min(in_features, out_features): |
| 504 | raise ValueError(f"LoRA rank {rank} must be less or equal than {min(in_features, out_features)}") |
| 505 | |
| 506 | self.down = nn.Linear(in_features, rank, bias=False) |
| 507 | self.up = nn.Linear(rank, out_features, bias=False) |
| 508 | |
| 509 | nn.init.normal_(self.down.weight, std=1 / rank) |
| 510 | nn.init.zeros_(self.up.weight) |
| 511 | |
| 512 | def forward(self, hidden_states): |
| 513 | orig_dtype = hidden_states.dtype |
| 514 | dtype = self.down.weight.dtype |
| 515 | |
| 516 | down_hidden_states = self.down(hidden_states.to(dtype)) |
| 517 | up_hidden_states = self.up(down_hidden_states) |
| 518 | |
| 519 | return up_hidden_states.to(orig_dtype) |
| 520 | |
| 521 | |
| 522 | class LoRAAttnProcessor(nn.Module): |