[AutoParallel] add intermediate api for tp_conv (#73063)

AndSonder · web-flow · commit d9d70c408013 · 2025-06-11T18:09:50.000+08:00
* support ring_conv forward and backward

* add intermediate api for tp-conv

* apply suggestions from code review

* fix the convergence issue

* recover third_party
diff --git a/python/paddle/distributed/__init__.py b/python/paddle/distributed/__init__.py
@@ -52,6 +52,7 @@
 from .auto_parallel.intermediate.pipeline_parallel import SplitPoint
 from .auto_parallel.intermediate.tensor_parallel import (
     ColWiseParallel,
+    ConvParallel,
     PrepareLayerInput,
     PrepareLayerOutput,
     RowWiseParallel,
@@ -217,4 +218,5 @@
     "set_mesh",
     "get_mesh",
     "to_distributed",
+    "ConvParallel",
 ]
diff --git a/python/paddle/distributed/auto_parallel/intermediate/tensor_parallel.py b/python/paddle/distributed/auto_parallel/intermediate/tensor_parallel.py
@@ -635,6 +635,180 @@ def apply(self, layer, process_mesh, shard_param_list):
         )
 
 
+class ConvParallel(PlanBase):
+    """
+    A strategy for enabling spatial parallelism on ``paddle.nn.Conv2D`` layers
+    by sharding the input tensor along its Width (W) dimension.
+
+    When this ``ConvParallel`` configuration is applied to a ``Conv2D`` layer,
+    the layer's input tensor will have its width dimension split across devices
+    in the model parallel group. This can help reduce memory usage from activations,
+    especially when dealing with inputs that have a large width.
+
+    To enable width-wise input sharding correctly, make sure your `Conv2D` layer
+    satisfies the following conditions along the width dimension:
+
+    - **Dilation** must be set to `1`.
+    - **If no width padding is used:**
+        - The input width must be evenly divisible by the stride width.
+        - The stride width must be equal to the kernel width.
+    - **If width padding is used:**
+        - The stride width must be `1`.
+        - The total input width must be at least half the kernel width.
+
+    Examples:
+        .. code-block:: python
+
+            >>> import paddle
+            >>> import paddle.nn as nn
+            >>> import paddle.distributed as dist
+
+            >>> class SimpleConvNet(nn.Layer):
+            ...     def __init__(self, data_format="NCHW"):
+            ...         super().__init__()
+            ...         self.conv1 = nn.Conv2D(
+            ...             3, 8, kernel_size=3, padding=1, data_format=data_format
+            ...         )
+            ...         self.relu = nn.ReLU()
+            ...     def forward(self, x):
+            ...         x = self.conv1(x)
+            ...         return self.relu(x)
+            ...
+            >>> # doctest: +REQUIRES(env:DISTRIBUTED)
+            >>> model = SimpleConvNet(data_format="NCHW")
+            >>> mp_config = {
+            ...    "parallelize_plan": {
+            ...        "conv1": dist.ConvParallel()
+            ...     }
+            ... }
+    """
+
+    def __init__(self) -> None:
+        super().__init__()
+
+    @staticmethod
+    def _is_supported(
+        input_size,
+        kernel_size,
+        stride,
+        padding,
+        dilation,
+        data_format,
+        mp_group_size,
+    ):
+        idx_w_input = -1
+        idx_w_kernel = -1
+
+        if data_format == "NCHW":
+            idx_w_input = 3
+            idx_w_kernel = 3
+        elif data_format == "NHWC":
+            idx_w_input = 2
+            idx_w_kernel = 3
+        else:
+            return False
+
+        if input_size[idx_w_input] % mp_group_size != 0:
+            return False
+
+        dilation_w = dilation[1]
+        padding_w = padding[1]
+        stride_w = stride[1]
+
+        input_w = input_size[idx_w_input]
+        kernel_w = kernel_size[idx_w_kernel]
+
+        if dilation_w != 1:
+            # RingConv2d only supports dilation=1.
+            # Larger dilation would require enlarged halo regions and more complex communication.
+            return False
+
+        if padding_w == 0:
+            # To avoid halo exchange when padding=0, we require:
+            # - input_w must be divisible by stride_w, so partitions align evenly across ranks.
+            # - stride_w == kernel_w, so each kernel operates on disjoint local regions.
+            if input_w % stride_w != 0:
+                return False
+            if stride_w != kernel_w:
+                return False
+
+        else:
+            # When padding > 0, halo exchange is needed.
+            # To simplify halo logic, we require:
+            # - stride_w == 1: ensures each output element is computed from overlapping input,
+            #   and no input region is skipped, simplifying halo construction.
+            # - kernel_w // 2 <= input_w: prevents the kernel from exceeding local input.
+            if stride_w != 1:
+                return False
+            if kernel_w // 2 > input_w:
+                return False
+
+        return True
+
+    def conv_parallel_start(self, process_mesh, data_format):
+        def start(layer, input, output=None):
+            if data_format == "NCHW":
+                shard_w_dim = 3
+            elif data_format == "NHWC":
+                shard_w_dim = 2
+            else:
+                raise ValueError(
+                    f"Unsupported data_format: {data_format}. "
+                    "Only NCHW and NHWC are supported."
+                )
+
+            if isinstance(input, tuple):
+                x = input[0]
+            else:
+                x = input
+
+            placements = x.placements
+            mp_index = process_mesh.dim_names.index('mp')
+            mp_group_size = process_mesh.get_dim_size('mp')
+
+            # Note(luchang): for intermediate api, when this ConvLayer is
+            # not supported, we just skip apply parallelization.
+            if not ConvParallel._is_supported(
+                x.shape,
+                layer.weight.shape,
+                layer._stride,
+                layer._updated_padding,
+                layer._dilation,
+                data_format,
+                mp_group_size,
+            ):
+                return input
+
+            if placements is None:
+                placements = [
+                    dist.Replicate() for _ in range(len(process_mesh.shape))
+                ]
+            if placements[mp_index] == dist.Shard(shard_w_dim):
+                return input
+
+            placements[mp_index] = dist.Shard(shard_w_dim)
+
+            if not x.is_dist():
+                x = dist.shard_tensor(x, process_mesh, placements)
+            else:
+                x = dist.reshard(x, process_mesh, placements)
+
+            if isinstance(input, tuple):
+                input = list(input)
+                input[0] = x
+                input = tuple(input)
+            else:
+                input = x
+            return input
+
+        return start
+
+    def apply(self, layer, process_mesh, shard_param_list):
+        layer.register_forward_pre_hook(
+            self.conv_parallel_start(process_mesh, layer._data_format)
+        )
+
+
 class TensorParallel(ParallelModel):
     def __init__(self, model, parallelize_plan=None):
         super().__init__(model)
