paddlex 2.0.0rc4__py3-none-any.whl → 3.0.0__py3-none-any.whl

paddlex/inference/models/common/vlm/conversion_utils.py

@@ -0,0 +1,99 @@
+# Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+import numpy as np
+import paddle
+
+
+def fuse_param_func():
+    def fn(fuse_params, is_qkv=False, num_heads=None, num_key_value_heads=None):
+        concat_fn = np.concatenate
+        split_fn = np.split
+        if isinstance(fuse_params[0], paddle.Tensor):
+            concat_fn = paddle.concat
+            split_fn = paddle.split
+
+        if is_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:
+            return concat_fn(fuse_params, axis=-1)
+
+    return fn
+
+
+def split_param_func():
+    def fn(
+        fused_param,
+        split_nums=2,
+        is_qkv=False,
+        num_heads=None,
+        num_key_value_heads=None,
+    ):
+        concat_fn = np.concatenate
+        split_fn = np.split
+        if isinstance(fused_param, paddle.Tensor):
+            concat_fn = paddle.concat
+            split_fn = paddle.split
+
+        if is_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}"
+            num_query_groups = num_heads // num_key_value_heads
+            q_list, k_list, v_list = [], [], []
+            split_heads = split_fn(
+                fused_param, num_heads + 2 * num_key_value_heads, axis=-1
+            )
+            for i in range(num_key_value_heads):
+                q_list += split_heads[
+                    i * (num_query_groups + 2) : (i + 1) * (num_query_groups + 2) - 2
+                ]
+                k_list.append(split_heads[(i + 1) * (num_query_groups + 2) - 2])
+                v_list.append(split_heads[(i + 1) * (num_query_groups + 2) - 1])
+            return (
+                concat_fn(q_list, axis=-1),
+                concat_fn(k_list, axis=-1),
+                concat_fn(v_list, axis=-1),
+            )
+        else:
+            return split_fn(fused_param, split_nums, axis=-1)
+
+    return fn
+
+
+def split_or_fuse_func(is_fuse=True):
+    return fuse_param_func() if is_fuse else split_param_func()

paddlex/inference/models/common/vlm/distributed.py

