(self, input_)
| 349 | self.register_parameter("bias", None) |
| 350 | |
| 351 | def forward(self, input_): |
| 352 | # Set up backprop all-reduce. |
| 353 | input_parallel = copy_to_tensor_model_parallel_region(input_) |
| 354 | # Matrix multiply. |
| 355 | |
| 356 | bias = self.bias if not self.skip_bias_add else None |
| 357 | output_parallel = F.linear(input_parallel, self.weight, bias) |
| 358 | if self.gather_output: |
| 359 | # All-gather across the partitions. |
| 360 | output = gather_from_tensor_model_parallel_region(output_parallel) |
| 361 | else: |
| 362 | output = output_parallel |
| 363 | output_bias = self.bias if self.skip_bias_add else None |
| 364 | return output, output_bias |
| 365 | |
| 366 | |
| 367 | class RowParallelLinear(torch.nn.Module): |
nothing calls this directly
no test coverage detected