chore: import upstream snapshot with attribution

test/auto_parallel/tp_conv.py

@@ -0,0 +1,520 @@
+# Copyright (c) 2025 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 random
+
+import numpy as np
+
+import paddle
+import paddle.distributed as dist
+from paddle import nn
+
+dist.init_parallel_env()
+
+
+class SimpleConvNet(nn.Layer):
+    def __init__(
+        self,
+        in_channel,
+        out_channel,
+        kernel_size,
+        padding,
+        bias_attr,
+        stride,
+        data_format="NCHW",
+    ):
+        super().__init__()
+        self.conv1 = nn.Conv2D(
+            in_channel,
+            out_channel,
+            kernel_size=kernel_size,
+            padding=padding,
+            data_format=data_format,
+            bias_attr=bias_attr,
+            stride=stride,
+        )
+        self.relu = nn.ReLU()
+
+    def forward(self, x):
+        x = self.conv1(x)
+        return self.relu(x)
+
+
+class TestTPConv:
+    def __init__(self):
+        self.rank = dist.get_rank()
+        self.world_size = dist.get_world_size()
+        self._tp_mesh = dist.ProcessMesh(
+            list(range(self.world_size)), dim_names=["tp"]
+        )
+
+    def set_seed(self, seed):
+        paddle.seed(seed)
+        np.random.seed(seed)
+        random.seed(seed)
+
+    def _test_intermediate(
+        self,
+        N,
+        C,
+        H,
+        W,
+        kernel_size,
+        padding,
+        bias_attr,
+        mesh,
+        test_name="conv_test",
+        dtype_str="float32",
+        data_format="NCHW",
+        stride=1,
+    ):
+        self.set_seed(2025)
+
+        dist.auto_parallel.set_mesh(mesh)
+
+        conv_layer = SimpleConvNet(
+            C,
+            C,
+            kernel_size=kernel_size,
+            padding=padding,
+            data_format=data_format,
+            bias_attr=bias_attr,
+            stride=stride,
+        )
+        original_weight = conv_layer.conv1.weight
+        conv_layer.conv1.weight = original_weight
+
+        if data_format == "NCHW":
+            input_tensor = paddle.randn([N, C, H, W])
+            shard_axis_input = 3
+        else:
+            input_tensor = paddle.randn([N, H, W, C])
+            shard_axis_input = 2
+
+        input_placements = [
+            dist.Replicate() for _ in range(len(mesh.dim_names))
+        ]
+        mp_dim_index = mesh.dim_names.index("mp")
+        input_placements[mp_dim_index] = dist.Shard(shard_axis_input)
+
+        sharded_input = dist.shard_tensor(input_tensor, mesh, input_placements)
+        output_ref = conv_layer(input_tensor)
+        loss_ref = output_ref.mean()
+        loss_ref.backward()
+        weight_grad_ref = conv_layer.conv1.weight.grad.clone()
+
+        if (
+            conv_layer.conv1.bias is not None
+            and conv_layer.conv1.bias.grad is not None
+        ):
+            bias_grad_ref = conv_layer.conv1.bias.grad.clone()
+
+        conv_layer.clear_gradients()
+        conv_layer.conv1.weight = original_weight
+
+        opt = paddle.optimizer.AdamW(
+            learning_rate=0.001, parameters=conv_layer.parameters()
+        )
+        mp_config = {"parallelize_plan": {"conv1": dist.ConvParallel()}}
+
+        parallel_config = {
+            "mp_config": mp_config,
+        }
+
+        dist_model, dist_opt = dist.parallelize(
+            conv_layer, opt, config=parallel_config
+        )
+
+        output_intermediate = dist_model(input_tensor)
+        loss_intermediate = paddle.mean(output_intermediate)
+        loss_intermediate.backward()
+        weight_grad_intermediate = dist_model.conv1.weight.grad.clone()
+
+        if (
+            dist_model.conv1.bias is not None
+            and dist_model.conv1.bias.grad is not None
+        ):
+            bias_grad_intermediate = dist_model.conv1.bias.grad.clone()
+
+        def compare_tensors(name, tensor1, tensor2):
+            np.testing.assert_allclose(
+                tensor1.numpy(), tensor2.numpy(), rtol=1e-8, atol=1e-8
+            )
+
+        def compare_grads(name, grad1, grad2):
+            np.testing.assert_allclose(
+                grad1.numpy(), grad2.numpy(), rtol=1e-6, atol=1e-6
+            )
+
+        if data_format == "NCHW":
+            w_size = output_ref.shape[-1]
+        else:
+            w_size = output_ref.shape[-2]
+
+        if dist.get_rank() == 0:
+            start_index = 0
+            end_index = w_size // 2
+        else:
+            start_index = w_size // 2
+            end_index = w_size
+
+        if data_format == "NCHW":
+            compare_tensors(
+                "output",
+                output_ref[:, :, :, start_index:end_index],
+                output_intermediate._local_value(),
+            )
+        else:
+            compare_tensors(
+                "output",
+                output_ref[:, :, start_index:end_index, :],
+                output_intermediate._local_value(),
+            )
+
+        compare_grads("w", weight_grad_ref, weight_grad_intermediate)
+
+        if (
+            conv_layer.conv1.bias is not None
+            and conv_layer.conv1.bias.grad is not None
+        ):
+            compare_grads("b", bias_grad_ref, bias_grad_intermediate)
+
+    def _test_conv_case(
+        self,
+        N,
+        C,
+        H,
+        W,
+        kernel_size,
+        padding,
+        bias_attr,
+        mesh,
+        test_name="conv_test",
+        dtype_str="float32",
+        data_format="NCHW",
+        stride=1,
+    ):
+        self.set_seed(2025)
+
+        conv_layer = nn.Conv2D(
+            C,
+            C,
+            kernel_size=kernel_size,
+            padding=padding,
+            bias_attr=bias_attr,
+            data_format=data_format,
+            stride=stride,
+        )
+        original_weight = conv_layer.weight
+        conv_layer.weight = original_weight
+
+        if data_format == "NCHW":
+            input_tensor = paddle.randn([N, C, H, W])
+            shard_axis_input = 3
+        else:
+            input_tensor = paddle.randn([N, H, W, C])
+            shard_axis_input = 2
+
+        output_ref = conv_layer(input_tensor)
+        loss_ref = output_ref.mean()
+        loss_ref.backward()
+        weight_grad_ref = conv_layer.weight.grad.clone()
+
+        if conv_layer.bias is not None and conv_layer.bias.grad is not None:
+            bias_grad_ref = conv_layer.bias.grad.clone()
+
+        conv_layer.clear_gradients()
+        conv_layer.weight = original_weight
+
+        rank = dist.get_rank()
+        input_placements = [
+            dist.Replicate() for _ in range(len(mesh.dim_names))
+        ]
+        mp_dim_index = mesh.dim_names.index("mp")
+        input_placements[mp_dim_index] = dist.Shard(shard_axis_input)
+
+        sharded_input = dist.shard_tensor(input_tensor, mesh, input_placements)
+
+        output_sharded = conv_layer(sharded_input)
+        loss_sharded = paddle.mean(output_sharded)
+        loss_sharded.backward()
+
+        weight_grad_shard = conv_layer.weight.grad.clone()
+        if conv_layer.bias is not None and conv_layer.bias.grad is not None:
+            bias_grad_shard = conv_layer.bias.grad.clone()
+
+        def compare_grads(name, grad1, grad2):
+            np.testing.assert_allclose(
+                grad1.numpy(), grad2.numpy(), rtol=1e-6, atol=1e-7
+            )
+
+        def compare_tensors(name, tensor1, tensor2):
+            np.testing.assert_allclose(
+                tensor1.numpy(), tensor2.numpy(), rtol=1e-8, atol=1e-8
+            )
+
+        if data_format == "NCHW":
+            w_size = output_ref.shape[-1]
+        else:
+            w_size = output_ref.shape[-2]
+
+        if dist.get_rank() == 0:
+            start_index = 0
+            end_index = w_size // 2
+        else:
+            start_index = w_size // 2
+            end_index = w_size
+
+        if data_format == "NCHW":
+            compare_tensors(
+                "output",
+                output_ref[:, :, :, start_index:end_index],
+                output_sharded._local_value(),
+            )
+        else:
+            compare_tensors(
+                "output",
+                output_ref[:, :, start_index:end_index, :],
+                output_sharded._local_value(),
+            )
+
+        compare_grads("w", weight_grad_ref, weight_grad_shard)
+
+        if conv_layer.bias is not None and conv_layer.bias.grad is not None:
+            compare_grads("b", bias_grad_ref, bias_grad_shard)
+
+    def run_test_cases(self):
+        mesh1 = dist.ProcessMesh([0, 1], dim_names=['mp'])
+
+        # ========= Case 1: padding > 0, stride = 1 =========
+        # Typical convolution with halo exchange required.
+        self._test_conv_case(
+            N=1,
+            C=10,
+            H=32,
+            W=32,
+            kernel_size=3,
+            padding=1,
+            bias_attr=True,
+            mesh=mesh1,
+        )
+        self._test_conv_case(
+            N=2,
+            C=8,
+            H=16,
+            W=32,
+            kernel_size=5,
+            padding=2,
+            bias_attr=False,
+            mesh=mesh1,
+        )
+        self._test_conv_case(
+            N=4,
+            C=6,
+            H=28,
+            W=28,
+            kernel_size=3,
+            padding=1,
+            bias_attr=True,
+            mesh=mesh1,
+        )
+
+        # NHWC format with padding > 0
+        self._test_conv_case(
+            N=2,
+            C=8,
+            H=16,
+            W=32,
+            kernel_size=3,
+            padding=1,
+            bias_attr=True,
+            mesh=mesh1,
+            data_format="NHWC",
+        )
+        self._test_conv_case(
+            N=4,
+            C=6,
+            H=28,
+            W=28,
+            kernel_size=5,
+            padding=2,
+            bias_attr=False,
+            mesh=mesh1,
+            data_format="NHWC",
+        )
+
+        # ========= Case 2: padding = 0, stride == kernel_size =========
+        # No halo exchange needed, input width must be divisible by stride.
+        self._test_conv_case(
+            N=1,
+            C=10,
+            H=32,
+            W=32,
+            kernel_size=1,
+            padding=0,
+            bias_attr=True,
+            mesh=mesh1,
+            stride=1,
+        )
+        self._test_conv_case(
+            N=4,
+            C=6,
+            H=32,
+            W=32,
+            kernel_size=2,
+            padding=0,
+            bias_attr=False,
+            mesh=mesh1,
+            stride=2,
+        )
+        self._test_conv_case(
+            N=2,
+            C=8,
+            H=16,
+            W=32,
+            kernel_size=4,
+            padding=0,
+            bias_attr=True,
+            mesh=mesh1,
+            stride=4,
+        )
+
+        # NHWC format with padding = 0
+        self._test_conv_case(
+            N=1,
+            C=10,
+            H=32,
+            W=32,
+            kernel_size=2,
+            padding=0,
+            bias_attr=True,
+            mesh=mesh1,
+            stride=2,
+            data_format="NHWC",
+        )
+        self._test_conv_case(
+            N=4,
+            C=6,
+            H=32,
+            W=32,
+            kernel_size=4,
+            padding=0,
+            bias_attr=False,
+            mesh=mesh1,
+            stride=4,
+            data_format="NHWC",
+        )
+
+        # ========= Case 3: 2D ProcessMesh (dp + tp) =========
+        mesh2 = dist.ProcessMesh([[0, 1]], dim_names=['dp', 'mp'])
+
+        # padding > 0
+        self._test_conv_case(
+            N=2,
+            C=8,
+            H=32,
+            W=32,
+            kernel_size=3,
+            padding=1,
+            bias_attr=True,
+            mesh=mesh2,
+        )
+        self._test_conv_case(
+            N=4,
+            C=6,
+            H=28,
+            W=28,
+            kernel_size=5,
+            padding=2,
+            bias_attr=False,
+            mesh=mesh2,
+        )
+
+        # padding = 0
+        self._test_conv_case(
+            N=2,
+            C=8,
+            H=16,
+            W=32,
+            kernel_size=1,
+            padding=0,
+            bias_attr=True,
+            mesh=mesh2,
+            stride=1,
+        )
+
+        # NHWC format, both padding > 0 and = 0
+        self._test_conv_case(
+            N=4,
+            C=6,
+            H=28,
+            W=28,
+            kernel_size=3,
+            padding=1,
+            bias_attr=True,
+            mesh=mesh2,
+            data_format="NHWC",
+        )
+        self._test_conv_case(
+            N=1,
+            C=10,
+            H=32,
+            W=32,
+            kernel_size=4,
+            padding=0,
+            bias_attr=True,
+            mesh=mesh2,
+            stride=4,
+            data_format="NHWC",
+        )
+
+        self._test_intermediate(
+            N=1,
+            C=10,
+            H=32,
+            W=32,
+            kernel_size=3,
+            padding=1,
+            bias_attr=True,
+            mesh=mesh1,
+        )
+        self._test_intermediate(
+            N=1,
+            C=10,
+            H=32,
+            W=32,
+            kernel_size=2,
+            padding=0,
+            bias_attr=True,
+            mesh=mesh1,
+            stride=2,
+            data_format="NHWC",
+        )
+        self._test_intermediate(
+            N=4,
+            C=6,
+            H=28,
+            W=28,
+            kernel_size=3,
+            padding=1,
+            bias_attr=True,
+            mesh=mesh2,
+            data_format="NHWC",
+        )
+
+
+if __name__ == '__main__':
+    tester = TestTPConv()
+    tester.run_test_cases()