Function fn

paddleformers/transformers/conversion_utils.py:620–670

fuse function for fusing weights (1) fuse_attention_qkv q => [q1,q2,q3,q4] k => [k1,k2,k3,k4] or [k1,k2] for GQA v => [v1,v2,v3,v4] or [v1,v2] for GQA fused weight => [q1,k1,v1,q2,k2,v2,q3,k3,v3,q4,k4,v4] or for GQA [q1,q2,k1,

(fuse_params, is_qkv=False, num_heads=None, num_key_value_heads=None)

def fuse_param_func():
    def fn(fuse_params, is_qkv=False, num_heads=None, num_key_value_heads=None):
        """fuse function for fusing weights

        (1) fuse_attention_qkv
            q => [q1,q2,q3,q4]
            k => [k1,k2,k3,k4] or [k1,k2] for GQA
            v => [v1,v2,v3,v4] or [v1,v2] for GQA
            fused weight => [q1,k1,v1,q2,k2,v2,q3,k3,v3,q4,k4,v4]
                 or for GQA [q1,q2,k1,v1,q3,q4,k2,v2]
        (2) fuse_attention_ffn
            directly fuse weights to 1 parts
            [gate_weight], [up_weight] => [gate_weight, up_weight]

        Args:
            fuse_params (_type_): to be fused weights
            is_qkv (bool, optional): for attention qkv weights. Defaults to False.
            num_heads (_type_, optional): query heads. Defaults to None.
            num_key_value_heads (_type_, optional): key and value heads. Defaults to None.

        Returns:
            _type_: fused weights
        """
        concat_fn = np.concatenate
        split_fn = np.split
        if isinstance(fuse_params[0], paddle.Tensor):
            concat_fn = paddle.cat
            split_fn = paddle.split

        if is_qkv:
            # fuse_attention_qkv
            assert num_heads, f"num_heads should be number of heads for Q, but got {num_heads}"
            assert (
                num_key_value_heads
            ), f"num_key_value_heads should be number of key_value_heads for K and V, but got {num_key_value_heads}"
            assert (
                len(fuse_params) == 3
            ), f"fuse_params length is not equal 3, it should be Q K V list. but got length {len(fuse_params)}"
            num_query_groups = num_heads // num_key_value_heads
            q_list = split_fn(fuse_params[0], num_heads, axis=-1)
            k_list = split_fn(fuse_params[1], num_key_value_heads, axis=-1)
            v_list = split_fn(fuse_params[2], num_key_value_heads, axis=-1)

            qkv_pairs = []
            for i in range(num_key_value_heads):
                qkv_pairs += q_list[i * num_query_groups : (i + 1) * num_query_groups]
                qkv_pairs.append(k_list[i])
                qkv_pairs.append(v_list[i])
            return concat_fn(qkv_pairs, axis=-1)
        else:
            # fuse_attention_ffn
            return concat_fn(fuse_params, axis=-1)

    return fn