chore: import upstream snapshot with attribution

test/auto_parallel/semi_auto_parallel_subgraph_embedding.py

@@ -0,0 +1,138 @@
+# 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 os
+
+import numpy as np
+
+import paddle
+import paddle.distributed as dist
+
+
+class Layer(paddle.nn.Layer):
+    def __init__(self, vocab_size, hidden_size):
+        super().__init__()
+        self.embedding = paddle.nn.Embedding(vocab_size, hidden_size)
+
+    def forward(self, x):
+        return self.embedding(x)
+
+
+class TestEmbeddingSubgraphSemiAutoParallel:
+    def __init__(self):
+        self._dtype = os.getenv("dtype")
+        self._backend = os.getenv("backend")
+        self._seed = eval(os.getenv("seed"))
+        self._mesh = dist.ProcessMesh([0, 1], dim_names=["x"])
+        self._batch_size = 17
+        self._seq_length = 23
+        self._vocab_size = 48
+        self._hidden_size = 16
+
+    def test_dp(self):
+        paddle.seed(self._seed)
+        np.random.seed(self._seed)
+        self._input = np.random.randint(
+            0, self._vocab_size, size=(self._batch_size, self._seq_length)
+        )
+        x = paddle.to_tensor(self._input)
+        paddle.seed(self._seed)
+        np.random.seed(self._seed)
+        layer = Layer(self._vocab_size, self._hidden_size)
+        desired_out = layer(x)
+        desired_out.backward()
+        desired_grad = layer.embedding.weight.grad
+
+        paddle.seed(self._seed)
+        np.random.seed(self._seed)
+        dist_x = dist.shard_tensor(x, self._mesh, placements=(dist.Shard(0),))
+        layer = Layer(self._vocab_size, self._hidden_size)
+        actual_out = layer(x)
+        actual_out.backward()
+        actual_grad = layer.embedding.weight.grad
+
+        np.testing.assert_allclose(actual_out, desired_out, rtol=1e-6, atol=0)
+        np.testing.assert_allclose(actual_grad, desired_grad, rtol=1e-6, atol=0)
+        # The threshold setting refers to Megatron-LM
+        assert (
+            np.max(np.abs(actual_out.numpy() - desired_out.numpy())) < 1.0e-12
+        ), (
+            f'embedding dp forward error. actual: {actual_out}, desired: {desired_out}'
+        )
+        assert (
+            np.max(np.abs(actual_grad.numpy() - desired_grad.numpy())) < 1.0e-12
+        ), (
+            f'embedding dp backward error. actual: {actual_out}, desired: {desired_out}'
+        )
+
+    def test_mp(self):
+        paddle.seed(self._seed)
+        np.random.seed(self._seed)
+        self._input = np.random.randint(
+            0, self._vocab_size, size=(self._batch_size, self._seq_length)
+        )
+
+        x = paddle.to_tensor(self._input)
+        paddle.seed(self._seed)
+        np.random.seed(self._seed)
+        layer = Layer(self._vocab_size, self._hidden_size)
+        desired_out = layer(x)
+        desired_out.backward()
+        desired_grad = layer.embedding.weight.grad
+
+        paddle.seed(self._seed)
+        np.random.seed(self._seed)
+        dist_x = dist.shard_tensor(
+            x, self._mesh, placements=(dist.Replicate(),)
+        )
+
+        def shard_fn(layer_name, layer, process_mesh):
+            if 'embedding' in layer_name:
+                layer.weight = dist.shard_tensor(
+                    layer.weight, process_mesh, (dist.Shard(1),)
+                )
+
+        layer = dist.shard_layer(
+            Layer(self._vocab_size, self._hidden_size), self._mesh, shard_fn
+        )
+        actual_out = layer(x)
+        actual_out.backward()
+        actual_grad = layer.embedding.weight.grad
+
+        # The threshold setting refers to Megatron-LM
+        assert (
+            np.max(np.abs(actual_out.numpy() - desired_out.numpy())) < 1.0e-12
+        ), (
+            f'embedding mp forward error. actual: {actual_out}, desired: {desired_out}'
+        )
+        assert (
+            np.max(np.abs(actual_grad.numpy() - desired_grad.numpy())) < 1.0e-12
+        ), (
+            f'embedding mp backward error. actual: {actual_out}, desired: {desired_out}'
+        )
+
+    def run_test_case(self):
+        if self._backend == "cpu":
+            paddle.set_device("cpu")
+        elif self._backend == "gpu":
+            paddle.set_device("gpu:" + str(dist.get_rank()))
+        else:
+            raise ValueError("Only support cpu or gpu backend.")
+
+        self.test_dp()
+        self.test_mp()
+
+
+if __name__ == '__main__':
+    TestEmbeddingSubgraphSemiAutoParallel().run_test_case()