@@ -0,0 +1,229 @@
+# Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import re
+from typing import Any, Union
+
+import numpy as np
+import paddle
+import paddle.distributed as distributed
+
+from .utils import device_guard
+
+world_size = distributed.get_world_size()
+
+
+def convert_file_size_to_int(size: Union[int, str]):
+    """
+    Converts a size expressed as a string with digits an unit (like `"5MB"`) to an integer (in bytes).
+    Args:
+        size (`int` or `str`): The size to convert. Will be directly returned if an `int`.
+    """
+    if isinstance(size, int):
+        return size
+    if size.upper().endswith("GIB"):
+        return int(size[:-3]) * (2**30)
+    if size.upper().endswith("MIB"):
+        return int(size[:-3]) * (2**20)
+    if size.upper().endswith("KIB"):
+        return int(size[:-3]) * (2**10)
+    if size.upper().endswith("GB"):
+        int_size = int(size[:-2]) * (10**9)
+        return int_size // 8 if size.endswith("b") else int_size
+    if size.upper().endswith("MB"):
+        int_size = int(size[:-2]) * (10**6)
+        return int_size // 8 if size.endswith("b") else int_size
+    if size.upper().endswith("KB"):
+        int_size = int(size[:-2]) * (10**3)
+        return int_size // 8 if size.endswith("b") else int_size
+    raise ValueError(
+        "`size` is not in a valid format. Use an integer followed by the unit, e.g., '5GB'."
+    )
+
+
+def reduce_tensor(tensor, buffer_size="32MiB"):
+    if tensor.dtype == paddle.int8:
+        numel = np.prod(tensor.shape)
+    else:
+        numel = int(paddle.numel(tensor).item())
+    buffer_size = convert_file_size_to_int(buffer_size)
+    tensor.reshape_([-1])
+
+    send_size = buffer_size // dtype_byte_size(tensor.dtype)
+
+    for x in range(0, numel, send_size):
+        part_tensor = tensor[x : min(numel, x + send_size)]
+        yield part_tensor, (x, min(numel, x + send_size))
+
+
+def dtype_byte_size(dtype):
+    """
+    Returns the size (in bytes) occupied by one parameter of type `dtype`.
+    """
+    if dtype == paddle.bool:
+        return 1 / 8
+    bit_search = re.search(r"[^\d](\d+)$", str(dtype))
+    if bit_search is None:
+        raise ValueError(f"`dtype` is not a valid dtype: {dtype}.")
+    bit_size = int(bit_search.groups()[0])
+    return bit_size // 8
+
+
+@paddle.no_grad()
+def distributed_gather(tensor: Any, dst: int = 0, group=None, offload=False) -> Any:
+    try:
+        if isinstance(tensor, (tuple, list)):
+            return type(tensor)(
+                distributed_gather(t, dst, group, offload) for t in tensor
+            )
+        if isinstance(tensor, dict):
+            return {
+                k: distributed_gather(v, dst, group, offload) for k, v in tensor.items()
+            }
+
+        output_tensors = None
+
+        is_dst = dst == distributed.get_rank(group=group)
+        if is_dst:
+            if offload:
+                output_tensors = [
+                    [] for _ in range(distributed.get_world_size(group=group))
+                ]
+            else:
+                output_tensors = [
+                    paddle.empty_like(tensor)
+                    for _ in range(distributed.get_world_size(group=group))
+                ]
+                output_tensors = [
+                    t if len(t.shape) > 0 else t[None] for t in output_tensors
+                ]
+
+        if offload:
+            origin_shape = tensor.shape
+            tensor.reshape_([-1])
+
+            for slice_tensor, index in reduce_tensor(tensor):
+                slice_output_tensors = None
+                if distributed.get_rank(group=group) == dst:
+                    slice_output_tensors = [
+                        paddle.empty_like(slice_tensor)
+                        for _ in range(distributed.get_world_size(group=group))
+                    ]
+                paddle.distributed.communication.stream.gather(
+                    slice_tensor,
+                    slice_output_tensors,
+                    dst=group.ranks[dst] if group else dst,
+                    group=group,
+                    sync_op=True,
+                    use_calc_stream=False,
+                )
+
+                if is_dst:
+                    for i in range(len(output_tensors)):
+                        output_tensors[i].append(slice_output_tensors[i].cpu().numpy())
+
+            tensor.reshape_(origin_shape)
+            if is_dst:
+                with device_guard("cpu"):
+                    new_output_tensors = []
+                    for x in output_tensors:
+                        t = np.concatenate(x)
+                        t = t.reshape(origin_shape)
+                        new_output_tensors.append(t)
+                    output_tensors = new_output_tensors
+
+        else:
+            paddle.distributed.communication.stream.gather(
+                tensor,
+                output_tensors,
+                dst=group.ranks[dst] if group else dst,
+                group=group,
+                sync_op=True,
+                use_calc_stream=False,
+            )
+
+        return output_tensors
+
+    except AssertionError:
+        raise AssertionError("Not currently using distributed training")
+
+
+@paddle.no_grad()
+def distributed_allgather(tensor: Any, group=None, offload=False):
+    """nested all gather function with offload
+
+    Args:
+        tensor (Any): the desired tensor, list of tensor, dict of tensor to allgather.
+        group (_type_, optional): the communication group. Defaults to None.
+        offload (bool, optional): If True, we offload the received tensor to cpu/(numpy). Defaults to False.
+
+    Raises:
+        AssertionError: Unexpected errors.
+
+    Returns:
+        tensor list: list of all gathered tensors
+    """
+    try:
+        if isinstance(tensor, (tuple, list)):
+            return type(tensor)(
+                distributed_allgather(t, group, offload) for t in tensor
+            )
+        if isinstance(tensor, dict):
+            return {
+                k: distributed_allgather(v, group, offload) for k, v in tensor.items()
+            }
+
+        output_tensors = []
+
+        if offload:
+            with device_guard("cpu"):
+                output_tensors = [
+                    paddle.empty_like(tensor)
+                    for _ in range(distributed.get_world_size(group))
+                ]
+        else:
+            output_tensors = [
+                paddle.empty_like(tensor)
+                for _ in range(distributed.get_world_size(group))
+            ]
+
+        output_tensors = [t if len(t.shape) > 0 else t[None] for t in output_tensors]
+
+        if offload:
+            origin_shape = tensor.shape
+            tensor.reshape_([-1])
+            for x in output_tensors:
+                x.reshape_([-1])
+
+            for slice_tensor, index in reduce_tensor(tensor):
+                slice_output_tensors = [
+                    paddle.empty_like(slice_tensor)
+                    for _ in range(distributed.get_world_size(group))
+                ]
+                distributed.all_gather(slice_output_tensors, slice_tensor, group=group)
+                for x, y in zip(slice_output_tensors, output_tensors):
+                    with device_guard("cpu"):
+                        y[index[0] : index[1]] = x.cpu()
+
+            tensor.reshape_(origin_shape)
+            for x in output_tensors:
+                x.reshape_(origin_shape)
+
+        else:
+            distributed.all_gather(output_tensors, tensor)
+
+        return output_tensors
+
+    except AssertionError:
+        raise AssertionError("Not currently using distributed training")

