Forward pass of the BitLora layer. Args: x (Tensor): The input tensor. Returns: Tensor: The output tensor.
(self, x: Tensor)
| 61 | self.rms_norm = SimpleRMSNorm(self.in_features) |
| 62 | |
| 63 | def forward(self, x: Tensor) -> Tensor: |
| 64 | """ |
| 65 | Forward pass of the BitLora layer. |
| 66 | |
| 67 | Args: |
| 68 | x (Tensor): The input tensor. |
| 69 | |
| 70 | Returns: |
| 71 | Tensor: The output tensor. |
| 72 | |
| 73 | """ |
| 74 | w = self.weight |
| 75 | |
| 76 | # Normalize the input tensor |
| 77 | x_norm = self.rms_norm(x) |
| 78 | |
| 79 | # Activation Quant |
| 80 | x_quant = activation_quant(x_norm) |
| 81 | |
| 82 | if not self.merged and self.rank > 0: |
| 83 | lora = self.lora_a @ self.lora_b |
| 84 | w = w + lora * self.scaling |
| 85 | |
| 86 | # w_quant, scale = weight_quant(w) |
| 87 | w_quant = weight_quant(w) |
| 88 | # scale = weight_quant(w) |
| 89 | scale = 1.0 |
| 90 | output = nn.functional.linear(x_quant, w_quant, self.bias) |
| 91 | return output * scale |
| 92 | |
| 93 | def merge(self): |
| 94 | """ |
nothing calls this directly
no test coverage detected