Method repartition

SwissArmyTransformer/sat/mpu/layers.py:242–281

(self)

        return output
    
    def repartition(self):
        assert self.output_size_per_partition == self.output_size
        self.output_size_per_partition = divide(self.output_size, get_model_parallel_world_size())
        mp_rank = get_model_parallel_rank()
        mp_size = get_model_parallel_world_size()
        self.original_weight = self.weight
        # weight is arranged as [stride0...stride1...stride2] * [input_size], extract non-contiguous parts
        strides = [1]*self.stride if isinstance(self.stride, int) else self.stride # int means equal number of qkv, or ratios
        assert self.weight.shape[0] % sum(strides) == 0, 'cannot divide weight evenly'
        factor = self.weight.shape[0] // sum(strides)
        # decompose weight according to strides
        strided_weights, _acm = [], 0
        for i in range(len(strides)):
            strided_weights.append(self.weight[_acm:_acm+factor*strides[i], :].detach())
            _acm += factor*strides[i]
        new_weight = torch.cat([
            strided_weight[
                (strided_weight.shape[0]//mp_size)*mp_rank:
                (strided_weight.shape[0]//mp_size)*(mp_rank+1)
                ]
            for strided_weight in strided_weights
        ], dim=0).contiguous().view(self.output_size_per_partition, self.input_size)
        self.weight = torch.nn.Parameter(new_weight)
        del self.original_weight
        if self.bias is not None and self.bias.numel() != 0:
            self.original_bias = self.bias
            # decompose bias according to strides
            strided_biases, _acm = [], 0
            for i in range(len(strides)):
                strided_biases.append(self.bias[_acm:_acm+factor*strides[i]].detach())
                _acm += factor*strides[i]
            new_bias = torch.cat([
                strided_bias[
                    (strided_bias.shape[0]//mp_size)*mp_rank:
                    (strided_bias.shape[0]//mp_size)*(mp_rank+1)
                    ]
                for strided_bias in strided_biases
            ], dim=0).contiguous().view(self.output_size_per_partition)
            self.bias = torch.nn.Parameter(new_bias)
            del self.original_bias

    def partition(self, new_model_parallel_size=None, full_weight=None):
        assert self.output_size_per_partition == self.output_size or full_weight is not None