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

# 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

environment_variables = {
    "NVIDIA_TF32_OVERRIDE": "0",
    "FLAGS_embedding_deterministic": "1",
    "FLAGS_cudnn_deterministic": "1",
}
for k, v in environment_variables.items():
    os.environ[k] = v
import unittest
from typing import Optional, Tuple

import paddle
import paddle.device
import paddle.nn as nn
import paddle.nn.functional as F
from paddle.distributed.fleet.recompute import recompute as original_recompute

from paddlenlp.trainer.training_args import TrainingArguments
from paddlenlp.transformers.refined_recompute import no_recompute as rr_no_recompute
from paddlenlp.transformers.refined_recompute import recompute as rr_recompute
from paddlenlp.utils.import_utils import is_paddle_cuda_available

ACT2FN = {
    "relu": F.relu,
    "gelu": F.gelu,
    "tanh": F.tanh,
    "sigmoid": F.sigmoid,
}
dtype = paddle.float16


class PyLayerMatmul(paddle.autograd.PyLayer):
    @staticmethod
    def forward(ctx, a, b):
        ctx.save_for_backward(a, b)
        return a @ b

    @staticmethod
    def backward(ctx, dy):
        a, b = ctx.saved_tensor()
        if hasattr(a, "main_grad"):
            a.main_grad.add_(paddle.ones_like(a.main_grad))
        if hasattr(b, "main_grad"):
            b.main_grad.add_(paddle.ones_like(b.main_grad))
        grad_a = paddle.matmul(dy, b, transpose_y=True)
        grad_b = paddle.matmul(a, dy, transpose_x=True)
        return grad_a, grad_b


pylayer_matmul = PyLayerMatmul.apply


class BertConfig:
    def __init__(
        self,
        vocab_size: int = 30522,
        hidden_size: int = 768,
        num_hidden_layers: int = 4,
        num_attention_heads: int = 12,
        intermediate_size: int = 3072,
        hidden_act: str = "gelu",
        hidden_dropout_prob: float = 0.0,
        attention_probs_dropout_prob: float = 0.0,
        max_position_embeddings: int = 1024,
        type_vocab_size: int = 2,
        initializer_range: float = 0.2,
        pad_token_id: int = 0,
        pool_act: str = "tanh",
        layer_norm_eps: float = 1e-12,
        output_attentions: bool = False,
        output_hidden_states: bool = False,
        num_labels=2,
        recompute=False,
        use_rr_recompute=False,
        recompute_use_reentrant=False,
        **kwargs
    ):
        self.pad_token_id = pad_token_id
        self.vocab_size = vocab_size
        self.hidden_size = hidden_size
        self.num_hidden_layers = num_hidden_layers
        self.num_attention_heads = num_attention_heads
        self.intermediate_size = intermediate_size
        self.hidden_act = hidden_act
        self.hidden_dropout_prob = hidden_dropout_prob
        self.attention_probs_dropout_prob = attention_probs_dropout_prob
        self.max_position_embeddings = max_position_embeddings
        self.type_vocab_size = type_vocab_size
        self.initializer_range = initializer_range
        self.pool_act = pool_act
        self.layer_norm_eps = layer_norm_eps
        self.output_attentions = output_attentions
        self.output_hidden_states = output_hidden_states
        self.num_labels = num_labels
        self.recompute = recompute
        self.use_rr_recompute = use_rr_recompute
        self.recompute_use_reentrant = recompute_use_reentrant


class BertEmbeddings(nn.Layer):
    """Construct the embeddings from word, position and token_type embeddings."""

    def __init__(self, config):
        super().__init__()
        self.word_embeddings = nn.Embedding(config.vocab_size, config.hidden_size, padding_idx=config.pad_token_id)
        self.position_embeddings = nn.Embedding(config.max_position_embeddings, config.hidden_size)
        self.token_type_embeddings = nn.Embedding(config.type_vocab_size, config.hidden_size)

        self.LayerNorm = nn.LayerNorm(config.hidden_size, epsilon=config.layer_norm_eps)
        self.dropout = nn.Dropout(config.hidden_dropout_prob)

        self.register_buffer(
            "position_ids", paddle.arange(config.max_position_embeddings, dtype="int64").reshape((1, -1))
        )

    def forward(
        self,
        input_ids: Optional[paddle.Tensor] = None,
        token_type_ids: Optional[paddle.Tensor] = None,
        position_ids: Optional[paddle.Tensor] = None,
    ) -> paddle.Tensor:
        input_shape = input_ids.shape
        seq_length = input_ids.shape[1]

        if position_ids is None:
            position_ids = self.position_ids[:, :seq_length]

        if token_type_ids is None:
            token_type_ids = paddle.zeros(input_shape, dtype=paddle.int64)

        inputs_embeds = self.word_embeddings(input_ids)
        token_type_embeddings = self.token_type_embeddings(token_type_ids)
        position_embeddings = self.position_embeddings(position_ids)
        embeddings = inputs_embeds + token_type_embeddings + position_embeddings
        embeddings = self.LayerNorm(embeddings)
        embeddings = self.dropout(embeddings)
        return embeddings


