Linear layer with row parallelism. The linear layer is defined as Y = XA + b. A is parallelized along its first dimension and X along its second dimension as: - - | A_1 | | . | A = | . | X = [X_1, ..., X_p] | .
| 160 | |
| 161 | |
| 162 | class LoraRowParallelLinear(RowParallelLinear): |
| 163 | """Linear layer with row parallelism. |
| 164 | |
| 165 | The linear layer is defined as Y = XA + b. A is parallelized along |
| 166 | its first dimension and X along its second dimension as: |
| 167 | - - |
| 168 | | A_1 | |
| 169 | | . | |
| 170 | A = | . | X = [X_1, ..., X_p] |
| 171 | | . | |
| 172 | | A_p | |
| 173 | - - |
| 174 | Arguments: |
| 175 | in_features: first dimension of matrix A. |
| 176 | out_features: second dimension of matrix A. |
| 177 | bias: If true, add bias. Note that bias is not parallelized. |
| 178 | input_is_parallel: If true, we assume that the input is already |
| 179 | split across the GPUs and we do not split |
| 180 | again. |
| 181 | init_method: method to initialize weights. Note that bias is always set |
| 182 | to zero. |
| 183 | stride: For the strided linear layers. |
| 184 | keep_master_weight_for_test: This was added for testing and should be |
| 185 | set to False. It returns the master weights |
| 186 | used for initialization. |
| 187 | """ |
| 188 | |
| 189 | def __init__( |
| 190 | self, |
| 191 | in_features: int, |
| 192 | out_features: int, |
| 193 | bias: bool = True, |
| 194 | input_is_parallel: bool = False, |
| 195 | init_method: Callable[[torch.Tensor], torch.Tensor] = init.xavier_normal_, |
| 196 | stride: int = 1, |
| 197 | keep_master_weight_for_test: bool = False, |
| 198 | lora_rank = 0 |
| 199 | ): |
| 200 | nn.Module.__init__(self) |
| 201 | |
| 202 | # Keep input parameters |
| 203 | self.in_features = in_features |
| 204 | self.out_features = out_features |
| 205 | self.input_is_parallel = input_is_parallel |
| 206 | # Divide the weight matrix along the last dimension. |
| 207 | world_size = get_model_parallel_world_size() |
| 208 | self.input_size_per_partition = divide_and_check_no_remainder(in_features, world_size) |
| 209 | |
| 210 | # Parameters. |
| 211 | # Note: torch.nn.functional.linear performs XA^T + b and as a result |
| 212 | # we allocate the transpose. |
| 213 | self.weight = Parameter(torch.Tensor(self.out_features, self.input_size_per_partition)) |
| 214 | if bias: |
| 215 | self.bias = Parameter(torch.Tensor(self.out_features)) |
| 216 | # Always initialize bias to zero. |
| 217 | with torch.no_grad(): |
| 218 | self.bias.zero_() |
| 219 | else: |