Method split

fastdeploy/model_executor/layers/normalization.py:162–183

Split the input tensor across tensor parallel dimension. Args: x (paddle.Tensor): Input tensor to be split. Returns: paddle.Tensor: Splitted tensor.

(self, x)

        self.weight.set_value(weight_tensor.astype(self._norm_weight_dtype))

    def split(self, x):
        """
        Split the input tensor across tensor parallel dimension.

        Args:
            x (paddle.Tensor): Input tensor to be split.

        Returns:
            paddle.Tensor: Splitted tensor.
        """
        token_num = x.shape[0]
        token_num_per_rank = (token_num + self.tp_size - 1) // self.tp_size
        # AllGather will hang when the data shapes on multi-ranks are different!
        start_offset = self.tp_rank * token_num_per_rank
        end_offset = (self.tp_rank + 1) * token_num_per_rank
        if start_offset >= token_num:
            start_offset = token_num
        if end_offset > token_num:
            end_offset = token_num
        part_x = paddle.zeros(shape=[token_num_per_rank, x.shape[1]], dtype=x.dtype)
        part_x[: (end_offset - start_offset), :] = x[start_offset:end_offset, :]
        return part_x

    def allgather(self, out, token_num):
        """