chore: import upstream snapshot with attribution

tests/transformers/test_segment_parallel_utils.py

@@ -0,0 +1,172 @@
+# Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+# Copyright 2020 The HuggingFace Team. 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 unittest
+
+import numpy as np
+import paddle
+import paddle.distributed as dist
+from paddle.distributed.communication.group import _get_global_group
+
+from paddlenlp.transformers.segment_parallel_utils import (
+    ReshardLayer,
+    split_inputs_sequence_dim,
+)
+
+
+def prepare_data(batch_major=True, dim_size=4, batch_size=2, seq_len=2, num_head=2, h=4):
+    assert dim_size == 3 or dim_size == 4, f"dim_size should be 3 or 4, but {dim_size}"
+    batch_size = batch_size
+    seq_len = seq_len
+    h = h
+    num_head = num_head
+    sep = dist.get_world_size()
+    # assert sep == 2, f"sep should be 2, but {sep}"
+    num_elem = batch_size * seq_len // sep * num_head * h // num_head
+    local_rank = dist.get_rank()
+    input_data_list = []
+    split_tensor_list = []
+    if dim_size == 4:
+        shape = (
+            [batch_size, seq_len, num_head // sep, h // num_head]
+            if batch_major
+            else [seq_len, batch_size, num_head // sep, h // num_head]
+        )
+        split_axis = 1 if batch_major else 0
+        concat_axis = 2
+    else:
+        shape = [batch_size, seq_len // sep, h] if batch_major else [seq_len // sep, batch_size, h]
+        split_axis = 2
+        concat_axis = 1 if batch_major else 0
+    for rank in range(sep):
+        t = paddle.to_tensor(
+            np.reshape(np.arange(rank * num_elem, (rank + 1) * num_elem) + 1, shape), dtype=paddle.float32
+        )
+        input_data_list.append(t)
+        split_tensor_list.append(paddle.split(t, sep, axis=split_axis))
+    input_data = input_data_list[dist.get_rank()]
+    expected_output_data = [t[local_rank] for t in split_tensor_list]
+    expected_output_data = paddle.concat(expected_output_data, axis=concat_axis)
+    return input_data, expected_output_data
+
+
+def run_forward_backward(x, y_grad, split_axis=0, concat_axis=2):
+    x = x.detach()
+    x.stop_gradient = False
+    reshard_layer = ReshardLayer(sep_group=_get_global_group())
+    y = reshard_layer(
+        x,
+        split_axis=split_axis,
+        concat_axis=concat_axis,
+    )
+    paddle.autograd.backward([y], [y_grad], True)
+    return y, x.grad
+
+
+def should_test(sep_degree):
+    if sep_degree <= 1:
+        print(f"sep degree should greater than 1, but is {sep_degree}, skip this test")
+    return sep_degree > 1
+
+
+class TestReshardLayer(unittest.TestCase):
+    def setUp(self):
+        dist.init_parallel_env()
+
+    def test_split_inputs(self):
+        batch_size = 8
+        seq_len = 4096
+        sep = dist.get_world_size()
+        sep_rank = dist.get_rank()
+        if not should_test(sep):
+            return
+
+        inputs_ids = paddle.randint(low=0, high=65535, shape=(batch_size, seq_len))
+        labels = paddle.randint(low=0, high=2, shape=(batch_size, seq_len))
+        inputs = {"inputs_ids": inputs_ids, "labels": labels}
+
+        splited_inputs = split_inputs_sequence_dim(inputs, sep_rank, sep)
+        expected_local_inputs = {}
+        for k, v in inputs.items():
+            expected_local_inputs[k] = paddle.split(v, sep, axis=1)[sep_rank]
+            assert k in splited_inputs
+            np.testing.assert_equal(expected_local_inputs[k].numpy(), splited_inputs[k].numpy())
+
+    def test_reshard(self):
+        # [s / sep, b, h] -> [s, b, h / sep]
+        seq_len = 16
+        sep = dist.get_world_size()
+        if not should_test(sep):
+            return
+        assert seq_len % sep == 0, f"seq_len should be divisible by sep, seq_len:{seq_len}, sep:{sep}"
+
+        def check_equal(input_data, expected_output_data, split_axis=0, concat_axis=2):
+            bin_bout_output_grad = expected_output_data
+            bin_bout_output, bin_bout_input_grad = run_forward_backward(
+                input_data,
+                bin_bout_output_grad,
+                split_axis=split_axis,
+                concat_axis=concat_axis,
+            )
+            np.testing.assert_equal(bin_bout_output.numpy(), expected_output_data.numpy())
+            np.testing.assert_equal(bin_bout_input_grad.numpy(), input_data.numpy())
+
+        dim_size = 3
+        for batch_major in [True, False]:
+            for dim_size in [3, 4]:
+                for batch_size in [1, 2]:
+                    for seq_len in [4096, 4096 * 2, 4096 * 4]:
+                        for num_head in [32, 64]:
+                            for hidden_size in [num_head * 16, num_head * 32]:
+                                if dim_size == 3:
+                                    # check reshard before flash attn, shaped: [b, s, h]
+                                    input_data, expected_output_data = prepare_data(
+                                        batch_major=batch_major,
+                                        dim_size=dim_size,
+                                        batch_size=batch_size,
+                                        seq_len=seq_len,
+                                        num_head=num_head,
+                                        h=hidden_size,
+                                    )
+                                    split_axis_for_num_head = 2
+                                    concat_axis_for_seq = 1 if batch_major else 0
+                                    check_equal(
+                                        input_data,
+                                        expected_output_data,
+                                        split_axis=split_axis_for_num_head,
+                                        concat_axis=concat_axis_for_seq,
+                                    )
+                                elif dim_size == 4:
+                                    # check reshard after flash attn, shaped: [b, s, num_head, head_dim]
+                                    input_data, expected_output_data = prepare_data(
+                                        batch_major=batch_major,
+                                        dim_size=dim_size,
+                                        batch_size=batch_size,
+                                        seq_len=seq_len,
+                                        num_head=num_head,
+                                        h=hidden_size,
+                                    )
+                                    split_axis_for_seq = 1 if batch_major else 0
+                                    concat_axis_for_num_head = 2
+                                    check_equal(
+                                        input_data,
+                                        expected_output_data,
+                                        split_axis=split_axis_for_seq,
+                                        concat_axis=concat_axis_for_num_head,
+                                    )
+
+
+if __name__ == "__main__":
+    unittest.main()