(
self,
input_size,
output_size,
bias=True,
input_is_parallel=False,
init_method=init.xavier_normal_,
stride=1,
keep_master_weight_for_test=False,
skip_bias_add=False,
params_dtype=None,
skip_init=False,
device=None,
)
| 395 | """ |
| 396 | |
| 397 | def __init__( |
| 398 | self, |
| 399 | input_size, |
| 400 | output_size, |
| 401 | bias=True, |
| 402 | input_is_parallel=False, |
| 403 | init_method=init.xavier_normal_, |
| 404 | stride=1, |
| 405 | keep_master_weight_for_test=False, |
| 406 | skip_bias_add=False, |
| 407 | params_dtype=None, |
| 408 | skip_init=False, |
| 409 | device=None, |
| 410 | ): |
| 411 | super(RowParallelLinear, self).__init__() |
| 412 | |
| 413 | # Keep input parameters |
| 414 | self.input_size = input_size |
| 415 | self.output_size = output_size |
| 416 | self.input_is_parallel = input_is_parallel |
| 417 | # Divide the weight matrix along the last dimension. |
| 418 | world_size = get_tensor_model_parallel_world_size() |
| 419 | self.input_size_per_partition = divide(input_size, world_size) |
| 420 | self.skip_bias_add = skip_bias_add |
| 421 | self.params_dtype = params_dtype |
| 422 | self.device = device |
| 423 | |
| 424 | # Parameters. |
| 425 | # Note: torch.nn.functional.linear performs XA^T + b and as a result |
| 426 | # we allocate the transpose. |
| 427 | # Initialize weight. |
| 428 | args = get_args() |
| 429 | if not skip_init: |
| 430 | if args.use_cpu_initialization: |
| 431 | self.weight = Parameter( |
| 432 | torch.empty( |
| 433 | self.output_size, |
| 434 | self.input_size_per_partition, |
| 435 | dtype=self.params_dtype if self.params_dtype is not None else args.params_dtype, |
| 436 | ) |
| 437 | ) |
| 438 | self.master_weight = _initialize_affine_weight_cpu( |
| 439 | self.weight, |
| 440 | self.output_size, |
| 441 | self.input_size, |
| 442 | self.input_size_per_partition, |
| 443 | 1, |
| 444 | init_method, |
| 445 | stride=stride, |
| 446 | return_master_weight=keep_master_weight_for_test, |
| 447 | ) |
| 448 | else: |
| 449 | self.weight = Parameter( |
| 450 | torch.empty( |
| 451 | self.output_size, |
| 452 | self.input_size_per_partition, |
| 453 | device=self.device if self.device is not None else torch.cuda.current_device(), |
| 454 | dtype=self.params_dtype if self.params_dtype is not None else args.params_dtype, |
no test coverage detected