class BertSelfAttention(nn.Layer):
    def __init__(self, config):
        super().__init__()
        self.config = config
        self.num_attention_heads = config.num_attention_heads
        self.attention_head_size = config.hidden_size // config.num_attention_heads
        self.all_head_size = self.num_attention_heads * self.attention_head_size

        self.query = nn.Linear(config.hidden_size, self.all_head_size)
        self.key = nn.Linear(config.hidden_size, self.all_head_size)
        self.value = nn.Linear(config.hidden_size, self.all_head_size)

        self.dropout = nn.Dropout(config.attention_probs_dropout_prob)

    def forward(
        self,
        hidden_states: paddle.Tensor,
        attention_mask: Optional[paddle.Tensor] = None,
        output_attentions: Optional[bool] = False,
    ) -> Tuple[paddle.Tensor]:

        reshape_fn = lambda x: x.reshape([0, 0, -1, self.attention_head_size])
        # compute q,k,v
        query_layer = reshape_fn(self.query(hidden_states))
        key_layer = reshape_fn(self.key(hidden_states))
        value_layer = reshape_fn(self.value(hidden_states))

        context_layer = rr_no_recompute(
            F.scaled_dot_product_attention,
            query=query_layer,
            key=key_layer,
            value=value_layer,
            is_causal=True,
            enable=self.config.use_rr_recompute and self.config.recompute,
        )

        new_context_layer_shape = context_layer.shape[:-2] + [
            self.all_head_size,
        ]
        context_layer = context_layer.reshape(new_context_layer_shape)

        outputs = (context_layer, None) if output_attentions else (context_layer,)

        return outputs


class BertSelfOutput(nn.Layer):
    def __init__(self, config):
        super().__init__()
        self.config = config
        self.dense = nn.Linear(config.hidden_size, config.hidden_size)
        self.LayerNorm = nn.LayerNorm(config.hidden_size, epsilon=config.layer_norm_eps)
        self.dropout = nn.Dropout(config.hidden_dropout_prob)

    def forward(self, hidden_states: paddle.Tensor, input_tensor: paddle.Tensor) -> paddle.Tensor:
        hidden_states = rr_no_recompute(
            self.dense, hidden_states, enable=self.config.use_rr_recompute and self.config.recompute
        )
        hidden_states = self.dropout(hidden_states)

        hidden_states = rr_no_recompute(
            self.LayerNorm, hidden_states + input_tensor, enable=self.config.use_rr_recompute and self.config.recompute
        )
        return hidden_states


class BertAttention(nn.Layer):
    def __init__(self, config):
        super().__init__()
        self.config = config
        self.self = BertSelfAttention(config)
        self.output = BertSelfOutput(config)

    def forward(
        self,
        hidden_states: paddle.Tensor,
        attention_mask: Optional[paddle.Tensor] = None,
        output_attentions: Optional[bool] = False,
    ) -> Tuple[paddle.Tensor]:
        self_outputs = self.self(
            hidden_states,
            attention_mask,
            output_attentions,
        )
        attention_output = self.output(self_outputs[0], hidden_states)
        outputs = (attention_output,) + self_outputs[1:]  # add attentions if we output them
        return outputs


class BertIntermediate(nn.Layer):
    def __init__(self, config):
        super().__init__()
        self.config = config
        self.dense = nn.Linear(config.hidden_size, config.intermediate_size, bias_attr=False)
        self.dense.weight.main_grad = paddle.zeros_like(self.dense.weight).cast("float32")
        if isinstance(config.hidden_act, str):
            self.intermediate_act_fn = ACT2FN[config.hidden_act]
        else:
            self.intermediate_act_fn = config.hidden_act

    def forward(self, hidden_states: paddle.Tensor) -> paddle.Tensor:
        def pylayer_dense(hidden_states):
            return pylayer_matmul(hidden_states, self.dense.weight)

        hidden_states = rr_no_recompute(
            pylayer_dense, hidden_states, enable=self.config.use_rr_recompute and self.config.recompute
        )
        hidden_states = self.intermediate_act_fn(hidden_states)

        return hidden_states