diff --git a/python/paddle/distributed/auto_parallel/ring_conv.py b/python/paddle/distributed/auto_parallel/ring_conv.py
@@ -85,7 +85,7 @@ def _get_conv_tp_group(x_mesh, x_placements, data_format):
     )
 
 
-def _ring_send_recv_construct(
+def _ring_conv_halo_exchange(
     local_input_tensor,
     halo_width_to_receive_from_left,
     halo_width_to_receive_from_right,
@@ -183,10 +183,10 @@ def _ring_send_recv_construct(
             axis=width_dim_idx,
         )
 
-    return reconstructed_tensor
+    return reconstructed_tensor.contiguous()
 
 
-def _ring_send_recv_aggregate(
+def _ring_conv_halo_aggregate(
     local_gradient_tensor,
     halo_width_send_left,
     halo_width_send_right,
@@ -314,7 +314,7 @@ def _ring_send_recv_aggregate(
         ]
         target_slice_left.add_(buffer_for_gradient_from_left)
 
-    return processed_gradient_tensor
+    return processed_gradient_tensor.contiguous()
 
 
 class RingConv2d(paddle.autograd.PyLayer):
@@ -489,7 +489,7 @@ def forward(
 
             # step 1: reconstruct the local input tensor including halo regions via ring communication
             # `x` is updated here, now including halo data received from neighboring ranks.
-            x = _ring_send_recv_construct(
+            x = _ring_conv_halo_exchange(
                 x,
                 left_halo_width,
                 right_halo_width,
@@ -570,7 +570,7 @@ def forward(
             final_local_results, x_mesh, x_placements
         )
 
-        return final_local_results
+        return final_local_results.contiguous()
 
     @staticmethod
     def backward(ctx, grad_out):
@@ -627,7 +627,7 @@ def backward(ctx, grad_out):
             output_width_dim_idx = ctx.output_width_dim_idx
 
             # Step 1: Reconstruct `in_tensor_augmented` (original input to local conv in forward)
-            in_tensor_augmented = _ring_send_recv_construct(
+            in_tensor_augmented = _ring_conv_halo_exchange(
                 x,
                 left_halo_width,
                 right_halo_width,
@@ -671,7 +671,7 @@ def backward(ctx, grad_out):
 
             # Step 4: Aggregate "halo" regions for grad_input
             if not x_stop_gradient:
-                grad_x = _ring_send_recv_aggregate(
+                grad_x = _ring_conv_halo_aggregate(
                     grad_x_augmented,
                     left_halo_width,
                     right_halo_width,
@@ -709,6 +709,7 @@ def backward(ctx, grad_out):
         grad_weight = dist.auto_parallel.api.dtensor_from_local(
             grad_weight, weight_mesh, weight_placements
         )
+        # do allreduce to get right grad_weight
         grad_weight = dist.reshard(
             grad_weight,
             weight_mesh,
@@ -719,11 +720,12 @@ def backward(ctx, grad_out):
             grad_bias = dist.auto_parallel.api.dtensor_from_local(
                 grad_bias, bias_mesh, bias_placements
             )
-        grad_weight = dist.reshard(
-            grad_weight,
-            weight_mesh,
-            [dist.Replicate() for _ in range(len(weight_placements))],
-        )
+            # do allreduce to get right grad_bias
+            grad_bias = dist.reshard(
+                grad_bias,
+                weight_mesh,
+                [dist.Replicate() for _ in range(len(bias_placements))],
+            )
 
         if x_stop_gradient:
             grad_x = None
diff --git a/python/paddle/nn/layer/conv.py b/python/paddle/nn/layer/conv.py
@@ -741,6 +741,8 @@ def forward(self, x: Tensor) -> Tensor:
                 data_format=self._data_format,
             )
 
+        # Note(luchang): If the input tensor is sharded along the spatial width
+        # dimension (W), this indicates spatially parallel convolution is being used.
         if (
             in_dygraph_mode()
             and x.is_dist()
diff --git a/test/auto_parallel/tp_conv.py b/test/auto_parallel/tp_conv.py

Original file line number	Diff line number	Diff line change
`@@ -741,6 +741,8 @@ def forward(self, x: Tensor) -> Tensor:`
`741`	`741`	`data_format=self._data_format,`
`742`	`742`	`)`
`743`	`743`
	`744`	`+ # Note(luchang): If the input tensor is sharded along the spatial width`
	`745`	`+ # dimension (W), this indicates spatially parallel convolution is being used.`
`744`	`746`	`if (`
`745`	`747`	`in_dygraph_mode()`
`746`	`748`	`and x.is_dist()`