Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / core / shard / ops / parallel_conv3d.py: 68%
93 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-07-06 05:41 +0800
1# Copyright 2026 Huawei Technologies Co., Ltd
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7# http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14# ============================================================================
15"""
16Distributed implementation for Conv3d operator.
17"""
19from typing import Callable, Optional, Tuple
21from hyper_parallel.core.dtensor.layout import Layout
22from .parallel_ops import DistributedOp
25def _normalize_conv3d_args(input_tensor, weight, bias=None, stride=1, padding=0, dilation=1, groups=1):
26 return (input_tensor, weight, bias, stride, padding, dilation, groups), {}
29class Conv3dDistributedOp(DistributedOp):
30 """
31 Distributed implementation for torch.nn.functional.conv3d.
32 Supports Data Parallel, Tensor Parallel (Column/Row), and Spatial Parallel.
33 """
35 def __init__(self, op_name):
36 super().__init__(op_name)
37 self._allow_partial_inputs = False
39 def _validate_row_parallelism(self, in_map, w_map, groups):
40 """
41 Validate constraints for Row Parallelism.
42 """
43 # 1. Handle Groups Constraint for Row Parallelism
44 if groups > 1:
45 if in_map[1] != "None" or w_map[1] != "None":
46 # Row Parallelism with groups > 1 requires advanced group-wise communication
47 raise ValueError(f"For {self.op_name}, Sharding on C_in with groups > 1 is not supported.")
49 # 2. Check Row Parallelism (Sharding on Channel In)
50 # Input: (N, C_in, D, H, W), Weight: (C_out, C_in/groups, kD, kH, kW)
51 if in_map[1] != "None":
52 if in_map[1] != w_map[1]:
53 raise ValueError(f"For {self.op_name}, Input C_in and Weight C_in must be sharded on the same axis.")
55 def _validate_column_parallelism(self, w_layout, b_layout, groups):
56 """
57 Validate constraints for Column Parallelism.
58 """
59 w_map = w_layout.alias_tensor_map
60 w_map_0 = w_map[0]
62 if w_map_0 != "None":
63 # Check bias alignment
64 if b_layout is not None:
65 b_map = b_layout.alias_tensor_map
66 b_map_0 = b_map[0]
67 if w_map_0 != b_map_0:
68 raise ValueError(
69 f"For {self.op_name}, Weight C_out and Bias C_out must be sharded on the same axis."
70 )
72 # Check groups divisibility for Column Parallelism
73 if groups > 1:
74 dev_num = 1
75 axes = w_map_0 if isinstance(w_map_0, tuple) else (w_map_0,)
76 for axis_name in axes:
77 dev_num *= w_layout.mesh.get_device_num_along_axis(axis_name)
79 if groups % dev_num != 0:
80 raise ValueError(
81 f"For {self.op_name}, groups ({groups}) "
82 f"must be divisible by tp_size ({dev_num})."
83 )
85 def preprocess(self, args: tuple, kwargs: dict) -> tuple:
86 """
87 Preprocess arguments for Conv3d operator.
89 Args:
90 args (tuple): Conv3d positional arguments.
91 kwargs (dict): Conv3d keyword arguments.
93 Returns:
94 tuple: (local_args, local_kwargs, cache_values)
95 """
96 args, _ = _normalize_conv3d_args(*args, **kwargs)
97 input_tensor, weight, bias, stride, padding, dilation, groups = args
98 local_args = (
99 input_tensor.to_local(),
100 weight.to_local(),
101 bias.to_local() if hasattr(bias, '_layout') else bias,
102 stride,
103 padding,
104 dilation,
105 groups,
106 )
107 local_kwargs = {}
108 cache_values = [
109 input_tensor.layout,
110 weight.layout,
111 bias.layout if hasattr(bias, '_layout') else None,
112 stride,
113 padding,
114 dilation,
115 groups,
116 ]
117 return local_args, local_kwargs, cache_values
119 def infer_layout(self, cache_values: list) -> Tuple[tuple, None]: # pylint: disable=W0221
120 """
121 Infer output layout for Conv3d operator.
123 Rules:
124 1. Input and weight must not have Partial status.
125 2. Input and weight must both be 5D.
126 3. Input C_in and weight C_in sharding must match for row parallelism.
127 4. Sharding C_in with groups > 1 is not supported.
128 5. Bias C_out sharding must match weight C_out sharding.
129 6. Output layout inherits N/D/H/W sharding from input and C_out sharding from weight.
130 7. Row parallelism marks output as Partial('sum') on the C_in mesh axis.
132 Args:
133 cache_values (list): [input_layout, weight_layout, bias_layout_or_None,
134 stride, padding, dilation, groups]
136 Returns:
137 tuple: ((output_layout,), None)
139 Raises:
140 ValueError: If layouts are missing, partial, malformed, or violate Conv3d
141 sharding constraints.
142 """
144 in_layout, w_layout, b_layout = cache_values[0], cache_values[1], cache_values[2]
145 groups = cache_values[6]
147 if not in_layout or not w_layout:
148 raise ValueError(f"For {self.op_name}, Requires at least input and weight layouts.")
150 self._check_partial_inputs([in_layout, w_layout])
152 if b_layout is not None:
153 self._check_partial_inputs([b_layout])
155 if in_layout.mesh_shape != w_layout.mesh_shape:
156 raise ValueError(
157 f"For {self.op_name}, input and weight must have the same mesh_shape, "
158 f"but got input: {in_layout.mesh_shape} and weight: {w_layout.mesh_shape}"
159 )
160 if b_layout is not None and b_layout.mesh_shape != in_layout.mesh_shape:
161 raise ValueError(
162 f"For {self.op_name}, bias and input must have the same mesh_shape, "
163 f"but got bias: {b_layout.mesh_shape} and input: {in_layout.mesh_shape}"
164 )
166 in_map = in_layout.alias_tensor_map
167 w_map = w_layout.alias_tensor_map
169 # Validate dimensions
170 if len(in_map) != 5 or len(w_map) != 5:
171 raise ValueError(f"For {self.op_name}, Input and weight must be 5D.")
173 # Delegate validation to helper methods to reduce cyclomatic complexity
174 self._validate_row_parallelism(in_map, w_map, groups)
175 self._validate_column_parallelism(w_layout, b_layout, groups)
177 # Construct Output Map (N, C_out, D_out, H_out, W_out)
178 out_map = [
179 in_map[0], # N
180 w_map[0], # C_out
181 in_map[2], # D
182 in_map[3], # H
183 in_map[4] # W
184 ]
186 # Build Layout
187 output_layout = Layout(
188 mesh_shape=in_layout.mesh_shape,
189 alias_name=in_layout.alias_name,
190 rank_list=in_layout.rank_list,
191 )
192 output_layout = output_layout(*tuple(out_map))
194 # Set Partial status for Row Parallelism
195 if in_map[1] != "None":
196 axes = in_map[1] if isinstance(in_map[1], tuple) else (in_map[1],)
197 for axis in axes:
198 output_layout.set_partial_by_dev_axis(axis, "sum")
200 return (output_layout,), None
202 def get_expand_impl(self, func: Optional[Callable], infer_result: tuple, # pylint: disable=W0221
203 cache_values: list) -> Optional[Callable]:
204 """
205 Get expand implementation for the operator.
206 Intercepts the execution to handle Grouped Convolution with Column Parallelism.
207 """
208 w_layout = cache_values[1]
209 w_map = w_layout.alias_tensor_map
210 w_map_0 = w_map[0]
212 # If Weight is NOT sharded on C_out (dim=0), native conv3d works fine.
213 if w_map_0 == "None":
214 return None
216 parsed_groups = cache_values[6]
217 if parsed_groups == 1:
218 return None
220 mesh = w_layout.mesh
221 axes = w_map_0 if isinstance(w_map_0, tuple) else (w_map_0,)
222 dev_num = 1
223 local_rank = 0
224 for axis_name in axes:
225 axis_size = mesh.get_device_num_along_axis(axis_name)
226 dev_num *= axis_size
227 local_rank = local_rank * axis_size + mesh.get_local_rank(axis_name)
229 # Pre-calculate local groups and group boundaries for the current device ahead of time.
230 # This hoisting optimization avoids redundant calculations during every forward pass.
231 local_groups = parsed_groups // dev_num
232 start_group = local_rank * local_groups
233 end_group = start_group + local_groups
235 def distributed_conv3d_impl(input_tensor, weight_tensor, bias=None, stride=1, padding=0, dilation=1, groups=1):
236 # --- Handling Groups > 1 with Column Parallelism ---
237 # Calculate the input channel chunk size
238 c_in = input_tensor.shape[1]
239 c_in_per_group = c_in // groups
241 # Map the pre-calculated groups to the actual input channels
242 # Uses start_group and end_group captured from the outer scope
243 start_channel = start_group * c_in_per_group
244 end_channel = end_group * c_in_per_group
246 # Slice the replicated input to match the local groups
247 sliced_input = input_tensor[:, start_channel:end_channel, ...]
249 # Execute native conv3d with the sliced input and adjusted local groups
250 return func(sliced_input, weight_tensor, bias, stride, padding, dilation, local_groups)
252 return distributed_conv3d_impl