class BertOutput(nn.Layer):
    def __init__(self, config):
        super().__init__()
        self.config = config
        self.dense = nn.Linear(config.intermediate_size, config.hidden_size)
        self.LayerNorm = nn.LayerNorm(config.hidden_size, epsilon=config.layer_norm_eps)
        self.dropout = nn.Dropout(config.hidden_dropout_prob)

    def forward(self, hidden_states: paddle.Tensor, input_tensor: paddle.Tensor) -> paddle.Tensor:
        def custom_dense(hidden_states, weight, bias=None):
            return F.linear(hidden_states, weight, bias)

        bias = self.dense.bias * 1.1
        hidden_states = rr_no_recompute(
            custom_dense,
            hidden_states,
            weight=self.dense.weight,
            bias=bias,
            enable=self.config.use_rr_recompute and self.config.recompute,
            keys_ignore_to_save=["bias"],
        )
        hidden_states = self.dropout(hidden_states)
        hidden_states = self.LayerNorm(hidden_states + input_tensor)
        return hidden_states


class BertLayer(nn.Layer):
    def __init__(self, config):
        super().__init__()
        self.config = config
        self.seq_len_dim = 1
        self.attention = BertAttention(config)
        self.intermediate = BertIntermediate(config)
        self.output = BertOutput(config)

    def forward(
        self,
        hidden_states: paddle.Tensor,
        attention_mask: Optional[paddle.Tensor] = None,
        output_attentions: Optional[bool] = False,
    ) -> Tuple[paddle.Tensor]:
        # self attn
        self_attention_outputs = self.attention(
            hidden_states,
            attention_mask,
            output_attentions=output_attentions,
        )
        attention_output = self_attention_outputs[0]

        outputs = self_attention_outputs[1:]  # add self attentions if we output attention weights

        # ffn
        intermediate_output = self.intermediate(attention_output)
        layer_output = self.output(intermediate_output, attention_output)

        outputs = (layer_output,) + outputs

        return outputs


class BertEncoder(nn.Layer):
    def __init__(self, config):
        super().__init__()
        self.config = config
        self.layer = nn.LayerList([BertLayer(config) for _ in range(config.num_hidden_layers)])

    def forward(
        self,
        hidden_states: paddle.Tensor,
        attention_mask: Optional[paddle.Tensor] = None,
        output_attentions: Optional[bool] = False,
        output_hidden_states: Optional[bool] = False,
    ) -> Tuple[paddle.Tensor]:
        all_hidden_states = () if output_hidden_states else None
        all_self_attentions = () if output_attentions else None

        for layer_module in self.layer:
            # add hidden_states
            if output_hidden_states:
                all_hidden_states = all_hidden_states + (hidden_states,)

            if self.training and self.config.recompute:
                recompute_function = rr_recompute if self.config.use_rr_recompute else original_recompute
                layer_outputs = recompute_function(
                    layer_module,
                    hidden_states,
                    attention_mask,
                    output_attentions,
                    use_reentrant=self.config.recompute_use_reentrant,
                )
            else:
                layer_outputs = layer_module(
                    hidden_states,
                    attention_mask,
                    output_attentions,
                )
            hidden_states = layer_outputs[0]

            # add self attn
            if output_attentions:
                all_self_attentions = all_self_attentions + (layer_outputs[1],)

        if output_hidden_states:
            all_hidden_states = all_hidden_states + (hidden_states,)

        return tuple(
            v
            for v in [
                hidden_states,
                all_hidden_states,
                all_self_attentions,
            ]
            if v is not None
        )


class BertPreTrainedModel(nn.Layer):
    def _init_weights(self, module):
        """Initialize the weights"""
        pass


class BertModel(BertPreTrainedModel):
    def __init__(self, config):
        super().__init__()
        self.config = config
        self.embeddings = BertEmbeddings(config)
        self.encoder = BertEncoder(config)

    def forward(
        self,
        input_ids: Optional[paddle.Tensor] = None,
        attention_mask: Optional[paddle.Tensor] = None,
        token_type_ids: Optional[paddle.Tensor] = None,
        position_ids: Optional[paddle.Tensor] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
    ) -> Tuple[paddle.Tensor]:
        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
        output_hidden_states = (
            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
        )

        if token_type_ids is None:
            token_type_ids = paddle.zeros(input_ids.shape, dtype=paddle.int64)

        embedding_output = self.embeddings(
            input_ids=input_ids,
            position_ids=position_ids,
            token_type_ids=token_type_ids,
        )
        encoder_outputs = self.encoder(
            embedding_output,
            attention_mask=attention_mask,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
        )
        return encoder_outputs


