paddlepaddle--paddlenlp/tests/transformers/test_conversion_common.py

# Copyright (c) 2023 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 copy
import glob
import os
import tempfile
import unittest

import paddle

input_ids = paddle.to_tensor([[0, 345, 232, 328, 740, 140, 1695, 69, 6078, 1588, 2]])


def prepare_default_config(config):
    config = copy.deepcopy(config)
    config.hidden_size = 512
    config.num_layers = 2
    config.num_hidden_layers = 2
    config.num_attention_heads = 16
    config.num_key_value_heads = 16
    config.intermediate_size = config.hidden_size
    config.word_embed_proj_dim = 512
    return config


def prepare_split_config(config):
    config = prepare_default_config(config)
    config = copy.deepcopy(config)
    config.fuse_attention_qkv = False
    config.fuse_attention_ffn = False
    return config


def prepare_fuse_config(config):
    config = prepare_default_config(config)
    config = copy.deepcopy(config)
    config.fuse_attention_qkv = True
    config.fuse_attention_ffn = True
    return config


def common_test_load(model_class, model_first, config_second, tempdir):
    model_first.eval()
    with paddle.no_grad():
        first = model_first(input_ids)[0]

    model_second = model_class.from_pretrained(tempdir, config=config_second)
    model_second.eval()
    with paddle.no_grad():
        second = model_second(input_ids)[0]

    assert paddle.allclose(paddle.mean(first), paddle.mean(second), atol=1e-5)
    # assert paddle.allclose(first, second, atol=1e-4)

    files = glob.glob(tempdir + "/*")
    for f in files:
        os.remove(f)


def common_test_save_and_load(config_first, config_second, model_class):
    model_first = model_class.from_config(config_first)

    with tempfile.TemporaryDirectory() as tempdir:
        # test load pdparams: model.pdparams
        model_first.save_pretrained(save_dir=tempdir)
        common_test_load(model_class, model_first, config_second, tempdir)

        # test load shard pdparams: model-001-0f-008.pdparams
        model_first.save_pretrained(tempdir, max_shard_size="5MB")
        common_test_load(model_class, model_first, config_second, tempdir)

        # test save safetensors: model.safetensors
        model_first.save_pretrained(tempdir, safe_serialization=True)
        common_test_load(model_class, model_first, config_second, tempdir)

        # test load shard safetensors: model-001-0f-008.safetensors
        model_first.save_pretrained(tempdir, max_shard_size="5MB", safe_serialization=True)
        common_test_load(model_class, model_first, config_second, tempdir)


def _test_split_to_fuse(config_class, model_class):
    config = config_class()

    config_split = prepare_split_config(config)
    config_fuse = prepare_fuse_config(config)

    # Test from splitted weights to fused weight
    common_test_save_and_load(config_split, config_fuse, model_class)


def _test_fuse_to_split(config_class, model_class):
    config = config_class()

    config_split = prepare_split_config(config)
    config_fuse = prepare_fuse_config(config)

    # Test from fused weight to splitted weights
    common_test_save_and_load(config_fuse, config_split, model_class)