paddlex/inference/models/common/vlm/flash_attn_utils.py

@@ -0,0 +1,119 @@
+# Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+import paddle
+
+
+def is_flash_attn_available():
+    try:
+        import os
+
+        if "npu" in paddle.get_device():  # NOTE: flash attn has not been tested yet
+            for lib in os.listdir(os.getenv("CUSTOM_DEVICE_ROOT")):
+                if lib.endswith(".so"):
+                    paddle.utils.cpp_extension.extension_utils.load_op_meta_info_and_register_op(
+                        lib
+                    )
+            from paddle.base import core
+
+            def flash_attention_npu(
+                query,
+                key,
+                value,
+                dropout=0.0,
+                causal=False,
+                return_softmax=False,
+                *,
+                fixed_seed_offset=None,
+                rng_name="",
+                training=True,
+                name=None,
+                attn_mask=None,
+                is_varlen=False,
+                batch_size=None,
+                seq_length=None,
+            ):
+
+                is_triangle_upper_mask = True
+
+                if is_varlen:
+                    if len(query.shape) == 4:
+                        B, S, N, D = query.shape
+                        query = query.reshape([B * S, N, D])
+                        key = key.reshape([B * S, N, D])
+                        value = value.reshape([B * S, N, D])
+                    else:
+                        assert batch_size is not None
+                        assert seq_length is not None
+                        B = batch_size
+                        S = seq_length
+                    actual_seq_q_len = actual_seq_kv_len = list(range(S, B * S + S, S))
+                else:
+                    actual_seq_q_len = actual_seq_kv_len = []
+
+                out = core.eager._run_custom_op(
+                    "flash_attention_npu",
+                    query,
+                    key,
+                    value,
+                    fixed_seed_offset,
+                    attn_mask,
+                    actual_seq_q_len,
+                    actual_seq_kv_len,
+                    dropout,
+                    causal,
+                    return_softmax,
+                    not training,
+                    is_triangle_upper_mask,
+                    is_varlen,
+                )[0]
+
+                return out
+
+            q = paddle.rand((1, 4, 2, 8)).astype("bfloat16")
+            _ = flash_attention_npu(q, q, q, 0.9, False, False)
+            paddle.nn.functional.flash_attention_npu = flash_attention_npu
+            return True
+        q = paddle.rand((1, 4, 2, 8)).astype("bfloat16")
+        _ = paddle.nn.functional.flash_attention.flash_attention(
+            q, q, q, 0.9, False, False
+        )
+        return True
+    except:
+        return False
+
+
+HAS_FLASH_ATTN = is_flash_attn_available()
+
+
+def has_flash_attn_func():
+    if HAS_FLASH_ATTN:
+        try:
+            if "npu" in paddle.get_device():
+                flash_attn_func_npu = paddle.nn.functional.flash_attention_npu
+                return flash_attn_func_npu, flash_attn_func_npu
+            else:
+                from paddle.nn.functional.flash_attention import (
+                    flash_attention as flash_attn_func,
+                )
+                from paddle.nn.functional.flash_attention import (
+                    flash_attn_unpadded as flash_attn_varlen_func,
+                )
+
+                return flash_attn_func, flash_attn_varlen_func
+        except:
+            return None, None
+    else:
+        return None, None