class BertRefinedRecomputeTest(unittest.TestCase):
    def no_pp_fwd_bwd(
        self,
        recompute=False,
        use_rr_recompute=False,
        recompute_use_reentrant=False,
        num_hidden_layers=4,
        shape=[2, 64],
    ):
        paddle.set_default_dtype(dtype)
        paddle.seed(42)
        config = BertConfig(
            num_hidden_layers=num_hidden_layers,
            recompute=recompute,
            use_rr_recompute=use_rr_recompute,
            recompute_use_reentrant=recompute_use_reentrant,
        )
        model = BertModel(config)
        model.train()
        input_ids = paddle.randint(10, config.vocab_size, shape=shape)
        gpu_mem_used_before = paddle.device.cuda.memory_allocated()
        outputs = model(input_ids=input_ids)[0]
        gpu_mem_used_after = paddle.device.cuda.memory_allocated()
        outputs.sum().backward()

        # div = 1024**3 # GB
        div = 1  # KB
        return (
            model,
            round((gpu_mem_used_after - gpu_mem_used_before) / div, 2),
            round(paddle.device.cuda.max_memory_allocated() / div, 2),
        )

    @unittest.skipIf(not is_paddle_cuda_available(), "refined-recompute only support on gpu")
    def test_refined_recompute(self):
        raw_dtype = paddle.get_default_dtype()

        model1, mem_usage_forward1, max_mem_usage_forward1 = self.no_pp_fwd_bwd(
            recompute=True, use_rr_recompute=False
        )  # with recompute
        model2, mem_usage_forward2, max_mem_usage_forward2 = self.no_pp_fwd_bwd(
            recompute=True, use_rr_recompute=True
        )  # with rr recompute
        model3, mem_usage_forward3, max_mem_usage_forward3 = self.no_pp_fwd_bwd(
            recompute=False, use_rr_recompute=False
        )  # without recompute

        name_list = [n for n, _ in model1.named_parameters()]

        for param1, param2, name in zip(model1.parameters(), model3.parameters(), name_list):
            # test main grad
            if "intermediate.dense.weight" in name:
                self.assertTrue(param1.main_grad.sum().item() > 0)
                self.assertTrue(param2.main_grad.sum().item() > 0)
            self.assertTrue(paddle.equal_all(param1.grad.cast("float32"), param2.grad.cast("float32")))

        for param1, param2, name in zip(model2.parameters(), model3.parameters(), name_list):
            # test main grad
            if "intermediate.dense.weight" in name:
                self.assertTrue(param1.main_grad.sum().item() > 0)
                self.assertTrue(param2.main_grad.sum().item() > 0)
            self.assertTrue(paddle.equal_all(param1.grad.cast("float32"), param2.grad.cast("float32")))

        # self.assertTrue(mem_usage_forward1 < mem_usage_forward2 < mem_usage_forward3)
        # self.assertTrue(max_mem_usage_forward1 < max_mem_usage_forward2 < max_mem_usage_forward3)

        del model1, model2, model3
        paddle.device.cuda.empty_cache()
        paddle.set_default_dtype(raw_dtype)

    def pp_fwd_bwd(
        self,
        recompute=False,
        use_rr_recompute=False,
        recompute_use_reentrant=False,
        num_iter=4,
        shape=[2, 64],
    ):
        paddle.set_default_dtype(dtype)
        paddle.seed(42)
        config = BertConfig(
            num_hidden_layers=1,
            recompute=recompute,
            use_rr_recompute=use_rr_recompute,
            recompute_use_reentrant=recompute_use_reentrant,
        )
        layer = BertLayer(config)
        layer.train()

        x = paddle.randn([*shape, config.hidden_size])
        x.stop_gradient = False
        x_copy = x

        if layer.training and config.recompute:
            recompute_function = rr_recompute if config.use_rr_recompute else original_recompute
            for _ in range(num_iter):
                x = recompute_function(layer, x, use_reentrant=config.recompute_use_reentrant)[0]
        else:
            for _ in range(num_iter):
                x = layer(x)[0]

        x.sum().backward()

        return x_copy.grad, layer

    @unittest.skipIf(not is_paddle_cuda_available(), "refined-recompute-pp only support on gpu")
    def test_refined_recompute_pp(self):
        paddle.set_device("gpu")
        raw_dtype = paddle.get_default_dtype()
        grad1, layer1 = self.pp_fwd_bwd(recompute=True, use_rr_recompute=False)
        grad2, layer2 = self.pp_fwd_bwd(recompute=True, use_rr_recompute=True)
        grad3, layer3 = self.pp_fwd_bwd(recompute=False, use_rr_recompute=False)

        name_list = [n for n, _ in layer1.named_parameters()]

        for param1, param2, name in zip(layer1.parameters(), layer3.parameters(), name_list):
            # test main grad
            if "intermediate.dense.weight" in name:
                self.assertTrue(param1.main_grad.sum().item() > 0)
                self.assertTrue(param2.main_grad.sum().item() > 0)
            self.assertTrue(paddle.equal_all(param1.grad.cast("float32"), param2.grad.cast("float32")))

        self.assertTrue(paddle.equal_all(grad1.cast("float32"), grad3.cast("float32")))
        for param1, param2, name in zip(layer2.parameters(), layer3.parameters(), name_list):
            # test main grad
            if "intermediate.dense.weight" in name:
                self.assertTrue(param1.main_grad.sum().item() > 0)
                self.assertTrue(param2.main_grad.sum().item() > 0)
            self.assertTrue(paddle.equal_all(param1.grad.cast("float32"), param2.grad.cast("float32")))

        self.assertTrue(paddle.equal_all(grad2.cast("float32"), grad3.cast("float32")))

        del grad1, grad2, grad3
        del layer1, layer2, layer3
        paddle.device.cuda.empty_cache()
        paddle.set_default_dtype(raw_dtype)


