Method __init__

SwissArmyTransformer/sat/mpu/layers.py:190–225

(self, input_size, output_size, bias=True, gather_output=True,
                 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'))

                                     used for initialization.
    """
    def __init__(self, input_size, output_size, bias=True, gather_output=True,
                 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')):
        super(ColumnParallelLinear, self).__init__()

        # Keep input parameters
        self.stride = stride
        self.input_size = input_size
        self.output_size = output_size
        self.gather_output = gather_output
        # Divide the weight matrix along the last dimension.
        world_size = get_model_parallel_world_size()
        self.output_size_per_partition = divide(output_size, world_size)

        # Parameters.
        # Note: torch.nn.functional.linear performs XA^T + b and as a result
        # we allocate the transpose.
        self.weight = Parameter(torch.empty(self.output_size_per_partition,
                                             self.input_size, dtype=params_dtype,
                                             device=device))
        self.weight.model_parallel = True
        if bias:
            self.bias = Parameter(torch.empty(self.output_size_per_partition,dtype=params_dtype, device=device))
            self.bias.model_parallel = True
            # Always initialize bias to zero.
            with torch.no_grad():
                self.bias.zero_()
        else:
            self.register_parameter('bias', None)

        # Initialize weight.
        if not skip_init:
            self.master_weight = _initialize_affine_weight(
                self.weight, self.output_size, self.input_size,
                self.output_size_per_partition, 0, init_method,
                stride=self.stride, return_master_weight=keep_master_weight_for_test, module=module, name=name, self=self)

    def forward(self, input_):
        # Set up backprop all-reduce, and don't change the input.