(self, input_size, output_size, bias=True,
input_is_parallel=False,
init_method=unscaled_init_method(0.02), stride=1,
keep_master_weight_for_test=False, params_dtype=torch.float, module=None, name=None, skip_init=False, device=torch.device('cpu'), final_bias=True)
| 383 | used for initialization. |
| 384 | """ |
| 385 | def __init__(self, input_size, output_size, bias=True, |
| 386 | input_is_parallel=False, |
| 387 | init_method=unscaled_init_method(0.02), stride=1, |
| 388 | keep_master_weight_for_test=False, params_dtype=torch.float, module=None, name=None, skip_init=False, device=torch.device('cpu'), final_bias=True): |
| 389 | super(RowParallelLinear, self).__init__() |
| 390 | |
| 391 | # Keep input parameters |
| 392 | self.input_size = input_size |
| 393 | self.output_size = output_size |
| 394 | self.input_is_parallel = input_is_parallel |
| 395 | # Divide the weight matrix along the last dimension. |
| 396 | world_size = get_model_parallel_world_size() |
| 397 | self.input_size_per_partition = divide(input_size, world_size) |
| 398 | self.final_bias = final_bias |
| 399 | |
| 400 | # Parameters. |
| 401 | # Note: torch.nn.functional.linear performs XA^T + b and as a result |
| 402 | # we allocate the transpose. |
| 403 | self.weight = Parameter(torch.empty(self.output_size, |
| 404 | self.input_size_per_partition, dtype=params_dtype, device=device)) |
| 405 | self.weight.model_parallel = True |
| 406 | if bias: |
| 407 | self.bias = Parameter(torch.empty(self.output_size, dtype=params_dtype, device=device)) |
| 408 | # Always initialize bias to zero. |
| 409 | with torch.no_grad(): |
| 410 | self.bias.zero_() |
| 411 | else: |
| 412 | self.register_parameter('bias', None) |
| 413 | |
| 414 | # Initialize weight. |
| 415 | if not skip_init: |
| 416 | self.master_weight = _initialize_affine_weight( |
| 417 | self.weight, self.output_size, self.input_size, |
| 418 | self.input_size_per_partition, 1, init_method, |
| 419 | stride=stride, return_master_weight=keep_master_weight_for_test, module=module, name=name, self=self) |
| 420 | |
| 421 | def forward(self, input_): |
| 422 | # Split the input vector along the last dimension. |
no test coverage detected