class TestRefinedRecomputeModel(unittest.TestCase):
    def setUp(self):
        self.args = TrainingArguments(
            output_dir="./",
            do_train=True,
            max_steps=100,
            tensor_parallel_degree=1,
            pipeline_parallel_degree=1,
            refined_recompute="attention_column_ln:1,attention_row_ln:2,flash_attn:-1,mlp_column_ln:2,mlp_row_ln:-1",
        )

    @unittest.skipIf(not is_paddle_cuda_available(), "refined-recompute-pp only support on gpu")
    def test_llama_refined_recompute(self):
        paddle.set_device("gpu")
        from paddlenlp.transformers.llama import LlamaConfig, LlamaModel

        llama_model = "__internal_testing__/tiny-random-llama"
        config = LlamaConfig.from_pretrained(llama_model)
        config.recompute = True
        config.recompute_granularity = "full"
        config.recompute_use_reentrant = False
        config.sequence_parallel = False
        config.use_flash_attention = True
        config.refined_recompute = self.args.refined_recompute
        model = LlamaModel.from_config(config=config, dtype="bfloat16")
        input_ids = paddle.randint(0, 100, shape=[1, 1024], dtype="int64")
        output = model(input_ids)
        output[0].mean().backward()

    @unittest.skipIf(not is_paddle_cuda_available(), "refined-recompute-pp only support on gpu")
    def test_qwen_refined_recompute(self):
        paddle.set_device("gpu")
        from paddlenlp.transformers.qwen import QWenConfig, QWenModel

        llama_model = "__internal_testing__/tiny-random-qwen"
        config = QWenConfig.from_pretrained(llama_model)
        config.recompute = True
        config.recompute_granularity = "full"
        config.recompute_use_reentrant = False
        config.sequence_parallel = False
        config.use_flash_attention = True
        config.refined_recompute = self.args.refined_recompute
        config.seq_length = 1024
        model = QWenModel.from_config(config=config, dtype="bfloat16")
        input_ids = paddle.randint(0, 100, shape=[1, 1024], dtype="int64")
        output = model(input_ids)
        output[0].mean().backward()

    @unittest.skipIf(not is_paddle_cuda_available(), "refined-recompute-pp only support on gpu")
    def test_qwen2_refined_recompute(self):
        paddle.set_device("gpu")
        from paddlenlp.transformers.qwen2 import Qwen2Config, Qwen2Model

        llama_model = "__internal_testing__/tiny-random-qwen2"
        config = Qwen2Config.from_pretrained(llama_model)
        config.recompute = True
        config.recompute_granularity = "full"
        config.recompute_use_reentrant = False
        config.sequence_parallel = False
        config.use_flash_attention = True
        config.refined_recompute = self.args.refined_recompute
        model = Qwen2Model.from_config(config=config, dtype="bfloat16")
        input_ids = paddle.randint(0, 100, shape=[1, 1024], dtype="int64")
        output = model(input_ids)
        output[0].mean().backward()