paddlex/inference/models/common/vlm/fusion_ops.py

@@ -0,0 +1,205 @@
+# Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import os
+
+import paddle
+import paddle.nn.functional as F
+
+try:
+    from paddle.incubate.nn.functional import fused_rotary_position_embedding
+except ImportError:
+    fused_rotary_position_embedding = None
+
+try:
+    from paddle.incubate.nn.functional import swiglu
+except ImportError:
+
+    def swiglu(x, y=None):
+        if y is None:
+            x, y = paddle.chunk(x, chunks=2, axis=-1)
+        return F.silu(x) * y
+
+
+from paddle.utils import try_import
+
+
+def get_env_device():
+    """
+    Return the device name of running environment.
+    """
+    if paddle.is_compiled_with_cuda():
+        return "gpu"
+    elif "npu" in paddle.device.get_all_custom_device_type():
+        return "npu"
+    elif "mlu" in paddle.device.get_all_custom_device_type():
+        return "mlu"
+    elif "gcu" in paddle.device.get_all_custom_device_type():
+        return "gcu"
+    elif "intel_hpu" in paddle.device.get_all_custom_device_type():
+        return "intel_hpu"
+    elif paddle.is_compiled_with_rocm():
+        return "rocm"
+    elif paddle.is_compiled_with_xpu():
+        return "xpu"
+    return "cpu"
+
+
+try:
+    from paddle.incubate.nn.functional import fused_rotary_position_embedding
+except ImportError:
+    fused_rotary_position_embedding = None
+try:
+    if get_env_device() in ["npu", "mlu", "gcu"]:
+        from paddle.base import core
+
+        for lib in os.listdir(os.getenv("CUSTOM_DEVICE_ROOT")):
+            if lib.endswith(".so"):
+                paddle.utils.cpp_extension.extension_utils.load_op_meta_info_and_register_op(
+                    lib
+                )
+    from paddle.nn.functional.flash_attention import flash_attention
+except:
+    flash_attention = None
+
+
+def fusion_rope(
+    query_states,
+    key_states,
+    value_states,
+    hidden_states,
+    position_ids,
+    past_key_value,
+    rotary_emb,
+    context_parallel_degree=-1,
+):
+    if get_env_device() not in ["gcu", "intel_hpu"]:
+        assert past_key_value is None, "fuse rotary not support cache kv for now"
+    batch_size, seq_length, num_heads, head_dim = query_states.shape
+    _, kv_seq_len, num_key_value_heads, _ = key_states.shape
+    if context_parallel_degree > 1:
+        assert (
+            get_env_device() == "gpu"
+        ), "context parallel only support cuda device for now"
+        kv_seq_len *= context_parallel_degree
+    if get_env_device() not in ["gcu", "intel_hpu"]:
+        cos, sin = rotary_emb(value_states, seq_len=kv_seq_len)
+    if get_env_device() == "npu":
+        query_states = core.eager._run_custom_op("fused_rope", query_states, cos, sin)[
+            0
+        ]
+        key_states = core.eager._run_custom_op("fused_rope", key_states, cos, sin)[0]
+    elif get_env_device() == "intel_hpu":
+        if past_key_value is not None:
+            kv_seq_len += past_key_value[0].shape[-3]
+        cos, sin = rotary_emb(value_states, seq_len=kv_seq_len)
+        cos = cos.squeeze().unsqueeze(0).unsqueeze(0)
+        sin = sin.squeeze().unsqueeze(0).unsqueeze(0)
+        query_states, _, _ = (
+            paddle.incubate.nn.functional.fused_rotary_position_embedding(
+                paddle.transpose(query_states, [0, 2, 1, 3]),
+                None,
+                None,
+                sin=sin,
+                cos=cos,
+                position_ids=position_ids,
+            )
+        )
+        key_states, _, _ = (
+            paddle.incubate.nn.functional.fused_rotary_position_embedding(
+                paddle.transpose(key_states, [0, 2, 1, 3]),
+                None,
+                None,
+                sin=sin,
+                cos=cos,
+                position_ids=position_ids,
+            )
+        )
+        query_states = paddle.transpose(query_states, [0, 2, 1, 3])
+        key_states = paddle.transpose(key_states, [0, 2, 1, 3])
+    elif get_env_device() == "gcu":
+        cos_sin = rotary_emb.get_fused_cos_sin(value_states, seq_len=kv_seq_len)
+        query_states, key_states = core.eager._run_custom_op(
+            "fused_rotary_embedding_gcu",
+            query_states,
+            key_states,
+            cos_sin,
+            position_ids,
+            True,
+        )
+    else:
+        # paddle version > 2.6 or develop support q and k/v with different num_heads
+        paddle_version = float(paddle.__version__[:3])
+        if ((paddle_version != 0.0) and (paddle_version <= 2.6)) and (
+            num_heads != num_key_value_heads
+        ):
+            query_states, _, _ = fused_rotary_position_embedding(
+                query_states,
+                None,
+                None,
+                sin=sin,
+                cos=cos,
+                position_ids=position_ids,
+                use_neox_rotary_style=False,
+            )
+            key_states, _, _ = fused_rotary_position_embedding(
+                key_states,
+                None,
+                None,
+                sin=sin,
+                cos=cos,
+                position_ids=position_ids,
+                use_neox_rotary_style=False,
+            )
+        else:
+            query_states, key_states, _ = fused_rotary_position_embedding(
+                query_states,
+                key_states,
+                v=None,
+                sin=sin,
+                cos=cos,
+                position_ids=position_ids,
+                use_neox_rotary_style=False,
+            )
+    return query_states, key_states
+
+
+def rms_norm_fused(x_in, w, eps, use_fast_ln=False):
+    if use_fast_ln:
+        fast_ln = try_import("fast_ln")
+        return fast_ln.fast_rms_norm(x_in, w, eps)[0]
+    else:
+        fused_ln = try_import("fused_ln")
+        return fused_ln.fused_rms_norm(x_in, w, eps)[0]
+
+
+def fusion_rms_norm(hidden_states, weight, variance_epsilon, use_fast_ln=False):
+    if get_env_device() == "npu":
+        return core.eager._run_custom_op(
+            "rms_norm_npu", hidden_states, weight, variance_epsilon
+        )[0]
+    if get_env_device() == "mlu":
+        return core.eager._run_custom_op(
+            "rms_norm_mlu", hidden_states, weight, variance_epsilon
+        )[0]
+    elif get_env_device() == "gcu":
+        return core.eager._run_custom_op(
+            "rms_norm_gcu", hidden_states, weight, variance_epsilon
+        )[0]
+    elif get_env_device() == "intel_hpu":
+        return paddle.incubate.nn.functional.fused_rms_norm(
+            hidden_states, weight, None, variance_epsilon, hidden_states.dim() - 1
+        )[0]
+
+    return rms_norm_fused(hidden_states, weight, variance_epsilon, use_fast_ln)

paddlex/inference/models/common/vlm/generation/__init__.py

@@ -0,0 +1,34 @@
+# Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from .configuration_utils import GenerationConfig
+from .logits_process import (
+    ForcedBOSTokenLogitsProcessor,
+    ForcedEOSTokenLogitsProcessor,
+    HammingDiversityLogitsProcessor,
+    LogitsProcessor,
+    LogitsProcessorList,
+    MinLengthLogitsProcessor,
+    RepetitionPenaltyLogitsProcessor,
+    TopKProcess,
+    TopPProcess,
+)
+from .stopping_criteria import (
+    MaxLengthCriteria,
+    MaxTimeCriteria,
+    StoppingCriteria,
+    StoppingCriteriaList,
+    validate_stopping_criteria,
+)
+from .utils import GenerationMixin