def _test_fast_ffn():
    from functools import partial

    import paddle
    from paddle import nn

    from paddlenlp.transformers import PretrainedModel
    from paddlenlp.transformers.configuration_utils import PretrainedConfig

    class TestConfig(PretrainedConfig):
        def __init__(self, fast_ffn_state=False, convert_fast_ffn=False):
            self.fast_ffn_state = fast_ffn_state
            self.convert_fast_ffn = convert_fast_ffn
            super().__init__()

    class TestMLP(nn.Layer):
        def __init__(self, config):
            super().__init__()
            self.config = config
            self.hidden_size = config.hidden_size
            self.gate_up_fused_proj = nn.Linear(self.hidden_size, self.hidden_size * 2, bias_attr=True)

        def forward(self, hidden_state):
            hidden_state = self.gate_up_fused_proj(hidden_state)
            if self.config.use_fast_ffn:
                x, y = paddle.chunk(hidden_state, chunks=2, axis=-1)
            else:
                x, y = hidden_state[..., ::2], hidden_state[..., 1::2]

            return nn.functional.silu(x) * y

    class TestPretrainedModel(PretrainedModel):
        config_class = TestConfig

        @classmethod
        def _get_fuse_or_split_param_mappings(cls, config: TestConfig, is_fuse=False):

            #  user defined function to get convert param mappings
            def convert_fast_ffn_fn(fuse_params, convert_fast_ffn=False):
                import numpy as np

                concat_fn = np.concatenate
                if isinstance(fuse_params[0], paddle.Tensor):
                    concat_fn = paddle.concat

                if convert_fast_ffn:
                    # fast_ffn
                    first = fuse_params[0][..., ::2]
                    second = fuse_params[0][..., 1::2]
                    return concat_fn([first, second], axis=-1)

            fn = convert_fast_ffn_fn

            convert_fast_ffn_keys = (
                "layers.0.gate_up_fused_proj.weight",
                "layers.0.gate_up_fused_proj.weight",
            )
            convert_fast_ffn_bias_keys = (
                "layers.0.gate_up_fused_proj.bias",
                "layers.0.gate_up_fused_proj.bias",
            )
            fast_ffn_state = getattr(config, "fast_ffn_state", False)
            convert_fast_ffn = getattr(config, "convert_fast_ffn", False)
            convert_fast_ffn &= not fast_ffn_state

            final_actions = {}
            if is_fuse:
                # for_get_fuse_or_split_param_mappings, is_fuse have two conditions, true and false,
                # to fit partial fuse or split conditions, is_fuse will called twice(True and False).
                # thus, for this func, we only use one condition.

                # use_fast_ffn only in one condition
                # convert when use_fast_ffn is False
                if convert_fast_ffn:
                    for i in range(config.num_hidden_layers):
                        for keys in [convert_fast_ffn_keys, convert_fast_ffn_bias_keys]:
                            keys = tuple([key.replace("layers.0.", f"layers.{i}.") for key in keys])
                            final_actions[keys] = partial(fn, convert_fast_ffn=convert_fast_ffn)
            return final_actions

        def _init_weights(self, layer):
            if isinstance(layer, (nn.Linear, nn.Embedding)):
                if isinstance(layer.weight, paddle.Tensor):
                    layer.weight.set_value(paddle.tensor.normal(mean=0.0, std=1.0, shape=layer.weight.shape))
                if hasattr(layer, "bias") and isinstance(layer.bias, paddle.Tensor):
                    layer.bias.set_value(paddle.tensor.normal(mean=0.0, std=1.0, shape=layer.bias.shape))

    class TestModel(TestPretrainedModel):
        def __init__(self, config):
            super().__init__(config)
            self.layers = nn.LayerList([TestMLP(config=config) for i in range(config.num_hidden_layers)])

        def forward(self, hidden_state):
            for idx, (decoder_layer) in enumerate(self.layers):
                hidden_state = decoder_layer(hidden_state)
            return hidden_state

    class TestForCausalLM(TestPretrainedModel):
        def __init__(self, config):
            super().__init__(config)
            self.config = config
            self.embedding_layer = nn.Embedding(65535, self.config.hidden_size)
            self.test = TestModel(config=config)

        def forward(self, input_ids):
            hidden_state = self.embedding_layer(input_ids)
            return self.test(hidden_state)

    config = TestConfig()
    config = prepare_default_config(config)
    config_no_fast_ffn = copy.deepcopy(config)
    config_fast_ffn = copy.deepcopy(config)

    config_no_fast_ffn.use_fast_ffn = False

    config_fast_ffn.use_fast_ffn = True
    config_fast_ffn.fast_ffn_state = False
    config_fast_ffn.convert_fast_ffn = True

    common_test_save_and_load(config_no_fast_ffn, config_fast_ffn, TestForCausalLM)


from paddlenlp.transformers import (
    GPTConfig,
    GPTForCausalLM,
    LlamaConfig,
    LlamaForCausalLM,
    OPTConfig,
    OPTForCausalLM,
)


class TestFuseOrSplit(unittest.TestCase):
    def test_model_split_to_fuse(self):
        _test_split_to_fuse(LlamaConfig, LlamaForCausalLM)
        _test_split_to_fuse(GPTConfig, GPTForCausalLM)
        _test_split_to_fuse(OPTConfig, OPTForCausalLM)

    def test_model_fuse_to_split(self):
        _test_fuse_to_split(LlamaConfig, LlamaForCausalLM)
        _test_fuse_to_split(GPTConfig, GPTForCausalLM)
        _test_fuse_to_split(OPTConfig, OPTForCausalLM)

    def test_model_convert_fast_ffn(self):
        _test_fast_ffn()