已合并
migrate distributed ops to the new dispatch flow #736
rongyue创建于 6月1日
migrate distributed ops to the new dispatch flow #736
已合并
共 27 个文件变更+1859-1066
| @@ -278,6 +278,7 @@ class OpDispatcher: | |||
| 278 | "normal_", "uniform_", "bernoulli", "bernoulli_", | 278 | "normal_", "uniform_", "bernoulli", "bernoulli_", |
| 279 | "native_dropout", "rand", "rand_like", "randn", | 279 | "native_dropout", "rand", "rand_like", "randn", |
| 280 | "randn_like", "randint_like", "kaiming_uniform_", | 280 | "randn_like", "randint_like", "kaiming_uniform_", |
| 281 | + "multinomial", | ||
| 281 | } | 282 | } |
| 282 | # Only mint random op support | 283 | # Only mint random op support |
| 283 | # MindSpore use the actual kernel name. | 284 | # MindSpore use the actual kernel name. |
| @@ -438,14 +439,15 @@ class OpDispatcher: | |||
| 438 | 439 | ||
| 439 | return input_layouts, extra_args, input_args, input_kwargs, cache_key_values | 440 | return input_layouts, extra_args, input_args, input_kwargs, cache_key_values |
| 440 | 441 | ||
| 441 | - def _with_layout_infer(self, func: callable, *args, **kwargs) -> Tensor: | 442 | + def _with_layout_infer(self, func: callable, *args, _packed_call=None, **kwargs) -> Tensor: |
| 442 | - """_with_layout_infer""" | 443 | + """_with_layout_infer |
| 444 | + | ||
| 445 | + NOTE: aclop packed args are already normalized upstream by | ||
| 446 | + ``_dispatch_layout_infer`` via ``_normalize_aclop_args``. The | ||
| 447 | + ``_packed_call`` kwarg carries the ``(prim, name)`` tuple when the | ||
| 448 | + op uses aclop format, or ``None`` otherwise. | ||
| 449 | + """ | ||
| 443 | func_name = platform.get_op_name(func) | 450 | func_name = platform.get_op_name(func) |
| 444 | - packed_call = None | ||
| 445 | - if(func_name in self.unpack_ops and len(args) == 3 and | ||
| 446 | - isinstance(args[1], str) and isinstance(args[2],(tuple,list))): | ||
| 447 | - packed_call = (args[0], args[1]) | ||
| 448 | - args = tuple(args[2]) | ||
| 449 | 451 | ||
| 450 | input_layouts, extra_args, input_args, input_kwargs, cache_key_values = \ | 452 | input_layouts, extra_args, input_args, input_kwargs, cache_key_values = \ |
| 451 | OpDispatcher._process_args_and_kwargs(args, kwargs) | 453 | OpDispatcher._process_args_and_kwargs(args, kwargs) |
| @@ -469,10 +471,7 @@ class OpDispatcher: | |||
| 469 | if op_impl is None: | 471 | if op_impl is None: |
| 470 | op_impl = func | 472 | op_impl = func |
| 471 | 473 | ||
| 472 | - if packed_call is not None: | 474 | + py_output = OpDispatcher._call_op_impl(op_impl, _packed_call, input_args, input_kwargs) |
| 473 | - py_output = op_impl(packed_call[0], packed_call[1], tuple(input_args), **input_kwargs) | ||
| 474 | - else: | ||
| 475 | - py_output = op_impl(*input_args, **input_kwargs) | ||
| 476 | 475 | ||
| 477 | if isinstance(py_output, (tuple, list)): | 476 | if isinstance(py_output, (tuple, list)): |
| 478 | output = () | 477 | output = () |
| @@ -530,8 +529,14 @@ class OpDispatcher: | |||
| 530 | 529 | ||
| 531 | return DTensor.from_local(py_output, output_layout.mesh, output_layout.alias_placements) | 530 | return DTensor.from_local(py_output, output_layout.mesh, output_layout.alias_placements) |
| 532 | 531 | ||
| 533 | - def _with_layout_infer_with_tuple_expand(self, func: callable, *args, **kwargs) -> Tensor: | 532 | + def _with_layout_infer_with_tuple_expand(self, func: callable, *args, _packed_call=None, **kwargs) -> Tensor: |
| 534 | - """_with_layout_infer_with_tuple_expand""" | 533 | + """_with_layout_infer_with_tuple_expand |
| 534 | + | ||
| 535 | + NOTE: aclop packed args are already normalized upstream by | ||
| 536 | + ``_dispatch_layout_infer`` via ``_normalize_aclop_args``. The | ||
| 537 | + ``_packed_call`` kwarg carries the ``(prim, name)`` tuple when the | ||
| 538 | + op uses aclop format, or ``None`` otherwise. | ||
| 539 | + """ | ||
| 535 | expanded_args = [] | 540 | expanded_args = [] |
| 536 | input_args = [] | 541 | input_args = [] |
| 537 | for arg in args: | 542 | for arg in args: |
| @@ -571,7 +576,7 @@ class OpDispatcher: | |||
| 571 | if op_impl is None: | 576 | if op_impl is None: |
| 572 | op_impl = func | 577 | op_impl = func |
| 573 | 578 | ||
| 574 | - py_output = op_impl(*input_args, **input_kwargs) | 579 | + py_output = OpDispatcher._call_op_impl(op_impl, _packed_call, input_args, input_kwargs) |
| 575 | return distribute_op.wrap_output(py_output, output_layout) | 580 | return distribute_op.wrap_output(py_output, output_layout) |
| 576 | 581 | ||
| 577 | @staticmethod | 582 | @staticmethod |
| @@ -704,15 +709,15 @@ class OpDispatcher: | |||
| 704 | 709 | ||
| 705 | return input_layouts, input_shapes, extra_args, input_args, input_kwargs, cache_key_values | 710 | return input_layouts, input_shapes, extra_args, input_args, input_kwargs, cache_key_values |
| 706 | 711 | ||
| 707 | - def _with_layout_infer_with_shape(self, func: callable, *args, **kwargs) -> Tensor: | 712 | + def _with_layout_infer_with_shape(self, func: callable, *args, _packed_call=None, **kwargs) -> Tensor: |
| 708 | - """_with_layout_infer_with_shape""" | 713 | + """_with_layout_infer_with_shape |
| 714 | + | ||
| 715 | + NOTE: aclop packed args are already normalized upstream by | ||
| 716 | + ``_dispatch_layout_infer`` via ``_normalize_aclop_args``. The | ||
| 717 | + ``_packed_call`` kwarg carries the ``(prim, name)`` tuple when the | ||
| 718 | + op uses aclop format, or ``None`` otherwise. | ||
| 719 | + """ | ||
| 709 | func_name = platform.get_op_name(func) | 720 | func_name = platform.get_op_name(func) |
| 710 | - packed_call = None | ||
| 711 | - # Packed fallback args for some ops (e.g. Mod: (prim_obj, "Mod", (x, y))). | ||
| 712 | - if (func_name in self.unpack_ops and len(args) == 3 and | ||
| 713 | - isinstance(args[1], str) and isinstance(args[2], (tuple, list))): | ||
| 714 | - packed_call = (args[0], args[1]) | ||
| 715 | - args = tuple(args[2]) | ||
| 716 | 721 | ||
| 717 | (input_layouts, input_shapes, extra_args, input_args, | 722 | (input_layouts, input_shapes, extra_args, input_args, |
| 718 | input_kwargs, cache_key_values) = OpDispatcher._process_args_and_kwargs_with_shape(args, kwargs) | 723 | input_kwargs, cache_key_values) = OpDispatcher._process_args_and_kwargs_with_shape(args, kwargs) |
| @@ -738,10 +743,7 @@ class OpDispatcher: | |||
| 738 | if op_impl is None: | 743 | if op_impl is None: |
| 739 | op_impl = func | 744 | op_impl = func |
| 740 | 745 | ||
| 741 | - if packed_call is not None: | 746 | + py_output = OpDispatcher._call_op_impl(op_impl, _packed_call, input_args, input_kwargs) |
| 742 | - py_output = op_impl(packed_call[0], packed_call[1], tuple(input_args), **input_kwargs) | ||
| 743 | - else: | ||
| 744 | - py_output = op_impl(*input_args, **input_kwargs) | ||
| 745 | 747 | ||
| 746 | # set output layout | 748 | # set output layout |
| 747 | if isinstance(py_output, (tuple, list)): | 749 | if isinstance(py_output, (tuple, list)): |
| @@ -868,7 +870,7 @@ class OpDispatcher: | |||
| 868 | local_kwargs = {k: v.to_local() if isinstance(v, DTensor) else v for k, v in kwargs.items()} | 870 | local_kwargs = {k: v.to_local() if isinstance(v, DTensor) else v for k, v in kwargs.items()} |
| 869 | first_local_arg = first_arg.to_local() | 871 | first_local_arg = first_arg.to_local() |
| 870 | 872 | ||
| 871 | - if self._rng_tracker is None and is_rng_supported_mesh(): | 873 | + if self._rng_tracker is None and is_rng_supported_mesh(first_arg.device_mesh): |
| 872 | self._rng_tracker = OffsetBasedRNGTracker() | 874 | self._rng_tracker = OffsetBasedRNGTracker() |
| 873 | 875 | ||
| 874 | maybe_user_generator = local_kwargs.pop("generator", None) | 876 | maybe_user_generator = local_kwargs.pop("generator", None) |
| @@ -897,6 +899,8 @@ class OpDispatcher: | |||
| 897 | local_args = _apply_shard_offset_to_rng_args(local_args, offset_incr) | 899 | local_args = _apply_shard_offset_to_rng_args(local_args, offset_incr) |
| 898 | local_results = op_call(*local_args, **local_kwargs) | 900 | local_results = op_call(*local_args, **local_kwargs) |
| 899 | else: | 901 | else: |
| 902 | + if maybe_user_generator is not None: | ||
| 903 | + local_kwargs["generator"] = maybe_user_generator | ||
| 900 | local_results = op_call(*local_args, **local_kwargs) | 904 | local_results = op_call(*local_args, **local_kwargs) |
| 901 | 905 | ||
| 902 | return self._wrap_random_result(op_name, local_results, first_arg, args, kwargs) | 906 | return self._wrap_random_result(op_name, local_results, first_arg, args, kwargs) |
| @@ -1138,6 +1142,58 @@ class OpDispatcher: | |||
| 1138 | f"(not recommended for large vocabulary), use logits.full_tensor() explicitly." | 1142 | f"(not recommended for large vocabulary), use logits.full_tensor() explicitly." |
| 1139 | ) | 1143 | ) |
| 1140 | 1144 | ||
| 1145 | + @staticmethod | ||
| 1146 | + def _normalize_aclop_args(op_name: str, unpack_ops: list, args: tuple) -> tuple: | ||
| 1147 | + """ | ||
| 1148 | + Normalize aclop-packed arguments for MindSpore backend operators. | ||
| 1149 | + | ||
| 1150 | + NOTE: This handles MindSpore aclop operators whose kernel signature packs | ||
| 1151 | + arguments as ``(prim, op_name_str, (real_arg0, real_arg1, ...))``. The | ||
| 1152 | + ``prim`` and ``op_name_str`` are preserved as ``packed_call`` for the | ||
| 1153 | + final kernel invocation, while the real tensor arguments are extracted | ||
| 1154 | + for layout inference and preprocessing. | ||
| 1155 | + | ||
| 1156 | + **aclop is planned for deprecation.** Once aclop is fully removed, this | ||
| 1157 | + normalization and the associated ``unpack_ops`` list can be deleted. | ||
| 1158 | + | ||
| 1159 | + Args: | ||
| 1160 | + op_name (str): Canonical operator name. | ||
| 1161 | + unpack_ops (list): List of op names that may use aclop packed format. | ||
| 1162 | + args (tuple): Raw positional arguments from the op call. | ||
| 1163 | + | ||
| 1164 | + Returns: | ||
| 1165 | + tuple: ``(packed_call, normalized_args)`` | ||
| 1166 | + - **packed_call**: ``(prim, op_name_str)`` tuple for kernel | ||
| 1167 | + invocation, or ``None`` if no unpacking was performed. | ||
| 1168 | + - **normalized_args**: The real tensor arguments (unpacked if | ||
| 1169 | + the packed format was detected, otherwise the original args). | ||
| 1170 | + """ | ||
| 1171 | + if op_name in unpack_ops and len(args) == 3 and \ | ||
| 1172 | + isinstance(args[1], str) and isinstance(args[2], (tuple, list)): | ||
| 1173 | + return (args[0], args[1]), tuple(args[2]) | ||
| 1174 | + return None, args | ||
| 1175 | + | ||
| 1176 | + @staticmethod | ||
| 1177 | + def _call_op_impl(op_impl: callable, packed_call, args, kwargs: dict): | ||
| 1178 | + """Invoke *op_impl* with optional aclop packed-call wrapping. | ||
| 1179 | + | ||
| 1180 | + When *packed_call* is not ``None`` the MindSpore aclop kernel expects | ||
| 1181 | + ``(prim, op_name, (arg0, arg1, ...))``. Otherwise *args* are spread | ||
| 1182 | + as positional arguments in the usual way. | ||
| 1183 | + | ||
| 1184 | + Args: | ||
| 1185 | + op_impl: The op implementation callable. | ||
| 1186 | + packed_call: ``(prim, op_name)`` tuple or ``None``. | ||
| 1187 | + args: Local tensor arguments (list or tuple). | ||
| 1188 | + kwargs: Keyword arguments dict. | ||
| 1189 | + | ||
| 1190 | + Returns: | ||
| 1191 | + Result of the *op_impl* invocation. | ||
| 1192 | + """ | ||
| 1193 | + if packed_call is not None: | ||
| 1194 | + return op_impl(packed_call[0], packed_call[1], tuple(args), **kwargs) | ||
| 1195 | + return op_impl(*args, **kwargs) | ||
| 1196 | + | ||
| 1141 | def _dispatch_layout_infer( | 1197 | def _dispatch_layout_infer( |
| 1142 | self, op_name: str, op_call: callable, args: tuple, kwargs: dict | 1198 | self, op_name: str, op_call: callable, args: tuple, kwargs: dict |
| 1143 | ): | 1199 | ): |
| @@ -1189,25 +1245,39 @@ class OpDispatcher: | |||
| 1189 | cache_manager = LayoutCacheManager.get_instance() | 1245 | cache_manager = LayoutCacheManager.get_instance() |
| 1190 | distribute_op = cache_manager.distributed_op(op_name) | 1246 | distribute_op = cache_manager.distributed_op(op_name) |
| 1191 | 1247 | ||
| 1248 | + # Normalize aclop-packed args before any per-op processing. | ||
| 1249 | + # This allows all distributed op classes (ElementWiseDistributedOp, | ||
| 1250 | + # GatherNdDistributedOp, etc.) to receive clean unpacked args, | ||
| 1251 | + # avoiding duplicated unpack logic in each class's preprocess. | ||
| 1252 | + packed_call, args = self._normalize_aclop_args(op_name, getattr(self, 'unpack_ops', []), args) | ||
| 1253 | + | ||
| 1192 | result = distribute_op.preprocess(args, kwargs) | 1254 | result = distribute_op.preprocess(args, kwargs) |
| 1193 | if result is not None: | 1255 | if result is not None: |
| 1194 | - return self._dispatch_new(op_call, distribute_op, result) | 1256 | + return self._dispatch_new(op_call, distribute_op, packed_call, result) |
| 1195 | 1257 | ||
| 1196 | suffix = self.layout_infer_ops[op_name].get('infer_layout_suffix', '') | 1258 | suffix = self.layout_infer_ops[op_name].get('infer_layout_suffix', '') |
| 1197 | if not suffix: | 1259 | if not suffix: |
| 1198 | - return self._with_layout_infer(op_call, *args, **kwargs) | 1260 | + return self._with_layout_infer(op_call, *args, _packed_call=packed_call, **kwargs) |
| 1261 | + | ||
| 1262 | + if suffix == 'WithShape': | ||
| 1263 | + return self._with_layout_infer_with_shape(op_call, *args, _packed_call=packed_call, **kwargs) | ||
| 1264 | + | ||
| 1265 | + if suffix == 'WithTupleExpand': | ||
| 1266 | + return self._with_layout_infer_with_tuple_expand(op_call, *args, _packed_call=packed_call, **kwargs) | ||
| 1199 | 1267 | ||
| 1200 | handler_name = self._suffix_dispatch.get(suffix) | 1268 | handler_name = self._suffix_dispatch.get(suffix) |
| 1201 | if handler_name is None: | 1269 | if handler_name is None: |
| 1202 | raise RuntimeError(f"Operator {op_name} specified wrong suffix in parallel yaml.") | 1270 | raise RuntimeError(f"Operator {op_name} specified wrong suffix in parallel yaml.") |
| 1203 | return getattr(self, handler_name)(op_call, *args, **kwargs) | 1271 | return getattr(self, handler_name)(op_call, *args, **kwargs) |
| 1204 | 1272 | ||
| 1205 | - def _dispatch_new(self, func, distribute_op, result) -> Tensor: | 1273 | + def _dispatch_new(self, func, distribute_op, packed_call, result) -> Tensor: |
| 1206 | """New dispatch flow using preprocess result. | 1274 | """New dispatch flow using preprocess result. |
| 1207 | 1275 | ||
| 1208 | Args: | 1276 | Args: |
| 1209 | func: Original function. | 1277 | func: Original function. |
| 1210 | distribute_op: Distributed operation instance. | 1278 | distribute_op: Distributed operation instance. |
| 1279 | + packed_call: (prim, op_name_str) tuple for aclop kernel invocation, | ||
| 1280 | + or None for regular ops. | ||
| 1211 | result: Preprocessed result (local_args, local_kwargs, cache_values). | 1281 | result: Preprocessed result (local_args, local_kwargs, cache_values). |
| 1212 | 1282 | ||
| 1213 | Returns: | 1283 | Returns: |
| @@ -1229,7 +1299,7 @@ class OpDispatcher: | |||
| 1229 | op_layout_cache[cache_key] = (infer_result, op_impl) | 1299 | op_layout_cache[cache_key] = (infer_result, op_impl) |
| 1230 | output_layouts, _ = infer_result | 1300 | output_layouts, _ = infer_result |
| 1231 | op_impl = func if op_impl is None else op_impl | 1301 | op_impl = func if op_impl is None else op_impl |
| 1232 | - py_output = op_impl(*local_args, **local_kwargs) | 1302 | + py_output = OpDispatcher._call_op_impl(op_impl, packed_call, local_args, local_kwargs) |
| 1233 | return distribute_op.wrap_output(py_output, output_layouts) | 1303 | return distribute_op.wrap_output(py_output, output_layouts) |
| 1234 | 1304 | ||
| 1235 | def dispatch(self, op_call: callable, args: tuple, kwargs: dict) -> object: | 1305 | def dispatch(self, op_call: callable, args: tuple, kwargs: dict) -> object: |
| @@ -16,82 +16,95 @@ | |||
| 16 | Distributed implementation for Cumsum operator. | 16 | Distributed implementation for Cumsum operator. |
| 17 | """ | 17 | """ |
| 18 | 18 | ||
| 19 | -from hyper_parallel.core.dtensor.layout import Layout | 19 | +import copy |
| 20 | +from typing import Tuple | ||
| 21 | + | ||
| 20 | from .parallel_ops import DistributedOp | 22 | from .parallel_ops import DistributedOp |
| 21 | 23 | ||
| 22 | 24 | ||
| 25 | +def _normalize_cumsum_args(x, dim, dtype=None): | ||
| 26 | + return (x, dim), {'dtype': dtype} | ||
| 27 | + | ||
| 28 | + | ||
| 23 | class CumsumDistributedOp(DistributedOp): | 29 | class CumsumDistributedOp(DistributedOp): |
| 24 | """Distributed implementation for torch.cumsum.""" | 30 | """Distributed implementation for torch.cumsum.""" |
| 31 | + _MS_PRIMITIVE_OP_NAMES = frozenset({'CumsumExt'}) | ||
| 25 | 32 | ||
| 26 | - def infer_layout(self, layouts, extra_args=None): | 33 | + def preprocess(self, args: tuple, kwargs: dict) -> tuple: |
| 27 | """ | 34 | """ |
| 28 | - Infer output layout for torch.cumsum | 35 | + Preprocess arguments for Cumsum operator. |
| 29 | - | ||
| 30 | - PyTorch semantics: | ||
| 31 | - - Computes cumulative sum along dimension `dim` | ||
| 32 | - - Output shape is identical to input shape | ||
| 33 | - - Operation is sequential along `dim`: each element depends on all preceding elements in that dimension | ||
| 34 | - | ||
| 35 | - Critical sharding constraint: | ||
| 36 | - - The dimension `dim` MUST be unsharded (-1 in tensor_map) | ||
| 37 | 36 | ||
| 38 | Args: | 37 | Args: |
| 39 | - layouts (tuple): Layouts of inputs. Expected: | 38 | + args (tuple): Input arguments, first element is the input tensor. |
| 40 | - layouts[0] (Layout): Input tensor layout (required). | 39 | + kwargs (dict): Keyword arguments (dim, dtype). |
| 41 | - extra_args (tuple): Should contain 'dim'. Expected: | ||
| 42 | - extra_args[0] (int): Dimension to perform cumsum over (required). | ||
| 43 | 40 | ||
| 44 | Returns: | 41 | Returns: |
| 45 | - Layout: Output tensor layout (identical to input layout after validation). | 42 | + tuple: (local_args, local_kwargs, cache_values) |
| 46 | """ | 43 | """ |
| 47 | - if not layouts or layouts[0] is None: | 44 | + args, kwargs = _normalize_cumsum_args(*args, **kwargs) |
| 48 | - raise ValueError( | 45 | + input_tensor = args[0] |
| 49 | - f"Operation {self.op_name}: cumsum requires a valid input tensor layout." | 46 | + dim = args[1] |
| 50 | - ) | 47 | + dtype = kwargs['dtype'] |
| 51 | - input_layout = layouts[0] | 48 | + |
| 52 | - in_tensor_map = input_layout.tensor_map | 49 | + local_input = input_tensor.to_local() |
| 53 | - input_ndim = len(in_tensor_map) | 50 | + if self.op_name in self._MS_PRIMITIVE_OP_NAMES: |
| 51 | + local_args = (local_input, dim, dtype) | ||
| 52 | + local_kwargs = {} | ||
| 53 | + else: | ||
| 54 | + local_args = (local_input, dim) | ||
| 55 | + local_kwargs = {} | ||
| 56 | + if dtype is not None: | ||
| 57 | + local_kwargs['dtype'] = dtype | ||
| 58 | + | ||
| 59 | + cache_values = [input_tensor.layout, dim] | ||
| 60 | + return local_args, local_kwargs, cache_values | ||
| 61 | + | ||
| 62 | + def infer_layout(self, cache_values: list) -> Tuple[tuple, None]: | ||
| 63 | + """ | ||
| 64 | + Infer output layout for Cumsum operator. | ||
| 65 | + | ||
| 66 | + Rules: | ||
| 67 | + 1. Input must not have Partial status. | ||
| 68 | + 2. dim must be an integer within the valid range [-ndim, ndim-1]. | ||
| 69 | + 3. The cumsum dimension must not be sharded, including StridedShard mappings. | ||
| 70 | + 4. Output layout is identical to the input layout. | ||
| 71 | + | ||
| 72 | + Args: | ||
| 73 | + cache_values (list): [input_layout, dim] | ||
| 74 | + | ||
| 75 | + Returns: | ||
| 76 | + tuple: ((output_layout,), None) | ||
| 77 | + | ||
| 78 | + Raises: | ||
| 79 | + ValueError: If input has Partial status, dim is not an int, dim is out of range, | ||
| 80 | + or the cumsum dimension is sharded. | ||
| 81 | + """ | ||
| 82 | + layout = cache_values[0] | ||
| 83 | + dim = cache_values[1] | ||
| 84 | + | ||
| 85 | + if not self._allow_partial_inputs: | ||
| 86 | + self._check_partial_inputs([layout]) | ||
| 54 | 87 | ||
| 55 | - if not extra_args or extra_args[0] is None: | ||
| 56 | - raise ValueError( | ||
| 57 | - f"Operation {self.op_name}: cumsum requires 'dim' parameter in extra_args." | ||
| 58 | - ) | ||
| 59 | - dim = extra_args[0] | ||
| 60 | if not isinstance(dim, int): | 88 | if not isinstance(dim, int): |
| 61 | raise ValueError( | 89 | raise ValueError( |
| 62 | - f"Operation {self.op_name}: 'dim' must be an integer, got {type(dim)}." | 90 | + f"For {self.op_name}, dimension should be int, but got {type(dim)}" |
| 63 | ) | 91 | ) |
| 64 | 92 | ||
| 65 | - # Normalize negative dimensions | 93 | + alias_map = layout.alias_tensor_map |
H | |||
| 66 | - normalized_dim = dim if dim >= 0 else dim + input_ndim | 94 | + ndim = len(alias_map) |
| 67 | - if normalized_dim < 0 or normalized_dim >= input_ndim: | 95 | + if dim < -ndim or dim >= ndim: |
| 68 | raise ValueError( | 96 | raise ValueError( |
| 69 | - f"Operation {self.op_name}: Dimension {dim} out of range for " | 97 | + f"For {self.op_name}, dimension out of range " |
| 70 | - f"{input_ndim}-dimensional input tensor." | 98 | + f"(expected to be in range of [{-ndim}, {ndim - 1}], but got {dim})" |
| 99 | + ) | ||
| 100 | + if dim < 0: | ||
| 101 | + dim += ndim | ||
| 102 | + | ||
| 103 | + mapping = alias_map[dim] | ||
| 104 | + if mapping != "None": | ||
| 105 | + raise ValueError( | ||
| 106 | + f"For {self.op_name}, cumsum along a sharded dimension " | ||
| 107 | + f"(dim {dim} mapped to {mapping}) is not supported." | ||
| 71 | ) | 108 | ) |
| 72 | 109 | ||
| 73 | - # cumsum dimension must be unsharded | 110 | + return ((copy.deepcopy(layout),), None) |
| 74 | - if in_tensor_map[normalized_dim] != -1: | ||
| 75 | - raise ValueError( | ||
| 76 | - f"Operation {self.op_name}: Cannot perform sharding on normalized dimension {normalized_dim}, " | ||
| 77 | - f"but found sharding assignment: {in_tensor_map[normalized_dim]}" | ||
| 78 | - ) | ||
| 79 | - | ||
| 80 | - mesh_shape = input_layout.mesh_shape | ||
| 81 | - alias_name = input_layout.alias_name | ||
| 82 | - rank_list = input_layout.rank_list | ||
| 83 | - | ||
| 84 | - # Create output layout | ||
| 85 | - def idx_to_alias(idx, aliases): | ||
| 86 | - if idx == -1: | ||
| 87 | - return "None" | ||
| 88 | - return aliases[len(aliases) - idx - 1] | ||
| 89 | - output_map = tuple(idx_to_alias(idx, alias_name) for idx in in_tensor_map) | ||
| 90 | - | ||
| 91 | - output_layout = Layout( | ||
| 92 | - mesh_shape=mesh_shape, | ||
| 93 | - alias_name=alias_name, | ||
| 94 | - rank_list=rank_list | ||
| 95 | - ) | ||
| 96 | - output_layout = output_layout(*output_map) | ||
| 97 | - return output_layout | ||
| @@ -17,9 +17,44 @@ Distributed implementation for Element-wise operator. | |||
| 17 | """ | 17 | """ |
| 18 | 18 | ||
| 19 | import copy | 19 | import copy |
| 20 | +from typing import Callable, Optional, Tuple | ||
| 21 | + | ||
| 20 | from .parallel_ops import DistributedOp | 22 | from .parallel_ops import DistributedOp |
| 21 | 23 | ||
| 22 | 24 | ||
| 25 | +def _unwrap_local_value(value): | ||
| 26 | + """Convert DTensor-like values to local tensors while preserving containers.""" | ||
| 27 | + if hasattr(value, "_layout"): | ||
| 28 | + return value.to_local() | ||
| 29 | + if isinstance(value, tuple): | ||
| 30 | + return tuple(_unwrap_local_value(item) for item in value) | ||
| 31 | + if isinstance(value, list): | ||
| 32 | + return [_unwrap_local_value(item) for item in value] | ||
| 33 | + return value | ||
| 34 | + | ||
| 35 | + | ||
| 36 | +def _collect_layout_and_shape(value): | ||
| 37 | + """Collect layout and shape from one argument for layout inference cache.""" | ||
| 38 | + layout = value.layout if hasattr(value, "_layout") else None | ||
| 39 | + shape = value.shape if hasattr(value, "shape") else None | ||
| 40 | + return layout, shape | ||
| 41 | + | ||
| 42 | + | ||
| 43 | +def _build_elementwise_cache_values(args, kwargs): | ||
| 44 | + """Build cache_values from the real element-wise input arguments.""" | ||
| 45 | + input_layouts = [] | ||
| 46 | + input_shapes = [] | ||
| 47 | + for value in args: | ||
| 48 | + layout, shape = _collect_layout_and_shape(value) | ||
| 49 | + input_layouts.append(layout) | ||
| 50 | + input_shapes.append(shape) | ||
| 51 | + for value in kwargs.values(): | ||
| 52 | + layout, shape = _collect_layout_and_shape(value) | ||
| 53 | + input_layouts.append(layout) | ||
| 54 | + input_shapes.append(shape) | ||
| 55 | + return [*input_layouts, input_shapes] | ||
| 56 | + | ||
| 57 | + | ||
| 23 | class ElementWiseDistributedOp(DistributedOp): | 58 | class ElementWiseDistributedOp(DistributedOp): |
| 24 | """ | 59 | """ |
| 25 | Base class for distributed element-wise operators. | 60 | Base class for distributed element-wise operators. |
| @@ -31,28 +66,51 @@ class ElementWiseDistributedOp(DistributedOp): | |||
| 31 | op_name (str): Name of the operator to register. | 66 | op_name (str): Name of the operator to register. |
| 32 | """ | 67 | """ |
| 33 | 68 | ||
| 34 | - def infer_layout(self, layouts, extra_args=None): | 69 | + def preprocess(self, args: tuple, kwargs: dict) -> tuple: |
| 70 | + """ | ||
| 71 | + Preprocess arguments for element-wise operators. | ||
| 72 | + | ||
| 73 | + NOTE: aclop packed-args normalization (for MindSpore aclop operators | ||
| 74 | + like Mod, StopGradient that pack args as ``(prim, name, (real_args...))``) | ||
| 75 | + is handled upstream in ``OpDispatcher._dispatch_layout_infer`` via | ||
| 76 | + ``_normalize_aclop_args``. This method receives clean unpacked args. | ||
| 77 | + | ||
| 78 | + Args: | ||
| 79 | + args (tuple): Positional arguments passed to the operator. | ||
| 80 | + kwargs (dict): Keyword arguments passed to the operator. | ||
| 81 | + | ||
| 82 | + Returns: | ||
| 83 | + tuple: (local_args, local_kwargs, cache_values) | ||
| 84 | + """ | ||
| 85 | + local_kwargs = {key: _unwrap_local_value(value) for key, value in kwargs.items()} | ||
| 86 | + | ||
| 87 | + local_args = tuple(_unwrap_local_value(arg) for arg in args) | ||
| 88 | + cache_values = _build_elementwise_cache_values(args, kwargs) | ||
H 这个分布式算子比较基础，迁移起来影响比较大，基础原则是保证存量功能正确不回退就行，能支持现在注册列表的所有算子。 从代码来看，我们后续还是要区分单输入（不需要shape），双输入（涉及广播，需要shape），各类输入可能涉及属性值等，目前统一缓存了shape也没能为他。后续可以按计划进一步优化。   | |||
| 89 | + return local_args, local_kwargs, cache_values | ||
| 90 | + | ||
| 91 | + def infer_layout(self, cache_values: list) -> Tuple[tuple, None]: | ||
| 35 | """ | 92 | """ |
| 36 | Infer output layouts for element-wise operations with broadcasting support. | 93 | Infer output layouts for element-wise operations with broadcasting support. |
| 37 | 94 | ||
| 38 | - For element-wise operations: | 95 | + Rules: |
| 39 | - - Supports broadcasting following NumPy broadcasting rules | 96 | + 1. Inputs must not have Partial status unless the operator explicitly allows it. |
| 40 | - - All inputs must have compatible shapes for broadcasting | 97 | + 2. Input shapes must be broadcast-compatible when shape information is available. |
| 41 | - - Output will have the broadcasted shape and appropriate sharding strategy | 98 | + 3. Broadcasting dimensions cannot be sharded. |
| 42 | - - Handles both simple and complex sharding patterns (including tuple-type tensor_maps) | 99 | + 4. Non-broadcast sharding patterns must be compatible across inputs. |
| 100 | + 5. Output layout uses the merged sharding strategy and merged Partial status. | ||
| 43 | 101 | ||
| 44 | Args: | 102 | Args: |
| 45 | - layouts (tuple): Tuple of layouts for input tensors | 103 | + cache_values (list): [input_layout_0, ..., input_layout_n, input_shapes]. |
| 46 | - extra_args: Extra arguments for the operation. It can be: | ||
| 47 | - - dict containing 'input_shapes' | ||
| 48 | - - list/tuple where the last element is input_shapes (WithShape path) | ||
| 49 | 104 | ||
| 50 | Returns: | 105 | Returns: |
| 51 | - Layout: Layout for output tensor with merged sharding strategy. | 106 | + tuple: ((output_layout,), None) |
| 52 | 107 | ||
| 53 | Raises: | 108 | Raises: |
| 54 | ValueError: If input layouts are not compatible for broadcasting. | 109 | ValueError: If input layouts are not compatible for broadcasting. |
| 55 | """ | 110 | """ |
| 111 | + layouts = tuple(cache_values[:-1]) | ||
| 112 | + input_shapes = cache_values[-1] if cache_values else None | ||
| 113 | + | ||
| 56 | if not layouts: | 114 | if not layouts: |
| 57 | return None | 115 | return None |
| 58 | 116 | ||
| @@ -61,23 +119,19 @@ class ElementWiseDistributedOp(DistributedOp): | |||
| 61 | if not valid_layouts: | 119 | if not valid_layouts: |
| 62 | return None | 120 | return None |
| 63 | 121 | ||
| 64 | - # Check partial inputs - ElementWiseDistributedOp does not support partial by default | ||
| 65 | - # This check is performed after basic layout validation | ||
| 66 | if not self._allow_partial_inputs: | 122 | if not self._allow_partial_inputs: |
| 67 | self._check_partial_inputs(layouts) | 123 | self._check_partial_inputs(layouts) |
| 68 | 124 | ||
| 69 | if len(valid_layouts) == 1: | 125 | if len(valid_layouts) == 1: |
| 70 | - return valid_layouts[0] | 126 | + return ((copy.deepcopy(valid_layouts[0]),), None) |
| 71 | - | ||
| 72 | - input_shapes = self._extract_input_shapes(extra_args) | ||
| 73 | 127 | ||
| 74 | if not input_shapes: | 128 | if not input_shapes: |
| 75 | - return self._handle_no_input_shapes(valid_layouts) | 129 | + return ((self._handle_no_input_shapes(valid_layouts),), None) |
| 76 | 130 | ||
| 77 | aligned_layouts, aligned_shapes = self._align_layouts_and_shapes(layouts, input_shapes) | 131 | aligned_layouts, aligned_shapes = self._align_layouts_and_shapes(layouts, input_shapes) |
| 78 | 132 | ||
| 79 | if len(aligned_layouts) <= 1 or len(aligned_layouts) != len(aligned_shapes): | 133 | if len(aligned_layouts) <= 1 or len(aligned_layouts) != len(aligned_shapes): |
| 80 | - return valid_layouts[0] | 134 | + return ((copy.deepcopy(valid_layouts[0]),), None) |
| 81 | 135 | ||
| 82 | output_shape = self._compute_output_shape(aligned_shapes) | 136 | output_shape = self._compute_output_shape(aligned_shapes) |
| 83 | merged_tensor_map, merged_partial = self._merge_all_layouts( | 137 | merged_tensor_map, merged_partial = self._merge_all_layouts( |
| @@ -89,20 +143,22 @@ class ElementWiseDistributedOp(DistributedOp): | |||
| 89 | 143 | ||
| 90 | self._check_all_inputs_broadcasts_and_partial(aligned_layouts, aligned_shapes, output_shape) | 144 | self._check_all_inputs_broadcasts_and_partial(aligned_layouts, aligned_shapes, output_shape) |
| 91 | 145 | ||
| 92 | - return self._create_output_layout(aligned_layouts[0], merged_tensor_map, merged_partial) | 146 | + output_layout = self._create_output_layout(aligned_layouts[0], merged_tensor_map, merged_partial) |
| 147 | + return ((output_layout,), None) | ||
| 93 | 148 | ||
| 94 | def _handle_no_input_shapes(self, valid_layouts): | 149 | def _handle_no_input_shapes(self, valid_layouts): |
| 95 | """ | 150 | """ |
| 96 | Handle the case when input shapes are not available. | 151 | Handle the case when input shapes are not available. |
| 97 | """ | 152 | """ |
| 98 | first_layout = valid_layouts[0] | 153 | first_layout = valid_layouts[0] |
| 154 | + first_alias_map = first_layout.alias_tensor_map | ||
| 99 | for layout in valid_layouts[1:]: | 155 | for layout in valid_layouts[1:]: |
| 100 | - if layout.tensor_map != first_layout.tensor_map: | 156 | + if layout.alias_tensor_map != first_alias_map: |
| 101 | raise ValueError( | 157 | raise ValueError( |
| 102 | f"For {self.op_name}, cannot infer layout without shapes: " | 158 | f"For {self.op_name}, cannot infer layout without shapes: " |
| 103 | - f"mismatched tensor_map {first_layout.tensor_map} vs {layout.tensor_map}." | 159 | + f"mismatched alias_tensor_map {first_alias_map} vs {layout.alias_tensor_map}." |
| 104 | ) | 160 | ) |
| 105 | - return first_layout | 161 | + return copy.deepcopy(first_layout) |
| 106 | 162 | ||
| 107 | def _align_layouts_and_shapes(self, layouts, input_shapes): | 163 | def _align_layouts_and_shapes(self, layouts, input_shapes): |
| 108 | """ | 164 | """ |
| @@ -169,24 +225,6 @@ class ElementWiseDistributedOp(DistributedOp): | |||
| 169 | 225 | ||
| 170 | return merged_tensor_map, merged_partial | 226 | return merged_tensor_map, merged_partial |
| 171 | 227 | ||
| 172 | - def _extract_input_shapes(self, extra_args): | ||
| 173 | - """ | ||
| 174 | - Extract input_shapes from extra_args. | ||
| 175 | - | ||
| 176 | - Compatible with: | ||
| 177 | - - dict: {"input_shapes": [...]} | ||
| 178 | - - list/tuple (WithShape dispatcher): extra_args = [..., input_shapes] | ||
| 179 | - """ | ||
| 180 | - if isinstance(extra_args, dict): | ||
| 181 | - return extra_args.get("input_shapes", None) | ||
| 182 | - | ||
| 183 | - if isinstance(extra_args, (list, tuple)) and extra_args: | ||
| 184 | - maybe_shapes = extra_args[-1] | ||
| 185 | - if isinstance(maybe_shapes, (list, tuple)): | ||
| 186 | - return maybe_shapes | ||
| 187 | - | ||
| 188 | - return None | ||
| 189 | - | ||
| 190 | def _merge_partial_status(self, partial1, partial2, merged_tensor_map, tensor_map1, tensor_map2, layouts): | 228 | def _merge_partial_status(self, partial1, partial2, merged_tensor_map, tensor_map1, tensor_map2, layouts): |
| 191 | """ | 229 | """ |
| 192 | Merge partial status from two inputs. | 230 | Merge partial status from two inputs. |
| @@ -313,61 +351,6 @@ class ElementWiseDistributedOp(DistributedOp): | |||
| 313 | f"Should be without Partial status for broadcasting without communication" | 351 | f"Should be without Partial status for broadcasting without communication" |
| 314 | ) | 352 | ) |
| 315 | 353 | ||
| 316 | - def _merge_tensor_maps_without_shape(self, layout1, layout2): | ||
| 317 | - """ | ||
| 318 | - Merge tensor_maps without shape information (for broadcasting scenarios). | ||
| 319 | - | ||
| 320 | - Merging rules without shape: | ||
| 321 | - - If both dimensions are not sharded: use -1 | ||
| 322 | - - If one is sharded and one is not: use the sharded one (assume broadcasting) | ||
| 323 | - - If both are sharded: they must be identical, otherwise raise error | ||
| 324 | - | ||
| 325 | - Args: | ||
| 326 | - layout1: Layout of the first input | ||
| 327 | - layout2: Layout of the second input | ||
| 328 | - | ||
| 329 | - Returns: | ||
| 330 | - tuple: Merged tensor_map | ||
| 331 | - | ||
| 332 | - Raises: | ||
| 333 | - ValueError: If sharding strategies conflict | ||
| 334 | - """ | ||
| 335 | - map1 = layout1.tensor_map if layout1.tensor_map else tuple() | ||
| 336 | - map2 = layout2.tensor_map if layout2.tensor_map else tuple() | ||
| 337 | - | ||
| 338 | - # Align ranks by padding with -1 | ||
| 339 | - max_len = max(len(map1), len(map2)) | ||
| 340 | - padded_map1 = (-1,) * (max_len - len(map1)) + map1 | ||
| 341 | - padded_map2 = (-1,) * (max_len - len(map2)) + map2 | ||
| 342 | - | ||
| 343 | - merged_map = [] | ||
| 344 | - for i, (m1, m2) in enumerate(zip(padded_map1, padded_map2)): | ||
| 345 | - m1_axis = self._normalize_tensor_map_element(m1) | ||
| 346 | - m2_axis = self._normalize_tensor_map_element(m2) | ||
| 347 | - | ||
| 348 | - m1_axis_for_compare = frozenset(m1_axis) | ||
| 349 | - m2_axis_for_compare = frozenset(m2_axis) | ||
| 350 | - | ||
| 351 | - m1_is_sharded = bool(m1_axis) | ||
| 352 | - m2_is_sharded = bool(m2_axis) | ||
| 353 | - | ||
| 354 | - if not m1_is_sharded and not m2_is_sharded: | ||
| 355 | - merged_map.append(-1) | ||
| 356 | - elif not m1_is_sharded: | ||
| 357 | - merged_map.append(self._denormalize_tensor_map_element(m2_axis)) | ||
| 358 | - elif not m2_is_sharded: | ||
| 359 | - merged_map.append(self._denormalize_tensor_map_element(m1_axis)) | ||
| 360 | - else: | ||
| 361 | - if m1_axis_for_compare != m2_axis_for_compare: | ||
| 362 | - raise ValueError( | ||
| 363 | - f"For {self.op_name}, inputs should have same sharding pattern, " | ||
| 364 | - f"but got confilcting sharding at dimension {i}, " | ||
| 365 | - f"input1 shaded on {m1_axis} and input2 shaded on {m2_axis}." | ||
| 366 | - ) | ||
| 367 | - merged_map.append(self._denormalize_tensor_map_element(m1_axis)) | ||
| 368 | - | ||
| 369 | - return tuple(merged_map) | ||
| 370 | - | ||
| 371 | def _broadcast_shapes(self, shape1, shape2): | 354 | def _broadcast_shapes(self, shape1, shape2): |
| 372 | """ | 355 | """ |
| 373 | Calculate the broadcasted shape of two shapes according to broadcasting rules. | 356 | Calculate the broadcasted shape of two shapes according to broadcasting rules. |
| @@ -621,33 +604,65 @@ class AddDistributedOp(ElementWiseWithPartialDistributedOp): | |||
| 621 | which is useful for operations like gradient accumulation where partial | 604 | which is useful for operations like gradient accumulation where partial |
| 622 | results need to be preserved through the computation graph. | 605 | results need to be preserved through the computation graph. |
| 623 | """ | 606 | """ |
| 624 | - def get_expand_impl(self, func, infer_result, layouts, extra_args=None): | 607 | + def infer_layout(self, cache_values: list) -> Tuple[tuple, None]: |
| 608 | + """ | ||
| 609 | + Infer output layout for Add operator. | ||
| 610 | + | ||
| 611 | + Rules: | ||
| 612 | + 1. Follow element-wise broadcasting and sharding merge rules. | ||
| 613 | + 2. Partial status is allowed and propagated. | ||
| 614 | + 3. Propagated Partial status must be "sum" or None. | ||
| 615 | + | ||
| 616 | + Args: | ||
| 617 | + cache_values (list): [input_layout_0, ..., input_layout_n, input_shapes]. | ||
| 618 | + | ||
| 619 | + Returns: | ||
| 620 | + tuple: ((output_layout,), None) | ||
| 621 | + | ||
| 622 | + Raises: | ||
| 623 | + ValueError: If propagated Partial status is not "sum" or None. | ||
| 624 | + """ | ||
| 625 | + infer_result = super().infer_layout(cache_values) | ||
| 626 | + if infer_result is None: | ||
| 627 | + return infer_result | ||
| 628 | + | ||
| 629 | + output_layout = infer_result[0][0] | ||
| 630 | + for i, partial_type in enumerate(output_layout.partial): | ||
| 631 | + if partial_type is not None and partial_type != "sum": | ||
| 632 | + raise ValueError( | ||
| 633 | + f"For {self.op_name}, inputs partial status should be 'sum' or None, " | ||
| 634 | + f"but got {partial_type} at index {i}." | ||
| 635 | + ) | ||
| 636 | + | ||
| 637 | + return infer_result | ||
| 638 | + | ||
| 639 | + def get_expand_impl(self, func: Optional[Callable], infer_result: tuple, | ||
| 640 | + cache_values: list) -> Optional[Callable]: | ||
| 625 | """ | 641 | """ |
| 626 | Get expand implementation for the operator | 642 | Get expand implementation for the operator |
| 627 | """ | 643 | """ |
| 644 | + layouts = tuple(cache_values[:-1]) | ||
| 628 | x1_layout = layouts[0] | 645 | x1_layout = layouts[0] |
| 629 | x2_layout = layouts[1] | 646 | x2_layout = layouts[1] |
| 630 | - x1_partial = x1_layout.is_partial() if x1_layout is not None else None | 647 | + x1_partial = x1_layout.is_partial() if x1_layout is not None else False |
| 631 | - x2_partial = x2_layout.is_partial() if x2_layout is not None else None | 648 | + x2_partial = x2_layout.is_partial() if x2_layout is not None else False |
| 632 | 649 | ||
| 633 | - if x1_partial != x2_partial: | 650 | + if x1_partial == x2_partial: |
| 634 | - scaling_factor = 1 | 651 | + return None |
| 635 | - for i, partial_type in enumerate(infer_result.partial): | ||
| 636 | - if partial_type == "sum": | ||
| 637 | - scaling_factor *= infer_result.mesh_shape[i] | ||
| 638 | - elif partial_type is not None: | ||
| 639 | - raise ValueError( | ||
| 640 | - f"For {self.op_name}, inputs partial status should be 'sum' or None, " | ||
| 641 | - f"but got {partial_type} at index {i}." | ||
| 642 | - ) | ||
| 643 | 652 | ||
| 644 | - # use expand_impl only when one of x1 and x2 is with partial placement. | 653 | + output_layout = infer_result[0][0] |
| 645 | - def expand_impl1(x1, x2): | 654 | + scaling_factor = 1 |
| 646 | - add_out = func(x1 / scaling_factor, x2) | 655 | + for i, partial_type in enumerate(output_layout.partial): |
| 647 | - return add_out | 656 | + if partial_type == "sum": |
| 657 | + scaling_factor *= output_layout.mesh_shape[i] | ||
| 648 | 658 | ||
| 649 | - def expand_impl2(x1, x2): | 659 | + # use expand_impl only when one of x1 and x2 is with partial placement. |
| 650 | - add_out = func(x1, x2 / scaling_factor) | 660 | + def _expand_impl1(x1, x2, *args, **kwargs): |
| 651 | - return add_out | 661 | + add_out = func(x1 / scaling_factor, x2, *args, **kwargs) |
| 652 | - return expand_impl1 if not x1_partial else expand_impl2 | 662 | + return add_out |
| 653 | - return None | 663 | + |
| 664 | + def _expand_impl2(x1, x2, *args, **kwargs): | ||
| 665 | + add_out = func(x1, x2 / scaling_factor, *args, **kwargs) | ||
| 666 | + return add_out | ||
| 667 | + | ||
| 668 | + return _expand_impl1 if not x1_partial else _expand_impl2 | ||
| @@ -12,37 +12,109 @@ | |||
| 12 | # See the License for the specific language governing permissions and | 12 | # See the License for the specific language governing permissions and |
| 13 | # limitations under the License. | 13 | # limitations under the License. |
| 14 | # ============================================================================ | 14 | # ============================================================================ |
| 15 | -""" | 15 | +"""Distributed implementation for the embedding operator.""" |
| 16 | -Distributed implementation for the embedding operator. | 16 | + |
| 17 | -This module handles Tensor Parallelism for Embedding layers (Column/Row Parallel). | 17 | +from typing import Callable, Optional, Tuple |
| 18 | -""" | ||
| 19 | 18 | ||
| 20 | from hyper_parallel.core.dtensor.layout import Layout | 19 | from hyper_parallel.core.dtensor.layout import Layout |
| 21 | from .parallel_ops import DistributedOp | 20 | from .parallel_ops import DistributedOp |
| 22 | 21 | ||
| 23 | 22 | ||
| 23 | +def _normalize_embedding_args(input_tensor, weight_tensor, padding_idx=None, max_norm=None, | ||
| 24 | + norm_type=2.0, scale_grad_by_freq=False, sparse=False): | ||
| 25 | + return ( | ||
| 26 | + input_tensor, | ||
| 27 | + weight_tensor, | ||
| 28 | + padding_idx, | ||
| 29 | + max_norm, | ||
| 30 | + norm_type, | ||
| 31 | + scale_grad_by_freq, | ||
| 32 | + sparse, | ||
| 33 | + ), {} | ||
| 34 | + | ||
| 35 | + | ||
| 24 | class EmbeddingDistributedOp(DistributedOp): | 36 | class EmbeddingDistributedOp(DistributedOp): |
| 25 | """ | 37 | """ |
| 26 | Distributed implementation for embedding operators. | 38 | Distributed implementation for embedding operators. |
| 27 | Supports Column Parallelism (CP) and Row Parallelism (RP). | 39 | Supports Column Parallelism (CP) and Row Parallelism (RP). |
| 28 | """ | 40 | """ |
| 41 | + _MS_PRIMITIVE_OP_NAMES = frozenset({'Embedding'}) | ||
| 29 | 42 | ||
| 30 | - def infer_layout(self, layouts, extra_args=None): | 43 | + def preprocess(self, args: tuple, kwargs: dict) -> tuple: |
| 31 | """ | 44 | """ |
| 32 | - Infers the output layout (sharding state) based on input and weight layouts. | 45 | + Preprocess arguments for Embedding operator. |
| 46 | + | ||
| 47 | + Args: | ||
| 48 | + args (tuple): Input arguments containing input and weight tensors. | ||
| 49 | + kwargs (dict): Keyword arguments for embedding options. | ||
| 50 | + | ||
| 51 | + Returns: | ||
| 52 | + tuple: (local_args, local_kwargs, cache_values) where cache_values is | ||
| 53 | + [input_layout, weight_layout]. | ||
| 33 | """ | 54 | """ |
| 34 | - if not layouts or len(layouts) < 2: | 55 | + args, _ = _normalize_embedding_args(*args, **kwargs) |
| 56 | + input_tensor, weight_tensor = args[0], args[1] | ||
| 57 | + if self.op_name in self._MS_PRIMITIVE_OP_NAMES: | ||
| 58 | + local_args = (input_tensor.to_local(), weight_tensor.to_local()) + args[2:6] | ||
| 59 | + else: | ||
| 60 | + local_args = (input_tensor.to_local(), weight_tensor.to_local()) + args[2:] | ||
| 61 | + local_kwargs = {} | ||
| 62 | + cache_values = [input_tensor.layout, weight_tensor.layout] | ||
| 63 | + return local_args, local_kwargs, cache_values | ||
| 64 | + | ||
| 65 | + def infer_layout(self, cache_values: list) -> Tuple[tuple, None]: | ||
| 66 | + """ | ||
| 67 | + Infer output layout for Embedding operator. | ||
| 68 | + | ||
| 69 | + Rules: | ||
| 70 | + 1. Input and weight must not have Partial status. | ||
| 71 | + 2. Input and weight must share the same mesh_shape. | ||
| 72 | + 3. Weight must be 2D [vocab_size, embedding_dim]. | ||
| 73 | + 4. Output shape is [*input_shape, embedding_dim], preserving input sharding | ||
| 74 | + and inheriting weight embedding-dimension sharding. | ||
| 75 | + 5. If weight vocab dimension is sharded, output carries Partial('sum') on | ||
| 76 | + that vocab sharding axis. | ||
| 77 | + | ||
| 78 | + Args: | ||
| 79 | + cache_values (list): [input_layout, weight_layout]. | ||
| 80 | + | ||
| 81 | + Returns: | ||
| 82 | + tuple: ((output_layout,), None) | ||
| 83 | + | ||
| 84 | + Raises: | ||
| 85 | + ValueError: If layouts are missing, partial, incompatible, or weight is not 2D. | ||
| 86 | + """ | ||
| 87 | + if len(cache_values) != 2: | ||
| 35 | raise ValueError( | 88 | raise ValueError( |
| 36 | - f"Operation {self.op_name}: requires both input and weight layouts." | 89 | + f"For {self.op_name}, cache_values length should be 2, but got {len(cache_values)}" |
| 90 | + ) | ||
| 91 | + | ||
| 92 | + input_layout = cache_values[0] | ||
| 93 | + weight_layout = cache_values[1] | ||
| 94 | + if not input_layout or not weight_layout: | ||
| 95 | + raise ValueError( | ||
| 96 | + f"For {self.op_name}, requires both input and weight layouts." | ||
| 97 | + ) | ||
| 98 | + | ||
| 99 | + self._check_partial_inputs([input_layout, weight_layout]) | ||
| 100 | + | ||
| 101 | + if input_layout.mesh_shape != weight_layout.mesh_shape: | ||
| 102 | + raise ValueError( | ||
| 103 | + f"For {self.op_name}, input and weight must have the same mesh_shape, " | ||
| 104 | + f"but got input: {input_layout.mesh_shape} and weight: {weight_layout.mesh_shape}" | ||
| 37 | ) | 105 | ) |
| 38 | 106 | ||
| 39 | - input_layout = layouts[0] | ||
| 40 | - weight_layout = layouts[1] | ||
| 41 | weight_tensor_map = weight_layout.tensor_map | 107 | weight_tensor_map = weight_layout.tensor_map |
| 108 | + if len(weight_tensor_map) != 2: | ||
| 109 | + raise ValueError( | ||
| 110 | + f"For {self.op_name}, weight should be 2D [vocab_size, embedding_dim], " | ||
| 111 | + f"but got {len(weight_tensor_map)}D" | ||
| 112 | + ) | ||
| 42 | 113 | ||
| 43 | # weight_tensor_map: [vocab_size_dim, embed_dim_dim] | 114 | # weight_tensor_map: [vocab_size_dim, embed_dim_dim] |
| 44 | - w_shard_vocab_axis = weight_tensor_map[0] | ||
| 45 | w_shard_embed_axis = weight_tensor_map[1] | 115 | w_shard_embed_axis = weight_tensor_map[1] |
| 116 | + weight_alias_map = weight_layout.alias_tensor_map | ||
| 117 | + w_shard_vocab_alias = weight_alias_map[0] | ||
| 46 | 118 | ||
| 47 | # Output shape is [*input_shape, embed_dim] | 119 | # Output shape is [*input_shape, embed_dim] |
| 48 | output_tensor_map = list(input_layout.tensor_map) | 120 | output_tensor_map = list(input_layout.tensor_map) |
| @@ -56,18 +128,17 @@ class EmbeddingDistributedOp(DistributedOp): | |||
| 56 | output_layout.set_tensor_map(tuple(output_tensor_map)) | 128 | output_layout.set_tensor_map(tuple(output_tensor_map)) |
| 57 | 129 | ||
| 58 | # If vocab is sharded (Row Parallelism), output is in Partial Sum state | 130 | # If vocab is sharded (Row Parallelism), output is in Partial Sum state |
| 59 | - if w_shard_vocab_axis != -1: | 131 | + if w_shard_vocab_alias != "None": |
| 60 | # pylint: disable=protected-access | 132 | # pylint: disable=protected-access |
| 61 | output_layout._partial = list(input_layout.partial) | 133 | output_layout._partial = list(input_layout.partial) |
| 62 | - vocab_axis_name = input_layout.alias_name[len(input_layout.alias_name) - 1 - w_shard_vocab_axis] | 134 | + output_layout.set_partial_by_dev_axis(w_shard_vocab_alias, 'sum') |
| 63 | - output_layout.set_partial_by_dev_axis(vocab_axis_name, 'sum') | ||
| 64 | # pylint: disable=protected-access | 135 | # pylint: disable=protected-access |
| 65 | output_layout._alias_tensor_map = output_layout._build_readable_tensor_map() | 136 | output_layout._alias_tensor_map = output_layout._build_readable_tensor_map() |
| 66 | # pylint: disable=protected-access | 137 | # pylint: disable=protected-access |
| 67 | output_layout.tensor_map_to_placement() | 138 | output_layout.tensor_map_to_placement() |
| 68 | output_layout.update_compact_str() | 139 | output_layout.update_compact_str() |
| 69 | 140 | ||
| 70 | - return output_layout | 141 | + return ((output_layout,), None) |
| 71 | 142 | ||
| 72 | def _parse_params(self, args, kwargs): | 143 | def _parse_params(self, args, kwargs): |
| 73 | """ | 144 | """ |
| @@ -122,17 +193,20 @@ class EmbeddingDistributedOp(DistributedOp): | |||
| 122 | 193 | ||
| 123 | return local_input, mask_int | 194 | return local_input, mask_int |
| 124 | 195 | ||
| 125 | - def get_expand_impl(self, func, infer_result, layouts, extra_args=None): | 196 | + def get_expand_impl(self, func: Optional[Callable], infer_result: tuple, |
| 197 | + cache_values: list) -> Optional[Callable]: | ||
| 126 | """ | 198 | """ |
| 127 | Returns the execution implementation wrapper. | 199 | Returns the execution implementation wrapper. |
| 128 | Helper functions are used to keep Cyclomatic Complexity (CCN) low. | 200 | Helper functions are used to keep Cyclomatic Complexity (CCN) low. |
| 129 | """ | 201 | """ |
| 130 | - weight_layout = layouts[1] | 202 | + weight_layout = cache_values[1] |
| 131 | w_shard_vocab_axis = weight_layout.tensor_map[0] | 203 | w_shard_vocab_axis = weight_layout.tensor_map[0] |
| 132 | - w_shard_embed_axis = weight_layout.tensor_map[1] | 204 | + weight_alias_map = weight_layout.alias_tensor_map |
| 205 | + w_shard_vocab_alias = weight_alias_map[0] | ||
| 206 | + w_shard_embed_alias = weight_alias_map[1] | ||
| 133 | 207 | ||
| 134 | # Use native implementation if no weight sharding is applied | 208 | # Use native implementation if no weight sharding is applied |
| 135 | - if w_shard_vocab_axis == -1 and w_shard_embed_axis == -1: | 209 | + if w_shard_vocab_alias == "None" and w_shard_embed_alias == "None": |
| 136 | return None | 210 | return None |
| 137 | 211 | ||
| 138 | def distributed_embedding_impl(*args, **kwargs): | 212 | def distributed_embedding_impl(*args, **kwargs): |
| @@ -144,18 +218,24 @@ class EmbeddingDistributedOp(DistributedOp): | |||
| 144 | 218 | ||
| 145 | # Check for max_norm with specific error messages for CP and RP | 219 | # Check for max_norm with specific error messages for CP and RP |
| 146 | if max_norm is not None: | 220 | if max_norm is not None: |
| 147 | - if w_shard_embed_axis != -1: | 221 | + if w_shard_embed_alias != "None": |
| 148 | - raise ValueError("Column-Parallel Embedding does not support `max_norm` parameter.") | 222 | + raise ValueError( |
| 149 | - if w_shard_vocab_axis != -1: | 223 | + f"For {self.op_name}, Column-Parallel Embedding does not support `max_norm` parameter." |
| 150 | - raise ValueError("Row-Parallel Embedding does not support `max_norm` parameter.") | 224 | + ) |
| 225 | + if w_shard_vocab_alias != "None": | ||
| 226 | + raise ValueError( | ||
| 227 | + f"For {self.op_name}, Row-Parallel Embedding does not support `max_norm` parameter." | ||
| 228 | + ) | ||
| 151 | 229 | ||
| 152 | # Check for scale_grad_by_freq with RP | 230 | # Check for scale_grad_by_freq with RP |
| 153 | - if scale_grad and w_shard_vocab_axis != -1: | 231 | + if scale_grad and w_shard_vocab_alias != "None": |
| 154 | - raise ValueError("Row-Parallel Embedding does not support `scale_grad_by_freq=True`.") | 232 | + raise ValueError( |
| 233 | + f"For {self.op_name}, Row-Parallel Embedding does not support `scale_grad_by_freq=True`." | ||
| 234 | + ) | ||
| 155 | 235 | ||
| 156 | # 2. Row Parallel Processing | 236 | # 2. Row Parallel Processing |
| 157 | input_mask_int = None | 237 | input_mask_int = None |
| 158 | - if w_shard_vocab_axis != -1: | 238 | + if w_shard_vocab_alias != "None": |
| 159 | is_args_pad = len(args) > 2 | 239 | is_args_pad = len(args) > 2 |
| 160 | mapped_input, input_mask_int = self._handle_rp_input( | 240 | mapped_input, input_mask_int = self._handle_rp_input( |
| 161 | input_tensor, weight_tensor, weight_layout, w_shard_vocab_axis, | 241 | input_tensor, weight_tensor, weight_layout, w_shard_vocab_axis, |
| @@ -167,7 +247,7 @@ class EmbeddingDistributedOp(DistributedOp): | |||
| 167 | output = func(*new_args, **new_kwargs) | 247 | output = func(*new_args, **new_kwargs) |
| 168 | 248 | ||
| 169 | # 4. Erase invalid partial embeddings (Row-Parallel only) | 249 | # 4. Erase invalid partial embeddings (Row-Parallel only) |
| 170 | - if w_shard_vocab_axis != -1 and input_mask_int is not None: | 250 | + if w_shard_vocab_alias != "None" and input_mask_int is not None: |
| 171 | expanded_mask = input_mask_int[..., None].to(output.dtype) | 251 | expanded_mask = input_mask_int[..., None].to(output.dtype) |
| 172 | output = output * expanded_mask | 252 | output = output * expanded_mask |
| 173 | 253 | ||
| @@ -15,45 +15,92 @@ | |||
| 15 | """ | 15 | """ |
| 16 | Distributed implementation for ExpandDims operator. | 16 | Distributed implementation for ExpandDims operator. |
| 17 | """ | 17 | """ |
| 18 | +from typing import Tuple | ||
| 19 | + | ||
| 18 | from hyper_parallel.core.dtensor.layout import Layout | 20 | from hyper_parallel.core.dtensor.layout import Layout |
| 19 | from .parallel_ops import DistributedOp | 21 | from .parallel_ops import DistributedOp |
| 20 | 22 | ||
| 21 | 23 | ||
| 24 | +def _normalize_expand_dims_args(x, axis=None, dim=None): | ||
| 25 | + if axis is None: | ||
| 26 | + axis = dim | ||
| 27 | + return (x, axis), {} | ||
| 28 | + | ||
| 29 | + | ||
| 22 | class ExpandDimsDistributedOp(DistributedOp): | 30 | class ExpandDimsDistributedOp(DistributedOp): |
| 23 | """Distributed implementation for ExpandDims operator.""" | 31 | """Distributed implementation for ExpandDims operator.""" |
| 24 | 32 | ||
| 25 | - def infer_layout(self, layouts, extra_args=None): | 33 | + def preprocess(self, args: tuple, kwargs: dict) -> tuple: |
| 26 | """ | 34 | """ |
| 27 | - Infer output layout for ExpandDims. | 35 | + Preprocess arguments for ExpandDims operator. |
| 28 | 36 | ||
| 29 | Args: | 37 | Args: |
| 30 | - layouts (tuple): Tuple containing input layout. | 38 | + args (tuple): Input arguments, first element is the input tensor. |
| 31 | - extra_args: axis parameter (int or list/tuple containing int). | 39 | + kwargs (dict): Keyword arguments (axis or dim). |
| 32 | 40 | ||
| 33 | Returns: | 41 | Returns: |
| 34 | - Layout: Output layout with inserted dimension. | 42 | + tuple: (local_args, local_kwargs, cache_values) |
| 35 | """ | 43 | """ |
| 36 | - if not layouts: | 44 | + args, _ = _normalize_expand_dims_args(*args, **kwargs) |
| 37 | - raise ValueError(f"{self.__class__.__name__} requires at least one input layout") | 45 | + input_tensor, axis = args[0], args[1] |
| 46 | + local_args = (input_tensor.to_local(), axis) | ||
| 47 | + cache_values = [input_tensor.layout, axis] | ||
| 48 | + return local_args, {}, cache_values | ||
| 38 | 49 | ||
| 39 | - x_layout = layouts[0] | 50 | + def infer_layout(self, cache_values: list) -> Tuple[tuple, None]: |
| 51 | + """ | ||
| 52 | + Infer output layout for ExpandDims operator. | ||
| 53 | + | ||
| 54 | + Rules: | ||
| 55 | + 1. Input must not have Partial status. | ||
| 56 | + 2. axis must be an integer within the valid range [-(rank + 1), rank]. | ||
| 57 | + 3. The inserted dimension is replicated. | ||
| 58 | + 4. Existing input dimension mappings are shifted and otherwise preserved. | ||
| 59 | + | ||
| 60 | + Args: | ||
| 61 | + cache_values (list): [input_layout, axis] | ||
| 62 | + | ||
| 63 | + Returns: | ||
| 64 | + tuple: ((output_layout,), None) | ||
| 65 | + | ||
| 66 | + Raises: | ||
| 67 | + ValueError: If input has Partial status, input layout is missing, | ||
| 68 | + axis is missing or invalid, or axis is out of range. | ||
| 69 | + """ | ||
| 70 | + if not cache_values: | ||
| 71 | + raise ValueError( | ||
| 72 | + f"For {self.op_name}, cache_values should contain input layout, " | ||
| 73 | + f"but got empty cache_values." | ||
| 74 | + ) | ||
| 75 | + | ||
| 76 | + x_layout = cache_values[0] | ||
| 77 | + if not self._allow_partial_inputs: | ||
| 78 | + self._check_partial_inputs([x_layout]) | ||
| 40 | 79 | ||
| 41 | if x_layout.mesh_shape is None: | 80 | if x_layout.mesh_shape is None: |
| 42 | - raise ValueError("Input layout cannot be None.") | 81 | + raise ValueError( |
| 82 | + f"For {self.op_name}, input layout mesh_shape should not be None, " | ||
| 83 | + f"but got None." | ||
| 84 | + ) | ||
| 43 | 85 | ||
| 44 | - if not extra_args: | 86 | + axis = cache_values[1] if len(cache_values) > 1 else None |
| 45 | - raise ValueError(f"{self.op_name}: axis parameter is required") | ||
| 46 | 87 | ||
| 47 | - axis = extra_args[0] if isinstance(extra_args, (tuple, list)) else extra_args | 88 | + if axis is None: |
| 89 | + raise ValueError(f"For {self.op_name}, axis parameter is required.") | ||
| 90 | + if not isinstance(axis, int): | ||
| 91 | + raise ValueError( | ||
| 92 | + f"For {self.op_name}, axis should be int, but got {type(axis)}." | ||
| 93 | + ) | ||
| 48 | 94 | ||
| 49 | in_rank = len(x_layout.alias_tensor_map) | 95 | in_rank = len(x_layout.alias_tensor_map) |
| 96 | + original_axis = axis | ||
| 50 | if axis < 0: | 97 | if axis < 0: |
| 51 | axis = axis + in_rank + 1 | 98 | axis = axis + in_rank + 1 |
| 52 | 99 | ||
| 53 | if axis < 0 or axis > in_rank: | 100 | if axis < 0 or axis > in_rank: |
| 54 | raise ValueError( | 101 | raise ValueError( |
| 55 | - f"{self.op_name}: axis {axis} out of range for input rank {in_rank}. " | 102 | + f"For {self.op_name}, axis {original_axis} out of range for input rank {in_rank}. " |
| 56 | - f"Valid range is [{-in_rank-1}, {in_rank}]" | 103 | + f"Valid range is [{-in_rank - 1}, {in_rank}]." |
| 57 | ) | 104 | ) |
| 58 | 105 | ||
| 59 | x_map = list(x_layout.alias_tensor_map) | 106 | x_map = list(x_layout.alias_tensor_map) |
| @@ -66,9 +113,10 @@ class ExpandDimsDistributedOp(DistributedOp): | |||
| 66 | ) | 113 | ) |
| 67 | output_layout = output_layout(*x_map) | 114 | output_layout = output_layout(*x_map) |
| 68 | 115 | ||
| 69 | - for i, partial_op in enumerate(x_layout.partial): | 116 | + if self._allow_partial_inputs: |
| 70 | - if partial_op is not None: | 117 | + for i, partial_op in enumerate(x_layout.partial): |
| 71 | - dev_axis_name = x_layout.alias_name[i] | 118 | + if partial_op is not None: |
| 72 | - output_layout.set_partial_by_dev_axis(dev_axis_name, partial_op) | 119 | + dev_axis_name = x_layout.alias_name[i] |
| 120 | + output_layout.set_partial_by_dev_axis(dev_axis_name, partial_op) | ||
| 73 | 121 | ||
| 74 | - return output_layout | 122 | + return ((output_layout,), None) |
| @@ -15,96 +15,108 @@ | |||
| 15 | """ | 15 | """ |
| 16 | Distributed implementation for Flatten operator. | 16 | Distributed implementation for Flatten operator. |
| 17 | """ | 17 | """ |
| 18 | +from typing import Tuple | ||
| 19 | + | ||
| 18 | from hyper_parallel.core.shard.ops.parallel_reshape import ReshapeDistributedOp | 20 | from hyper_parallel.core.shard.ops.parallel_reshape import ReshapeDistributedOp |
| 19 | 21 | ||
| 20 | 22 | ||
| 23 | +def _normalize_flatten_args(x, start_dim=0, end_dim=-1): | ||
| 24 | + return (x, start_dim, end_dim), {} | ||
| 25 | + | ||
| 26 | + | ||
| 21 | class FlattenDistributedOp(ReshapeDistributedOp): | 27 | class FlattenDistributedOp(ReshapeDistributedOp): |
| 22 | """Distributed implementation for torch.flatten.""" | 28 | """Distributed implementation for torch.flatten.""" |
| 23 | 29 | ||
| 24 | - def infer_layout(self, layouts, extra_args=None): | 30 | + def preprocess(self, args: tuple, kwargs: dict) -> tuple: |
| 25 | """ | 31 | """ |
| 26 | - Infer output layout for torch.flatten. | 32 | + Preprocess arguments for Flatten operator. |
| 27 | - | ||
| 28 | - PyTorch semantics: | ||
| 29 | - - flatten(input, start_dim=0, end_dim=-1) | ||
| 30 | - - Flattens the input tensor starting from `start_dim` to `end_dim`. | ||
| 31 | 33 | ||
| 32 | Args: | 34 | Args: |
| 33 | - layouts (tuple): Layouts of inputs. | 35 | + args (tuple): Input tensor followed by optional start_dim and end_dim. |
| 34 | - extra_args (tuple): Contains scalar arguments (start_dim, end_dim) | 36 | + kwargs (dict): Optional keyword arguments. |
| 35 | - and the input shapes list appended by WithShape suffix. | ||
| 36 | 37 | ||
| 37 | Returns: | 38 | Returns: |
| 38 | - Layout: Output tensor layout. | 39 | + tuple: (local_args, local_kwargs, cache_values) |
| 39 | """ | 40 | """ |
| 40 | - if not layouts or layouts[0] is None: | 41 | + args, _ = _normalize_flatten_args(*args, **kwargs) |
| 42 | + input_tensor, start_dim, end_dim = args[0], args[1], args[2] | ||
| 43 | + local_args = (input_tensor.to_local(), start_dim, end_dim) | ||
| 44 | + local_kwargs = {} | ||
| 45 | + cache_values = [input_tensor.layout, start_dim, end_dim, tuple(input_tensor.shape)] | ||
| 46 | + return local_args, local_kwargs, cache_values | ||
| 47 | + | ||
| 48 | + def infer_layout(self, cache_values: list) -> Tuple[tuple, None]: | ||
| 49 | + """ | ||
| 50 | + Infer output layout for Flatten operator. | ||
| 51 | + | ||
| 52 | + Rules: | ||
| 53 | + 1. Partial input is allowed and preserved by reshape layout inference. | ||
| 54 | + 2. input_shape must be provided and match the input rank. | ||
| 55 | + 3. start_dim and end_dim must be integers within the valid input rank. | ||
| 56 | + 4. If start_dim >= end_dim after normalization, the output layout is the | ||
| 57 | + input layout. | ||
| 58 | + 5. Otherwise, dimensions from start_dim through end_dim are merged using | ||
| 59 | + reshape-compatible sharding rules. | ||
| 60 | + | ||
| 61 | + Args: | ||
| 62 | + cache_values (list): [input_layout, start_dim, end_dim, input_shape]. | ||
| 63 | + | ||
| 64 | + Returns: | ||
| 65 | + tuple: ((output_layout,), None) | ||
| 66 | + | ||
| 67 | + Raises: | ||
| 68 | + ValueError: If cache_values are invalid, dimensions are out of range, | ||
| 69 | + or the flatten would change sharded slices incompatibly. | ||
| 70 | + """ | ||
| 71 | + input_layout, start_dim, end_dim, input_shape = ( | ||
| 72 | + cache_values[0], cache_values[1], cache_values[2], cache_values[3] | ||
| 73 | + ) | ||
| 74 | + if input_layout is None: | ||
| 41 | raise ValueError( | 75 | raise ValueError( |
| 42 | - f"Operation {self.op_name}: flatten requires a valid input tensor layout." | 76 | + f"For {self.op_name}, flatten requires a valid input tensor layout." |
| 77 | + ) | ||
| 78 | + if not isinstance(input_shape, (list, tuple)): | ||
| 79 | + raise ValueError( | ||
| 80 | + f"For {self.op_name}, input_shape should be list or tuple, " | ||
| 81 | + f"but got {type(input_shape)}." | ||
| 82 | + ) | ||
| 83 | + if len(input_shape) != len(input_layout.tensor_map): | ||
| 84 | + raise ValueError( | ||
| 85 | + f"For {self.op_name}, input shape rank should match layout rank, " | ||
| 86 | + f"but got {len(input_shape)} and {len(input_layout.tensor_map)}." | ||
| 87 | + ) | ||
| 88 | + if not isinstance(start_dim, int) or not isinstance(end_dim, int): | ||
| 89 | + raise ValueError( | ||
| 90 | + f"For {self.op_name}, start_dim and end_dim should be int, " | ||
| 91 | + f"but got {type(start_dim)} and {type(end_dim)}." | ||
| 43 | ) | 92 | ) |
| 44 | - | ||
| 45 | - input_layout = layouts[0] | ||
| 46 | - # WithShape suffix appends a list of input_shapes as the last item in extra_args | ||
| 47 | - input_shapes = extra_args[-1] | ||
| 48 | - input_shape = input_shapes[0] | ||
| 49 | - | ||
| 50 | - # PyTorch flatten defaults: start_dim=0, end_dim=-1 | ||
| 51 | - start_dim = 0 | ||
| 52 | - end_dim = -1 | ||
| 53 | - | ||
| 54 | - # Parse scalar args (start_dim, end_dim) from extra_args | ||
| 55 | - num_scalar_args = len(extra_args) - 1 | ||
| 56 | - if num_scalar_args > 0: | ||
| 57 | - start_dim = extra_args[0] | ||
| 58 | - if num_scalar_args > 1: | ||
| 59 | - end_dim = extra_args[1] | ||
| 60 | 93 | ||
| 61 | ndim = len(input_shape) | 94 | ndim = len(input_shape) |
| 62 | 95 | ||
| 63 | - # Handle 0-D tensor case specifically to avoid parent class indexing errors | ||
| 64 | if ndim == 0: | 96 | if ndim == 0: |
| 65 | - out_layout = layouts[0].__class__.from_device_mesh(input_layout.mesh) | 97 | + out_layout = input_layout.__class__.from_device_mesh(input_layout.mesh) |
| 66 | out_layout.set_placements(input_layout.placements) | 98 | out_layout.set_placements(input_layout.placements) |
| 67 | - out_layout.placement_to_tensor_map(1) # Flattened 0-D becomes 1-D of shape (1,) | 99 | + out_layout.placement_to_tensor_map(1) |
| 68 | - return out_layout | 100 | + return ((out_layout,), None) |
| 69 | 101 | ||
| 70 | - # Handle negative dimensions | ||
| 71 | if start_dim < 0: | 102 | if start_dim < 0: |
| 72 | start_dim += ndim | 103 | start_dim += ndim |
| 73 | if end_dim < 0: | 104 | if end_dim < 0: |
| 74 | end_dim += ndim | 105 | end_dim += ndim |
| 75 | 106 | ||
| 76 | - # Validate dimension bounds | ||
| 77 | if start_dim < 0 or start_dim >= ndim or end_dim < 0 or end_dim >= ndim: | 107 | if start_dim < 0 or start_dim >= ndim or end_dim < 0 or end_dim >= ndim: |
| 78 | - raise ValueError(f"Dimension out of range for flatten: start_dim={start_dim}, end_dim={end_dim}") | 108 | + raise ValueError( |
| 109 | + f"For {self.op_name}, dimension out of range " | ||
| 110 | + f"(start_dim={start_dim}, end_dim={end_dim}, ndim={ndim})." | ||
| 111 | + ) | ||
| 79 | 112 | ||
| 80 | - | ||
| 81 | - # If start_dim > end_dim, PyTorch returns the tensor unchanged (identity operation). | ||
| 82 | - # This also covers the case where start_dim == end_dim (e.g., flatten(1, 1)) | ||
| 83 | - # because the loop `range(start_dim, end_dim + 1)` would still execute, but | ||
| 84 | - # effectively for a single dimension, resulting in `dst_shape == input_shape`. | ||
| 85 | - # However, the underlying `_merge_unshared_axis` logic in ReshapeDistributedOp | ||
| 86 | - # is not designed for an identity mapping and introduces inconsistencies. | ||
| 87 | - # Thus, for an identity flatten, we should simply return the original layout. | ||
| 88 | if start_dim >= end_dim: | 113 | if start_dim >= end_dim: |
| 89 | - # If flattening a single dimension or a range that results in no actual flattening, | 114 | + return ((input_layout,), None) |
| 90 | - # the layout should remain unchanged. | ||
| 91 | - return input_layout # Directly return the input layout | ||
| 92 | 115 | ||
| 93 | - | ||
| 94 | - # Calculate the target flattened shape | ||
| 95 | flattened_size = 1 | 116 | flattened_size = 1 |
| 96 | for i in range(start_dim, end_dim + 1): | 117 | for i in range(start_dim, end_dim + 1): |
| 97 | flattened_size *= input_shape[i] | 118 | flattened_size *= input_shape[i] |
| 98 | dst_shape = list(input_shape[:start_dim]) + [flattened_size] + list(input_shape[end_dim + 1:]) | 119 | dst_shape = list(input_shape[:start_dim]) + [flattened_size] + list(input_shape[end_dim + 1:]) |
| 99 | 120 | ||
| 100 | - # Prepare args for ReshapeDistributedOp.infer_layout | 121 | + out_layout, _ = self._infer_reshape_layout(input_layout, dst_shape, input_shape) |
| 101 | - # Since self.op_name is 'flatten' (not in ["reshape", "view"]), | 122 | + return ((out_layout,), None) |
| 102 | - # the parent class will treat this as the MindSpore Reshape branch, | ||
| 103 | - # which expects: extra_args = [dst_shape, input_shape] | ||
| 104 | - mock_extra_args = [dst_shape, input_shape] | ||
| 105 | - | ||
| 106 | - # Reuse parent class logic to safely resolve Shard tensor mapping merges | ||
| 107 | - out_layout, _ = super().infer_layout(layouts, mock_extra_args) | ||
| 108 | - | ||
| 109 | - # Returning purely the Layout object, which is what the WithShape suffix expects | ||
| 110 | - return out_layout | ||
| @@ -386,6 +386,11 @@ class GatherNdDistributedOp(DistributedOp): | |||
| 386 | """ | 386 | """ |
| 387 | Preprocess arguments for GatherNd operator. | 387 | Preprocess arguments for GatherNd operator. |
| 388 | 388 | ||
| 389 | + NOTE: aclop packed-args normalization (for MindSpore aclop operators | ||
| 390 | + that pack args as ``(prim, name, (real_args...))``) is handled | ||
| 391 | + upstream in ``OpDispatcher._dispatch_layout_infer`` via | ||
| 392 | + ``_normalize_aclop_args``. This method receives clean unpacked args. | ||
| 393 | + | ||
| 389 | Args: | 394 | Args: |
| 390 | args (tuple): Input arguments (input, indices). | 395 | args (tuple): Input arguments (input, indices). |
| 391 | kwargs (dict): Keyword arguments. | 396 | kwargs (dict): Keyword arguments. |
| @@ -393,19 +398,11 @@ class GatherNdDistributedOp(DistributedOp): | |||
| 393 | Returns: | 398 | Returns: |
| 394 | tuple: (local_args, local_kwargs, cache_values) | 399 | tuple: (local_args, local_kwargs, cache_values) |
| 395 | """ | 400 | """ |
| 396 | - packed_call = None | ||
| 397 | - if len(args) == 3 and isinstance(args[1], str) and isinstance(args[2], (tuple, list)): | ||
| 398 | - packed_call = (args[0], args[1]) | ||
| 399 | - args = tuple(args[2]) | ||
| 400 | - | ||
| 401 | args, _ = _normalize_gathernd_args(*args, **kwargs) | 401 | args, _ = _normalize_gathernd_args(*args, **kwargs) |
| 402 | input_tensor, indices = args[0], args[1] | 402 | input_tensor, indices = args[0], args[1] |
| 403 | local_input = input_tensor.to_local() if hasattr(input_tensor, "_layout") else input_tensor | 403 | local_input = input_tensor.to_local() if hasattr(input_tensor, "_layout") else input_tensor |
| 404 | local_indices = indices.to_local() | 404 | local_indices = indices.to_local() |
| 405 | - if packed_call is not None: | 405 | + local_args = (local_input, local_indices) |
| 406 | - local_args = (packed_call[0], packed_call[1], (local_input, local_indices)) | ||
| 407 | - else: | ||
| 408 | - local_args = (local_input, local_indices) | ||
| 409 | local_kwargs = {} | 406 | local_kwargs = {} |
| 410 | cache_values = [ | 407 | cache_values = [ |
| 411 | input_tensor.layout if hasattr(input_tensor, "_layout") else None, | 408 | input_tensor.layout if hasattr(input_tensor, "_layout") else None, |
| @@ -1,86 +0,0 @@ | |||
| 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 | -""" | ||
| 16 | -Distributed implementation for Multinomial operator. | ||
| 17 | -""" | ||
| 18 | - | ||
| 19 | -from hyper_parallel.core.dtensor.layout import Layout | ||
| 20 | -from .parallel_ops import DistributedOp | ||
| 21 | - | ||
| 22 | - | ||
| 23 | -class MultinomialDistributedOp(DistributedOp): | ||
| 24 | - """Distributed implementation for Multinomial operator.""" | ||
| 25 | - | ||
| 26 | - def infer_layout(self, layouts, extra_args=None): | ||
| 27 | - """ | ||
| 28 | - Infer output layout for Multinomial operator. | ||
| 29 | - | ||
| 30 | - Args: | ||
| 31 | - layouts (tuple): Layouts of input tensor. | ||
| 32 | - extra_args (tuple): Arguments for the operator (num_samples, replacement, generator). | ||
| 33 | - Note: logic assumes num_samples affects shape but not the sharding map pattern. | ||
| 34 | - | ||
| 35 | - Returns: | ||
| 36 | - Layout: Layout for output tensor. | ||
| 37 | - """ | ||
| 38 | - # Check partial inputs | ||
| 39 | - if not self._allow_partial_inputs: | ||
| 40 | - self._check_partial_inputs(layouts) | ||
| 41 | - layout = layouts[0] | ||
| 42 | - in_tensor_map = layout.alias_tensor_map | ||
| 43 | - ndim = len(in_tensor_map) | ||
| 44 | - | ||
| 45 | - # PyTorch multinomial supports 1D or 2D tensors | ||
| 46 | - if ndim not in (1, 2): | ||
| 47 | - raise ValueError(f"For 'multinomial', input dimension must be 1 or 2, but got {ndim}.") | ||
| 48 | - | ||
| 49 | - # The last dimension (probability dimension) must NOT be sharded. | ||
| 50 | - # Multinomial sampling requires the full probability distribution (cumulative sum, normalization) | ||
| 51 | - # to be present on the device. If it's sharded, we cannot sample correctly without communication. | ||
| 52 | - # In alias_tensor_map, "None" indicates the dimension is NOT sharded. | ||
| 53 | - prob_dim_map = in_tensor_map[-1] | ||
| 54 | - | ||
| 55 | - # Check if the last dimension is sharded (i.e., not "None") | ||
| 56 | - if prob_dim_map != "None": | ||
| 57 | - raise ValueError( | ||
| 58 | - f"For 'multinomial', the last dimension (probability category dimension) " | ||
| 59 | - f"must not be sharded (Replicated). Got layout map: {in_tensor_map}. " | ||
| 60 | - f"Please redistribute the tensor to replicate the last dimension before calling multinomial." | ||
| 61 | - ) | ||
| 62 | - out_tensor_map = () | ||
| 63 | - if ndim == 1: | ||
| 64 | - # Input: (C,) -> Probability distribution | ||
| 65 | - # Output: (num_samples,) -> Sampled indices | ||
| 66 | - # Since the input distribution C is replicated (checked above), the samples generated | ||
| 67 | - # from it have no spatial correlation to shards, so the output is also Replicated ("None"). | ||
| 68 | - out_tensor_map = ("None",) | ||
| 69 | - else: | ||
| 70 | - # Input: (N, C) -> Batch of distributions | ||
| 71 | - # Output: (N, num_samples) -> Batch of sampled indices | ||
| 72 | - # | ||
| 73 | - # Rule: | ||
| 74 | - # 1. Preserve sharding on the Batch dimension (dim 0). If input is data parallel, output is too. | ||
| 75 | - # 2. The new dimension (num_samples) is created locally and is not sharded ("None"). | ||
| 76 | - batch_dim_map = in_tensor_map[0] | ||
| 77 | - out_tensor_map = (batch_dim_map, "None") | ||
| 78 | - | ||
| 79 | - output_layout = Layout( | ||
| 80 | - mesh_shape=layout.mesh_shape, | ||
| 81 | - alias_name=layout.alias_name, | ||
| 82 | - rank_list=layout.rank_list | ||
| 83 | - ) | ||
| 84 | - | ||
| 85 | - # Create output layout using the inferred tensor map aliases | ||
| 86 | - return output_layout(*out_tensor_map) | ||
| @@ -16,79 +16,166 @@ | |||
| 16 | """ | 16 | """ |
| 17 | Distributed implementation for RepeatInterleave operator. | 17 | Distributed implementation for RepeatInterleave operator. |
| 18 | """ | 18 | """ |
| 19 | +import copy | ||
| 20 | +from typing import Tuple | ||
| 21 | + | ||
| 19 | from hyper_parallel.core.dtensor.layout import Layout | 22 | from hyper_parallel.core.dtensor.layout import Layout |
| 20 | from .parallel_ops import DistributedOp | 23 | from .parallel_ops import DistributedOp |
| 21 | 24 | ||
| 22 | 25 | ||
| 26 | +def _normalize_repeat_interleave_args(x, repeats, dim=None, *, output_size=None): | ||
| 27 | + """Normalize positional and keyword arguments into a canonical form. | ||
| 28 | + | ||
| 29 | + Args: | ||
| 30 | + x: Input tensor. | ||
| 31 | + repeats: Number of repetitions (int or Tensor). | ||
| 32 | + dim: Dimension along which to repeat values. None means flatten first. | ||
| 33 | + output_size: Total output size (keyword-only in torch). | ||
| 34 | + | ||
| 35 | + Returns: | ||
| 36 | + tuple: (positional_args_tuple, kwargs_dict) | ||
| 37 | + """ | ||
| 38 | + kwargs = {} | ||
| 39 | + if output_size is not None: | ||
| 40 | + kwargs['output_size'] = output_size | ||
| 41 | + return (x, repeats, dim), kwargs | ||
| 42 | + | ||
| 43 | + | ||
| 23 | class RepeatInterleaveDistributedOp(DistributedOp): | 44 | class RepeatInterleaveDistributedOp(DistributedOp): |
| 24 | - """Distributed implementation for torch.repeat_interleave.""" | 45 | + """Distributed implementation for torch.repeat_interleave. |
| 25 | 46 | ||
| 26 | - def infer_layout(self, layouts, extra_args=None): | 47 | + Sharding constraints: |
| 27 | - """ | 48 | + - When dim is specified: the repeat dimension must be replicated. |
| 28 | - Infer output layout for RepeatInterleave operator. | 49 | + - When dim is None (flatten mode): only the first dimension may be sharded. |
| 29 | 50 | ||
| 30 | - RepeatInterleave: output = repeat_interleave(input, repeats, dim) | 51 | + Output layout: |
| 31 | - | 52 | + - When dim is specified: same as input layout. |
| 32 | - Rules: | 53 | + - When dim is None: 1-D layout preserving sharding on dim 0. |
| 33 | - 1. dim = None if not specified. | 54 | + """ |
| 34 | - 2. The dimension `dim` MUST be unsharded to ensure global repeat_interleave correctness. | 55 | + |
| 35 | - 3. Output layout usually same as input, but shape changes. | 56 | + @staticmethod |
| 57 | + def _validate_input_layouts(input_layout, dim, op_name: str) -> None: | ||
| 58 | + """Validate sharding constraints for repeat_interleave. | ||
| 36 | 59 | ||
| 37 | Args: | 60 | Args: |
| 38 | - layouts (tuple): Layouts of inputs. Expected: | 61 | + input_layout: Layout of the input tensor. |
| 39 | - layouts[0] (Layout): Input tensor layout (required). | 62 | + dim: The repeat dimension, or None for flatten mode. |
| 40 | - extra_args (tuple, optional): Contains repeats and dim. Expected: | 63 | + op_name: Operator name used in error messages. |
| 41 | - extra_args[0] (int or Tensor): Number of repeats or Tensor | ||
| 42 | - extra_args[1] (int, optional): Dimension to repeat. Defaults to -1. | ||
| 43 | 64 | ||
| 44 | - Returns: | 65 | + Raises: |
| 45 | - tuple: Layouts for values. | 66 | + ValueError: If dim is out of range or the repeat dimension is sharded. |
| 46 | """ | 67 | """ |
| 47 | - if not layouts or layouts[0] is None: | 68 | + in_tensor_map = input_layout.alias_tensor_map |
| 48 | - raise ValueError("repeat_interleave requires a valid input tensor layout.") | 69 | + ndim = len(in_tensor_map) |
| 49 | 70 | ||
| 50 | - input_layout = layouts[0] | 71 | + if dim is not None: |
| 51 | - in_tensor_map = input_layout.tensor_map | 72 | + actual_dim = dim if dim >= 0 else ndim + dim |
| 52 | - dim = None # The dimension along which to repeat values. | 73 | + if not 0 <= actual_dim < ndim: |
| 53 | - # By default, use the flattened input array, and return a flat output array. | ||
| 54 | - if len(extra_args) >= 2 and extra_args[1] is not None: | ||
| 55 | - dim = extra_args[1] | ||
| 56 | - if dim is None: | ||
| 57 | - sharded_dims = [i for i, shard in enumerate(in_tensor_map) if shard != -1] | ||
| 58 | - if not sharded_dims: # not shard | ||
| 59 | - output_tensor_map = [-1] | ||
| 60 | - # Only can shard on the first dimension. | ||
| 61 | - elif sharded_dims == [0] and in_tensor_map[0] != -1: | ||
| 62 | - output_tensor_map = [in_tensor_map[0]] | ||
| 63 | - else: | ||
| 64 | - # Other dims must NOT be sharded. | ||
| 65 | raise ValueError( | 74 | raise ValueError( |
| 66 | - f"Operation {self.op_name}: Cannot flatten tensor when dim=None." | 75 | + f"For {op_name}, dimension should be in [0, {ndim}), but got {dim}." |
| 76 | + ) | ||
| 77 | + mapping = in_tensor_map[actual_dim] | ||
| 78 | + if isinstance(mapping, (list, tuple)): | ||
| 79 | + is_sharded = any(axis != "None" for axis in mapping) | ||
| 80 | + else: | ||
| 81 | + is_sharded = mapping != "None" | ||
| 82 | + if is_sharded: | ||
| 83 | + raise ValueError( | ||
| 84 | + f"For {op_name}, the repeat dimension should be replicated, " | ||
| 85 | + f"but got dim={dim} mapped to {mapping}." | ||
| 67 | ) | 86 | ) |
| 68 | 87 | ||
| 69 | - def idx_to_alias(idx, aliases): | 88 | + def preprocess(self, args: tuple, kwargs: dict) -> tuple: |
| 70 | - if idx == -1: | 89 | + """Preprocess arguments for RepeatInterleave operator. |
| 71 | - return "None" | 90 | + |
| 72 | - return aliases[len(aliases) - idx - 1] | 91 | + Extracts local tensors from DTensor inputs and builds cache_values |
| 73 | - output_map = tuple(idx_to_alias(idx, input_layout.alias_name) for idx in output_tensor_map) | 92 | + containing only the information needed for layout inference. |
| 93 | + | ||
| 94 | + Args: | ||
| 95 | + args: Positional arguments (input, repeats, dim). | ||
| 96 | + kwargs: Keyword arguments (output_size). | ||
| 97 | + | ||
| 98 | + Returns: | ||
| 99 | + tuple: (local_args, local_kwargs, cache_values) | ||
| 100 | + """ | ||
| 101 | + args, kwargs = _normalize_repeat_interleave_args(*args, **kwargs) | ||
| 102 | + input_tensor = args[0] | ||
| 103 | + repeats = args[1] | ||
| 104 | + dim = args[2] | ||
| 105 | + output_size = kwargs.get('output_size', None) | ||
| 106 | + | ||
| 107 | + local_input = input_tensor.to_local() | ||
| 108 | + | ||
| 109 | + # repeats can be int or Tensor; handle DTensor defensively | ||
| 110 | + if hasattr(repeats, 'to_local'): | ||
| 111 | + local_repeats = repeats.to_local() | ||
| 112 | + else: | ||
| 113 | + local_repeats = repeats | ||
| 114 | + | ||
| 115 | + local_args = (local_input, local_repeats, dim) | ||
| 116 | + local_kwargs = {} | ||
| 117 | + if output_size is not None: | ||
| 118 | + local_kwargs['output_size'] = output_size | ||
| 119 | + | ||
| 120 | + cache_values = [input_tensor.layout, dim] | ||
| 121 | + return local_args, local_kwargs, cache_values | ||
| 122 | + | ||
| 123 | + def infer_layout(self, cache_values: list) -> Tuple[tuple, None]: | ||
| 124 | + """Infer output layout for RepeatInterleave operator. | ||
| 125 | + | ||
| 126 | + Rules: | ||
| 127 | + 1. Input must not have Partial status. | ||
| 128 | + 2. When dim is specified: the repeat dimension must be replicated. | ||
| 129 | + Output layout is a deep copy of the input layout. | ||
| 130 | + 3. When dim is None (flatten mode): only dim 0 may be sharded. | ||
| 131 | + Output is a 1-D layout preserving the mesh axis on dim 0. | ||
| 132 | + | ||
| 133 | + Args: | ||
| 134 | + cache_values: [input_layout, dim] where dim is the repeat dimension | ||
| 135 | + or None for flatten mode. | ||
| 136 | + | ||
| 137 | + Returns: | ||
| 138 | + tuple: ((output_layout,), None) | ||
| 139 | + | ||
| 140 | + Raises: | ||
| 141 | + ValueError: If input has Partial status, dim is out of range, | ||
| 142 | + the repeat dimension is sharded, or flatten is attempted on | ||
| 143 | + a tensor sharded on a non-first dimension. | ||
| 144 | + """ | ||
| 145 | + input_layout = cache_values[0] | ||
| 146 | + dim = cache_values[1] | ||
| 147 | + | ||
| 148 | + self._check_partial_inputs([input_layout]) | ||
| 149 | + self._validate_input_layouts(input_layout, dim, self.op_name) | ||
| 150 | + | ||
| 151 | + in_tensor_map = input_layout.alias_tensor_map | ||
| 152 | + | ||
| 153 | + if dim is None: | ||
| 154 | + # Flatten mode: output is 1-D. | ||
| 155 | + sharded_dims = [] | ||
| 156 | + for i, shard in enumerate(in_tensor_map): | ||
| 157 | + if isinstance(shard, (list, tuple)): | ||
| 158 | + is_sharded = any(axis != "None" for axis in shard) | ||
| 159 | + else: | ||
| 160 | + is_sharded = shard != "None" | ||
| 161 | + if is_sharded: | ||
| 162 | + sharded_dims.append(i) | ||
| 163 | + if not sharded_dims: | ||
| 164 | + output_tensor_map = ("None",) | ||
| 165 | + elif sharded_dims == [0]: | ||
| 166 | + output_tensor_map = (in_tensor_map[0],) | ||
| 167 | + else: | ||
| 168 | + raise ValueError( | ||
| 169 | + f"For {self.op_name}, sharded dims in flatten mode should be [] or [0], " | ||
| 170 | + f"but got {sharded_dims}." | ||
| 171 | + ) | ||
| 74 | 172 | ||
| 75 | output_layout = Layout( | 173 | output_layout = Layout( |
| 76 | mesh_shape=input_layout.mesh_shape, | 174 | mesh_shape=input_layout.mesh_shape, |
| 77 | alias_name=input_layout.alias_name, | 175 | alias_name=input_layout.alias_name, |
| 78 | rank_list=input_layout.rank_list | 176 | rank_list=input_layout.rank_list |
| 79 | ) | 177 | ) |
| 178 | + return (output_layout(*output_tensor_map),), None | ||
| 80 | 179 | ||
| 81 | - return output_layout(*output_map) | 180 | + # dim specified: output layout = deep copy of input layout |
| 82 | - input_dim = len(in_tensor_map) | 181 | + return (copy.deepcopy(input_layout),), None |
| 83 | - if dim < 0: | ||
| 84 | - dim = input_dim + dim | ||
| 85 | - # Check if dimension is within valid range | ||
| 86 | - if not 0 <= dim < input_dim: | ||
| 87 | - raise ValueError(f"Dimension out of range (expected to be in [0, {input_dim}), but got {dim}).") | ||
| 88 | - # The chosen dim must NOT be sharded | ||
| 89 | - if in_tensor_map[dim] != -1: | ||
| 90 | - raise ValueError( | ||
| 91 | - f"Operation {self.op_name}: Cannot perform sharding on params along the chosen dim" | ||
| 92 | - ) | ||
| 93 | - # Output layout same as input layout (shape change does not affect sharding pattern) | ||
| 94 | - return input_layout | ||
| @@ -16,6 +16,8 @@ | |||
| 16 | Distributed implementation for Reshape operator. | 16 | Distributed implementation for Reshape operator. |
| 17 | """ | 17 | """ |
| 18 | 18 | ||
| 19 | +from typing import Callable, Optional, Tuple | ||
| 20 | + | ||
| 19 | from hyper_parallel.core.dtensor.layout import Layout | 21 | from hyper_parallel.core.dtensor.layout import Layout |
| 20 | from hyper_parallel.platform import get_platform | 22 | from hyper_parallel.platform import get_platform |
| 21 | from .parallel_ops import DistributedOp | 23 | from .parallel_ops import DistributedOp |
| @@ -23,6 +25,21 @@ platform = get_platform() | |||
| 23 | Tensor = platform.Tensor | 25 | Tensor = platform.Tensor |
| 24 | 26 | ||
| 25 | 27 | ||
| 28 | +def _normalize_reshape_args(x, *shape, **kwargs): | ||
| 29 | + """Normalize reshape/view arguments into positional args and empty kwargs.""" | ||
| 30 | + unexpected_kwargs = set(kwargs) - {'shape'} | ||
| 31 | + if unexpected_kwargs: | ||
| 32 | + unexpected = next(iter(unexpected_kwargs)) | ||
| 33 | + raise TypeError(f"reshape got an unexpected keyword argument '{unexpected}'.") | ||
| 34 | + if shape and 'shape' in kwargs: | ||
| 35 | + raise TypeError("reshape got shape from both args and kwargs.") | ||
| 36 | + if not shape and 'shape' in kwargs: | ||
| 37 | + shape = (kwargs['shape'],) | ||
| 38 | + if not shape: | ||
| 39 | + raise TypeError("reshape missing required shape argument.") | ||
| 40 | + return (x,) + shape, {} | ||
| 41 | + | ||
| 42 | + | ||
| 26 | def _filter_none_split_tensor_map(tensor_map, mesh_shape): | 43 | def _filter_none_split_tensor_map(tensor_map, mesh_shape): |
| 27 | """ | 44 | """ |
| 28 | Filter out the elements in tensor_map where the size of the corresponding dimension in device_matrix is 1. | 45 | Filter out the elements in tensor_map where the size of the corresponding dimension in device_matrix is 1. |
| @@ -152,47 +169,6 @@ class ReshapeDistributedOp(DistributedOp): | |||
| 152 | local_dst_shape.append(dst_shape[idx] // x_mesh_shape[-1 - map_id] if dst_shape[idx] > 0 else -1) | 169 | local_dst_shape.append(dst_shape[idx] // x_mesh_shape[-1 - map_id] if dst_shape[idx] > 0 else -1) |
| 153 | return output_map, local_dst_shape | 170 | return output_map, local_dst_shape |
| 154 | 171 | ||
| 155 | - def _parse_shape_args(self, extra_args): | ||
| 156 | - """Parse shape arguments from extra_args. | ||
| 157 | - | ||
| 158 | - Args: | ||
| 159 | - extra_args: Extra arguments containing shape info | ||
| 160 | - | ||
| 161 | - Returns: | ||
| 162 | - tuple: (dst_shape, input_shape) | ||
| 163 | - """ | ||
| 164 | - if self.op_name in ["reshape", "view"]: | ||
| 165 | - return self._parse_torch_shape_args(extra_args) | ||
| 166 | - return self._parse_mindspore_shape_args(extra_args) | ||
| 167 | - | ||
| 168 | - def _parse_torch_shape_args(self, extra_args): | ||
| 169 | - """Parse PyTorch style shape arguments.""" | ||
| 170 | - if len(extra_args) < 2: | ||
| 171 | - raise ValueError(f"{self.op_name} requires output shape and input shape.") | ||
| 172 | - | ||
| 173 | - input_shape = extra_args[-1] | ||
| 174 | - shape_args = extra_args[:-1] | ||
| 175 | - | ||
| 176 | - if len(shape_args) == 1: | ||
| 177 | - first_arg = shape_args[0] | ||
| 178 | - if isinstance(first_arg, (list, tuple)): | ||
| 179 | - dst_shape = first_arg | ||
| 180 | - elif isinstance(first_arg, Tensor): | ||
| 181 | - dst_shape = first_arg.tolist() | ||
| 182 | - else: | ||
| 183 | - dst_shape = shape_args | ||
| 184 | - else: | ||
| 185 | - dst_shape = shape_args | ||
| 186 | - | ||
| 187 | - return dst_shape, input_shape | ||
| 188 | - | ||
| 189 | - def _parse_mindspore_shape_args(self, extra_args): | ||
| 190 | - """Parse MindSpore style shape arguments.""" | ||
| 191 | - if len(extra_args) != 2: | ||
| 192 | - raise ValueError("Reshape requires output shape and input shape.") | ||
| 193 | - | ||
| 194 | - return extra_args[0], extra_args[1] | ||
| 195 | - | ||
| 196 | def _normalize_shape(self, dst_shape): | 172 | def _normalize_shape(self, dst_shape): |
| 197 | """Normalize dst_shape to list format.""" | 173 | """Normalize dst_shape to list format.""" |
| 198 | if isinstance(dst_shape, Tensor): | 174 | if isinstance(dst_shape, Tensor): |
| @@ -265,25 +241,29 @@ class ReshapeDistributedOp(DistributedOp): | |||
| 265 | if partial_op is not None and i < len(out_layout.alias_name): | 241 | if partial_op is not None and i < len(out_layout.alias_name): |
| 266 | out_layout.set_partial_by_dev_axis(out_layout.alias_name[i], partial_op) | 242 | out_layout.set_partial_by_dev_axis(out_layout.alias_name[i], partial_op) |
| 267 | 243 | ||
| 268 | - def infer_layout(self, layouts, extra_args=None): | 244 | + def preprocess(self, args: tuple, kwargs: dict) -> tuple: |
| 269 | """ | 245 | """ |
| 270 | - Infer output layout for reshape operator. | 246 | + Preprocess arguments for Reshape operator. |
| 271 | - | ||
| 272 | - For reshape operations, data slice on each device after reshape should be same as data slice before reshape. | ||
| 273 | 247 | ||
| 274 | Args: | 248 | Args: |
| 275 | - layouts (Layout): Layout of input x | 249 | + args (tuple): Input tensor followed by target shape arguments. |
| 276 | - extra_args: | 250 | + kwargs (dict): Keyword arguments. |
| 277 | - For MindSpore Reshape: (destination shape, original shape) | ||
| 278 | - For PyTorch reshape/view: (shape_arg1, shape_arg2, ..., original shape) or (shape_tuple, original shape) | ||
| 279 | 251 | ||
| 280 | Returns: | 252 | Returns: |
| 281 | - tuple: Layout for output tensor | 253 | + tuple: (local_args, local_kwargs, cache_values) |
| 282 | """ | 254 | """ |
| 283 | - x_layout = layouts[0] | 255 | + args, _ = _normalize_reshape_args(*args, **kwargs) |
| 284 | - x_dict = x_layout.to_dict() | 256 | + input_tensor = args[0] |
| 257 | + dst_shape = args[1:] if len(args) > 2 else args[1] | ||
| 285 | 258 | ||
| 286 | - dst_shape, input_shape = self._parse_shape_args(extra_args) | 259 | + local_args = (input_tensor.to_local(), dst_shape) |
| 260 | + local_kwargs = {} | ||
| 261 | + cache_values = [input_tensor.layout, dst_shape, tuple(input_tensor.shape)] | ||
| 262 | + return local_args, local_kwargs, cache_values | ||
| 263 | + | ||
| 264 | + def _infer_reshape_layout(self, x_layout, dst_shape, input_shape): | ||
| 265 | + """Infer reshape output layout and local destination shape.""" | ||
| 266 | + x_dict = x_layout.to_dict() | ||
| 287 | dst_shape = self._normalize_shape(dst_shape) | 267 | dst_shape = self._normalize_shape(dst_shape) |
| 288 | 268 | ||
| 289 | x_map = _filter_none_split_tensor_map(x_dict["tensor_map"], x_dict["mesh_shape"]) | 269 | x_map = _filter_none_split_tensor_map(x_dict["tensor_map"], x_dict["mesh_shape"]) |
| @@ -307,3 +287,53 @@ class ReshapeDistributedOp(DistributedOp): | |||
| 307 | self._apply_partial_status(x_layout, out_layout) | 287 | self._apply_partial_status(x_layout, out_layout) |
| 308 | 288 | ||
| 309 | return out_layout, local_dst_shape | 289 | return out_layout, local_dst_shape |
| 290 | + | ||
| 291 | + def infer_layout(self, cache_values: list) -> Tuple[tuple, list]: | ||
| 292 | + """ | ||
| 293 | + Infer output layout for Reshape operator. | ||
| 294 | + | ||
| 295 | + Rules: | ||
| 296 | + 1. Partial input is allowed and preserved on the output layout. | ||
| 297 | + 2. Target shape must be a Tensor, tuple, or list. | ||
| 298 | + 3. Input and output total element counts must match after resolving one dynamic axis. | ||
| 299 | + 4. Reshape must preserve each device's local data slice; sharded axes can only be | ||
| 300 | + split or merged when the shard boundary remains valid. | ||
| 301 | + 5. Output Partial status follows the input Partial status. | ||
| 302 | + | ||
| 303 | + Args: | ||
| 304 | + cache_values (list): [input_layout, dst_shape, input_shape]. | ||
| 305 | + | ||
| 306 | + Returns: | ||
| 307 | + tuple: ((output_layout,), local_dst_shape) | ||
| 308 | + | ||
| 309 | + Raises: | ||
| 310 | + ValueError: If target shape is invalid or the reshape would change sharded slices. | ||
| 311 | + """ | ||
| 312 | + if len(cache_values) != 3: | ||
| 313 | + raise ValueError( | ||
| 314 | + f"For {self.op_name}, cache_values length should be 3, but got {len(cache_values)}" | ||
| 315 | + ) | ||
| 316 | + | ||
| 317 | + x_layout, dst_shape, input_shape = cache_values[0], cache_values[1], cache_values[2] | ||
| 318 | + if x_layout is None: | ||
| 319 | + raise ValueError(f"For {self.op_name}, reshape requires a valid input tensor layout.") | ||
| 320 | + | ||
| 321 | + out_layout, local_dst_shape = self._infer_reshape_layout(x_layout, dst_shape, input_shape) | ||
| 322 | + return ((out_layout,), local_dst_shape) | ||
| 323 | + | ||
| 324 | + def get_expand_impl(self, func: Optional[Callable], infer_result: tuple, | ||
| 325 | + cache_values: list) -> Optional[Callable]: | ||
| 326 | + """Return a closure that calls reshape/view with the inferred local target shape.""" | ||
| 327 | + del cache_values | ||
| 328 | + if func is None: | ||
| 329 | + return None | ||
| 330 | + | ||
| 331 | + local_dst_shape = infer_result[1] | ||
| 332 | + if local_dst_shape is None: | ||
| 333 | + return None | ||
| 334 | + | ||
| 335 | + def expand_impl(x: object, shape: object) -> object: | ||
| 336 | + del shape | ||
| 337 | + return func(x, local_dst_shape) | ||
| 338 | + | ||
| 339 | + return expand_impl | ||
| @@ -16,29 +16,51 @@ | |||
| 16 | Distributed implementation for Slice operator. | 16 | Distributed implementation for Slice operator. |
| 17 | """ | 17 | """ |
| 18 | # pylint: disable=E0402 | 18 | # pylint: disable=E0402 |
| 19 | +from typing import Callable, Optional, Tuple | ||
| 20 | + | ||
| 19 | from .parallel_ops import DistributedOp | 21 | from .parallel_ops import DistributedOp |
| 20 | 22 | ||
| 21 | 23 | ||
| 24 | +def _normalize_slice_args(x, begin, end): | ||
| 25 | + return (x, begin, end), {} | ||
| 26 | + | ||
| 27 | + | ||
| 22 | class SliceDistributedOp(DistributedOp): | 28 | class SliceDistributedOp(DistributedOp): |
| 23 | - """Distributed implementation for MatMul operator.""" | 29 | + """Distributed implementation for Slice operator.""" |
| 30 | + | ||
| 31 | + def preprocess(self, args: tuple, kwargs: dict) -> tuple: | ||
| 32 | + """ | ||
| 33 | + Preprocess arguments for Slice operator. | ||
| 34 | + | ||
| 35 | + Args: | ||
| 36 | + args (tuple): Input arguments containing x, begin and end. | ||
| 37 | + kwargs (dict): Keyword arguments. | ||
| 38 | + | ||
| 39 | + Returns: | ||
| 40 | + tuple: (local_args, local_kwargs, cache_values) | ||
| 41 | + """ | ||
| 42 | + args, _ = _normalize_slice_args(*args, **kwargs) | ||
| 43 | + input_tensor, begin, end = args | ||
| 44 | + local_args = (input_tensor.to_local(), begin, end) | ||
| 45 | + local_kwargs = {} | ||
| 46 | + cache_values = [input_tensor.layout, begin, end, input_tensor.shape] | ||
| 47 | + return local_args, local_kwargs, cache_values | ||
| 24 | 48 | ||
| 25 | def _is_shard_dim(self, layout): | 49 | def _is_shard_dim(self, layout): |
| 26 | """return the shard num in each dim""" | 50 | """return the shard num in each dim""" |
| 27 | shard_dim = [] | 51 | shard_dim = [] |
| 28 | - dev_mat = layout.mesh_shape | 52 | + for axis_name in layout.alias_tensor_map: |
| 29 | - tensor_map = layout.tensor_map | 53 | + if axis_name == "None": |
| 30 | - for dev_idx in tensor_map: | 54 | + shard_dim.append(1) |
| 31 | - if isinstance(dev_idx, (tuple, list)): | 55 | + continue |
| 56 | + if isinstance(axis_name, (tuple, list)): | ||
| 32 | shard_num = 1 | 57 | shard_num = 1 |
| 33 | - for idx in dev_idx: | 58 | + for axis in axis_name: |
| 34 | - if idx != -1 and dev_mat[len(dev_mat) - idx - 1] != 1: | 59 | + if axis != "None": |
| 35 | - shard_num *= dev_mat[len(dev_mat) - idx - 1] | 60 | + shard_num *= layout.mesh.get_device_num_along_axis(axis) |
| 36 | shard_dim.append(shard_num) | 61 | shard_dim.append(shard_num) |
| 37 | - else: | 62 | + continue |
| 38 | - if dev_idx != -1 and dev_mat[len(dev_mat) - dev_idx - 1] != 1: | 63 | + shard_dim.append(layout.mesh.get_device_num_along_axis(axis_name)) |
| 39 | - shard_dim.append(dev_mat[len(dev_mat) - dev_idx - 1]) | ||
| 40 | - else: | ||
| 41 | - shard_dim.append(1) | ||
| 42 | return shard_dim | 64 | return shard_dim |
| 43 | 65 | ||
| 44 | def _check_layout(self, layout, begin, end, shape): | 66 | def _check_layout(self, layout, begin, end, shape): |
| @@ -54,23 +76,69 @@ class SliceDistributedOp(DistributedOp): | |||
| 54 | f"the begin is {begin}, the end is {end}, the shape is {shape}, layout is {layout.to_dict()}") | 76 | f"the begin is {begin}, the end is {end}, the shape is {shape}, layout is {layout.to_dict()}") |
| 55 | return shard_dim | 77 | return shard_dim |
| 56 | 78 | ||
| 57 | - def infer_layout(self, layouts, extra_args=None): | 79 | + def infer_layout(self, cache_values: list) -> Tuple[tuple, tuple]: |
| 58 | """ | 80 | """ |
| 59 | - Infer output layout for slice operator. The shard dim must be fully fetched. | 81 | + Infer output layout for Slice operator. |
| 82 | + | ||
| 83 | + Rules: | ||
| 84 | + 1. Input must not have Partial status. | ||
| 85 | + 2. begin, end and global_shape must have the same rank as input layout. | ||
| 86 | + 3. Any sharded dimension must be fully fetched. | ||
| 87 | + 4. Output layout is identical to the input layout. | ||
| 60 | 88 | ||
| 61 | Args: | 89 | Args: |
| 62 | - layouts (Layout): Layout of input x | 90 | + cache_values (list): [input_layout, begin, end, global_shape]. |
| 63 | - extra_args: (begin, end, global shape) | ||
| 64 | 91 | ||
| 65 | Returns: | 92 | Returns: |
| 66 | - layout (Layout): the out layout | 93 | + tuple: ((output_layout,), (new_begin, new_end)). |
| 67 | - new_begin (tuple): begin after modification | 94 | + |
| 68 | - new_end (tuple): end after modification | 95 | + Raises: |
| 96 | + ValueError: If input has Partial status, arguments rank mismatch, or a sharded | ||
| 97 | + dimension is not fully fetched. | ||
| 69 | """ | 98 | """ |
| 70 | - begin = extra_args[0] | 99 | + layout, begin, end, global_shape = cache_values |
| 71 | - end = extra_args[1] | 100 | + self._check_partial_inputs([layout]) |
| 72 | - global_shape = extra_args[2] | 101 | + |
| 73 | - shard_dim = self._check_layout(layouts, begin, end, global_shape) | 102 | + if len(begin) != len(end) or len(begin) != len(global_shape): |
| 103 | + raise ValueError( | ||
| 104 | + f"For {self.op_name}, begin, end and global_shape must have the same length, " | ||
| 105 | + f"but got begin: {len(begin)}, end: {len(end)}, global_shape: {len(global_shape)}" | ||
| 106 | + ) | ||
| 107 | + if len(begin) != len(layout.alias_tensor_map): | ||
| 108 | + raise ValueError( | ||
| 109 | + f"For {self.op_name}, slice arguments rank must match input layout rank, " | ||
| 110 | + f"but got args rank: {len(begin)} and layout rank: {len(layout.alias_tensor_map)}" | ||
| 111 | + ) | ||
| 112 | + | ||
| 113 | + shard_dim = self._check_layout((layout,), begin, end, global_shape) | ||
| 74 | new_begin = tuple(begin[i] // shard_dim[i] for i in range(len(begin))) | 114 | new_begin = tuple(begin[i] // shard_dim[i] for i in range(len(begin))) |
| 75 | new_end = tuple(end[i] // shard_dim[i] for i in range(len(end))) | 115 | new_end = tuple(end[i] // shard_dim[i] for i in range(len(end))) |
| 76 | - return layouts[0], new_begin, new_end | 116 | + return ((layout,), (new_begin, new_end)) |
| 117 | + | ||
| 118 | + def get_expand_impl(self, func: Optional[Callable], infer_result: tuple, | ||
| 119 | + cache_values: list) -> Optional[Callable]: | ||
| 120 | + """ | ||
| 121 | + Return a custom Slice implementation when local begin/end need adjustment. | ||
| 122 | + | ||
| 123 | + Args: | ||
| 124 | + func: Original operator callable. | ||
| 125 | + infer_result (tuple): ((output_layout,), (new_begin, new_end)) from infer_layout. | ||
| 126 | + cache_values (list): [input_layout, begin, end, global_shape]. | ||
| 127 | + | ||
| 128 | + Returns: | ||
| 129 | + callable | None: expand_impl closure when local slice bounds differ, else None. | ||
| 130 | + """ | ||
| 131 | + if func is None: | ||
| 132 | + return None | ||
| 133 | + | ||
| 134 | + begin = cache_values[1] | ||
| 135 | + end = cache_values[2] | ||
| 136 | + new_begin, new_end = infer_result[1] | ||
| 137 | + if begin == new_begin and end == new_end: | ||
| 138 | + return None | ||
| 139 | + | ||
| 140 | + def expand_impl(input_tensor: object, *_unused_args: object) -> object: | ||
| 141 | + """Call Slice with local slice bounds.""" | ||
| 142 | + return func(input_tensor, new_begin, new_end) | ||
| 143 | + | ||
| 144 | + return expand_impl | ||
| @@ -13,77 +13,89 @@ | |||
| 13 | # limitations under the License. | 13 | # limitations under the License. |
| 14 | # ============================================================================ | 14 | # ============================================================================ |
| 15 | """ | 15 | """ |
| 16 | -Distributed implementation for ExpandDims operator. | 16 | +Distributed implementation for Squeeze operator. |
| 17 | """ | 17 | """ |
| 18 | +from copy import deepcopy | ||
| 19 | +from typing import Tuple | ||
| 20 | + | ||
| 18 | from hyper_parallel.core.dtensor.layout import Layout | 21 | from hyper_parallel.core.dtensor.layout import Layout |
| 19 | from .parallel_ops import DistributedOp | 22 | from .parallel_ops import DistributedOp |
| 20 | 23 | ||
| 21 | 24 | ||
| 25 | +def _normalize_squeeze_args(x, axis=None): | ||
| 26 | + return (x, axis), {} | ||
| 27 | + | ||
| 28 | + | ||
| 22 | class SqueezeDistributedOp(DistributedOp): | 29 | class SqueezeDistributedOp(DistributedOp): |
| 23 | """Distributed implementation for Squeeze operator.""" | 30 | """Distributed implementation for Squeeze operator.""" |
| 24 | 31 | ||
| 25 | - def infer_layout(self, layouts, extra_args=None): | 32 | + def preprocess(self, args: tuple, kwargs: dict) -> tuple: |
| 33 | + args, _ = _normalize_squeeze_args(*args, **kwargs) | ||
| 34 | + input_tensor, axis = args[0], args[1] | ||
| 35 | + | ||
| 36 | + if axis is None: | ||
| 37 | + local_args = (input_tensor.to_local(),) | ||
| 38 | + else: | ||
| 39 | + local_args = (input_tensor.to_local(), axis) | ||
| 40 | + | ||
| 41 | + input_shape = getattr(input_tensor, "shape", None) | ||
| 42 | + cache_values = [input_tensor.layout, axis, input_shape] | ||
| 43 | + return local_args, {}, cache_values | ||
| 44 | + | ||
| 45 | + def infer_layout(self, cache_values: list) -> Tuple[tuple, None]: | ||
| 26 | """ | 46 | """ |
| 27 | - Infer output layout for Squeeze. | 47 | + Infer output layout for Squeeze operator. |
| 48 | + | ||
| 49 | + Rules: | ||
| 50 | + 1. Input must not have Partial status. | ||
| 51 | + 2. input_shape rank must match input layout rank. | ||
| 52 | + 3. If axis is None, only size-1 replicated dimensions are removed. | ||
| 53 | + 4. If axis is specified, every axis must be in range, have size 1, and | ||
| 54 | + must not be sharded. | ||
| 55 | + 5. Output layout removes the squeezed dimensions and preserves the | ||
| 56 | + remaining input dimension mappings. | ||
| 28 | 57 | ||
| 29 | Args: | 58 | Args: |
| 30 | - layouts (tuple): Tuple containing input layout. | 59 | + cache_values (list): [input_layout, axis, input_shape] |
| 31 | - extra_args: Extra arguments containing axis and input_shapes. | ||
| 32 | - Can be dict or list/tuple where last element is input_shapes. | ||
| 33 | 60 | ||
| 34 | Returns: | 61 | Returns: |
| 35 | - Layout: Output layout with squeezed dimensions removed. | 62 | + tuple: ((output_layout,), None) |
| 63 | + | ||
| 64 | + Raises: | ||
| 65 | + ValueError: If input has Partial status, input_shape is missing, axis | ||
| 66 | + is invalid, or a requested squeeze dimension is sharded. | ||
| 36 | """ | 67 | """ |
| 37 | - if not layouts: | 68 | + if not cache_values: |
| 38 | raise ValueError( | 69 | raise ValueError( |
| 39 | - f"For {self.op_name}, layouts should contain at least one input layout, " | 70 | + f"For {self.op_name}, cache_values should contain input layout, " |
| 40 | - f"but got empty layouts." | 71 | + f"but got empty cache_values." |
| 41 | ) | 72 | ) |
| 42 | 73 | ||
| 43 | - x_layout = layouts[0] | 74 | + x_layout = cache_values[0] |
| 75 | + if not self._allow_partial_inputs: | ||
| 76 | + self._check_partial_inputs([x_layout]) | ||
| 77 | + | ||
| 44 | if x_layout.mesh_shape is None: | 78 | if x_layout.mesh_shape is None: |
| 45 | raise ValueError( | 79 | raise ValueError( |
| 46 | f"For {self.op_name}, input layout mesh_shape should not be None, " | 80 | f"For {self.op_name}, input layout mesh_shape should not be None, " |
| 47 | f"but got None." | 81 | f"but got None." |
| 48 | ) | 82 | ) |
| 49 | 83 | ||
| 50 | - axis, input_shape = self._extract_args(extra_args) | 84 | + axis = cache_values[1] if len(cache_values) > 1 else None |
| 85 | + input_shape = cache_values[2] if len(cache_values) > 2 else None | ||
| 51 | if input_shape is None: | 86 | if input_shape is None: |
| 52 | raise ValueError( | 87 | raise ValueError( |
| 53 | - f"For {self.op_name}, input_shapes should be provided in extra_args, " | 88 | + f"For {self.op_name}, input_shape should be provided in cache_values, " |
| 54 | f"but got None." | 89 | f"but got None." |
| 55 | ) | 90 | ) |
| 56 | - | 91 | + if not isinstance(input_shape, (list, tuple)): |
| 57 | - return self._compute_squeeze_layout(x_layout, axis, input_shape) | ||
| 58 | - | ||
| 59 | - def _extract_args(self, extra_args): | ||
| 60 | - """Extract axis and input_shape from extra_args.""" | ||
| 61 | - if isinstance(extra_args, dict): | ||
| 62 | - input_shapes = extra_args.get("input_shapes", None) | ||
| 63 | - axis = extra_args.get("axis", None) | ||
| 64 | - elif isinstance(extra_args, (list, tuple)) and extra_args: | ||
| 65 | - # Last element is input_shapes | ||
| 66 | - input_shapes = extra_args[-1] | ||
| 67 | - if not isinstance(input_shapes, (list, tuple)): | ||
| 68 | - raise ValueError( | ||
| 69 | - f"For {self.op_name}, input_shapes should be list or tuple, " | ||
| 70 | - f"but got {type(input_shapes)}." | ||
| 71 | - ) | ||
| 72 | - # First element is axis (if available) | ||
| 73 | - axis = extra_args[0] if len(extra_args) > 1 else None | ||
| 74 | - else: | ||
| 75 | raise ValueError( | 92 | raise ValueError( |
| 76 | - f"For {self.op_name}, extra_args should be dict or list/tuple, " | 93 | + f"For {self.op_name}, input_shape should be list or tuple, " |
| 77 | - f"but got {type(extra_args)}." | 94 | + f"but got {type(input_shape)}." |
| 78 | ) | 95 | ) |
| 79 | 96 | ||
| 80 | - # Get input shape (first element of input_shapes) | 97 | + output_layout = self._compute_squeeze_layout(x_layout, axis, input_shape) |
| 81 | - if input_shapes: | 98 | + return ((output_layout,), None) |
| 82 | - input_shape = input_shapes[0] if isinstance(input_shapes[0], (list, tuple)) else input_shapes | ||
| 83 | - else: | ||
| 84 | - input_shape = None | ||
| 85 | - | ||
| 86 | - return axis, input_shape | ||
| 87 | 99 | ||
| 88 | def _compute_squeeze_layout(self, x_layout, axis, input_shape): | 100 | def _compute_squeeze_layout(self, x_layout, axis, input_shape): |
| 89 | """Compute the squeezed layout.""" | 101 | """Compute the squeezed layout.""" |
| @@ -108,14 +120,7 @@ class SqueezeDistributedOp(DistributedOp): | |||
| 108 | f"but got {axis}." | 120 | f"but got {axis}." |
| 109 | ) | 121 | ) |
| 110 | 122 | ||
| 111 | - # Return scalar layout | 123 | + return deepcopy(x_layout) |
| 112 | - output_layout = Layout( | ||
| 113 | - mesh_shape=x_layout.mesh_shape, | ||
| 114 | - alias_name=x_layout.alias_name, | ||
| 115 | - rank_list=x_layout.rank_list | ||
| 116 | - ) | ||
| 117 | - output_layout = output_layout() | ||
| 118 | - return output_layout | ||
| 119 | 124 | ||
| 120 | def _validate_input_shape(self, x_layout, input_shape): | 125 | def _validate_input_shape(self, x_layout, input_shape): |
| 121 | """Validate that input shape matches layout rank.""" | 126 | """Validate that input shape matches layout rank.""" |
| @@ -150,6 +155,19 @@ class SqueezeDistributedOp(DistributedOp): | |||
| 150 | # Convert axis to list if it's a single integer | 155 | # Convert axis to list if it's a single integer |
| 151 | if isinstance(axis, int): | 156 | if isinstance(axis, int): |
| 152 | axis = [axis] | 157 | axis = [axis] |
| 158 | + elif isinstance(axis, tuple): | ||
| 159 | + axis = list(axis) | ||
| 160 | + elif not isinstance(axis, list): | ||
| 161 | + raise ValueError( | ||
| 162 | + f"For {self.op_name}, axis should be int, list or tuple, " | ||
| 163 | + f"but got {type(axis)}." | ||
| 164 | + ) | ||
| 165 | + | ||
| 166 | + if not all(isinstance(ax, int) for ax in axis): | ||
| 167 | + raise ValueError( | ||
| 168 | + f"For {self.op_name}, every axis value should be int, " | ||
| 169 | + f"but got {axis}." | ||
| 170 | + ) | ||
| 153 | 171 | ||
| 154 | # Convert negative indices to positive | 172 | # Convert negative indices to positive |
| 155 | axis = [ax if ax >= 0 else ax + in_rank for ax in axis] | 173 | axis = [ax if ax >= 0 else ax + in_rank for ax in axis] |
| @@ -191,6 +209,9 @@ class SqueezeDistributedOp(DistributedOp): | |||
| 191 | 209 | ||
| 192 | def _create_output_layout(self, x_layout, dims_to_squeeze): | 210 | def _create_output_layout(self, x_layout, dims_to_squeeze): |
| 193 | """Create output layout after squeezing dimensions.""" | 211 | """Create output layout after squeezing dimensions.""" |
| 212 | + if not dims_to_squeeze: | ||
| 213 | + return deepcopy(x_layout) | ||
| 214 | + | ||
| 194 | # Get current alias tensor map | 215 | # Get current alias tensor map |
| 195 | x_map = list(x_layout.alias_tensor_map) | 216 | x_map = list(x_layout.alias_tensor_map) |
| 196 | 217 | ||
| @@ -16,33 +16,76 @@ | |||
| 16 | Element-wise distributed operator implementation. | 16 | Element-wise distributed operator implementation. |
| 17 | """ | 17 | """ |
| 18 | 18 | ||
| 19 | +import copy | ||
| 20 | +from typing import Tuple | ||
| 21 | + | ||
| 19 | from .parallel_ops import DistributedOp | 22 | from .parallel_ops import DistributedOp |
| 20 | 23 | ||
| 21 | 24 | ||
| 25 | +def _unwrap_local_value(value): | ||
| 26 | + """Convert DTensor-like values to local tensors while preserving scalar slots.""" | ||
| 27 | + return value.to_local() if hasattr(value, "_layout") else value | ||
| 28 | + | ||
| 29 | + | ||
| 22 | class TupleElementWiseDistributedOp(DistributedOp): | 30 | class TupleElementWiseDistributedOp(DistributedOp): |
| 23 | """ | 31 | """ |
| 24 | Distributed implementation for tuple element-wise operators. | 32 | Distributed implementation for tuple element-wise operators. |
| 25 | 33 | ||
| 26 | Inherits from DistributedOp and provides element-wise specific implementations. | 34 | Inherits from DistributedOp and provides element-wise specific implementations. |
| 27 | """ | 35 | """ |
| 28 | - def infer_layout(self, layouts, extra_args=None): | 36 | + |
| 37 | + def preprocess(self, args: tuple, kwargs: dict) -> tuple: | ||
| 38 | + """ | ||
| 39 | + Preprocess arguments for tuple element-wise operators. | ||
| 40 | + | ||
| 41 | + Args: | ||
| 42 | + args (tuple): Positional arguments passed to the operator. | ||
| 43 | + kwargs (dict): Keyword arguments passed to the operator. | ||
| 44 | + | ||
| 45 | + Returns: | ||
| 46 | + tuple: (local_args, local_kwargs, cache_values) | ||
| 47 | + """ | ||
| 48 | + expanded_args = [] | ||
| 49 | + local_args = [] | ||
| 50 | + for arg in args: | ||
| 51 | + if isinstance(arg, (tuple, list)): | ||
| 52 | + expanded_args.extend(arg) | ||
| 53 | + local_args.append(tuple(_unwrap_local_value(item) for item in arg)) | ||
| 54 | + else: | ||
| 55 | + expanded_args.append(arg) | ||
| 56 | + local_args.append(_unwrap_local_value(arg)) | ||
| 57 | + | ||
| 58 | + local_kwargs = {key: _unwrap_local_value(value) for key, value in kwargs.items()} | ||
| 59 | + cache_values = [getattr(arg, "layout", None) for arg in expanded_args] | ||
| 60 | + cache_values.extend(getattr(value, "layout", None) for value in kwargs.values()) | ||
| 61 | + | ||
| 62 | + return tuple(local_args), local_kwargs, cache_values | ||
| 63 | + | ||
| 64 | + def infer_layout(self, cache_values: list) -> Tuple[tuple, None]: | ||
| 29 | """ | 65 | """ |
| 30 | Infer output layouts for element-wise operations. | 66 | Infer output layouts for element-wise operations. |
| 31 | 67 | ||
| 32 | - For element-wise operations, all inputs should have the same layout, | 68 | + Rules: |
| 33 | - and the output will have the same layout. | 69 | + 1. Inputs must not have Partial status. |
| 70 | + 2. Tuple/list positional arguments are inferred from their expanded elements. | ||
| 71 | + 3. Output layouts are identical to the expanded input layouts. | ||
| 34 | 72 | ||
| 35 | Args: | 73 | Args: |
| 36 | - primitive: Primitive instance | 74 | + cache_values (list): Expanded input layouts, using None for non-DTensor slots. |
| 37 | - layouts: Layouts of input tensors | ||
| 38 | 75 | ||
| 39 | Returns: | 76 | Returns: |
| 40 | - tuple: Layout for output tensor. | 77 | + tuple: (output_layouts, None) |
| 41 | 78 | ||
| 42 | Raises: | 79 | Raises: |
| 43 | - ValueError: If input layouts are not compatible. | 80 | + ValueError: If input has Partial status. |
| 44 | """ | 81 | """ |
| 45 | - if not layouts: | 82 | + if not cache_values: |
| 46 | return None | 83 | return None |
| 47 | 84 | ||
| 48 | - return layouts | 85 | + self._check_partial_inputs(cache_values) |
| 86 | + | ||
| 87 | + output_layouts = tuple( | ||
| 88 | + copy.deepcopy(layout) if layout is not None else None | ||
| 89 | + for layout in cache_values | ||
| 90 | + ) | ||
| 91 | + return output_layouts, None | ||
| @@ -1,4 +0,0 @@ | |||
| 1 | -multinomial: | ||
| 2 | - dist_op_name: _torch_multinomial_dist_op | ||
| 3 | - distributed_op_class: MultinomialDistributedOp | ||
| 4 | - distributed_op_file: parallel_multinomial | ||
| @@ -15,10 +15,10 @@ | |||
| 15 | """parallel_cumsum test""" | 15 | """parallel_cumsum test""" |
| 16 | import os | 16 | import os |
| 17 | import unittest | 17 | import unittest |
| 18 | -from unittest.mock import patch | 18 | +from unittest.mock import MagicMock, patch |
| 19 | import numpy as np | 19 | import numpy as np |
| 20 | 20 | ||
| 21 | -from hyper_parallel.core.dtensor.dtensor import _build_layout | 21 | +from hyper_parallel.core.dtensor.dtensor import _build_layout, _LAYOUT_CACHE |
| 22 | from hyper_parallel.core.dtensor.placement_types import Shard, Replicate | 22 | from hyper_parallel.core.dtensor.placement_types import Shard, Replicate |
| 23 | from hyper_parallel.core.shard.ops.parallel_cumsum import CumsumDistributedOp | 23 | from hyper_parallel.core.shard.ops.parallel_cumsum import CumsumDistributedOp |
| 24 | from hyper_parallel.core.dtensor.device_mesh import ( | 24 | from hyper_parallel.core.dtensor.device_mesh import ( |
| @@ -28,6 +28,7 @@ from hyper_parallel.core.dtensor.device_mesh import ( | |||
| 28 | from hyper_parallel.platform.platform import EXISTING_COMM_GROUPS | 28 | from hyper_parallel.platform.platform import EXISTING_COMM_GROUPS |
| 29 | 29 | ||
| 30 | op = CumsumDistributedOp("cumsum") | 30 | op = CumsumDistributedOp("cumsum") |
| 31 | +op_ms = CumsumDistributedOp("CumsumExt") | ||
| 31 | 32 | ||
| 32 | 33 | ||
| 33 | class TestParallelCumsum(unittest.TestCase): | 34 | class TestParallelCumsum(unittest.TestCase): |
| @@ -40,11 +41,13 @@ class TestParallelCumsum(unittest.TestCase): | |||
| 40 | """ | 41 | """ |
| 41 | EXISTING_COMM_GROUPS.clear() | 42 | EXISTING_COMM_GROUPS.clear() |
| 42 | _DEVICE_MESH_MAP.clear() | 43 | _DEVICE_MESH_MAP.clear() |
| 44 | + _LAYOUT_CACHE.clear() | ||
| 43 | 45 | ||
| 44 | def tearDown(self): | 46 | def tearDown(self): |
| 45 | """Clean up after each test method.""" | 47 | """Clean up after each test method.""" |
| 46 | EXISTING_COMM_GROUPS.clear() | 48 | EXISTING_COMM_GROUPS.clear() |
| 47 | _DEVICE_MESH_MAP.clear() | 49 | _DEVICE_MESH_MAP.clear() |
| 50 | + _LAYOUT_CACHE.clear() | ||
| 48 | 51 | ||
| 49 | def _setup_mock_platform(self, mock_platform, platform_type=None, world_size=8): | 52 | def _setup_mock_platform(self, mock_platform, platform_type=None, world_size=8): |
| 50 | """Configure common mock-platform attributes used across tests. | 53 | """Configure common mock-platform attributes used across tests. |
| @@ -83,19 +86,23 @@ class TestParallelCumsum(unittest.TestCase): | |||
| 83 | x_placements = (Shard(0), Replicate()) | 86 | x_placements = (Shard(0), Replicate()) |
| 84 | x_layout = _build_layout(mesh, x_placements, 2) | 87 | x_layout = _build_layout(mesh, x_placements, 2) |
| 85 | 88 | ||
| 86 | - output_layout = op.infer_layout((x_layout,), extra_args=(-1,)) | 89 | + cache_values = [x_layout, -1] |
| 90 | + output_layouts, extra_info = op.infer_layout(cache_values) | ||
| 91 | + output_layout = output_layouts[0] | ||
| 87 | 92 | ||
| 88 | expected_map = (1, -1) | 93 | expected_map = (1, -1) |
| 89 | assert output_layout.tensor_map == expected_map, ( | 94 | assert output_layout.tensor_map == expected_map, ( |
| 90 | f"Data parallel cumsum failed. Expected {expected_map}, " | 95 | f"Data parallel cumsum failed. Expected {expected_map}, " |
| 91 | f"got {output_layout.tensor_map}" | 96 | f"got {output_layout.tensor_map}" |
| 92 | ) | 97 | ) |
| 98 | + assert output_layout is not x_layout, "Cumsum output layout should be a deep copy" | ||
| 99 | + assert extra_info is None, f"Cumsum extra_info should be None, got {extra_info}" | ||
| 93 | 100 | ||
| 94 | # Since `get_expand_impl` is not overridden, it returns None by default. | 101 | # Since `get_expand_impl` is not overridden, it returns None by default. |
| 95 | # The same applies to other test classes, so it is unnecessary to test its return value. | 102 | # The same applies to other test classes, so it is unnecessary to test its return value. |
| 96 | - assert op.get_expand_impl(None, output_layout, (x_layout,), (-1,)) is None, ( | 103 | + assert op.get_expand_impl(None, (output_layouts, None), cache_values) is None, ( |
| 97 | f"get_expand_impl test failed. Expected None, " | 104 | f"get_expand_impl test failed. Expected None, " |
| 98 | - f"got {op.get_expand_impl(None, output_layout, (x_layout,), (-1,))}" | 105 | + f"got {op.get_expand_impl(None, (output_layouts, None), cache_values)}" |
| 99 | ) | 106 | ) |
| 100 | 107 | ||
| 101 | @patch("hyper_parallel.core.dtensor.device_mesh.platform") | 108 | @patch("hyper_parallel.core.dtensor.device_mesh.platform") |
| @@ -109,13 +116,16 @@ class TestParallelCumsum(unittest.TestCase): | |||
| 109 | x_placements = (Replicate(), Shard(1)) | 116 | x_placements = (Replicate(), Shard(1)) |
| 110 | x_layout = _build_layout(mesh, x_placements, 2) | 117 | x_layout = _build_layout(mesh, x_placements, 2) |
| 111 | 118 | ||
| 112 | - output_layout = op.infer_layout((x_layout,), extra_args=(0,)) | 119 | + cache_values = [x_layout, 0] |
| 120 | + output_layouts, extra_info = op.infer_layout(cache_values) | ||
| 121 | + output_layout = output_layouts[0] | ||
| 113 | 122 | ||
| 114 | expected_map = (-1, 0) | 123 | expected_map = (-1, 0) |
| 115 | assert output_layout.tensor_map == expected_map, ( | 124 | assert output_layout.tensor_map == expected_map, ( |
| 116 | f"Tensor parallel cumsum failed. Expected {expected_map}, " | 125 | f"Tensor parallel cumsum failed. Expected {expected_map}, " |
| 117 | f"got {output_layout.tensor_map}" | 126 | f"got {output_layout.tensor_map}" |
| 118 | ) | 127 | ) |
| 128 | + assert extra_info is None, f"Cumsum extra_info should be None, got {extra_info}" | ||
| 119 | 129 | ||
| 120 | @patch("hyper_parallel.core.dtensor.device_mesh.platform") | 130 | @patch("hyper_parallel.core.dtensor.device_mesh.platform") |
| 121 | def test_cumsum_layout_mixed_parallel(self, mock_platform): | 131 | def test_cumsum_layout_mixed_parallel(self, mock_platform): |
| @@ -128,13 +138,16 @@ class TestParallelCumsum(unittest.TestCase): | |||
| 128 | x_placements = (Shard(0), Replicate(), Shard(2)) | 138 | x_placements = (Shard(0), Replicate(), Shard(2)) |
| 129 | x_layout = _build_layout(mesh, x_placements, 3) | 139 | x_layout = _build_layout(mesh, x_placements, 3) |
| 130 | 140 | ||
| 131 | - output_layout = op.infer_layout((x_layout,), extra_args=(1,)) | 141 | + cache_values = [x_layout, 1] |
| 142 | + output_layouts, extra_info = op.infer_layout(cache_values) | ||
| 143 | + output_layout = output_layouts[0] | ||
| 132 | 144 | ||
| 133 | expected_map = (2, -1, 0) | 145 | expected_map = (2, -1, 0) |
| 134 | assert output_layout.tensor_map == expected_map, ( | 146 | assert output_layout.tensor_map == expected_map, ( |
| 135 | f"Mixed parallel cumsum failed. Expected {expected_map}, " | 147 | f"Mixed parallel cumsum failed. Expected {expected_map}, " |
| 136 | f"got {output_layout.tensor_map}" | 148 | f"got {output_layout.tensor_map}" |
| 137 | ) | 149 | ) |
| 150 | + assert extra_info is None, f"Cumsum extra_info should be None, got {extra_info}" | ||
| 138 | 151 | ||
| 139 | @patch("hyper_parallel.core.dtensor.device_mesh.platform") | 152 | @patch("hyper_parallel.core.dtensor.device_mesh.platform") |
| 140 | def test_cumsum_layout_negative_dim(self, mock_platform): | 153 | def test_cumsum_layout_negative_dim(self, mock_platform): |
| @@ -147,13 +160,16 @@ class TestParallelCumsum(unittest.TestCase): | |||
| 147 | x_placements = (Replicate(), Shard(1)) | 160 | x_placements = (Replicate(), Shard(1)) |
| 148 | x_layout = _build_layout(mesh, x_placements, 2) | 161 | x_layout = _build_layout(mesh, x_placements, 2) |
| 149 | 162 | ||
| 150 | - output_layout = op.infer_layout((x_layout,), extra_args=(-2,)) | 163 | + cache_values = [x_layout, -2] |
| 164 | + output_layouts, extra_info = op.infer_layout(cache_values) | ||
| 165 | + output_layout = output_layouts[0] | ||
| 151 | 166 | ||
| 152 | expected_map = (-1, 0) | 167 | expected_map = (-1, 0) |
| 153 | assert output_layout.tensor_map == expected_map, ( | 168 | assert output_layout.tensor_map == expected_map, ( |
| 154 | f"Negative dimension cumsum failed. Expected {expected_map}, " | 169 | f"Negative dimension cumsum failed. Expected {expected_map}, " |
| 155 | f"got {output_layout.tensor_map}" | 170 | f"got {output_layout.tensor_map}" |
| 156 | ) | 171 | ) |
| 172 | + assert extra_info is None, f"Cumsum extra_info should be None, got {extra_info}" | ||
| 157 | 173 | ||
| 158 | @patch("hyper_parallel.core.dtensor.device_mesh.platform") | 174 | @patch("hyper_parallel.core.dtensor.device_mesh.platform") |
| 159 | def test_cumsum_layout_invalid_sharding_on_cumsum_dim(self, mock_platform): | 175 | def test_cumsum_layout_invalid_sharding_on_cumsum_dim(self, mock_platform): |
| @@ -166,8 +182,8 @@ class TestParallelCumsum(unittest.TestCase): | |||
| 166 | x_placements = (Shard(0), Shard(1)) | 182 | x_placements = (Shard(0), Shard(1)) |
| 167 | x_layout = _build_layout(mesh, x_placements, 2) | 183 | x_layout = _build_layout(mesh, x_placements, 2) |
| 168 | 184 | ||
| 169 | - with self.assertRaisesRegex(ValueError, "Cannot perform sharding on normalized dimension 1"): | 185 | + with self.assertRaisesRegex(ValueError, "sharded dimension"): |
| 170 | - op.infer_layout((x_layout,), extra_args=(-1,)) | 186 | + op.infer_layout([x_layout, -1]) |
| 171 | 187 | ||
| 172 | @patch("hyper_parallel.core.dtensor.device_mesh.platform") | 188 | @patch("hyper_parallel.core.dtensor.device_mesh.platform") |
| 173 | def test_cumsum_layout_dim_out_of_range_positive(self, mock_platform): | 189 | def test_cumsum_layout_dim_out_of_range_positive(self, mock_platform): |
| @@ -180,8 +196,8 @@ class TestParallelCumsum(unittest.TestCase): | |||
| 180 | x_placements = (Replicate(), Replicate()) | 196 | x_placements = (Replicate(), Replicate()) |
| 181 | x_layout = _build_layout(mesh, x_placements, 2) | 197 | x_layout = _build_layout(mesh, x_placements, 2) |
| 182 | 198 | ||
| 183 | - with self.assertRaisesRegex(ValueError, "Dimension 2 out of range for 2-dimensional input tensor"): | 199 | + with self.assertRaisesRegex(ValueError, "dimension out of range"): |
| 184 | - op.infer_layout((x_layout,), extra_args=(2,)) | 200 | + op.infer_layout([x_layout, 2]) |
| 185 | 201 | ||
| 186 | @patch("hyper_parallel.core.dtensor.device_mesh.platform") | 202 | @patch("hyper_parallel.core.dtensor.device_mesh.platform") |
| 187 | def test_cumsum_layout_missing_dim_parameter(self, mock_platform): | 203 | def test_cumsum_layout_missing_dim_parameter(self, mock_platform): |
| @@ -194,11 +210,8 @@ class TestParallelCumsum(unittest.TestCase): | |||
| 194 | x_placements = (Replicate(), Replicate()) | 210 | x_placements = (Replicate(), Replicate()) |
| 195 | x_layout = _build_layout(mesh, x_placements, 2) | 211 | x_layout = _build_layout(mesh, x_placements, 2) |
| 196 | 212 | ||
| 197 | - with self.assertRaisesRegex(ValueError, "cumsum requires 'dim' parameter in extra_args"): | 213 | + with self.assertRaisesRegex(ValueError, "dimension should be int"): |
| 198 | - op.infer_layout((x_layout,), extra_args=None) | 214 | + op.infer_layout([x_layout, None]) |
| 199 | - | ||
| 200 | - with self.assertRaisesRegex(ValueError, "cumsum requires 'dim' parameter in extra_args"): | ||
| 201 | - op.infer_layout((x_layout,), extra_args=(None,)) | ||
| 202 | 215 | ||
| 203 | @patch("hyper_parallel.core.dtensor.device_mesh.platform") | 216 | @patch("hyper_parallel.core.dtensor.device_mesh.platform") |
| 204 | def test_cumsum_layout_invalid_dim_type(self, mock_platform): | 217 | def test_cumsum_layout_invalid_dim_type(self, mock_platform): |
| @@ -211,8 +224,81 @@ class TestParallelCumsum(unittest.TestCase): | |||
| 211 | x_placements = (Replicate(), Replicate()) | 224 | x_placements = (Replicate(), Replicate()) |
| 212 | x_layout = _build_layout(mesh, x_placements, 2) | 225 | x_layout = _build_layout(mesh, x_placements, 2) |
| 213 | 226 | ||
| 214 | - with self.assertRaisesRegex(ValueError, "'dim' must be an integer, got <class 'str'>"): | 227 | + with self.assertRaisesRegex(ValueError, "dimension should be int"): |
| 215 | - op.infer_layout((x_layout,), extra_args=("invalid",)) | 228 | + op.infer_layout([x_layout, "invalid"]) |
| 229 | + | ||
| 230 | + @patch("hyper_parallel.core.dtensor.device_mesh.platform") | ||
| 231 | + def test_cumsum_partial_input_raises_error(self, mock_platform): | ||
| 232 | + """ | ||
| 233 | + Feature: Cumsum with Partial input | ||
| 234 | + Description: Input layout has Partial status | ||
| 235 | + Expectation: ValueError raised | ||
| 236 | + """ | ||
| 237 | + mesh = self._make_2x4_mesh(mock_platform) | ||
| 238 | + x_placements = (Replicate(), Replicate()) | ||
| 239 | + x_layout = _build_layout(mesh, x_placements, 2) | ||
| 240 | + x_layout.set_partial_by_dev_axis("dp", "sum") | ||
| 241 | + | ||
| 242 | + with self.assertRaisesRegex(ValueError, "Partial status"): | ||
| 243 | + op.infer_layout([x_layout, -1]) | ||
| 244 | + | ||
| 245 | + @patch("hyper_parallel.core.dtensor.device_mesh.platform") | ||
| 246 | + def test_cumsum_preprocess_torch_dtype_in_kwargs(self, mock_platform): | ||
| 247 | + """ | ||
| 248 | + Feature: Cumsum preprocess for PyTorch | ||
| 249 | + Description: dtype is keyword-only parameter | ||
| 250 | + Expectation: local tensor and dim are positional; dtype is in kwargs | ||
| 251 | + """ | ||
| 252 | + mesh = self._make_2x4_mesh(mock_platform) | ||
| 253 | + x_layout = _build_layout(mesh, (Shard(0), Replicate()), 2) | ||
| 254 | + | ||
| 255 | + mock_tensor = MagicMock() | ||
| 256 | + mock_tensor.layout = x_layout | ||
| 257 | + mock_tensor.to_local.return_value = MagicMock() | ||
| 258 | + | ||
| 259 | + local_args, local_kwargs, cache_values = op.preprocess( | ||
| 260 | + (mock_tensor,), | ||
| 261 | + {'dim': -1, 'dtype': 'float32'} | ||
| 262 | + ) | ||
| 263 | + | ||
| 264 | + assert len(local_args) == 2, ( | ||
| 265 | + f"For PyTorch 'cumsum', local_args should be (tensor, dim), got {local_args}" | ||
| 266 | + ) | ||
| 267 | + assert local_args[1] == -1, f"dim should be preserved as -1, got {local_args[1]}" | ||
| 268 | + assert local_kwargs == {'dtype': 'float32'}, ( | ||
| 269 | + f"For PyTorch 'cumsum', dtype should be kwargs, got {local_kwargs}" | ||
| 270 | + ) | ||
| 271 | + assert cache_values == [x_layout, -1], f"Unexpected cache_values: {cache_values}" | ||
| 272 | + | ||
| 273 | + @patch("hyper_parallel.core.dtensor.device_mesh.platform") | ||
| 274 | + def test_cumsum_preprocess_mindspore_primitive_in_args(self, mock_platform): | ||
| 275 | + """ | ||
| 276 | + Feature: Cumsum preprocess for MindSpore Primitive | ||
| 277 | + Description: CumsumExt does not accept kwargs | ||
| 278 | + Expectation: dim and dtype are routed to positional args, including dtype=None | ||
| 279 | + """ | ||
| 280 | + mesh = self._make_2x4_mesh(mock_platform) | ||
| 281 | + x_layout = _build_layout(mesh, (Shard(0), Replicate()), 2) | ||
| 282 | + | ||
| 283 | + mock_tensor = MagicMock() | ||
| 284 | + mock_tensor.layout = x_layout | ||
| 285 | + mock_tensor.to_local.return_value = MagicMock() | ||
| 286 | + | ||
| 287 | + local_args, local_kwargs, cache_values = op_ms.preprocess( | ||
| 288 | + (mock_tensor,), | ||
| 289 | + {'dim': -1} | ||
| 290 | + ) | ||
| 291 | + | ||
| 292 | + assert not local_kwargs, ( | ||
| 293 | + f"For MindSpore 'CumsumExt', local_kwargs should be empty, got {local_kwargs}" | ||
| 294 | + ) | ||
| 295 | + assert len(local_args) == 3, ( | ||
| 296 | + f"For MindSpore 'CumsumExt', local_args should be (tensor, dim, dtype), got {local_args}" | ||
| 297 | + ) | ||
| 298 | + assert local_args[1:] == (-1, None), ( | ||
| 299 | + f"MindSpore positional dim/dtype mismatch, got {local_args[1:]}" | ||
| 300 | + ) | ||
| 301 | + assert cache_values == [x_layout, -1], f"Unexpected cache_values: {cache_values}" | ||
| 216 | 302 | ||
| 217 | 303 | ||
| 218 | if __name__ == "__main__": | 304 | if __name__ == "__main__": |
| @@ -13,9 +13,8 @@ | |||
| 13 | # limitations under the License. | 13 | # limitations under the License. |
| 14 | # ============================================================================ | 14 | # ============================================================================ |
| 15 | """Unit tests for element-wise distributed operators""" | 15 | """Unit tests for element-wise distributed operators""" |
| 16 | -import os | ||
| 17 | import unittest | 16 | import unittest |
| 18 | -from unittest.mock import patch | 17 | +from unittest.mock import MagicMock, patch |
| 19 | import copy | 18 | import copy |
| 20 | import numpy as np | 19 | import numpy as np |
| 21 | 20 | ||
| @@ -29,6 +28,31 @@ from hyper_parallel.core.dtensor.device_mesh import ( | |||
| 29 | from hyper_parallel.platform.platform import EXISTING_COMM_GROUPS | 28 | from hyper_parallel.platform.platform import EXISTING_COMM_GROUPS |
| 30 | 29 | ||
| 31 | 30 | ||
| 31 | +def _extract_test_input_shapes(extra_args): | ||
| 32 | + """Extract input shapes from extra_args dict.""" | ||
| 33 | + if isinstance(extra_args, dict): | ||
| 34 | + return extra_args.get("input_shapes", None) | ||
| 35 | + return None | ||
| 36 | + | ||
| 37 | + | ||
| 38 | +def _make_cache_values(layouts, extra_args=None): | ||
| 39 | + """Build cache_values for the new preprocess/infer_layout protocol.""" | ||
| 40 | + return [*layouts, _extract_test_input_shapes(extra_args)] | ||
| 41 | + | ||
| 42 | + | ||
| 43 | +def _infer_output_layout(op, layouts, extra_args=None): | ||
| 44 | + """Infer and return the single output layout.""" | ||
| 45 | + output_layouts, extra_info = op.infer_layout(_make_cache_values(layouts, extra_args)) | ||
| 46 | + assert extra_info is None | ||
| 47 | + return output_layouts[0] | ||
| 48 | + | ||
| 49 | + | ||
| 50 | +def _get_expand_impl(op, func, output_layout, layouts, extra_args=None): | ||
| 51 | + """Call get_expand_impl with new-protocol infer_result and cache_values.""" | ||
| 52 | + cache_values = _make_cache_values(layouts, extra_args) | ||
| 53 | + return op.get_expand_impl(func, ((output_layout,), None), cache_values) | ||
| 54 | + | ||
| 55 | + | ||
| 32 | class TestParallelElementwiseOps(unittest.TestCase): | 56 | class TestParallelElementwiseOps(unittest.TestCase): |
| 33 | """Unit tests for ElementWiseDistributedOp.""" | 57 | """Unit tests for ElementWiseDistributedOp.""" |
| 34 | def setUp(self): | 58 | def setUp(self): |
| @@ -39,6 +63,7 @@ class TestParallelElementwiseOps(unittest.TestCase): | |||
| 39 | """ | 63 | """ |
| 40 | EXISTING_COMM_GROUPS.clear() | 64 | EXISTING_COMM_GROUPS.clear() |
| 41 | _DEVICE_MESH_MAP.clear() | 65 | _DEVICE_MESH_MAP.clear() |
| 66 | + _LAYOUT_CACHE.clear() | ||
| 42 | self.op = ElementWiseDistributedOp("element_wise") | 67 | self.op = ElementWiseDistributedOp("element_wise") |
| 43 | self.op_with_partial = AddDistributedOp("element_wise_with_partial") | 68 | self.op_with_partial = AddDistributedOp("element_wise_with_partial") |
| 44 | 69 | ||
| @@ -46,6 +71,7 @@ class TestParallelElementwiseOps(unittest.TestCase): | |||
| 46 | """Clean up after each test method.""" | 71 | """Clean up after each test method.""" |
| 47 | EXISTING_COMM_GROUPS.clear() | 72 | EXISTING_COMM_GROUPS.clear() |
| 48 | _DEVICE_MESH_MAP.clear() | 73 | _DEVICE_MESH_MAP.clear() |
| 74 | + _LAYOUT_CACHE.clear() | ||
| 49 | 75 | ||
| 50 | def _setup_mock_platform(self, mock_platform, platform_type=None, world_size=8): | 76 | def _setup_mock_platform(self, mock_platform, platform_type=None, world_size=8): |
| 51 | """Configure common mock-platform attributes used across tests. | 77 | """Configure common mock-platform attributes used across tests. |
| @@ -79,12 +105,21 @@ class TestParallelElementwiseOps(unittest.TestCase): | |||
| 79 | op_ = self.op_with_partial | 105 | op_ = self.op_with_partial |
| 80 | else: | 106 | else: |
| 81 | op_ = self.op | 107 | op_ = self.op |
| 82 | - output_layout = op_.infer_layout((x_layout, y_layout), extra_args) | 108 | + |
| 109 | + cache_values = [ | ||
| 110 | + x_layout, | ||
| 111 | + y_layout, | ||
| 112 | + extra_args.get("input_shapes"), | ||
| 113 | + ] | ||
| 114 | + | ||
| 115 | + infer_result = op_.infer_layout(cache_values) | ||
| 116 | + output_layout = infer_result[0][0] | ||
| 117 | + | ||
| 83 | got_map = output_layout.tensor_map | 118 | got_map = output_layout.tensor_map |
| 84 | assert got_map == expected_map, ( | 119 | assert got_map == expected_map, ( |
| 85 | f"Element-wise failed. Expected {expected_map}, got {got_map}" | 120 | f"Element-wise failed. Expected {expected_map}, got {got_map}" |
| 86 | ) | 121 | ) |
| 87 | - impl = op_.get_expand_impl(None, output_layout, (x_layout, y_layout), extra_args) | 122 | + impl = op_.get_expand_impl(None, infer_result, cache_values) |
| 88 | if flag: | 123 | if flag: |
| 89 | assert impl is not None, ( | 124 | assert impl is not None, ( |
| 90 | f"get_expand_impl test failed. Expected non-None, got {impl}" | 125 | f"get_expand_impl test failed. Expected non-None, got {impl}" |
| @@ -101,12 +136,20 @@ class TestParallelElementwiseOps(unittest.TestCase): | |||
| 101 | op_ = self.op_with_partial | 136 | op_ = self.op_with_partial |
| 102 | else: | 137 | else: |
| 103 | op_ = self.op | 138 | op_ = self.op |
| 104 | - output_layout = op_.infer_layout((x_layout,), extra_args) | 139 | + |
| 140 | + cache_values = [ | ||
| 141 | + x_layout, | ||
| 142 | + extra_args.get("input_shapes"), | ||
| 143 | + ] | ||
| 144 | + | ||
| 145 | + infer_result = op_.infer_layout(cache_values) | ||
| 146 | + output_layout = infer_result[0][0] | ||
| 147 | + | ||
| 105 | got_map = output_layout.tensor_map | 148 | got_map = output_layout.tensor_map |
| 106 | assert got_map == expected_map, ( | 149 | assert got_map == expected_map, ( |
| 107 | f"Element-wise failed. Expected {expected_map}, got {got_map}" | 150 | f"Element-wise failed. Expected {expected_map}, got {got_map}" |
| 108 | ) | 151 | ) |
| 109 | - impl = op_.get_expand_impl(None, output_layout, (x_layout,), extra_args) | 152 | + impl = op_.get_expand_impl(None, infer_result, cache_values) |
| 110 | if flag: | 153 | if flag: |
| 111 | assert impl is not None, ( | 154 | assert impl is not None, ( |
| 112 | f"get_expand_impl test failed. Expected non-None, got {impl}" | 155 | f"get_expand_impl test failed. Expected non-None, got {impl}" |
| @@ -117,6 +160,37 @@ class TestParallelElementwiseOps(unittest.TestCase): | |||
| 117 | f"get_expand_impl test failed. Expected None, got {impl}" | 160 | f"get_expand_impl test failed. Expected None, got {impl}" |
| 118 | ) | 161 | ) |
| 119 | 162 | ||
| 163 | + def test_preprocess_unwraps_dtensors_and_builds_cache_values(self): | ||
| 164 | + """ | ||
| 165 | + Feature: Element-wise preprocess | ||
| 166 | + Description: DTensor-like inputs are converted to local tensors and cached with layout/shape. | ||
| 167 | + Expectation: local args/kwargs use local tensors and cache_values carries layouts first. | ||
| 168 | + """ | ||
| 169 | + x_layout = MagicMock() | ||
| 170 | + y_layout = MagicMock() | ||
| 171 | + x_local = MagicMock() | ||
| 172 | + y_local = MagicMock() | ||
| 173 | + | ||
| 174 | + x_tensor = MagicMock() | ||
| 175 | + x_tensor.layout = x_layout | ||
| 176 | + x_tensor.shape = (4, 8) | ||
| 177 | + x_tensor.to_local.return_value = x_local | ||
| 178 | + x_tensor._layout = x_layout | ||
| 179 | + | ||
| 180 | + y_tensor = MagicMock() | ||
| 181 | + y_tensor.layout = y_layout | ||
| 182 | + y_tensor.shape = (4, 8) | ||
| 183 | + y_tensor.to_local.return_value = y_local | ||
| 184 | + y_tensor._layout = y_layout | ||
| 185 | + | ||
| 186 | + local_args, local_kwargs, cache_values = self.op.preprocess((x_tensor,), {"other": y_tensor}) | ||
| 187 | + | ||
| 188 | + assert local_args == (x_local,) | ||
| 189 | + assert local_kwargs == {"other": y_local} | ||
| 190 | + assert cache_values[0] is x_layout | ||
| 191 | + assert cache_values[1] is y_layout | ||
| 192 | + assert cache_values[2] == [(4, 8), (4, 8)] | ||
| 193 | + | ||
| 120 | @patch("hyper_parallel.core.dtensor.device_mesh.platform") | 194 | @patch("hyper_parallel.core.dtensor.device_mesh.platform") |
| 121 | def test_single_input_partial_0(self, mock_platform): | 195 | def test_single_input_partial_0(self, mock_platform): |
| 122 | """ | 196 | """ |
| @@ -129,7 +203,7 @@ class TestParallelElementwiseOps(unittest.TestCase): | |||
| 129 | x_layout.set_partial_by_dev_axis("dp", "sum") | 203 | x_layout.set_partial_by_dev_axis("dp", "sum") |
| 130 | 204 | ||
| 131 | with self.assertRaisesRegex(ValueError, "has Partial status which is not allowed"): | 205 | with self.assertRaisesRegex(ValueError, "has Partial status which is not allowed"): |
| 132 | - self.op.infer_layout((x_layout,), extra_args={"input_shapes": [(4, 8, 16)]}) | 206 | + self.op.infer_layout(_make_cache_values((x_layout,), {"input_shapes": [(4, 8, 16)]})) |
| 133 | 207 | ||
| 134 | @patch("hyper_parallel.core.dtensor.device_mesh.platform") | 208 | @patch("hyper_parallel.core.dtensor.device_mesh.platform") |
| 135 | def test_single_input_replicate_1(self, mock_platform): | 209 | def test_single_input_replicate_1(self, mock_platform): |
| @@ -177,7 +251,9 @@ class TestParallelElementwiseOps(unittest.TestCase): | |||
| 177 | x_layout = _build_layout(mesh, (Replicate(), Replicate(), Replicate()), 3) | 251 | x_layout = _build_layout(mesh, (Replicate(), Replicate(), Replicate()), 3) |
| 178 | x_layout.set_partial_by_dev_axis("dp", "sum") | 252 | x_layout.set_partial_by_dev_axis("dp", "sum") |
| 179 | 253 | ||
| 180 | - output_layout = self.op_with_partial.infer_layout((x_layout,), extra_args={"input_shapes": [(4, 8, 16)]}) | 254 | + output_layout = _infer_output_layout( |
| 255 | + self.op_with_partial, (x_layout,), {"input_shapes": [(4, 8, 16)]} | ||
| 256 | + ) | ||
| 181 | 257 | ||
| 182 | assert output_layout.partial[mesh.axis_index("dp")] == "sum" | 258 | assert output_layout.partial[mesh.axis_index("dp")] == "sum" |
| 183 | 259 | ||
| @@ -249,7 +325,7 @@ class TestParallelElementwiseOps(unittest.TestCase): | |||
| 249 | extra_args = {"input_shapes": [(4, 8, 16), (4, 8, 16)]} | 325 | extra_args = {"input_shapes": [(4, 8, 16), (4, 8, 16)]} |
| 250 | 326 | ||
| 251 | with self.assertRaisesRegex(ValueError, "should have same sharding pattern"): | 327 | with self.assertRaisesRegex(ValueError, "should have same sharding pattern"): |
| 252 | - self.op.infer_layout((x_layout, y_layout), extra_args=extra_args) | 328 | + self.op.infer_layout(_make_cache_values((x_layout, y_layout), extra_args)) |
| 253 | 329 | ||
| 254 | @patch("hyper_parallel.core.dtensor.device_mesh.platform") | 330 | @patch("hyper_parallel.core.dtensor.device_mesh.platform") |
| 255 | def test_sharded_broadcasts_to_replicated_8(self, mock_platform): | 331 | def test_sharded_broadcasts_to_replicated_8(self, mock_platform): |
| @@ -265,7 +341,7 @@ class TestParallelElementwiseOps(unittest.TestCase): | |||
| 265 | extra_args = {"input_shapes": [(1, 8, 16), (4, 8, 16)]} | 341 | extra_args = {"input_shapes": [(1, 8, 16), (4, 8, 16)]} |
| 266 | 342 | ||
| 267 | with self.assertRaisesRegex(ValueError, "Broadcasting dimension cannot be sharded"): | 343 | with self.assertRaisesRegex(ValueError, "Broadcasting dimension cannot be sharded"): |
| 268 | - self.op.infer_layout((x_layout, y_layout), extra_args=extra_args) | 344 | + self.op.infer_layout(_make_cache_values((x_layout, y_layout), extra_args)) |
| 269 | 345 | ||
| 270 | @patch("hyper_parallel.core.dtensor.device_mesh.platform") | 346 | @patch("hyper_parallel.core.dtensor.device_mesh.platform") |
| 271 | def test_replicated_broadcasts_to_sharded_9(self, mock_platform): | 347 | def test_replicated_broadcasts_to_sharded_9(self, mock_platform): |
| @@ -354,11 +430,11 @@ class TestParallelElementwiseOps(unittest.TestCase): | |||
| 354 | y_layout._partial = [None] * len(y_layout._partial) | 430 | y_layout._partial = [None] * len(y_layout._partial) |
| 355 | 431 | ||
| 356 | extra_args = {"input_shapes": [(4, 8, 16), (4, 8, 16)]} | 432 | extra_args = {"input_shapes": [(4, 8, 16), (4, 8, 16)]} |
| 357 | - output_layout = self.op_with_partial.infer_layout((x_layout, y_layout), extra_args) | 433 | + output_layout = _infer_output_layout(self.op_with_partial, (x_layout, y_layout), extra_args) |
| 358 | 434 | ||
| 359 | assert output_layout.partial[mesh.axis_index("dp")] == "sum" | 435 | assert output_layout.partial[mesh.axis_index("dp")] == "sum" |
| 360 | 436 | ||
| 361 | - impl = self.op_with_partial.get_expand_impl(None, output_layout, (x_layout, y_layout), extra_args) | 437 | + impl = _get_expand_impl(self.op_with_partial, None, output_layout, (x_layout, y_layout), extra_args) |
| 362 | assert impl is not None, ( | 438 | assert impl is not None, ( |
| 363 | f"get_expand_impl test failed. Expected non-None, got {impl}" | 439 | f"get_expand_impl test failed. Expected non-None, got {impl}" |
| 364 | ) | 440 | ) |
| @@ -378,11 +454,11 @@ class TestParallelElementwiseOps(unittest.TestCase): | |||
| 378 | y_layout.set_partial_by_dev_axis("dp", "sum") | 454 | y_layout.set_partial_by_dev_axis("dp", "sum") |
| 379 | 455 | ||
| 380 | extra_args = {"input_shapes": [(4, 8, 16), (4, 8, 16)]} | 456 | extra_args = {"input_shapes": [(4, 8, 16), (4, 8, 16)]} |
| 381 | - output_layout = self.op_with_partial.infer_layout((x_layout, y_layout), extra_args) | 457 | + output_layout = _infer_output_layout(self.op_with_partial, (x_layout, y_layout), extra_args) |
| 382 | 458 | ||
| 383 | assert output_layout.partial[mesh.axis_index("dp")] == "sum" | 459 | assert output_layout.partial[mesh.axis_index("dp")] == "sum" |
| 384 | 460 | ||
| 385 | - impl = self.op_with_partial.get_expand_impl(None, output_layout, (x_layout, y_layout), extra_args) | 461 | + impl = _get_expand_impl(self.op_with_partial, None, output_layout, (x_layout, y_layout), extra_args) |
| 386 | assert impl is None, ( | 462 | assert impl is None, ( |
| 387 | f"get_expand_impl test failed. Expected None, got {impl}" | 463 | f"get_expand_impl test failed. Expected None, got {impl}" |
| 388 | ) | 464 | ) |
| @@ -403,7 +479,7 @@ class TestParallelElementwiseOps(unittest.TestCase): | |||
| 403 | extra_args = {"input_shapes": [(4, 8, 16), (4, 8, 16)]} | 479 | extra_args = {"input_shapes": [(4, 8, 16), (4, 8, 16)]} |
| 404 | 480 | ||
| 405 | with self.assertRaisesRegex(ValueError, "partial operations should be same"): | 481 | with self.assertRaisesRegex(ValueError, "partial operations should be same"): |
| 406 | - self.op_with_partial.infer_layout((x_layout, y_layout), extra_args=extra_args) | 482 | + self.op_with_partial.infer_layout(_make_cache_values((x_layout, y_layout), extra_args)) |
| 407 | 483 | ||
| 408 | @patch("hyper_parallel.core.dtensor.device_mesh.platform") | 484 | @patch("hyper_parallel.core.dtensor.device_mesh.platform") |
| 409 | def test_shard_with_partial_conflict_16(self, mock_platform): | 485 | def test_shard_with_partial_conflict_16(self, mock_platform): |
| @@ -420,7 +496,7 @@ class TestParallelElementwiseOps(unittest.TestCase): | |||
| 420 | extra_args = {"input_shapes": [(4, 8, 16), (4, 8, 16)]} | 496 | extra_args = {"input_shapes": [(4, 8, 16), (4, 8, 16)]} |
| 421 | 497 | ||
| 422 | with self.assertRaisesRegex(ValueError, "Shard and Partial should not coexist on same device axis"): | 498 | with self.assertRaisesRegex(ValueError, "Shard and Partial should not coexist on same device axis"): |
| 423 | - self.op_with_partial.infer_layout((x_layout, y_layout), extra_args=extra_args) | 499 | + self.op_with_partial.infer_layout(_make_cache_values((x_layout, y_layout), extra_args)) |
| 424 | 500 | ||
| 425 | @patch("hyper_parallel.core.dtensor.device_mesh.platform") | 501 | @patch("hyper_parallel.core.dtensor.device_mesh.platform") |
| 426 | def test_partial_broadcasts_to_sharded_17(self, mock_platform): | 502 | def test_partial_broadcasts_to_sharded_17(self, mock_platform): |
| @@ -437,7 +513,7 @@ class TestParallelElementwiseOps(unittest.TestCase): | |||
| 437 | extra_args = {"input_shapes": [(1, 8, 16), (4, 8, 16)]} | 513 | extra_args = {"input_shapes": [(1, 8, 16), (4, 8, 16)]} |
| 438 | 514 | ||
| 439 | with self.assertRaisesRegex(ValueError, "Shard and Partial should not coexist on same device axis"): | 515 | with self.assertRaisesRegex(ValueError, "Shard and Partial should not coexist on same device axis"): |
| 440 | - self.op_with_partial.infer_layout((x_layout, y_layout), extra_args=extra_args) | 516 | + self.op_with_partial.infer_layout(_make_cache_values((x_layout, y_layout), extra_args)) |
| 441 | 517 | ||
| 442 | @patch("hyper_parallel.core.dtensor.device_mesh.platform") | 518 | @patch("hyper_parallel.core.dtensor.device_mesh.platform") |
| 443 | def test_multiple_inputs_18(self, mock_platform): | 519 | def test_multiple_inputs_18(self, mock_platform): |
| @@ -454,15 +530,15 @@ class TestParallelElementwiseOps(unittest.TestCase): | |||
| 454 | extra_args = {"input_shapes": [(4, 8, 16), (1, 8, 16), (1, 8, 16)]} | 530 | extra_args = {"input_shapes": [(4, 8, 16), (1, 8, 16), (1, 8, 16)]} |
| 455 | expected_map = (2, -1, -1) | 531 | expected_map = (2, -1, -1) |
| 456 | 532 | ||
| 457 | - output_layout = self.op.infer_layout((x_layout, y_layout, z_layout), extra_args=extra_args) | 533 | + output_layout = _infer_output_layout(self.op, (x_layout, y_layout, z_layout), extra_args) |
| 458 | got_map = output_layout.tensor_map | 534 | got_map = output_layout.tensor_map |
| 459 | assert got_map == expected_map | 535 | assert got_map == expected_map |
| 460 | 536 | ||
| 461 | - assert self.op.get_expand_impl(None, output_layout, | 537 | + assert _get_expand_impl(self.op, None, output_layout, |
| 462 | - (x_layout, y_layout, z_layout), extra_args) is None, ( | 538 | + (x_layout, y_layout, z_layout), extra_args) is None, ( |
| 463 | f"get_expand_impl test failed. Expected None, " | 539 | f"get_expand_impl test failed. Expected None, " |
| 464 | - f"""got {self.op.get_expand_impl(None, output_layout, | 540 | + f"""got {_get_expand_impl(self.op, None, output_layout, |
| 465 | - (x_layout, y_layout, z_layout), extra_args)}""" | 541 | + (x_layout, y_layout, z_layout), extra_args)}""" |
| 466 | ) | 542 | ) |
| 467 | 543 | ||
| 468 | @patch("hyper_parallel.core.dtensor.device_mesh.platform") | 544 | @patch("hyper_parallel.core.dtensor.device_mesh.platform") |
| @@ -480,7 +556,7 @@ class TestParallelElementwiseOps(unittest.TestCase): | |||
| 480 | extra_args = {"input_shapes": [(4, 8, 16), (4, 8, 16)]} | 556 | extra_args = {"input_shapes": [(4, 8, 16), (4, 8, 16)]} |
| 481 | 557 | ||
| 482 | with self.assertRaisesRegex(ValueError, "Shard and Partial should not coexist on same device axis"): | 558 | with self.assertRaisesRegex(ValueError, "Shard and Partial should not coexist on same device axis"): |
| 483 | - self.op_with_partial.infer_layout((x_layout, y_layout), extra_args=extra_args) | 559 | + self.op_with_partial.infer_layout(_make_cache_values((x_layout, y_layout), extra_args)) |
| 484 | 560 | ||
| 485 | @patch("hyper_parallel.core.dtensor.device_mesh.platform") | 561 | @patch("hyper_parallel.core.dtensor.device_mesh.platform") |
| 486 | def test_incompatible_broadcast_shapes_20(self, mock_platform): | 562 | def test_incompatible_broadcast_shapes_20(self, mock_platform): |
| @@ -496,7 +572,7 @@ class TestParallelElementwiseOps(unittest.TestCase): | |||
| 496 | extra_args = {"input_shapes": [(4, 8, 16), (4, 7, 16)]} | 572 | extra_args = {"input_shapes": [(4, 8, 16), (4, 7, 16)]} |
| 497 | 573 | ||
| 498 | with self.assertRaisesRegex(ValueError, "cannot be broadcast"): | 574 | with self.assertRaisesRegex(ValueError, "cannot be broadcast"): |
| 499 | - self.op_with_partial.infer_layout((x_layout, y_layout), extra_args=extra_args) | 575 | + self.op_with_partial.infer_layout(_make_cache_values((x_layout, y_layout), extra_args)) |
| 500 | 576 | ||
| 501 | @patch("hyper_parallel.core.dtensor.device_mesh.platform") | 577 | @patch("hyper_parallel.core.dtensor.device_mesh.platform") |
| 502 | def test_1d_to_3d_broadcast_21(self, mock_platform): | 578 | def test_1d_to_3d_broadcast_21(self, mock_platform): |
| @@ -528,7 +604,7 @@ class TestParallelElementwiseOps(unittest.TestCase): | |||
| 528 | y_layout = _build_layout(mesh, (Replicate(), Replicate(), Replicate()), 3) | 604 | y_layout = _build_layout(mesh, (Replicate(), Replicate(), Replicate()), 3) |
| 529 | 605 | ||
| 530 | with self.assertRaisesRegex(ValueError, "cannot infer layout without shapes"): | 606 | with self.assertRaisesRegex(ValueError, "cannot infer layout without shapes"): |
| 531 | - self.op.infer_layout((x_layout, y_layout), extra_args={}) | 607 | + self.op.infer_layout(_make_cache_values((x_layout, y_layout), {})) |
| 532 | 608 | ||
| 533 | 609 | ||
| 534 | class TestParallelArithmetic(unittest.TestCase): | 610 | class TestParallelArithmetic(unittest.TestCase): |
| @@ -538,11 +614,13 @@ class TestParallelArithmetic(unittest.TestCase): | |||
| 538 | """Set up test fixtures before each test method.""" | 614 | """Set up test fixtures before each test method.""" |
| 539 | EXISTING_COMM_GROUPS.clear() | 615 | EXISTING_COMM_GROUPS.clear() |
| 540 | _DEVICE_MESH_MAP.clear() | 616 | _DEVICE_MESH_MAP.clear() |
| 617 | + _LAYOUT_CACHE.clear() | ||
| 541 | 618 | ||
| 542 | def tearDown(self): | 619 | def tearDown(self): |
| 543 | """Clean up after each test method.""" | 620 | """Clean up after each test method.""" |
| 544 | EXISTING_COMM_GROUPS.clear() | 621 | EXISTING_COMM_GROUPS.clear() |
| 545 | _DEVICE_MESH_MAP.clear() | 622 | _DEVICE_MESH_MAP.clear() |
| 623 | + _LAYOUT_CACHE.clear() | ||
| 546 | 624 | ||
| 547 | def _setup_mock_platform(self, mock_platform, platform_type=None, world_size=8): | 625 | def _setup_mock_platform(self, mock_platform, platform_type=None, world_size=8): |
| 548 | """Configure common mock-platform attributes used across tests.""" | 626 | """Configure common mock-platform attributes used across tests.""" |
| @@ -580,7 +658,7 @@ class TestParallelArithmetic(unittest.TestCase): | |||
| 580 | w_layout = _build_layout(mesh, w_placements, 2) | 658 | w_layout = _build_layout(mesh, w_placements, 2) |
| 581 | 659 | ||
| 582 | extra_args = {"input_shapes": [(4, 16), (4, 16)]} | 660 | extra_args = {"input_shapes": [(4, 16), (4, 16)]} |
| 583 | - output_layout = op.infer_layout((x_layout, w_layout), (extra_args)) | 661 | + output_layout = _infer_output_layout(op, (x_layout, w_layout), extra_args) |
| 584 | expected_map = ("dp", "mp") | 662 | expected_map = ("dp", "mp") |
| 585 | assert output_layout.alias_tensor_map == expected_map, ( | 663 | assert output_layout.alias_tensor_map == expected_map, ( |
| 586 | f"Data Parallel test failed. Expected {expected_map}, " | 664 | f"Data Parallel test failed. Expected {expected_map}, " |
| @@ -589,9 +667,10 @@ class TestParallelArithmetic(unittest.TestCase): | |||
| 589 | 667 | ||
| 590 | # Since `get_expand_impl` is not overridden, it returns None by default. | 668 | # Since `get_expand_impl` is not overridden, it returns None by default. |
| 591 | # The same applies to other test classes, so it is unnecessary to test its return value. | 669 | # The same applies to other test classes, so it is unnecessary to test its return value. |
| 592 | - assert op.get_expand_impl(None, output_layout, (x_layout, None, None), (1, True)) is None, ( | 670 | + impl = _get_expand_impl(op, None, output_layout, (x_layout, w_layout), extra_args) |
| 671 | + assert impl is None, ( | ||
| 593 | f"get_expand_impl test failed. Expected None, " | 672 | f"get_expand_impl test failed. Expected None, " |
| 594 | - f"got {op.get_expand_impl(None, output_layout, (x_layout, None, None), (1, True))}" | 673 | + f"got {impl}" |
| 595 | ) | 674 | ) |
| 596 | 675 | ||
| 597 | @patch("hyper_parallel.core.dtensor.device_mesh.platform") | 676 | @patch("hyper_parallel.core.dtensor.device_mesh.platform") |
| @@ -609,7 +688,7 @@ class TestParallelArithmetic(unittest.TestCase): | |||
| 609 | w_placements = (Shard(0), Replicate()) | 688 | w_placements = (Shard(0), Replicate()) |
| 610 | w_layout = _build_layout(mesh, w_placements, 2) | 689 | w_layout = _build_layout(mesh, w_placements, 2) |
| 611 | extra_args = {"input_shapes": [(4, 16), (4, 16)]} | 690 | extra_args = {"input_shapes": [(4, 16), (4, 16)]} |
| 612 | - output_layout = op.infer_layout((x_layout, w_layout), (extra_args)) | 691 | + output_layout = _infer_output_layout(op, (x_layout, w_layout), extra_args) |
| 613 | expected_map = ("dp", "mp") | 692 | expected_map = ("dp", "mp") |
| 614 | assert output_layout.alias_tensor_map == expected_map, ( | 693 | assert output_layout.alias_tensor_map == expected_map, ( |
| 615 | f"Data Parallel test failed. Expected {expected_map}, " | 694 | f"Data Parallel test failed. Expected {expected_map}, " |
| @@ -631,7 +710,7 @@ class TestParallelArithmetic(unittest.TestCase): | |||
| 631 | w_placements = (Shard(0), Replicate()) | 710 | w_placements = (Shard(0), Replicate()) |
| 632 | w_layout = _build_layout(mesh, w_placements, 2) | 711 | w_layout = _build_layout(mesh, w_placements, 2) |
| 633 | extra_args = {"input_shapes": [(4, 16), (4, 16)]} | 712 | extra_args = {"input_shapes": [(4, 16), (4, 16)]} |
| 634 | - output_layout = op.infer_layout((x_layout, w_layout), extra_args) | 713 | + output_layout = _infer_output_layout(op, (x_layout, w_layout), extra_args) |
| 635 | expected_map = ("dp", "mp") | 714 | expected_map = ("dp", "mp") |
| 636 | assert output_layout.alias_tensor_map == expected_map, ( | 715 | assert output_layout.alias_tensor_map == expected_map, ( |
| 637 | f"Data Parallel test failed. Expected {expected_map}, " | 716 | f"Data Parallel test failed. Expected {expected_map}, " |
| @@ -653,7 +732,7 @@ class TestParallelArithmetic(unittest.TestCase): | |||
| 653 | w_placements = (Shard(0), Replicate()) | 732 | w_placements = (Shard(0), Replicate()) |
| 654 | w_layout = _build_layout(mesh, w_placements, 2) | 733 | w_layout = _build_layout(mesh, w_placements, 2) |
| 655 | extra_args = {"input_shapes": [(4, 16), (4, 16)]} | 734 | extra_args = {"input_shapes": [(4, 16), (4, 16)]} |
| 656 | - output_layout = op.infer_layout((x_layout, w_layout), extra_args) | 735 | + output_layout = _infer_output_layout(op, (x_layout, w_layout), extra_args) |
| 657 | expected_map = ("dp", "mp") | 736 | expected_map = ("dp", "mp") |
| 658 | assert output_layout.alias_tensor_map == expected_map, ( | 737 | assert output_layout.alias_tensor_map == expected_map, ( |
| 659 | f"Data Parallel test failed. Expected {expected_map}, " | 738 | f"Data Parallel test failed. Expected {expected_map}, " |
| @@ -675,7 +754,7 @@ class TestParallelArithmetic(unittest.TestCase): | |||
| 675 | w_placements = (Shard(0), Replicate()) | 754 | w_placements = (Shard(0), Replicate()) |
| 676 | w_layout = _build_layout(mesh, w_placements, 2) | 755 | w_layout = _build_layout(mesh, w_placements, 2) |
| 677 | extra_args = {"input_shapes": [(4, 16), (4, 16)]} | 756 | extra_args = {"input_shapes": [(4, 16), (4, 16)]} |
| 678 | - output_layout = op.infer_layout((x_layout, w_layout), extra_args) | 757 | + output_layout = _infer_output_layout(op, (x_layout, w_layout), extra_args) |
| 679 | expected_map = ("dp", "mp") | 758 | expected_map = ("dp", "mp") |
| 680 | assert output_layout.alias_tensor_map == expected_map, ( | 759 | assert output_layout.alias_tensor_map == expected_map, ( |
| 681 | f"Data Parallel test failed. Expected {expected_map}, " | 760 | f"Data Parallel test failed. Expected {expected_map}, " |
| @@ -728,7 +807,9 @@ class TestParallelZerosLike(unittest.TestCase): | |||
| 728 | placements = (Shard(0), Replicate()) | 807 | placements = (Shard(0), Replicate()) |
| 729 | x_layout = _build_layout(mesh, placements, 2) | 808 | x_layout = _build_layout(mesh, placements, 2) |
| 730 | 809 | ||
| 731 | - output_layout = self.op.infer_layout((x_layout,), {}) | 810 | + cache_values = [x_layout, None] |
| 811 | + infer_result = self.op.infer_layout(cache_values) | ||
| 812 | + output_layout = infer_result[0][0] | ||
| 732 | 813 | ||
| 733 | expected_map = (1, -1) | 814 | expected_map = (1, -1) |
| 734 | assert output_layout.tensor_map == expected_map, ( | 815 | assert output_layout.tensor_map == expected_map, ( |
| @@ -738,9 +819,8 @@ class TestParallelZerosLike(unittest.TestCase): | |||
| 738 | 819 | ||
| 739 | # Since `get_expand_impl` is not overridden, it returns None by default. | 820 | # Since `get_expand_impl` is not overridden, it returns None by default. |
| 740 | # The same applies to other test cases, so it is unnecessary to test its return value. | 821 | # The same applies to other test cases, so it is unnecessary to test its return value. |
| 741 | - assert self.op.get_expand_impl(None, output_layout, (x_layout,), {}) is None, ( | 822 | + assert self.op.get_expand_impl(None, infer_result, cache_values) is None, ( |
| 742 | - f"get_expand_impl should return None, " | 823 | + "get_expand_impl should return None" |
| 743 | - f"got={self.op.get_expand_impl(None, output_layout, (x_layout,), {})}" | ||
| 744 | ) | 824 | ) |
| 745 | 825 | ||
| 746 | @patch("hyper_parallel.core.dtensor.device_mesh.platform") | 826 | @patch("hyper_parallel.core.dtensor.device_mesh.platform") |
| @@ -754,7 +834,9 @@ class TestParallelZerosLike(unittest.TestCase): | |||
| 754 | placements = (Replicate(), Replicate()) | 834 | placements = (Replicate(), Replicate()) |
| 755 | x_layout = _build_layout(mesh, placements, 2) | 835 | x_layout = _build_layout(mesh, placements, 2) |
| 756 | 836 | ||
| 757 | - output_layout = self.op.infer_layout((x_layout,), {}) | 837 | + cache_values = [x_layout, None] |
| 838 | + infer_result = self.op.infer_layout(cache_values) | ||
| 839 | + output_layout = infer_result[0][0] | ||
| 758 | 840 | ||
| 759 | expected_map = (-1, -1) | 841 | expected_map = (-1, -1) |
| 760 | assert output_layout.tensor_map == expected_map, ( | 842 | assert output_layout.tensor_map == expected_map, ( |
| @@ -774,7 +856,7 @@ class TestParallelZerosLike(unittest.TestCase): | |||
| 774 | x_layout.set_partial_by_dev_axis("dp", "sum") | 856 | x_layout.set_partial_by_dev_axis("dp", "sum") |
| 775 | 857 | ||
| 776 | with self.assertRaisesRegex(ValueError, "has Partial status which is not allowed"): | 858 | with self.assertRaisesRegex(ValueError, "has Partial status which is not allowed"): |
| 777 | - self.op.infer_layout((x_layout,), {}) | 859 | + self.op.infer_layout([x_layout, None]) |
| 778 | 860 | ||
| 779 | 861 | ||
| 780 | if __name__ == "__main__": | 862 | if __name__ == "__main__": |
| @@ -18,7 +18,7 @@ import unittest | |||
| 18 | from unittest.mock import patch, MagicMock | 18 | from unittest.mock import patch, MagicMock |
| 19 | import numpy as np | 19 | import numpy as np |
| 20 | 20 | ||
| 21 | -from hyper_parallel.core.dtensor.dtensor import _build_layout | 21 | +from hyper_parallel.core.dtensor.dtensor import _build_layout, _LAYOUT_CACHE |
| 22 | from hyper_parallel.core.dtensor.placement_types import Shard, Replicate, Partial | 22 | from hyper_parallel.core.dtensor.placement_types import Shard, Replicate, Partial |
| 23 | from hyper_parallel.core.shard.ops.parallel_embedding import EmbeddingDistributedOp | 23 | from hyper_parallel.core.shard.ops.parallel_embedding import EmbeddingDistributedOp |
| 24 | from hyper_parallel.core.dtensor.device_mesh import ( | 24 | from hyper_parallel.core.dtensor.device_mesh import ( |
| @@ -41,11 +41,13 @@ class TestParallelEmbedding(unittest.TestCase): | |||
| 41 | """Set up test fixtures before each test method.""" | 41 | """Set up test fixtures before each test method.""" |
| 42 | EXISTING_COMM_GROUPS.clear() | 42 | EXISTING_COMM_GROUPS.clear() |
| 43 | _DEVICE_MESH_MAP.clear() | 43 | _DEVICE_MESH_MAP.clear() |
| 44 | + _LAYOUT_CACHE.clear() | ||
| 44 | 45 | ||
| 45 | def tearDown(self): | 46 | def tearDown(self): |
| 46 | """Clean up after each test method.""" | 47 | """Clean up after each test method.""" |
| 47 | EXISTING_COMM_GROUPS.clear() | 48 | EXISTING_COMM_GROUPS.clear() |
| 48 | _DEVICE_MESH_MAP.clear() | 49 | _DEVICE_MESH_MAP.clear() |
| 50 | + _LAYOUT_CACHE.clear() | ||
| 49 | 51 | ||
| 50 | def _setup_mock_platform(self, mock_platform, world_size=8): | 52 | def _setup_mock_platform(self, mock_platform, world_size=8): |
| 51 | """Configure common mock-platform attributes.""" | 53 | """Configure common mock-platform attributes.""" |
| @@ -79,13 +81,15 @@ class TestParallelEmbedding(unittest.TestCase): | |||
| 79 | weight_layout = _build_layout(mesh, (Replicate(), Replicate()), 2) | 81 | weight_layout = _build_layout(mesh, (Replicate(), Replicate()), 2) |
| 80 | 82 | ||
| 81 | for op in embedding_ops: | 83 | for op in embedding_ops: |
| 82 | - output_layout = op.infer_layout((input_layout, weight_layout)) | 84 | + cache_values = [input_layout, weight_layout] |
| 85 | + output_layouts, _ = op.infer_layout(cache_values) | ||
| 86 | + output_layout = output_layouts[0] | ||
| 83 | # Expected map: input(1, -1) + weight_embed(-1) = (1, -1, -1) | 87 | # Expected map: input(1, -1) + weight_embed(-1) = (1, -1, -1) |
| 84 | expected_map = (1, -1, -1) | 88 | expected_map = (1, -1, -1) |
| 85 | self.assertEqual(output_layout.tensor_map, expected_map, f"Op {op.op_name} failed") | 89 | self.assertEqual(output_layout.tensor_map, expected_map, f"Op {op.op_name} failed") |
| 86 | self.assertFalse(output_layout.is_partial()) | 90 | self.assertFalse(output_layout.is_partial()) |
| 87 | # Implementation should be native (None) for DP | 91 | # Implementation should be native (None) for DP |
| 88 | - self.assertIsNone(op.get_expand_impl(None, output_layout, (input_layout, weight_layout), None)) | 92 | + self.assertIsNone(op.get_expand_impl(None, (output_layouts, None), cache_values)) |
| 89 | 93 | ||
| 90 | @patch("hyper_parallel.core.dtensor.device_mesh.platform") | 94 | @patch("hyper_parallel.core.dtensor.device_mesh.platform") |
| 91 | def test_embedding_column_parallel(self, mock_platform): | 95 | def test_embedding_column_parallel(self, mock_platform): |
| @@ -101,13 +105,15 @@ class TestParallelEmbedding(unittest.TestCase): | |||
| 101 | weight_layout = _build_layout(mesh, (Replicate(), Shard(1)), 2) | 105 | weight_layout = _build_layout(mesh, (Replicate(), Shard(1)), 2) |
| 102 | 106 | ||
| 103 | for op in embedding_ops: | 107 | for op in embedding_ops: |
| 104 | - output_layout = op.infer_layout((input_layout, weight_layout)) | 108 | + cache_values = [input_layout, weight_layout] |
| 109 | + output_layouts, _ = op.infer_layout(cache_values) | ||
| 110 | + output_layout = output_layouts[0] | ||
| 105 | # Expected map: input(1, -1) + weight_embed(0) = (1, -1, 0) | 111 | # Expected map: input(1, -1) + weight_embed(0) = (1, -1, 0) |
| 106 | expected_map = (1, -1, 0) | 112 | expected_map = (1, -1, 0) |
| 107 | self.assertEqual(output_layout.tensor_map, expected_map, f"Op {op.op_name} failed") | 113 | self.assertEqual(output_layout.tensor_map, expected_map, f"Op {op.op_name} failed") |
| 108 | 114 | ||
| 109 | # CP requires wrapper to intercept max_norm | 115 | # CP requires wrapper to intercept max_norm |
| 110 | - impl = op.get_expand_impl(None, output_layout, (input_layout, weight_layout), None) | 116 | + impl = op.get_expand_impl(None, (output_layouts, None), cache_values) |
| 111 | self.assertTrue(callable(impl)) | 117 | self.assertTrue(callable(impl)) |
| 112 | with self.assertRaisesRegex(ValueError, "Column-Parallel.*does not support `max_norm`"): | 118 | with self.assertRaisesRegex(ValueError, "Column-Parallel.*does not support `max_norm`"): |
| 113 | impl(MagicMock(), MagicMock(), max_norm=1.0) | 119 | impl(MagicMock(), MagicMock(), max_norm=1.0) |
| @@ -125,7 +131,8 @@ class TestParallelEmbedding(unittest.TestCase): | |||
| 125 | weight_layout = _build_layout(mesh, (Replicate(), Shard(0)), 2) | 131 | weight_layout = _build_layout(mesh, (Replicate(), Shard(0)), 2) |
| 126 | 132 | ||
| 127 | for op in embedding_ops: | 133 | for op in embedding_ops: |
| 128 | - output_layout = op.infer_layout((input_layout, weight_layout)) | 134 | + output_layouts, _ = op.infer_layout([input_layout, weight_layout]) |
| 135 | + output_layout = output_layouts[0] | ||
| 129 | # Expected map: input(1, -1) + weight_embed(-1) = (1, -1, -1) | 136 | # Expected map: input(1, -1) + weight_embed(-1) = (1, -1, -1) |
| 130 | self.assertEqual(output_layout.tensor_map, (1, -1, -1)) | 137 | self.assertEqual(output_layout.tensor_map, (1, -1, -1)) |
| 131 | self.assertTrue(output_layout.is_partial()) | 138 | self.assertTrue(output_layout.is_partial()) |
| @@ -148,7 +155,8 @@ class TestParallelEmbedding(unittest.TestCase): | |||
| 148 | weight_layout = _build_layout(mesh, (Replicate(), Shard(0), Shard(1)), 2) | 155 | weight_layout = _build_layout(mesh, (Replicate(), Shard(0), Shard(1)), 2) |
| 149 | 156 | ||
| 150 | for op in embedding_ops: | 157 | for op in embedding_ops: |
| 151 | - output_layout = op.infer_layout((input_layout, weight_layout)) | 158 | + output_layouts, _ = op.infer_layout([input_layout, weight_layout]) |
| 159 | + output_layout = output_layouts[0] | ||
| 152 | # Expected map: input(2, -1) + weight_embed(0) = (2, -1, 0) | 160 | # Expected map: input(2, -1) + weight_embed(0) = (2, -1, 0) |
| 153 | expected_map = (2, -1, 0) | 161 | expected_map = (2, -1, 0) |
| 154 | self.assertEqual(output_layout.tensor_map, expected_map) | 162 | self.assertEqual(output_layout.tensor_map, expected_map) |
| @@ -169,7 +177,8 @@ class TestParallelEmbedding(unittest.TestCase): | |||
| 169 | weight_layout = _build_layout(mesh, (Replicate(), Replicate()), 2) | 177 | weight_layout = _build_layout(mesh, (Replicate(), Replicate()), 2) |
| 170 | 178 | ||
| 171 | for op in embedding_ops: | 179 | for op in embedding_ops: |
| 172 | - output_layout = op.infer_layout((input_layout, weight_layout)) | 180 | + output_layouts, _ = op.infer_layout([input_layout, weight_layout]) |
| 181 | + output_layout = output_layouts[0] | ||
| 173 | self.assertEqual(output_layout.tensor_map, (-1, 0, -1)) | 182 | self.assertEqual(output_layout.tensor_map, (-1, 0, -1)) |
| 174 | 183 | ||
| 175 | @patch("hyper_parallel.core.dtensor.device_mesh.platform") | 184 | @patch("hyper_parallel.core.dtensor.device_mesh.platform") |
| @@ -183,8 +192,8 @@ class TestParallelEmbedding(unittest.TestCase): | |||
| 183 | input_layout = _build_layout(mesh, (Replicate(), Replicate()), 2) | 192 | input_layout = _build_layout(mesh, (Replicate(), Replicate()), 2) |
| 184 | 193 | ||
| 185 | for op in embedding_ops: | 194 | for op in embedding_ops: |
| 186 | - with self.assertRaisesRegex(ValueError, "requires both input and weight layouts"): | 195 | + with self.assertRaisesRegex(ValueError, "cache_values length should be 2"): |
| 187 | - op.infer_layout((input_layout,)) | 196 | + op.infer_layout([input_layout]) |
| 188 | 197 | ||
| 189 | @patch("hyper_parallel.core.dtensor.device_mesh.platform") | 198 | @patch("hyper_parallel.core.dtensor.device_mesh.platform") |
| 190 | def test_embedding_partial_conflict_error(self, mock_platform): | 199 | def test_embedding_partial_conflict_error(self, mock_platform): |
| @@ -200,7 +209,7 @@ class TestParallelEmbedding(unittest.TestCase): | |||
| 200 | 209 | ||
| 201 | for op in embedding_ops: | 210 | for op in embedding_ops: |
| 202 | with self.assertRaisesRegex(ValueError, "Partial dim must be replicate"): | 211 | with self.assertRaisesRegex(ValueError, "Partial dim must be replicate"): |
| 203 | - op.infer_layout((input_layout, weight_layout)) | 212 | + op.infer_layout([input_layout, weight_layout]) |
| 204 | 213 | ||
| 205 | @patch("hyper_parallel.core.dtensor.device_mesh.platform") | 214 | @patch("hyper_parallel.core.dtensor.device_mesh.platform") |
| 206 | def test_embedding_row_parallel_params_validation(self, mock_platform): | 215 | def test_embedding_row_parallel_params_validation(self, mock_platform): |
| @@ -214,8 +223,9 @@ class TestParallelEmbedding(unittest.TestCase): | |||
| 214 | weight_layout = _build_layout(mesh, (Shard(0), Replicate()), 2) | 223 | weight_layout = _build_layout(mesh, (Shard(0), Replicate()), 2) |
| 215 | 224 | ||
| 216 | for op in embedding_ops: | 225 | for op in embedding_ops: |
| 217 | - output_layout = op.infer_layout((input_layout, weight_layout)) | 226 | + cache_values = [input_layout, weight_layout] |
| 218 | - impl = op.get_expand_impl(None, output_layout, (input_layout, weight_layout), None) | 227 | + output_layouts, _ = op.infer_layout(cache_values) |
| 228 | + impl = op.get_expand_impl(None, (output_layouts, None), cache_values) | ||
| 219 | 229 | ||
| 220 | with self.assertRaisesRegex(ValueError, "Row-Parallel.*does not support `scale_grad_by_freq=True`"): | 230 | with self.assertRaisesRegex(ValueError, "Row-Parallel.*does not support `scale_grad_by_freq=True`"): |
| 221 | impl(MagicMock(), MagicMock(), scale_grad_by_freq=True) | 231 | impl(MagicMock(), MagicMock(), scale_grad_by_freq=True) |
| @@ -234,7 +244,8 @@ class TestParallelEmbedding(unittest.TestCase): | |||
| 234 | weight_layout = _build_layout(mesh, (Shard(0), Shard(1)), 2) | 244 | weight_layout = _build_layout(mesh, (Shard(0), Shard(1)), 2) |
| 235 | 245 | ||
| 236 | for op in embedding_ops: | 246 | for op in embedding_ops: |
| 237 | - output_layout = op.infer_layout((input_layout, weight_layout)) | 247 | + output_layouts, _ = op.infer_layout([input_layout, weight_layout]) |
| 248 | + output_layout = output_layouts[0] | ||
| 238 | # Input map (-1, -1) + weight_embed(0) = (-1, -1, 0) | 249 | # Input map (-1, -1) + weight_embed(0) = (-1, -1, 0) |
| 239 | self.assertEqual(output_layout.tensor_map, (-1, -1, 0)) | 250 | self.assertEqual(output_layout.tensor_map, (-1, -1, 0)) |
| 240 | self.assertTrue(output_layout.is_partial()) | 251 | self.assertTrue(output_layout.is_partial()) |
| @@ -255,7 +266,8 @@ class TestParallelEmbedding(unittest.TestCase): | |||
| 255 | weight_layout = _build_layout(mesh, (Replicate(), Shard(1)), 2) | 266 | weight_layout = _build_layout(mesh, (Replicate(), Shard(1)), 2) |
| 256 | 267 | ||
| 257 | for op in embedding_ops: | 268 | for op in embedding_ops: |
| 258 | - output_layout = op.infer_layout((input_layout, weight_layout)) | 269 | + output_layouts, _ = op.infer_layout([input_layout, weight_layout]) |
| 270 | + output_layout = output_layouts[0] | ||
| 259 | # Expected map: (1, -1, -1, 0) | 271 | # Expected map: (1, -1, -1, 0) |
| 260 | self.assertEqual(output_layout.tensor_map, (1, -1, -1, 0)) | 272 | self.assertEqual(output_layout.tensor_map, (1, -1, -1, 0)) |
| 261 | 273 | ||
| @@ -272,10 +284,40 @@ class TestParallelEmbedding(unittest.TestCase): | |||
| 272 | weight_layout = _build_layout(mesh, (Replicate(), Replicate()), 2) | 284 | weight_layout = _build_layout(mesh, (Replicate(), Replicate()), 2) |
| 273 | 285 | ||
| 274 | for op in embedding_ops: | 286 | for op in embedding_ops: |
| 275 | - output_layout = op.infer_layout((input_layout, weight_layout)) | 287 | + output_layouts, _ = op.infer_layout([input_layout, weight_layout]) |
| 288 | + output_layout = output_layouts[0] | ||
| 276 | # Expected map: (1, 0, -1) | 289 | # Expected map: (1, 0, -1) |
| 277 | self.assertEqual(output_layout.tensor_map, (1, 0, -1)) | 290 | self.assertEqual(output_layout.tensor_map, (1, 0, -1)) |
| 278 | 291 | ||
| 292 | + def test_embedding_preprocess_routes_platform_args(self): | ||
| 293 | + """ | ||
| 294 | + Feature: New dispatch preprocessing. | ||
| 295 | + Description: Convert DTensor inputs to local tensors and cache layouts. | ||
| 296 | + Expectation: PyTorch embedding keeps sparse while MindSpore Embedding drops it. | ||
| 297 | + """ | ||
| 298 | + input_layout = MagicMock() | ||
| 299 | + weight_layout = MagicMock() | ||
| 300 | + input_tensor = MagicMock() | ||
| 301 | + input_tensor.layout = input_layout | ||
| 302 | + input_tensor.to_local.return_value = "local_input" | ||
| 303 | + weight_tensor = MagicMock() | ||
| 304 | + weight_tensor.layout = weight_layout | ||
| 305 | + weight_tensor.to_local.return_value = "local_weight" | ||
| 306 | + | ||
| 307 | + local_args, local_kwargs, cache_values = EmbeddingDistributedOp("embedding").preprocess( | ||
| 308 | + (input_tensor, weight_tensor), {} | ||
| 309 | + ) | ||
| 310 | + self.assertEqual(local_args, ("local_input", "local_weight", None, None, 2.0, False, False)) | ||
| 311 | + self.assertEqual(local_kwargs, {}) | ||
| 312 | + self.assertEqual(cache_values, [input_layout, weight_layout]) | ||
| 313 | + | ||
| 314 | + local_args, local_kwargs, cache_values = EmbeddingDistributedOp("Embedding").preprocess( | ||
| 315 | + (input_tensor, weight_tensor), {} | ||
| 316 | + ) | ||
| 317 | + self.assertEqual(local_args, ("local_input", "local_weight", None, None, 2.0, False)) | ||
| 318 | + self.assertEqual(local_kwargs, {}) | ||
| 319 | + self.assertEqual(cache_values, [input_layout, weight_layout]) | ||
| 320 | + | ||
| 279 | 321 | ||
| 280 | 322 | ||
| 281 | @patch("hyper_parallel.core.dtensor.device_mesh.platform") | 323 | @patch("hyper_parallel.core.dtensor.device_mesh.platform") |
| @@ -290,8 +332,9 @@ class TestParallelEmbedding(unittest.TestCase): | |||
| 290 | weight_layout = _build_layout(mesh, (Replicate(), Shard(1)), 2) # Column Parallel | 332 | weight_layout = _build_layout(mesh, (Replicate(), Shard(1)), 2) # Column Parallel |
| 291 | 333 | ||
| 292 | for op in embedding_ops: | 334 | for op in embedding_ops: |
| 293 | - output_layout = op.infer_layout((input_layout, weight_layout)) | 335 | + cache_values = [input_layout, weight_layout] |
| 294 | - impl = op.get_expand_impl(MagicMock(), output_layout, (input_layout, weight_layout), None) | 336 | + output_layouts, _ = op.infer_layout(cache_values) |
| 337 | + impl = op.get_expand_impl(MagicMock(), (output_layouts, None), cache_values) | ||
| 295 | 338 | ||
| 296 | # Should not raise ValueError for scale_grad_by_freq | 339 | # Should not raise ValueError for scale_grad_by_freq |
| 297 | try: | 340 | try: |
| @@ -310,10 +353,12 @@ class TestEmbeddingRowParallelImpl(unittest.TestCase): | |||
| 310 | def setUp(self): | 353 | def setUp(self): |
| 311 | EXISTING_COMM_GROUPS.clear() | 354 | EXISTING_COMM_GROUPS.clear() |
| 312 | _DEVICE_MESH_MAP.clear() | 355 | _DEVICE_MESH_MAP.clear() |
| 356 | + _LAYOUT_CACHE.clear() | ||
| 313 | 357 | ||
| 314 | def tearDown(self): | 358 | def tearDown(self): |
| 315 | EXISTING_COMM_GROUPS.clear() | 359 | EXISTING_COMM_GROUPS.clear() |
| 316 | _DEVICE_MESH_MAP.clear() | 360 | _DEVICE_MESH_MAP.clear() |
| 361 | + _LAYOUT_CACHE.clear() | ||
| 317 | 362 | ||
| 318 | @patch("hyper_parallel.core.dtensor.device_mesh.platform") | 363 | @patch("hyper_parallel.core.dtensor.device_mesh.platform") |
| 319 | def test_row_parallel_max_norm_raises(self, mock_platform): | 364 | def test_row_parallel_max_norm_raises(self, mock_platform): |
| @@ -330,9 +375,10 @@ class TestEmbeddingRowParallelImpl(unittest.TestCase): | |||
| 330 | weight_layout = _build_layout(mesh, (Shard(0),), 2) | 375 | weight_layout = _build_layout(mesh, (Shard(0),), 2) |
| 331 | 376 | ||
| 332 | for embedding_op in embedding_ops: | 377 | for embedding_op in embedding_ops: |
| 333 | - output_layout = embedding_op.infer_layout((input_layout, weight_layout)) | 378 | + cache_values = [input_layout, weight_layout] |
| 379 | + output_layouts, _ = embedding_op.infer_layout(cache_values) | ||
| 334 | impl = embedding_op.get_expand_impl( | 380 | impl = embedding_op.get_expand_impl( |
| 335 | - MagicMock(), output_layout, (input_layout, weight_layout), None | 381 | + MagicMock(), (output_layouts, None), cache_values |
| 336 | ) | 382 | ) |
| 337 | with self.assertRaisesRegex(ValueError, "Row-Parallel.*does not support.*max_norm"): | 383 | with self.assertRaisesRegex(ValueError, "Row-Parallel.*does not support.*max_norm"): |
| 338 | impl(MagicMock(), MagicMock(), max_norm=1.0) | 384 | impl(MagicMock(), MagicMock(), max_norm=1.0) |
| @@ -15,10 +15,10 @@ | |||
| 15 | """parallel_expand_dims test""" | 15 | """parallel_expand_dims test""" |
| 16 | import os | 16 | import os |
| 17 | import unittest | 17 | import unittest |
| 18 | -from unittest.mock import patch | 18 | +from unittest.mock import MagicMock, patch |
| 19 | import numpy as np | 19 | import numpy as np |
| 20 | 20 | ||
| 21 | -from hyper_parallel.core.dtensor.dtensor import _build_layout | 21 | +from hyper_parallel.core.dtensor.dtensor import _build_layout, _LAYOUT_CACHE |
| 22 | from hyper_parallel.core.dtensor.placement_types import Shard, Replicate | 22 | from hyper_parallel.core.dtensor.placement_types import Shard, Replicate |
| 23 | from hyper_parallel.core.shard.ops.parallel_expand_dims import ExpandDimsDistributedOp | 23 | from hyper_parallel.core.shard.ops.parallel_expand_dims import ExpandDimsDistributedOp |
| 24 | from hyper_parallel.core.dtensor.device_mesh import ( | 24 | from hyper_parallel.core.dtensor.device_mesh import ( |
| @@ -38,12 +38,14 @@ class TestParallelExpandDims(unittest.TestCase): | |||
| 38 | """ | 38 | """ |
| 39 | EXISTING_COMM_GROUPS.clear() | 39 | EXISTING_COMM_GROUPS.clear() |
| 40 | _DEVICE_MESH_MAP.clear() | 40 | _DEVICE_MESH_MAP.clear() |
| 41 | + _LAYOUT_CACHE.clear() | ||
| 41 | self.op = ExpandDimsDistributedOp("ExpandDims") | 42 | self.op = ExpandDimsDistributedOp("ExpandDims") |
| 42 | 43 | ||
| 43 | def tearDown(self): | 44 | def tearDown(self): |
| 44 | """Clean up after each test method.""" | 45 | """Clean up after each test method.""" |
| 45 | EXISTING_COMM_GROUPS.clear() | 46 | EXISTING_COMM_GROUPS.clear() |
| 46 | _DEVICE_MESH_MAP.clear() | 47 | _DEVICE_MESH_MAP.clear() |
| 48 | + _LAYOUT_CACHE.clear() | ||
| 47 | 49 | ||
| 48 | def _setup_mock_platform(self, mock_platform, platform_type=None, world_size=8): | 50 | def _setup_mock_platform(self, mock_platform, platform_type=None, world_size=8): |
| 49 | """Configure common mock-platform attributes used across tests. | 51 | """Configure common mock-platform attributes used across tests. |
| @@ -67,14 +69,18 @@ class TestParallelExpandDims(unittest.TestCase): | |||
| 67 | return init_device_mesh(device_type="npu", mesh_shape=(2, 2, 2), mesh_dim_names=mesh_dim_names) | 69 | return init_device_mesh(device_type="npu", mesh_shape=(2, 2, 2), mesh_dim_names=mesh_dim_names) |
| 68 | 70 | ||
| 69 | def _run_scenario(self, x_layout, expected_map, extra_args): | 71 | def _run_scenario(self, x_layout, expected_map, extra_args): |
| 70 | - """Infer layout of ExpandDims operator""" | 72 | + """Infer layout of ExpandDims operator.""" |
| 71 | - output_layout = self.op.infer_layout((x_layout,), extra_args) | 73 | + axis = extra_args[0] if isinstance(extra_args, (tuple, list)) else extra_args |
| 74 | + cache_values = [x_layout, axis] | ||
| 75 | + output_layouts, extra_info = self.op.infer_layout(cache_values) | ||
| 76 | + output_layout = output_layouts[0] | ||
| 72 | assert output_layout.tensor_map == expected_map, ( | 77 | assert output_layout.tensor_map == expected_map, ( |
| 73 | f"ExpandDims failed. Expected {expected_map}, got {output_layout.tensor_map}" | 78 | f"ExpandDims failed. Expected {expected_map}, got {output_layout.tensor_map}" |
| 74 | ) | 79 | ) |
| 75 | - assert self.op.get_expand_impl(None, output_layout, (x_layout,), extra_args) is None, ( | 80 | + assert extra_info is None, f"ExpandDims extra_info should be None, got {extra_info}" |
| 81 | + assert self.op.get_expand_impl(None, (output_layouts, None), cache_values) is None, ( | ||
| 76 | f"ExpandDims get_expand_impl failed. Expected None, " | 82 | f"ExpandDims get_expand_impl failed. Expected None, " |
| 77 | - f"got {self.op.get_expand_impl(None, output_layout, (x_layout,), extra_args)}" | 83 | + f"got {self.op.get_expand_impl(None, (output_layouts, None), cache_values)}" |
| 78 | ) | 84 | ) |
| 79 | 85 | ||
| 80 | @patch("hyper_parallel.core.dtensor.device_mesh.platform") | 86 | @patch("hyper_parallel.core.dtensor.device_mesh.platform") |
| @@ -206,7 +212,34 @@ class TestParallelExpandDims(unittest.TestCase): | |||
| 206 | x_layout = _build_layout(mesh, x_placements, 3) | 212 | x_layout = _build_layout(mesh, x_placements, 3) |
| 207 | 213 | ||
| 208 | with self.assertRaisesRegex(ValueError, "out of range for input rank"): | 214 | with self.assertRaisesRegex(ValueError, "out of range for input rank"): |
| 209 | - self.op.infer_layout((x_layout,), [5]) | 215 | + self.op.infer_layout([x_layout, 5]) |
| 216 | + | ||
| 217 | + @patch("hyper_parallel.core.dtensor.device_mesh.platform") | ||
| 218 | + def test_expanddims_partial_input_raises_error(self, mock_platform): | ||
| 219 | + """Verify ExpandDims rejects inputs with Partial status.""" | ||
| 220 | + mesh = self._make_2x2x2_mesh(mock_platform) | ||
| 221 | + x_layout = _build_layout(mesh, (Replicate(), Replicate(), Replicate()), 3) | ||
| 222 | + x_layout.set_partial_by_dev_axis("dp", "sum") | ||
| 223 | + | ||
| 224 | + with self.assertRaisesRegex(ValueError, "Partial status"): | ||
| 225 | + self.op.infer_layout([x_layout, 1]) | ||
| 226 | + | ||
| 227 | + @patch("hyper_parallel.core.dtensor.device_mesh.platform") | ||
| 228 | + def test_expanddims_preprocess(self, mock_platform): | ||
| 229 | + """Verify preprocess converts DTensor input to local args and cache values.""" | ||
| 230 | + mesh = self._make_2x2x2_mesh(mock_platform) | ||
| 231 | + x_layout = _build_layout(mesh, (Shard(0), Replicate(), Replicate()), 3) | ||
| 232 | + | ||
| 233 | + mock_tensor = MagicMock() | ||
| 234 | + mock_tensor.layout = x_layout | ||
| 235 | + local_tensor = MagicMock() | ||
| 236 | + mock_tensor.to_local.return_value = local_tensor | ||
| 237 | + | ||
| 238 | + local_args, local_kwargs, cache_values = self.op.preprocess((mock_tensor, 1), {}) | ||
| 239 | + | ||
| 240 | + assert local_args == (local_tensor, 1) | ||
| 241 | + assert not local_kwargs | ||
| 242 | + assert cache_values == [x_layout, 1] | ||
| 210 | 243 | ||
| 211 | 244 | ||
| 212 | class TestParallelExpandDimsView(unittest.TestCase): | 245 | class TestParallelExpandDimsView(unittest.TestCase): |
| @@ -216,12 +249,14 @@ class TestParallelExpandDimsView(unittest.TestCase): | |||
| 216 | """Set up test fixtures before each test method.""" | 249 | """Set up test fixtures before each test method.""" |
| 217 | EXISTING_COMM_GROUPS.clear() | 250 | EXISTING_COMM_GROUPS.clear() |
| 218 | _DEVICE_MESH_MAP.clear() | 251 | _DEVICE_MESH_MAP.clear() |
| 252 | + _LAYOUT_CACHE.clear() | ||
| 219 | self.op = ExpandDimsDistributedOp("ExpandDimsView") | 253 | self.op = ExpandDimsDistributedOp("ExpandDimsView") |
| 220 | 254 | ||
| 221 | def tearDown(self): | 255 | def tearDown(self): |
| 222 | """Clean up after each test method.""" | 256 | """Clean up after each test method.""" |
| 223 | EXISTING_COMM_GROUPS.clear() | 257 | EXISTING_COMM_GROUPS.clear() |
| 224 | _DEVICE_MESH_MAP.clear() | 258 | _DEVICE_MESH_MAP.clear() |
| 259 | + _LAYOUT_CACHE.clear() | ||
| 225 | 260 | ||
| 226 | def _setup_mock_platform(self, mock_platform, platform_type=None, world_size=8): | 261 | def _setup_mock_platform(self, mock_platform, platform_type=None, world_size=8): |
| 227 | """Configure common mock-platform attributes used across tests.""" | 262 | """Configure common mock-platform attributes used across tests.""" |
| @@ -240,13 +275,17 @@ class TestParallelExpandDimsView(unittest.TestCase): | |||
| 240 | 275 | ||
| 241 | def _run_scenario(self, x_layout, expected_map, extra_args): | 276 | def _run_scenario(self, x_layout, expected_map, extra_args): |
| 242 | """Infer layout of ExpandDimsView operator.""" | 277 | """Infer layout of ExpandDimsView operator.""" |
| 243 | - output_layout = self.op.infer_layout((x_layout,), extra_args) | 278 | + axis = extra_args[0] if isinstance(extra_args, (tuple, list)) else extra_args |
| 279 | + cache_values = [x_layout, axis] | ||
| 280 | + output_layouts, extra_info = self.op.infer_layout(cache_values) | ||
| 281 | + output_layout = output_layouts[0] | ||
| 244 | assert output_layout.tensor_map == expected_map, ( | 282 | assert output_layout.tensor_map == expected_map, ( |
| 245 | f"ExpandDimsView failed. Expected {expected_map}, got {output_layout.tensor_map}" | 283 | f"ExpandDimsView failed. Expected {expected_map}, got {output_layout.tensor_map}" |
| 246 | ) | 284 | ) |
| 247 | - assert self.op.get_expand_impl(None, output_layout, (x_layout,), extra_args) is None, ( | 285 | + assert extra_info is None, f"ExpandDimsView extra_info should be None, got {extra_info}" |
| 286 | + assert self.op.get_expand_impl(None, (output_layouts, None), cache_values) is None, ( | ||
| 248 | f"ExpandDimsView get_expand_impl failed. Expected None, " | 287 | f"ExpandDimsView get_expand_impl failed. Expected None, " |
| 249 | - f"got {self.op.get_expand_impl(None, output_layout, (x_layout,), extra_args)}" | 288 | + f"got {self.op.get_expand_impl(None, (output_layouts, None), cache_values)}" |
| 250 | ) | 289 | ) |
| 251 | 290 | ||
| 252 | @patch("hyper_parallel.core.dtensor.device_mesh.platform") | 291 | @patch("hyper_parallel.core.dtensor.device_mesh.platform") |
| @@ -378,7 +417,34 @@ class TestParallelExpandDimsView(unittest.TestCase): | |||
| 378 | x_layout = _build_layout(mesh, x_placements, 3) | 417 | x_layout = _build_layout(mesh, x_placements, 3) |
| 379 | 418 | ||
| 380 | with self.assertRaisesRegex(ValueError, "out of range for input rank"): | 419 | with self.assertRaisesRegex(ValueError, "out of range for input rank"): |
| 381 | - self.op.infer_layout((x_layout,), [5]) | 420 | + self.op.infer_layout([x_layout, 5]) |
| 421 | + | ||
| 422 | + @patch("hyper_parallel.core.dtensor.device_mesh.platform") | ||
| 423 | + def test_expanddims_partial_input_raises_error(self, mock_platform): | ||
| 424 | + """Verify ExpandDims rejects inputs with Partial status.""" | ||
| 425 | + mesh = self._make_2x2x2_mesh(mock_platform) | ||
| 426 | + x_layout = _build_layout(mesh, (Replicate(), Replicate(), Replicate()), 3) | ||
| 427 | + x_layout.set_partial_by_dev_axis("dp", "sum") | ||
| 428 | + | ||
| 429 | + with self.assertRaisesRegex(ValueError, "Partial status"): | ||
| 430 | + self.op.infer_layout([x_layout, 1]) | ||
| 431 | + | ||
| 432 | + @patch("hyper_parallel.core.dtensor.device_mesh.platform") | ||
| 433 | + def test_expanddims_preprocess(self, mock_platform): | ||
| 434 | + """Verify preprocess converts DTensor input to local args and cache values.""" | ||
| 435 | + mesh = self._make_2x2x2_mesh(mock_platform) | ||
| 436 | + x_layout = _build_layout(mesh, (Shard(0), Replicate(), Replicate()), 3) | ||
| 437 | + | ||
| 438 | + mock_tensor = MagicMock() | ||
| 439 | + mock_tensor.layout = x_layout | ||
| 440 | + local_tensor = MagicMock() | ||
| 441 | + mock_tensor.to_local.return_value = local_tensor | ||
| 442 | + | ||
| 443 | + local_args, local_kwargs, cache_values = self.op.preprocess((mock_tensor, 1), {}) | ||
| 444 | + | ||
| 445 | + assert local_args == (local_tensor, 1) | ||
| 446 | + assert not local_kwargs | ||
| 447 | + assert cache_values == [x_layout, 1] | ||
| 382 | 448 | ||
| 383 | 449 | ||
| 384 | if __name__ == "__main__": | 450 | if __name__ == "__main__": |
| @@ -15,10 +15,10 @@ | |||
| 15 | """parallel_flatten test""" | 15 | """parallel_flatten test""" |
| 16 | import os | 16 | import os |
| 17 | import unittest | 17 | import unittest |
| 18 | -from unittest.mock import patch | 18 | +from unittest.mock import MagicMock, patch |
| 19 | import numpy as np | 19 | import numpy as np |
| 20 | 20 | ||
| 21 | -from hyper_parallel.core.dtensor.dtensor import _build_layout | 21 | +from hyper_parallel.core.dtensor.dtensor import _build_layout, _LAYOUT_CACHE |
| 22 | from hyper_parallel.core.dtensor.placement_types import Shard, Replicate | 22 | from hyper_parallel.core.dtensor.placement_types import Shard, Replicate |
| 23 | from hyper_parallel.core.shard.ops.parallel_flatten import FlattenDistributedOp | 23 | from hyper_parallel.core.shard.ops.parallel_flatten import FlattenDistributedOp |
| 24 | from hyper_parallel.core.dtensor.device_mesh import ( | 24 | from hyper_parallel.core.dtensor.device_mesh import ( |
| @@ -40,11 +40,13 @@ class TestParallelFlatten(unittest.TestCase): | |||
| 40 | """ | 40 | """ |
| 41 | EXISTING_COMM_GROUPS.clear() | 41 | EXISTING_COMM_GROUPS.clear() |
| 42 | _DEVICE_MESH_MAP.clear() | 42 | _DEVICE_MESH_MAP.clear() |
| 43 | + _LAYOUT_CACHE.clear() | ||
| 43 | 44 | ||
| 44 | def tearDown(self): | 45 | def tearDown(self): |
| 45 | """Clean up after each test method.""" | 46 | """Clean up after each test method.""" |
| 46 | EXISTING_COMM_GROUPS.clear() | 47 | EXISTING_COMM_GROUPS.clear() |
| 47 | _DEVICE_MESH_MAP.clear() | 48 | _DEVICE_MESH_MAP.clear() |
| 49 | + _LAYOUT_CACHE.clear() | ||
| 48 | 50 | ||
| 49 | def _setup_mock_platform(self, mock_platform, platform_type=None, world_size=8): | 51 | def _setup_mock_platform(self, mock_platform, platform_type=None, world_size=8): |
| 50 | """Configure common mock-platform attributes used across tests. | 52 | """Configure common mock-platform attributes used across tests. |
| @@ -67,6 +69,33 @@ class TestParallelFlatten(unittest.TestCase): | |||
| 67 | self._setup_mock_platform(mock_platform, world_size=8) | 69 | self._setup_mock_platform(mock_platform, world_size=8) |
| 68 | return init_device_mesh(device_type="npu", mesh_shape=(2, 4), mesh_dim_names=("dp", "mp")) | 70 | return init_device_mesh(device_type="npu", mesh_shape=(2, 4), mesh_dim_names=("dp", "mp")) |
| 69 | 71 | ||
| 72 | + def _infer_flatten(self, x_layout, input_shape, start_dim=0, end_dim=-1): | ||
| 73 | + """Infer flatten layout using the new cache_values calling convention.""" | ||
| 74 | + output_layouts, extra_info = op.infer_layout([x_layout, start_dim, end_dim, input_shape]) | ||
| 75 | + assert extra_info is None, f"Flatten extra_info should be None, got {extra_info}" | ||
| 76 | + return output_layouts[0] | ||
| 77 | + | ||
| 78 | + @patch("hyper_parallel.core.dtensor.device_mesh.platform") | ||
| 79 | + def test_flatten_preprocess_builds_cache_values(self, mock_platform): | ||
| 80 | + """ | ||
| 81 | + Feature: FlattenDistributedOp preprocess. | ||
| 82 | + Description: Verify that preprocess converts DTensor input to local tensor and caches layout/shape metadata. | ||
| 83 | + Expectation: local args, kwargs, and cache values follow the new dispatch convention. | ||
| 84 | + """ | ||
| 85 | + mesh = self._make_2x4_mesh(mock_platform) | ||
| 86 | + x_layout = _build_layout(mesh, (Shard(0), Replicate()), 2) | ||
| 87 | + local_tensor = MagicMock() | ||
| 88 | + mock_tensor = MagicMock() | ||
| 89 | + mock_tensor.layout = x_layout | ||
| 90 | + mock_tensor.shape = (8, 4) | ||
| 91 | + mock_tensor.to_local.return_value = local_tensor | ||
| 92 | + | ||
| 93 | + local_args, local_kwargs, cache_values = op.preprocess((mock_tensor, 1), {}) | ||
| 94 | + | ||
| 95 | + assert local_args == (local_tensor, 1, -1) | ||
| 96 | + assert not local_kwargs | ||
| 97 | + assert cache_values == [x_layout, 1, -1, (8, 4)] | ||
| 98 | + | ||
| 70 | @patch("hyper_parallel.core.dtensor.device_mesh.platform") | 99 | @patch("hyper_parallel.core.dtensor.device_mesh.platform") |
| 71 | def test_flatten_layout_inference_all_dims(self, mock_platform): | 100 | def test_flatten_layout_inference_all_dims(self, mock_platform): |
| 72 | """ | 101 | """ |
| @@ -79,18 +108,19 @@ class TestParallelFlatten(unittest.TestCase): | |||
| 79 | x_placements = (Shard(0), Replicate(), Replicate()) | 108 | x_placements = (Shard(0), Replicate(), Replicate()) |
| 80 | x_layout = _build_layout(mesh, x_placements, 3) | 109 | x_layout = _build_layout(mesh, x_placements, 3) |
| 81 | 110 | ||
| 82 | - output_layout = op.infer_layout((x_layout,), extra_args=([(2, 4, 8)],)) | 111 | + output_layout = self._infer_flatten(x_layout, (2, 4, 8)) |
| 83 | 112 | ||
| 84 | expected_map = (1,) | 113 | expected_map = (1,) |
| 85 | assert output_layout.to_dict()["tensor_map"] == expected_map, ( | 114 | assert output_layout.to_dict()["tensor_map"] == expected_map, ( |
| 86 | f"Flatten all dims failed. Expected {expected_map}, got {output_layout.to_dict()['tensor_map']}" | 115 | f"Flatten all dims failed. Expected {expected_map}, got {output_layout.to_dict()['tensor_map']}" |
| 87 | ) | 116 | ) |
| 88 | 117 | ||
| 89 | - # Since `get_expand_impl` is not overridden, it returns None by default. | 118 | + def fake_flatten(x, start_dim=0, end_dim=-1): |
| 90 | - # The same applies to other test classes, so it is unnecessary to test its return value. | 119 | + return x, start_dim, end_dim |
| 91 | - assert op.get_expand_impl(None, output_layout, (x_layout,), ([(2, 4, 8)],)) is None, ( | 120 | + |
| 121 | + assert op.get_expand_impl(fake_flatten, ((output_layout,), None), [x_layout, 0, -1, (2, 4, 8)]) is None, ( | ||
| 92 | f"get_expand_impl test failed. Expected None, " | 122 | f"get_expand_impl test failed. Expected None, " |
| 93 | - f"got {op.get_expand_impl(None, output_layout, (x_layout,), ([(2, 4, 8)],))}" | 123 | + f"got {op.get_expand_impl(fake_flatten, ((output_layout,), None), [x_layout, 0, -1, (2, 4, 8)])}" |
| 94 | ) | 124 | ) |
| 95 | 125 | ||
| 96 | @patch("hyper_parallel.core.dtensor.device_mesh.platform") | 126 | @patch("hyper_parallel.core.dtensor.device_mesh.platform") |
| @@ -105,7 +135,7 @@ class TestParallelFlatten(unittest.TestCase): | |||
| 105 | x_placements = (Shard(0), Shard(1), Replicate()) | 135 | x_placements = (Shard(0), Shard(1), Replicate()) |
| 106 | x_layout = _build_layout(mesh, x_placements, 3) | 136 | x_layout = _build_layout(mesh, x_placements, 3) |
| 107 | 137 | ||
| 108 | - output_layout = op.infer_layout((x_layout,), extra_args=(1, [(2, 4, 8)])) | 138 | + output_layout = self._infer_flatten(x_layout, (2, 4, 8), start_dim=1) |
| 109 | 139 | ||
| 110 | expected_map = (1, 0) | 140 | expected_map = (1, 0) |
| 111 | assert output_layout.to_dict()["tensor_map"] == expected_map, ( | 141 | assert output_layout.to_dict()["tensor_map"] == expected_map, ( |
| @@ -123,7 +153,7 @@ class TestParallelFlatten(unittest.TestCase): | |||
| 123 | x_placements = () | 153 | x_placements = () |
| 124 | x_layout = _build_layout(mesh, x_placements, 0) | 154 | x_layout = _build_layout(mesh, x_placements, 0) |
| 125 | 155 | ||
| 126 | - output_layout = op.infer_layout((x_layout,), extra_args=([()],)) | 156 | + output_layout = self._infer_flatten(x_layout, ()) |
| 127 | 157 | ||
| 128 | expected_map = (-1,) | 158 | expected_map = (-1,) |
| 129 | assert output_layout.to_dict()["tensor_map"] == expected_map, ( | 159 | assert output_layout.to_dict()["tensor_map"] == expected_map, ( |
| @@ -141,7 +171,7 @@ class TestParallelFlatten(unittest.TestCase): | |||
| 141 | x_placements = (Shard(0), Replicate()) | 171 | x_placements = (Shard(0), Replicate()) |
| 142 | x_layout = _build_layout(mesh, x_placements, 2) | 172 | x_layout = _build_layout(mesh, x_placements, 2) |
| 143 | 173 | ||
| 144 | - output_layout = op.infer_layout((x_layout,), extra_args=(1, 0, [(2, 4)])) | 174 | + output_layout = self._infer_flatten(x_layout, (2, 4), start_dim=1, end_dim=0) |
| 145 | 175 | ||
| 146 | expected_map = (1, -1) | 176 | expected_map = (1, -1) |
| 147 | assert output_layout.to_dict()["tensor_map"] == expected_map, ( | 177 | assert output_layout.to_dict()["tensor_map"] == expected_map, ( |
| @@ -160,7 +190,7 @@ class TestParallelFlatten(unittest.TestCase): | |||
| 160 | x_layout = _build_layout(mesh, x_placements, 2) | 190 | x_layout = _build_layout(mesh, x_placements, 2) |
| 161 | 191 | ||
| 162 | with self.assertRaises(ValueError): | 192 | with self.assertRaises(ValueError): |
| 163 | - op.infer_layout((x_layout,), extra_args=(2, [(2, 4)])) | 193 | + op.infer_layout([x_layout, 2, -1, (2, 4)]) |
| 164 | 194 | ||
| 165 | 195 | ||
| 166 | if __name__ == "__main__": | 196 | if __name__ == "__main__": |
| @@ -439,19 +439,25 @@ class TestParallelGatherNd(unittest.TestCase): | |||
| 439 | 439 | ||
| 440 | @patch("hyper_parallel.core.dtensor.device_mesh.platform") | 440 | @patch("hyper_parallel.core.dtensor.device_mesh.platform") |
| 441 | def test_gathernd_preprocess_packed_call(self, mock_platform): | 441 | def test_gathernd_preprocess_packed_call(self, mock_platform): |
| 442 | - """Verify GatherNd preprocess preserves MindSpore packed fallback calls.""" | 442 | + """Verify GatherNd preprocess receives clean unpacked args. |
| 443 | + | ||
| 444 | + NOTE: aclop packed-args normalization now happens upstream in | ||
| 445 | + ``OpDispatcher._dispatch_layout_infer`` via ``_normalize_aclop_args``. | ||
| 446 | + The re-packing for the kernel call is handled by ``_dispatch_new``. | ||
| 447 | + This test validates that ``preprocess`` works correctly with | ||
| 448 | + the already-unpacked (input, indices) args. | ||
| 449 | + """ | ||
| 443 | mesh = self._make_1d_mesh(mock_platform, world_size=8) | 450 | mesh = self._make_1d_mesh(mock_platform, world_size=8) |
| 444 | input_layout = _build_layout(mesh, (Replicate(), Replicate()), 2) | 451 | input_layout = _build_layout(mesh, (Replicate(), Replicate()), 2) |
| 445 | indices_layout = _build_layout(mesh, (Shard(0), Replicate()), 2) | 452 | indices_layout = _build_layout(mesh, (Shard(0), Replicate()), 2) |
| 446 | input_tensor = _MockDTensor(input_layout, local_value="input", shape=(16, 64)) | 453 | input_tensor = _MockDTensor(input_layout, local_value="input", shape=(16, 64)) |
| 447 | indices = _MockDTensor(indices_layout, local_value="indices", shape=(16, 2)) | 454 | indices = _MockDTensor(indices_layout, local_value="indices", shape=(16, 2)) |
| 448 | - prim = object() | ||
| 449 | 455 | ||
| 450 | local_args, local_kwargs, cache_values = self.op.preprocess( | 456 | local_args, local_kwargs, cache_values = self.op.preprocess( |
| 451 | - (prim, "GatherNd", (input_tensor, indices)), {} | 457 | + (input_tensor, indices), {} |
| 452 | ) | 458 | ) |
| 453 | 459 | ||
| 454 | - assert local_args == (prim, "GatherNd", ("input", "indices")) | 460 | + assert local_args == ("input", "indices") |
| 455 | assert not local_kwargs | 461 | assert not local_kwargs |
| 456 | assert cache_values == [input_layout, indices_layout, (16, 64), (16, 2)] | 462 | assert cache_values == [input_layout, indices_layout, (16, 64), (16, 2)] |
| 457 | 463 | ||
| @@ -1,218 +0,0 @@ | |||
| 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 | -"""parallel_multinomial test""" | ||
| 16 | -import os | ||
| 17 | -import unittest | ||
| 18 | -from unittest.mock import patch | ||
| 19 | -import numpy as np | ||
| 20 | - | ||
| 21 | -from hyper_parallel.core.dtensor.dtensor import _build_layout | ||
| 22 | -from hyper_parallel.core.dtensor.placement_types import Shard, Replicate | ||
| 23 | -from hyper_parallel.core.shard.ops.parallel_multinomial import MultinomialDistributedOp | ||
| 24 | -from hyper_parallel.core.dtensor.device_mesh import ( | ||
| 25 | - init_device_mesh, | ||
| 26 | - _DEVICE_MESH_MAP | ||
| 27 | -) | ||
| 28 | -from hyper_parallel.platform.platform import EXISTING_COMM_GROUPS | ||
| 29 | - | ||
| 30 | -op = MultinomialDistributedOp("multinomial") | ||
| 31 | - | ||
| 32 | - | ||
| 33 | -class TestParallelMultinomial(unittest.TestCase): | ||
| 34 | - """Unit tests for MultinomialDistributedOp.""" | ||
| 35 | - def setUp(self): | ||
| 36 | - """Set up test fixtures before each test method. | ||
| 37 | - | ||
| 38 | - Clears global caches to ensure test isolation and initializes | ||
| 39 | - the platform for testing. | ||
| 40 | - """ | ||
| 41 | - EXISTING_COMM_GROUPS.clear() | ||
| 42 | - _DEVICE_MESH_MAP.clear() | ||
| 43 | - | ||
| 44 | - def tearDown(self): | ||
| 45 | - """Clean up after each test method.""" | ||
| 46 | - EXISTING_COMM_GROUPS.clear() | ||
| 47 | - _DEVICE_MESH_MAP.clear() | ||
| 48 | - | ||
| 49 | - def _setup_mock_platform(self, mock_platform, platform_type=None, world_size=8): | ||
| 50 | - """Configure common mock-platform attributes used across tests. | ||
| 51 | - | ||
| 52 | - Args: | ||
| 53 | - mock_platform: The MagicMock object injected by @patch. | ||
| 54 | - platform_type: Optional PlatformType to set on the mock. | ||
| 55 | - world_size: Value returned by mock_platform.get_world_size(). | ||
| 56 | - """ | ||
| 57 | - if platform_type is not None: | ||
| 58 | - mock_platform.platform_type = platform_type | ||
| 59 | - mock_platform.get_rank.return_value = 0 | ||
| 60 | - mock_platform.get_world_size.return_value = world_size | ||
| 61 | - mock_platform.tensor_to_numpy.side_effect = ( | ||
| 62 | - lambda t: t.numpy() if hasattr(t, "numpy") else np.array(t) | ||
| 63 | - ) | ||
| 64 | - | ||
| 65 | - def _make_2x4_mesh(self, mock_platform): | ||
| 66 | - """Set up mock and return a standard 2x4 (dp, mp) mesh via init_device_mesh.""" | ||
| 67 | - self._setup_mock_platform(mock_platform, world_size=8) | ||
| 68 | - return init_device_mesh(device_type="npu", mesh_shape=(2, 4), mesh_dim_names=("dp", "mp")) | ||
| 69 | - | ||
| 70 | - @patch("hyper_parallel.core.dtensor.device_mesh.platform") | ||
| 71 | - def test_multinomial_infer_layout_1d_replicated(self, mock_platform): | ||
| 72 | - """ | ||
| 73 | - Feature: Multinomial layout inference for 1D input | ||
| 74 | - Description: Input is a 1D probability vector (replicated). | ||
| 75 | - Expectation: Output layout should be 1D and Replicated ("None"). | ||
| 76 | - """ | ||
| 77 | - mesh = self._make_2x4_mesh(mock_platform) | ||
| 78 | - x_placements = (Replicate(), Replicate()) | ||
| 79 | - x_layout = _build_layout(mesh, x_placements, 1) | ||
| 80 | - extra_args = (10, True, None) | ||
| 81 | - | ||
| 82 | - output_layout = op.infer_layout((x_layout,), extra_args) | ||
| 83 | - | ||
| 84 | - expected_map = (-1,) | ||
| 85 | - assert output_layout.tensor_map == expected_map, ( | ||
| 86 | - f"1D inference failed. Expected {expected_map}, " | ||
| 87 | - f"got {output_layout.tensor_map}" | ||
| 88 | - ) | ||
| 89 | - | ||
| 90 | - # Since `get_expand_impl` is not overridden, it returns None by default. | ||
| 91 | - # The same applies to other test classes, so it is unnecessary to test its return value. | ||
| 92 | - assert op.get_expand_impl(None, output_layout, (x_layout,), extra_args) is None, ( | ||
| 93 | - f"get_expand_impl should return None, " | ||
| 94 | - f"but got {output_layout}" | ||
| 95 | - ) | ||
| 96 | - | ||
| 97 | - @patch("hyper_parallel.core.dtensor.device_mesh.platform") | ||
| 98 | - def test_multinomial_infer_layout_2d_data_parallel(self, mock_platform): | ||
| 99 | - """ | ||
| 100 | - Feature: Multinomial layout inference for 2D input (Data Parallel) | ||
| 101 | - Description: Input (N, C) is sharded on Batch dim (dim 0), C dim (dim 1) is replicated. | ||
| 102 | - Expectation: Output (N, num_samples) preserves sharding on dim 0, dim 1 is Replicated. | ||
| 103 | - """ | ||
| 104 | - mesh = self._make_2x4_mesh(mock_platform) | ||
| 105 | - x_placements = (Shard(0), Replicate()) | ||
| 106 | - x_layout = _build_layout(mesh, x_placements, 2) | ||
| 107 | - | ||
| 108 | - extra_args = (10, True, None) | ||
| 109 | - output_layout = op.infer_layout((x_layout,), extra_args=extra_args) | ||
| 110 | - | ||
| 111 | - expected_map = (1, -1) | ||
| 112 | - assert output_layout.tensor_map == expected_map, ( | ||
| 113 | - f"2D Data Parallel inference failed. Expected {expected_map}, " | ||
| 114 | - f"got {output_layout.tensor_map}" | ||
| 115 | - ) | ||
| 116 | - | ||
| 117 | - @patch("hyper_parallel.core.dtensor.device_mesh.platform") | ||
| 118 | - def test_multinomial_infer_layout_2d_fully_replicated(self, mock_platform): | ||
| 119 | - """ | ||
| 120 | - Feature: Multinomial layout inference for 2D input (Fully Replicated) | ||
| 121 | - Description: Input (N, C) is fully replicated. | ||
| 122 | - Expectation: Output (N, num_samples) is fully replicated. | ||
| 123 | - """ | ||
| 124 | - mesh = self._make_2x4_mesh(mock_platform) | ||
| 125 | - x_placements = (Replicate(), Replicate()) | ||
| 126 | - x_layout = _build_layout(mesh, x_placements, 2) | ||
| 127 | - | ||
| 128 | - extra_args = (10, True, None) | ||
| 129 | - output_layout = op.infer_layout((x_layout,), extra_args=extra_args) | ||
| 130 | - | ||
| 131 | - expected_map = (-1, -1) | ||
| 132 | - assert output_layout.tensor_map == expected_map, ( | ||
| 133 | - f"2D Replicated inference failed. Expected {expected_map}, " | ||
| 134 | - f"got {output_layout.tensor_map}" | ||
| 135 | - ) | ||
| 136 | - | ||
| 137 | - @patch("hyper_parallel.core.dtensor.device_mesh.platform") | ||
| 138 | - def test_multinomial_error_sharded_prob_dim(self, mock_platform): | ||
| 139 | - """ | ||
| 140 | - Feature: Error handling for sharded probability dimension | ||
| 141 | - Description: Attempt to run multinomial when the last dimension (probability classes) is sharded. | ||
| 142 | - Expectation: ValueError raised requiring redistribution. | ||
| 143 | - """ | ||
| 144 | - mesh = self._make_2x4_mesh(mock_platform) | ||
| 145 | - x_placements = (Replicate(), Shard(1)) | ||
| 146 | - x_layout = _build_layout(mesh, x_placements, 2) | ||
| 147 | - | ||
| 148 | - assert x_layout.alias_tensor_map == ("None", "mp") | ||
| 149 | - | ||
| 150 | - with self.assertRaisesRegex(ValueError, "must not be sharded"): | ||
| 151 | - op.infer_layout((x_layout,), extra_args=(10, True, None)) | ||
| 152 | - | ||
| 153 | - @patch("hyper_parallel.core.dtensor.device_mesh.platform") | ||
| 154 | - def test_multinomial_error_invalid_ndim_3d(self, mock_platform): | ||
| 155 | - """ | ||
| 156 | - Feature: Error handling for invalid input dimensions | ||
| 157 | - Description: Attempt to run multinomial on a 3D tensor. | ||
| 158 | - Expectation: ValueError raised. | ||
| 159 | - """ | ||
| 160 | - mesh = self._make_2x4_mesh(mock_platform) | ||
| 161 | - x_placements = (Replicate(), Replicate()) | ||
| 162 | - x_layout = _build_layout(mesh, x_placements, 3) | ||
| 163 | - | ||
| 164 | - with self.assertRaisesRegex(ValueError, "input dimension must be 1 or 2"): | ||
| 165 | - op.infer_layout((x_layout,), extra_args=(10, True, None)) | ||
| 166 | - | ||
| 167 | - @patch("hyper_parallel.core.dtensor.device_mesh.platform") | ||
| 168 | - def test_multinomial_error_invalid_ndim_0d(self, mock_platform): | ||
| 169 | - """ | ||
| 170 | - Feature: Error handling for 0D input | ||
| 171 | - Description: Attempt to run multinomial on a 0D tensor (scalar). | ||
| 172 | - Expectation: ValueError raised. | ||
| 173 | - """ | ||
| 174 | - mesh = self._make_2x4_mesh(mock_platform) | ||
| 175 | - x_placements = (Replicate(), Replicate()) | ||
| 176 | - x_layout = _build_layout(mesh, x_placements, 0) | ||
| 177 | - | ||
| 178 | - with self.assertRaisesRegex(ValueError, "input dimension must be 1 or 2"): | ||
| 179 | - op.infer_layout((x_layout,), extra_args=(10, True, None)) | ||
| 180 | - | ||
| 181 | - @patch("hyper_parallel.core.dtensor.device_mesh.platform") | ||
| 182 | - def test_multinomial_2d_with_replacement_false(self, mock_platform): | ||
| 183 | - """ | ||
| 184 | - Feature: Multinomial with replacement=False | ||
| 185 | - Description: Test 2D input with replacement=False. | ||
| 186 | - Expectation: Output layout follows same sharding rules. | ||
| 187 | - """ | ||
| 188 | - mesh = self._make_2x4_mesh(mock_platform) | ||
| 189 | - x_placements = (Shard(0), Replicate()) | ||
| 190 | - x_layout = _build_layout(mesh, x_placements, 2) | ||
| 191 | - | ||
| 192 | - extra_args = (5, False, None) | ||
| 193 | - output_layout = op.infer_layout((x_layout,), extra_args=extra_args) | ||
| 194 | - | ||
| 195 | - expected_map = (1, -1) | ||
| 196 | - assert output_layout.tensor_map == expected_map, ( | ||
| 197 | - f"2D with replacement=False inference failed. Expected {expected_map}, " | ||
| 198 | - f"got {output_layout.tensor_map}" | ||
| 199 | - ) | ||
| 200 | - | ||
| 201 | - @patch("hyper_parallel.core.dtensor.device_mesh.platform") | ||
| 202 | - def test_multinomial_partial_input(self, mock_platform): | ||
| 203 | - """ | ||
| 204 | - Feature: Multinomial with partial input | ||
| 205 | - Description: Input with partial state | ||
| 206 | - Expectation: Raise ValueError since _allow_partial_inputs is False | ||
| 207 | - """ | ||
| 208 | - mesh = self._make_2x4_mesh(mock_platform) | ||
| 209 | - x_placements = (Shard(0), Replicate()) | ||
| 210 | - x_layout = _build_layout(mesh, x_placements, 2) | ||
| 211 | - x_layout.set_partial_by_dev_axis("mp", "sum") | ||
| 212 | - | ||
| 213 | - with self.assertRaisesRegex(ValueError, "has Partial status which is not allowed"): | ||
| 214 | - op.infer_layout((x_layout,), extra_args=(10, True, None)) | ||
| 215 | - | ||
| 216 | - | ||
| 217 | -if __name__ == "__main__": | ||
| 218 | - unittest.main() | ||
| @@ -13,14 +13,17 @@ | |||
| 13 | # limitations under the License. | 13 | # limitations under the License. |
| 14 | # ============================================================================ | 14 | # ============================================================================ |
| 15 | """parallel_repeat_interleave test""" | 15 | """parallel_repeat_interleave test""" |
| 16 | -import os | ||
| 17 | import unittest | 16 | import unittest |
| 18 | -from unittest.mock import patch | 17 | +from unittest.mock import MagicMock, patch |
| 18 | + | ||
| 19 | import numpy as np | 19 | import numpy as np |
| 20 | 20 | ||
| 21 | -from hyper_parallel.core.dtensor.dtensor import _build_layout | 21 | +from hyper_parallel.core.dtensor.dtensor import _build_layout, _LAYOUT_CACHE |
| 22 | -from hyper_parallel.core.dtensor.placement_types import Shard, Replicate | 22 | +from hyper_parallel.core.dtensor.placement_types import Partial, Shard, Replicate |
| 23 | -from hyper_parallel.core.shard.ops.parallel_repeat_interleave import RepeatInterleaveDistributedOp | 23 | +from hyper_parallel.core.shard.ops.parallel_repeat_interleave import ( |
| 24 | + RepeatInterleaveDistributedOp, | ||
| 25 | + _normalize_repeat_interleave_args, | ||
| 26 | +) | ||
| 24 | from hyper_parallel.core.dtensor.device_mesh import ( | 27 | from hyper_parallel.core.dtensor.device_mesh import ( |
| 25 | init_device_mesh, | 28 | init_device_mesh, |
| 26 | _DEVICE_MESH_MAP | 29 | _DEVICE_MESH_MAP |
| @@ -32,15 +35,18 @@ op = RepeatInterleaveDistributedOp("repeat_interleave") | |||
| 32 | 35 | ||
| 33 | class TestParallelRepeatInterleave(unittest.TestCase): | 36 | class TestParallelRepeatInterleave(unittest.TestCase): |
| 34 | """Unit tests for RepeatInterleaveDistributedOp.""" | 37 | """Unit tests for RepeatInterleaveDistributedOp.""" |
| 35 | - def setUp(self): | 38 | + |
| 39 | + def setUp(self) -> None: | ||
| 36 | """Set up test fixtures before each test method.""" | 40 | """Set up test fixtures before each test method.""" |
| 37 | EXISTING_COMM_GROUPS.clear() | 41 | EXISTING_COMM_GROUPS.clear() |
| 38 | _DEVICE_MESH_MAP.clear() | 42 | _DEVICE_MESH_MAP.clear() |
| 43 | + _LAYOUT_CACHE.clear() | ||
| 39 | 44 | ||
| 40 | - def tearDown(self): | 45 | + def tearDown(self) -> None: |
| 41 | """Clean up after each test method.""" | 46 | """Clean up after each test method.""" |
| 42 | EXISTING_COMM_GROUPS.clear() | 47 | EXISTING_COMM_GROUPS.clear() |
| 43 | _DEVICE_MESH_MAP.clear() | 48 | _DEVICE_MESH_MAP.clear() |
| 49 | + _LAYOUT_CACHE.clear() | ||
| 44 | 50 | ||
| 45 | def _setup_mock_platform(self, mock_platform, platform_type=None, world_size=8): | 51 | def _setup_mock_platform(self, mock_platform, platform_type=None, world_size=8): |
| 46 | """Configure common mock-platform attributes used across tests.""" | 52 | """Configure common mock-platform attributes used across tests.""" |
| @@ -67,21 +73,22 @@ class TestParallelRepeatInterleave(unittest.TestCase): | |||
| 67 | mesh = self._make_2x4_mesh(mock_platform) | 73 | mesh = self._make_2x4_mesh(mock_platform) |
| 68 | x_placements = (Shard(0), Replicate()) | 74 | x_placements = (Shard(0), Replicate()) |
| 69 | x_layout = _build_layout(mesh, x_placements, 2) | 75 | x_layout = _build_layout(mesh, x_placements, 2) |
| 70 | - repeats = 2 | ||
| 71 | dim = 1 | 76 | dim = 1 |
| 72 | - output_layout = op.infer_layout((x_layout,), (repeats, dim)) | 77 | + cache_values = [x_layout, dim] |
| 78 | + result = op.infer_layout(cache_values) | ||
| 79 | + output_layout = result[0][0] | ||
| 73 | 80 | ||
| 74 | expected_map = (1, -1) | 81 | expected_map = (1, -1) |
| 75 | assert output_layout.to_dict()["tensor_map"] == expected_map, ( | 82 | assert output_layout.to_dict()["tensor_map"] == expected_map, ( |
| 76 | - f"Data Parallel with torch repeat_interleave test failed. Expected {expected_map}," | 83 | + f"Data Parallel with torch repeat_interleave test failed. Expected {expected_map}, " |
| 77 | - f" got {output_layout.to_dict()['tensor_map']}" | 84 | + f"got {output_layout.to_dict()['tensor_map']}" |
| 78 | ) | 85 | ) |
| 79 | 86 | ||
| 80 | # Since `get_expand_impl` is not overridden, it returns None by default. | 87 | # Since `get_expand_impl` is not overridden, it returns None by default. |
| 81 | # The same applies to other test classes, so it is unnecessary to test its return value. | 88 | # The same applies to other test classes, so it is unnecessary to test its return value. |
| 82 | - assert op.get_expand_impl(None, output_layout, (x_layout,), (repeats, dim)) is None, ( | 89 | + assert op.get_expand_impl(None, result, cache_values) is None, ( |
| 83 | f"get_expand_impl test failed. Expected None, " | 90 | f"get_expand_impl test failed. Expected None, " |
| 84 | - f"got {op.get_expand_impl(None, output_layout, (x_layout,), (repeats, dim))}" | 91 | + f"got {op.get_expand_impl(None, result, cache_values)}" |
| 85 | ) | 92 | ) |
| 86 | 93 | ||
| 87 | @patch("hyper_parallel.core.dtensor.device_mesh.platform") | 94 | @patch("hyper_parallel.core.dtensor.device_mesh.platform") |
| @@ -94,13 +101,14 @@ class TestParallelRepeatInterleave(unittest.TestCase): | |||
| 94 | mesh = self._make_2x4_mesh(mock_platform) | 101 | mesh = self._make_2x4_mesh(mock_platform) |
| 95 | x_placements = (Replicate(), Shard(0)) | 102 | x_placements = (Replicate(), Shard(0)) |
| 96 | x_layout = _build_layout(mesh, x_placements, 2) | 103 | x_layout = _build_layout(mesh, x_placements, 2) |
| 97 | - repeats = 2 | ||
| 98 | dim = 1 | 104 | dim = 1 |
| 99 | - output_layout = op.infer_layout((x_layout,), (repeats, dim)) | 105 | + cache_values = [x_layout, dim] |
| 106 | + result = op.infer_layout(cache_values) | ||
| 107 | + output_layout = result[0][0] | ||
| 100 | expected_map = (0, -1) | 108 | expected_map = (0, -1) |
| 101 | assert output_layout.to_dict()["tensor_map"] == expected_map, ( | 109 | assert output_layout.to_dict()["tensor_map"] == expected_map, ( |
| 102 | - f"Tensor Parallel with torch repeat_interleave test failed. Expected {expected_map}," | 110 | + f"Tensor Parallel with torch repeat_interleave test failed. Expected {expected_map}, " |
| 103 | - f" got {output_layout.to_dict()['tensor_map']}" | 111 | + f"got {output_layout.to_dict()['tensor_map']}" |
| 104 | ) | 112 | ) |
| 105 | 113 | ||
| 106 | @patch("hyper_parallel.core.dtensor.device_mesh.platform") | 114 | @patch("hyper_parallel.core.dtensor.device_mesh.platform") |
| @@ -113,14 +121,14 @@ class TestParallelRepeatInterleave(unittest.TestCase): | |||
| 113 | mesh = self._make_2x4_mesh(mock_platform) | 121 | mesh = self._make_2x4_mesh(mock_platform) |
| 114 | x_placements = (Shard(0), Replicate()) | 122 | x_placements = (Shard(0), Replicate()) |
| 115 | x_layout = _build_layout(mesh, x_placements, 2) | 123 | x_layout = _build_layout(mesh, x_placements, 2) |
| 116 | - repeats_tensor = [2, 1, 1, 1] | ||
| 117 | - | ||
| 118 | dim = 1 | 124 | dim = 1 |
| 119 | - output_layout = op.infer_layout((x_layout,), (repeats_tensor, dim)) | 125 | + cache_values = [x_layout, dim] |
| 126 | + result = op.infer_layout(cache_values) | ||
| 127 | + output_layout = result[0][0] | ||
| 120 | expected_map = (1, -1) | 128 | expected_map = (1, -1) |
| 121 | assert output_layout.to_dict()["tensor_map"] == expected_map, ( | 129 | assert output_layout.to_dict()["tensor_map"] == expected_map, ( |
| 122 | - f"Data Parallel with torch repeat_interleave test failed. Expected {expected_map}," | 130 | + f"Data Parallel with torch repeat_interleave test failed. Expected {expected_map}, " |
| 123 | - f" got {output_layout.to_dict()['tensor_map']}" | 131 | + f"got {output_layout.to_dict()['tensor_map']}" |
| 124 | ) | 132 | ) |
| 125 | 133 | ||
| 126 | @patch("hyper_parallel.core.dtensor.device_mesh.platform") | 134 | @patch("hyper_parallel.core.dtensor.device_mesh.platform") |
| @@ -133,14 +141,14 @@ class TestParallelRepeatInterleave(unittest.TestCase): | |||
| 133 | mesh = self._make_2x4_mesh(mock_platform) | 141 | mesh = self._make_2x4_mesh(mock_platform) |
| 134 | x_placements = (Replicate(), Shard(0)) | 142 | x_placements = (Replicate(), Shard(0)) |
| 135 | x_layout = _build_layout(mesh, x_placements, 2) | 143 | x_layout = _build_layout(mesh, x_placements, 2) |
| 136 | - repeats_tensor = [2, 1, 1, 1] | ||
| 137 | - | ||
| 138 | dim = 1 | 144 | dim = 1 |
| 139 | - output_layout = op.infer_layout((x_layout,), (repeats_tensor, dim)) | 145 | + cache_values = [x_layout, dim] |
| 146 | + result = op.infer_layout(cache_values) | ||
| 147 | + output_layout = result[0][0] | ||
| 140 | expected_map = (0, -1) | 148 | expected_map = (0, -1) |
| 141 | assert output_layout.to_dict()["tensor_map"] == expected_map, ( | 149 | assert output_layout.to_dict()["tensor_map"] == expected_map, ( |
| 142 | - f"Tensor Parallel with torch repeat_interleave test failed. Expected {expected_map}," | 150 | + f"Tensor Parallel with torch repeat_interleave test failed. Expected {expected_map}, " |
| 143 | - f" got {output_layout.to_dict()['tensor_map']}" | 151 | + f"got {output_layout.to_dict()['tensor_map']}" |
| 144 | ) | 152 | ) |
| 145 | 153 | ||
| 146 | @patch("hyper_parallel.core.dtensor.device_mesh.platform") | 154 | @patch("hyper_parallel.core.dtensor.device_mesh.platform") |
| @@ -153,12 +161,13 @@ class TestParallelRepeatInterleave(unittest.TestCase): | |||
| 153 | mesh = self._make_2x4_mesh(mock_platform) | 161 | mesh = self._make_2x4_mesh(mock_platform) |
| 154 | x_placements = (Shard(0), Replicate()) | 162 | x_placements = (Shard(0), Replicate()) |
| 155 | x_layout = _build_layout(mesh, x_placements, 2) | 163 | x_layout = _build_layout(mesh, x_placements, 2) |
| 156 | - repeats = 2 | 164 | + cache_values = [x_layout, None] |
| 157 | - output_layout = op.infer_layout((x_layout,), (repeats,)) | 165 | + result = op.infer_layout(cache_values) |
| 166 | + output_layout = result[0][0] | ||
| 158 | expected_map = (1,) | 167 | expected_map = (1,) |
| 159 | assert output_layout.to_dict()["tensor_map"] == expected_map, ( | 168 | assert output_layout.to_dict()["tensor_map"] == expected_map, ( |
| 160 | - f"Data Parallel with dim None repeat_interleave test failed. Expected {expected_map}," | 169 | + f"Data Parallel with dim None repeat_interleave test failed. Expected {expected_map}, " |
| 161 | - f" got {output_layout.to_dict()['tensor_map']}" | 170 | + f"got {output_layout.to_dict()['tensor_map']}" |
| 162 | ) | 171 | ) |
| 163 | 172 | ||
| 164 | @patch("hyper_parallel.core.dtensor.device_mesh.platform") | 173 | @patch("hyper_parallel.core.dtensor.device_mesh.platform") |
| @@ -171,15 +180,102 @@ class TestParallelRepeatInterleave(unittest.TestCase): | |||
| 171 | mesh = self._make_2x4_mesh(mock_platform) | 180 | mesh = self._make_2x4_mesh(mock_platform) |
| 172 | x_placements = (Replicate(), Shard(0)) | 181 | x_placements = (Replicate(), Shard(0)) |
| 173 | x_layout = _build_layout(mesh, x_placements, 2) | 182 | x_layout = _build_layout(mesh, x_placements, 2) |
| 174 | - | 183 | + cache_values = [x_layout, None] |
| 175 | - repeats = 2 | 184 | + result = op.infer_layout(cache_values) |
| 176 | - output_layout = op.infer_layout((x_layout,), (repeats,)) | 185 | + output_layout = result[0][0] |
| 177 | expected_map = (0,) | 186 | expected_map = (0,) |
| 178 | assert output_layout.to_dict()["tensor_map"] == expected_map, ( | 187 | assert output_layout.to_dict()["tensor_map"] == expected_map, ( |
| 179 | - f"Tensor Parallel dim None repeat_interleave test failed. Expected {expected_map}," | 188 | + f"Tensor Parallel dim None repeat_interleave test failed. Expected {expected_map}, " |
| 180 | - f" got {output_layout.to_dict()['tensor_map']}" | 189 | + f"got {output_layout.to_dict()['tensor_map']}" |
| 181 | ) | 190 | ) |
| 182 | 191 | ||
| 192 | + @patch("hyper_parallel.core.dtensor.device_mesh.platform") | ||
| 193 | + def test_torch_repeat_interleave_layout_hybrid_parallel(self, mock_platform): | ||
| 194 | + """ | ||
| 195 | + Feature: RepeatInterleave hybrid parallel | ||
| 196 | + Description: Hybrid scenario with multiple input dimensions sharded and repeated dim replicated | ||
| 197 | + Expectation: Success | ||
| 198 | + """ | ||
| 199 | + mesh = self._make_2x4_mesh(mock_platform) | ||
| 200 | + x_layout = _build_layout(mesh, ("dp", "None", "tp"), 3) | ||
| 201 | + dim = 1 | ||
| 202 | + cache_values = [x_layout, dim] | ||
| 203 | + result = op.infer_layout(cache_values) | ||
| 204 | + output_layout = result[0][0] | ||
| 205 | + expected_map = (1, -1, 0) | ||
| 206 | + assert output_layout.to_dict()["tensor_map"] == expected_map, ( | ||
| 207 | + f"Hybrid parallel repeat_interleave test failed. Expected {expected_map}, " | ||
| 208 | + f"got {output_layout.to_dict()['tensor_map']}" | ||
| 209 | + ) | ||
| 210 | + | ||
| 211 | + @patch("hyper_parallel.core.dtensor.device_mesh.platform") | ||
| 212 | + def test_torch_repeat_interleave_layout_all_replicated(self, mock_platform): | ||
| 213 | + """ | ||
| 214 | + Feature: RepeatInterleave all replicated | ||
| 215 | + Description: All-replicated input keeps replicated layout after repeat_interleave | ||
| 216 | + Expectation: Success | ||
| 217 | + """ | ||
| 218 | + mesh = self._make_2x4_mesh(mock_platform) | ||
| 219 | + x_placements = (Replicate(), Replicate()) | ||
| 220 | + x_layout = _build_layout(mesh, x_placements, 2) | ||
| 221 | + dim = 1 | ||
| 222 | + cache_values = [x_layout, dim] | ||
| 223 | + result = op.infer_layout(cache_values) | ||
| 224 | + output_layout = result[0][0] | ||
| 225 | + expected_map = (-1, -1) | ||
| 226 | + assert output_layout.to_dict()["tensor_map"] == expected_map, ( | ||
| 227 | + f"All replicated repeat_interleave test failed. Expected {expected_map}, " | ||
| 228 | + f"got {output_layout.to_dict()['tensor_map']}" | ||
| 229 | + ) | ||
| 230 | + | ||
| 231 | + @patch("hyper_parallel.core.dtensor.device_mesh.platform") | ||
| 232 | + def test_torch_repeat_interleave_layout_negative_dim(self, mock_platform): | ||
| 233 | + """ | ||
| 234 | + Feature: RepeatInterleave negative dim | ||
| 235 | + Description: Negative repeat dim is normalized before validating sharding | ||
| 236 | + Expectation: Success | ||
| 237 | + """ | ||
| 238 | + mesh = self._make_2x4_mesh(mock_platform) | ||
| 239 | + x_placements = (Shard(0), Replicate()) | ||
| 240 | + x_layout = _build_layout(mesh, x_placements, 2) | ||
| 241 | + dim = -1 | ||
| 242 | + cache_values = [x_layout, dim] | ||
| 243 | + result = op.infer_layout(cache_values) | ||
| 244 | + output_layout = result[0][0] | ||
| 245 | + expected_map = (1, -1) | ||
| 246 | + assert output_layout.to_dict()["tensor_map"] == expected_map, ( | ||
| 247 | + f"Negative dim repeat_interleave test failed. Expected {expected_map}, " | ||
| 248 | + f"got {output_layout.to_dict()['tensor_map']}" | ||
| 249 | + ) | ||
| 250 | + | ||
| 251 | + @patch("hyper_parallel.core.dtensor.device_mesh.platform") | ||
| 252 | + def test_torch_repeat_interleave_partial_input_error(self, mock_platform): | ||
| 253 | + """ | ||
| 254 | + Feature: RepeatInterleave partial input validation | ||
| 255 | + Description: Partial input layout should be rejected before layout inference | ||
| 256 | + Expectation: Success | ||
| 257 | + """ | ||
| 258 | + mesh = self._make_2x4_mesh(mock_platform) | ||
| 259 | + x_layout = _build_layout(mesh, (Partial(), Replicate()), 2) | ||
| 260 | + dim = 1 | ||
| 261 | + cache_values = [x_layout, dim] | ||
| 262 | + with self.assertRaisesRegex(ValueError, "status which is not allowed"): | ||
| 263 | + op.infer_layout(cache_values) | ||
| 264 | + | ||
| 265 | + @patch("hyper_parallel.core.dtensor.device_mesh.platform") | ||
| 266 | + def test_torch_repeat_interleave_strided_shard_repeat_dim_error(self, mock_platform): | ||
| 267 | + """ | ||
| 268 | + Feature: RepeatInterleave StridedShard repeat dim validation | ||
| 269 | + Description: Repeat dim using tuple alias map should be treated as sharded | ||
| 270 | + Expectation: Success | ||
| 271 | + """ | ||
| 272 | + mesh = self._make_2x4_mesh(mock_platform) | ||
| 273 | + x_layout = _build_layout(mesh, (("dp", "tp"), "None"), 2) | ||
| 274 | + dim = 0 | ||
| 275 | + cache_values = [x_layout, dim] | ||
| 276 | + with self.assertRaisesRegex(ValueError, "repeat dimension should be replicated"): | ||
| 277 | + op.infer_layout(cache_values) | ||
| 278 | + | ||
| 183 | @patch("hyper_parallel.core.dtensor.device_mesh.platform") | 279 | @patch("hyper_parallel.core.dtensor.device_mesh.platform") |
| 184 | def test_torch_repeat_interleave_layout_sharded_dim_error(self, mock_platform): | 280 | def test_torch_repeat_interleave_layout_sharded_dim_error(self, mock_platform): |
| 185 | """ | 281 | """ |
| @@ -190,11 +286,10 @@ class TestParallelRepeatInterleave(unittest.TestCase): | |||
| 190 | mesh = self._make_2x4_mesh(mock_platform) | 286 | mesh = self._make_2x4_mesh(mock_platform) |
| 191 | x_placements = (Shard(0), Shard(1)) | 287 | x_placements = (Shard(0), Shard(1)) |
| 192 | x_layout = _build_layout(mesh, x_placements, 2) | 288 | x_layout = _build_layout(mesh, x_placements, 2) |
| 193 | - | ||
| 194 | - repeats = 2 | ||
| 195 | dim = 1 | 289 | dim = 1 |
| 196 | - with self.assertRaisesRegex(ValueError, "Cannot perform sharding on params along the chosen dim"): | 290 | + cache_values = [x_layout, dim] |
| 197 | - op.infer_layout((x_layout,), (repeats, dim)) | 291 | + with self.assertRaisesRegex(ValueError, "repeat dimension should be replicated"): |
| 292 | + op.infer_layout(cache_values) | ||
| 198 | 293 | ||
| 199 | @patch("hyper_parallel.core.dtensor.device_mesh.platform") | 294 | @patch("hyper_parallel.core.dtensor.device_mesh.platform") |
| 200 | def test_torch_repeat_interleave_layout_error_dim_out_of_range(self, mock_platform): | 295 | def test_torch_repeat_interleave_layout_error_dim_out_of_range(self, mock_platform): |
| @@ -206,12 +301,98 @@ class TestParallelRepeatInterleave(unittest.TestCase): | |||
| 206 | mesh = self._make_2x4_mesh(mock_platform) | 301 | mesh = self._make_2x4_mesh(mock_platform) |
| 207 | x_placements = (Replicate(), Replicate()) | 302 | x_placements = (Replicate(), Replicate()) |
| 208 | x_layout = _build_layout(mesh, x_placements, 2) | 303 | x_layout = _build_layout(mesh, x_placements, 2) |
| 209 | - | ||
| 210 | - repeats = 2 | ||
| 211 | dim = 5 | 304 | dim = 5 |
| 305 | + cache_values = [x_layout, dim] | ||
| 306 | + with self.assertRaisesRegex(ValueError, "dimension should be in"): | ||
| 307 | + op.infer_layout(cache_values) | ||
| 212 | 308 | ||
| 213 | - with self.assertRaisesRegex(ValueError, "Dimension out of range"): | 309 | + def test_normalize_repeat_interleave_args_basic(self): |
| 214 | - op.infer_layout((x_layout,), extra_args=(repeats, dim)) | 310 | + """ |
| 311 | + Feature: Normalize repeat_interleave args | ||
| 312 | + Description: Verify normalize function returns correct structure | ||
| 313 | + Expectation: Success | ||
| 314 | + """ | ||
| 315 | + mock_input = MagicMock() | ||
| 316 | + args, kwargs = _normalize_repeat_interleave_args(mock_input, 3, 1) | ||
| 317 | + assert args == (mock_input, 3, 1), ( | ||
| 318 | + f"Normalize args mismatch: expected (mock_input, 3, 1), got {args}" | ||
| 319 | + ) | ||
| 320 | + assert not kwargs, ( | ||
| 321 | + f"Normalize kwargs mismatch: expected {{}}, got {kwargs}" | ||
| 322 | + ) | ||
| 323 | + | ||
| 324 | + def test_normalize_repeat_interleave_args_with_output_size(self): | ||
| 325 | + """ | ||
| 326 | + Feature: Normalize repeat_interleave args with output_size | ||
| 327 | + Description: Verify output_size goes into kwargs | ||
| 328 | + Expectation: Success | ||
| 329 | + """ | ||
| 330 | + mock_input = MagicMock() | ||
| 331 | + args, kwargs = _normalize_repeat_interleave_args(mock_input, 3, 1, output_size=10) | ||
| 332 | + assert args == (mock_input, 3, 1), ( | ||
| 333 | + f"Normalize args mismatch: expected (mock_input, 3, 1), got {args}" | ||
| 334 | + ) | ||
| 335 | + assert kwargs == {'output_size': 10}, ( | ||
| 336 | + f"Normalize kwargs mismatch: expected {{'output_size': 10}}, got {kwargs}" | ||
| 337 | + ) | ||
| 338 | + | ||
| 339 | + @patch("hyper_parallel.core.dtensor.device_mesh.platform") | ||
| 340 | + def test_preprocess_with_int_repeats(self, mock_platform): | ||
| 341 | + """ | ||
| 342 | + Feature: Preprocess with int repeats | ||
| 343 | + Description: Verify preprocess correctly handles int repeats | ||
| 344 | + Expectation: Success | ||
| 345 | + """ | ||
| 346 | + mesh = self._make_2x4_mesh(mock_platform) | ||
| 347 | + x_placements = (Shard(0), Replicate()) | ||
| 348 | + x_layout = _build_layout(mesh, x_placements, 2) | ||
| 349 | + | ||
| 350 | + mock_input = MagicMock() | ||
| 351 | + mock_input.layout = x_layout | ||
| 352 | + mock_local = MagicMock() | ||
| 353 | + mock_input.to_local.return_value = mock_local | ||
| 354 | + | ||
| 355 | + local_args, local_kwargs, cache_values = op.preprocess( | ||
| 356 | + (mock_input, 3, 1), {} | ||
| 357 | + ) | ||
| 358 | + | ||
| 359 | + assert local_args[0] is mock_local, ( | ||
| 360 | + f"local_args[0] should be local input: expected={mock_local}, got={local_args[0]}" | ||
| 361 | + ) | ||
| 362 | + assert local_args[1] == 3, f"local_args[1] should be 3, got {local_args[1]}" | ||
| 363 | + assert local_args[2] == 1, f"local_args[2] should be 1, got {local_args[2]}" | ||
| 364 | + assert not local_kwargs, f"local_kwargs should be empty, got {local_kwargs}" | ||
| 365 | + assert cache_values[0] is x_layout, ( | ||
| 366 | + f"cache_values[0] should be input layout: expected={x_layout}, got={cache_values[0]}" | ||
| 367 | + ) | ||
| 368 | + assert cache_values[1] == 1, f"cache_values[1] should be 1, got {cache_values[1]}" | ||
| 369 | + | ||
| 370 | + @patch("hyper_parallel.core.dtensor.device_mesh.platform") | ||
| 371 | + def test_preprocess_dim_none(self, mock_platform): | ||
| 372 | + """ | ||
| 373 | + Feature: Preprocess with dim=None | ||
| 374 | + Description: Verify preprocess correctly handles dim=None (flatten mode) | ||
| 375 | + Expectation: Success | ||
| 376 | + """ | ||
| 377 | + mesh = self._make_2x4_mesh(mock_platform) | ||
| 378 | + x_placements = (Shard(0), Replicate()) | ||
| 379 | + x_layout = _build_layout(mesh, x_placements, 2) | ||
| 380 | + | ||
| 381 | + mock_input = MagicMock() | ||
| 382 | + mock_input.layout = x_layout | ||
| 383 | + mock_local = MagicMock() | ||
| 384 | + mock_input.to_local.return_value = mock_local | ||
| 385 | + | ||
| 386 | + local_args, local_kwargs, cache_values = op.preprocess( | ||
| 387 | + (mock_input, 3, None), {} | ||
| 388 | + ) | ||
| 389 | + | ||
| 390 | + assert local_args[0] is mock_local, ( | ||
| 391 | + f"local_args[0] should be local input: expected={mock_local}, got={local_args[0]}" | ||
| 392 | + ) | ||
| 393 | + assert local_args[1] == 3, f"local_args[1] should be 3, got {local_args[1]}" | ||
| 394 | + assert local_args[2] is None, f"local_args[2] should be None, got {local_args[2]}" | ||
| 395 | + assert cache_values[1] is None, f"cache_values[1] should be None, got {cache_values[1]}" | ||
| 215 | 396 | ||
| 216 | 397 | ||
| 217 | if __name__ == "__main__": | 398 | if __name__ == "__main__": |
| @@ -15,10 +15,10 @@ | |||
| 15 | """parallel_reshape test""" | 15 | """parallel_reshape test""" |
| 16 | import os | 16 | import os |
| 17 | import unittest | 17 | import unittest |
| 18 | -from unittest.mock import patch | 18 | +from unittest.mock import MagicMock, patch |
| 19 | import numpy as np | 19 | import numpy as np |
| 20 | 20 | ||
| 21 | -from hyper_parallel.core.dtensor.dtensor import _build_layout | 21 | +from hyper_parallel.core.dtensor.dtensor import _build_layout, _LAYOUT_CACHE |
| 22 | from hyper_parallel.core.dtensor.placement_types import Shard, Replicate, Partial | 22 | from hyper_parallel.core.dtensor.placement_types import Shard, Replicate, Partial |
| 23 | from hyper_parallel.core.shard.ops.parallel_reshape import ReshapeDistributedOp | 23 | from hyper_parallel.core.shard.ops.parallel_reshape import ReshapeDistributedOp |
| 24 | from hyper_parallel.core.dtensor.device_mesh import ( | 24 | from hyper_parallel.core.dtensor.device_mesh import ( |
| @@ -42,11 +42,13 @@ class TestParallelReshape(unittest.TestCase): | |||
| 42 | """ | 42 | """ |
| 43 | EXISTING_COMM_GROUPS.clear() | 43 | EXISTING_COMM_GROUPS.clear() |
| 44 | _DEVICE_MESH_MAP.clear() | 44 | _DEVICE_MESH_MAP.clear() |
| 45 | + _LAYOUT_CACHE.clear() | ||
| 45 | 46 | ||
| 46 | def tearDown(self): | 47 | def tearDown(self): |
| 47 | """Clean up after each test method.""" | 48 | """Clean up after each test method.""" |
| 48 | EXISTING_COMM_GROUPS.clear() | 49 | EXISTING_COMM_GROUPS.clear() |
| 49 | _DEVICE_MESH_MAP.clear() | 50 | _DEVICE_MESH_MAP.clear() |
| 51 | + _LAYOUT_CACHE.clear() | ||
| 50 | 52 | ||
| 51 | def _setup_mock_platform(self, mock_platform, platform_type=None, world_size=8): | 53 | def _setup_mock_platform(self, mock_platform, platform_type=None, world_size=8): |
| 52 | """Configure common mock-platform attributes used across tests. | 54 | """Configure common mock-platform attributes used across tests. |
| @@ -79,6 +81,47 @@ class TestParallelReshape(unittest.TestCase): | |||
| 79 | self._setup_mock_platform(mock_platform, world_size=8) | 81 | self._setup_mock_platform(mock_platform, world_size=8) |
| 80 | return init_device_mesh(device_type="npu", mesh_shape=(8,), mesh_dim_names=(mesh_dim_name,)) | 82 | return init_device_mesh(device_type="npu", mesh_shape=(8,), mesh_dim_names=(mesh_dim_name,)) |
| 81 | 83 | ||
| 84 | + def _infer_reshape(self, reshape_op, x_layout, dst_shape, src_shape): | ||
| 85 | + """Infer reshape layout using the new cache_values calling convention.""" | ||
| 86 | + output_layouts, local_dst_shape = reshape_op.infer_layout([x_layout, dst_shape, src_shape]) | ||
| 87 | + return output_layouts[0], local_dst_shape | ||
| 88 | + | ||
| 89 | + @patch("hyper_parallel.core.dtensor.device_mesh.platform") | ||
| 90 | + def test_reshape_preprocess_builds_cache_values(self, mock_platform): | ||
| 91 | + """ | ||
| 92 | + Feature: ReshapeDistributedOp preprocess. | ||
| 93 | + Description: Verify that preprocess converts DTensor input to local tensor and caches layout/shape metadata. | ||
| 94 | + Expectation: local args, kwargs, and cache values follow the new dispatch convention. | ||
| 95 | + """ | ||
| 96 | + mesh = self._make_2x4_mesh(mock_platform) | ||
| 97 | + x_layout = _build_layout(mesh, (Shard(0), Replicate()), 2) | ||
| 98 | + local_tensor = MagicMock() | ||
| 99 | + mock_tensor = MagicMock() | ||
| 100 | + mock_tensor.layout = x_layout | ||
| 101 | + mock_tensor.shape = (8, 4) | ||
| 102 | + mock_tensor.to_local.return_value = local_tensor | ||
| 103 | + | ||
| 104 | + local_args, local_kwargs, cache_values = op.preprocess((mock_tensor, (4, 8)), {}) | ||
| 105 | + | ||
| 106 | + assert local_args == (local_tensor, (4, 8)) | ||
| 107 | + assert not local_kwargs | ||
| 108 | + assert cache_values == [x_layout, (4, 8), (8, 4)] | ||
| 109 | + | ||
| 110 | + def test_reshape_get_expand_impl_uses_local_shape(self): | ||
| 111 | + """ | ||
| 112 | + Feature: ReshapeDistributedOp get_expand_impl. | ||
| 113 | + Description: Verify that the wrapper replaces global target shape with inferred local target shape. | ||
| 114 | + Expectation: Underlying function receives local_dst_shape. | ||
| 115 | + """ | ||
| 116 | + infer_result = ((MagicMock(),), [2, 8]) | ||
| 117 | + | ||
| 118 | + def fake_reshape(x, shape): | ||
| 119 | + return x, shape | ||
| 120 | + | ||
| 121 | + impl = op.get_expand_impl(fake_reshape, infer_result, [MagicMock(), (4, 8), (8, 4)]) | ||
| 122 | + | ||
| 123 | + assert impl("local_x", (4, 8)) == ("local_x", [2, 8]) | ||
| 124 | + | ||
| 82 | @patch("hyper_parallel.core.dtensor.device_mesh.platform") | 125 | @patch("hyper_parallel.core.dtensor.device_mesh.platform") |
| 83 | def test_reshape_layout_not_change_sharded_axis(self, mock_platform): | 126 | def test_reshape_layout_not_change_sharded_axis(self, mock_platform): |
| 84 | """ | 127 | """ |
| @@ -91,7 +134,7 @@ class TestParallelReshape(unittest.TestCase): | |||
| 91 | src_shape = (1024, 512, 512) | 134 | src_shape = (1024, 512, 512) |
| 92 | dst_shape = (1024, 2, 256, 512) | 135 | dst_shape = (1024, 2, 256, 512) |
| 93 | 136 | ||
| 94 | - output_layout, local_dst_shape = op.infer_layout((x_layout,), (dst_shape, src_shape)) | 137 | + output_layout, local_dst_shape = self._infer_reshape(op, x_layout, dst_shape, src_shape) |
| 95 | expected_map = (1, -1, -1, -1) | 138 | expected_map = (1, -1, -1, -1) |
| 96 | assert output_layout.tensor_map == expected_map, ( | 139 | assert output_layout.tensor_map == expected_map, ( |
| 97 | f"Reshape do not change sharded axis failed. Expected {expected_map}, " | 140 | f"Reshape do not change sharded axis failed. Expected {expected_map}, " |
| @@ -103,11 +146,10 @@ class TestParallelReshape(unittest.TestCase): | |||
| 103 | f"got {local_dst_shape}" | 146 | f"got {local_dst_shape}" |
| 104 | ) | 147 | ) |
| 105 | 148 | ||
| 106 | - # Since `get_expand_impl` is not overridden, it returns None by default. | 149 | + # If func is None in a unit-level call, no runtime wrapper is needed. |
| 107 | - # The same applies to other test classes, so it is unnecessary to test its return value. | 150 | + assert op.get_expand_impl(None, ((output_layout,), local_dst_shape), [x_layout, dst_shape, src_shape]) is None, ( |
| 108 | - assert op.get_expand_impl(None, output_layout, (x_layout,), (dst_shape, src_shape)) is None, ( | ||
| 109 | f"get_expand_impl test failed. Expected None, " | 151 | f"get_expand_impl test failed. Expected None, " |
| 110 | - f"got {op.get_expand_impl(None, output_layout, (x_layout,), (dst_shape, src_shape))}" | 152 | + f"got {op.get_expand_impl(None, ((output_layout,), local_dst_shape), [x_layout, dst_shape, src_shape])}" |
| 111 | ) | 153 | ) |
| 112 | 154 | ||
| 113 | @patch("hyper_parallel.core.dtensor.device_mesh.platform") | 155 | @patch("hyper_parallel.core.dtensor.device_mesh.platform") |
| @@ -122,7 +164,7 @@ class TestParallelReshape(unittest.TestCase): | |||
| 122 | src_shape = (4, 4, 8) | 164 | src_shape = (4, 4, 8) |
| 123 | dst_shape = (16, 8) | 165 | dst_shape = (16, 8) |
| 124 | 166 | ||
| 125 | - output_layout, local_dst_shape = op.infer_layout((x_layout,), (dst_shape, src_shape)) | 167 | + output_layout, local_dst_shape = self._infer_reshape(op, x_layout, dst_shape, src_shape) |
| 126 | expected_map = (1, -1) | 168 | expected_map = (1, -1) |
| 127 | assert output_layout.tensor_map == expected_map, ( | 169 | assert output_layout.tensor_map == expected_map, ( |
| 128 | f"Reshape do not change sharded axis failed. Expected {expected_map}, " | 170 | f"Reshape do not change sharded axis failed. Expected {expected_map}, " |
| @@ -146,7 +188,7 @@ class TestParallelReshape(unittest.TestCase): | |||
| 146 | src_shape = (32, 128) | 188 | src_shape = (32, 128) |
| 147 | dst_shape = (4, 8, 128) | 189 | dst_shape = (4, 8, 128) |
| 148 | 190 | ||
| 149 | - output_layout, local_dst_shape = op.infer_layout((x_layout,), (dst_shape, src_shape)) | 191 | + output_layout, local_dst_shape = self._infer_reshape(op, x_layout, dst_shape, src_shape) |
| 150 | expected_map = (1, -1, -1) | 192 | expected_map = (1, -1, -1) |
| 151 | assert output_layout.tensor_map == expected_map, ( | 193 | assert output_layout.tensor_map == expected_map, ( |
| 152 | f"Reshape do not change sharded axis failed. Expected {expected_map}, " | 194 | f"Reshape do not change sharded axis failed. Expected {expected_map}, " |
| @@ -170,7 +212,7 @@ class TestParallelReshape(unittest.TestCase): | |||
| 170 | src_shape = (32, 6, 128, 28, 10) | 212 | src_shape = (32, 6, 128, 28, 10) |
| 171 | dst_shape = (4, 8, 2, 384, 280) | 213 | dst_shape = (4, 8, 2, 384, 280) |
| 172 | 214 | ||
| 173 | - output_layout, local_dst_shape = op.infer_layout((x_layout,), (dst_shape, src_shape)) | 215 | + output_layout, local_dst_shape = self._infer_reshape(op, x_layout, dst_shape, src_shape) |
| 174 | expected_map = (0, -1, 2, -1, 1) | 216 | expected_map = (0, -1, 2, -1, 1) |
| 175 | assert output_layout.tensor_map == expected_map, ( | 217 | assert output_layout.tensor_map == expected_map, ( |
| 176 | f"Reshape do not change sharded axis failed. Expected {expected_map}, " | 218 | f"Reshape do not change sharded axis failed. Expected {expected_map}, " |
| @@ -195,7 +237,7 @@ class TestParallelReshape(unittest.TestCase): | |||
| 195 | dst_shape = (4, 8, 12, 4) | 237 | dst_shape = (4, 8, 12, 4) |
| 196 | 238 | ||
| 197 | with self.assertRaises(ValueError): | 239 | with self.assertRaises(ValueError): |
| 198 | - _, _ = op.infer_layout((x_layout,), (dst_shape, src_shape)) | 240 | + _, _ = self._infer_reshape(op, x_layout, dst_shape, src_shape) |
| 199 | 241 | ||
| 200 | @patch("hyper_parallel.core.dtensor.device_mesh.platform") | 242 | @patch("hyper_parallel.core.dtensor.device_mesh.platform") |
| 201 | def test_reshape_layout_can_not_reshape2(self, mock_platform): | 243 | def test_reshape_layout_can_not_reshape2(self, mock_platform): |
| @@ -210,7 +252,7 @@ class TestParallelReshape(unittest.TestCase): | |||
| 210 | dst_shape = (4, 8, 2, 42) | 252 | dst_shape = (4, 8, 2, 42) |
| 211 | 253 | ||
| 212 | with self.assertRaises(ValueError): | 254 | with self.assertRaises(ValueError): |
| 213 | - _, _ = op.infer_layout((x_layout,), (dst_shape, src_shape)) | 255 | + _, _ = self._infer_reshape(op, x_layout, dst_shape, src_shape) |
| 214 | 256 | ||
| 215 | @patch("hyper_parallel.core.dtensor.device_mesh.platform") | 257 | @patch("hyper_parallel.core.dtensor.device_mesh.platform") |
| 216 | def test_reshape_layout_dynamic_shape1(self, mock_platform): | 258 | def test_reshape_layout_dynamic_shape1(self, mock_platform): |
| @@ -224,7 +266,7 @@ class TestParallelReshape(unittest.TestCase): | |||
| 224 | src_shape = (1024, -1, 256, 512) | 266 | src_shape = (1024, -1, 256, 512) |
| 225 | dst_shape = (1024, -1, 512) | 267 | dst_shape = (1024, -1, 512) |
| 226 | 268 | ||
| 227 | - output_layout, local_dst_shape = op.infer_layout((x_layout,), (dst_shape, src_shape)) | 269 | + output_layout, local_dst_shape = self._infer_reshape(op, x_layout, dst_shape, src_shape) |
| 228 | expected_map = (1, -1, -1) | 270 | expected_map = (1, -1, -1) |
| 229 | assert output_layout.tensor_map == expected_map, ( | 271 | assert output_layout.tensor_map == expected_map, ( |
| 230 | f"Reshape do not change sharded axis failed. Expected {expected_map}, " | 272 | f"Reshape do not change sharded axis failed. Expected {expected_map}, " |
| @@ -248,7 +290,7 @@ class TestParallelReshape(unittest.TestCase): | |||
| 248 | src_shape = (1024, -1, 512) | 290 | src_shape = (1024, -1, 512) |
| 249 | dst_shape = (1024, -1, 256, 512) | 291 | dst_shape = (1024, -1, 256, 512) |
| 250 | 292 | ||
| 251 | - output_layout, local_dst_shape = op.infer_layout((x_layout,), (dst_shape, src_shape)) | 293 | + output_layout, local_dst_shape = self._infer_reshape(op, x_layout, dst_shape, src_shape) |
| 252 | expected_map = (1, -1, -1, -1) | 294 | expected_map = (1, -1, -1, -1) |
| 253 | assert output_layout.tensor_map == expected_map, ( | 295 | assert output_layout.tensor_map == expected_map, ( |
| 254 | f"Reshape do not change sharded axis failed. Expected {expected_map}, " | 296 | f"Reshape do not change sharded axis failed. Expected {expected_map}, " |
| @@ -272,7 +314,7 @@ class TestParallelReshape(unittest.TestCase): | |||
| 272 | src_shape = (-1, 256, 512) | 314 | src_shape = (-1, 256, 512) |
| 273 | dst_shape = (-1, 512) | 315 | dst_shape = (-1, 512) |
| 274 | 316 | ||
| 275 | - output_layout, local_dst_shape = op.infer_layout((x_layout,), (dst_shape, src_shape)) | 317 | + output_layout, local_dst_shape = self._infer_reshape(op, x_layout, dst_shape, src_shape) |
| 276 | expected_map = (1, -1) | 318 | expected_map = (1, -1) |
| 277 | assert output_layout.tensor_map == expected_map, ( | 319 | assert output_layout.tensor_map == expected_map, ( |
| 278 | f"Reshape do not change sharded axis failed. Expected {expected_map}, " | 320 | f"Reshape do not change sharded axis failed. Expected {expected_map}, " |
| @@ -297,7 +339,7 @@ class TestParallelReshape(unittest.TestCase): | |||
| 297 | dst_shape = (-1, 256, 512) | 339 | dst_shape = (-1, 256, 512) |
| 298 | 340 | ||
| 299 | with self.assertRaises(ValueError): | 341 | with self.assertRaises(ValueError): |
| 300 | - _, _ = op.infer_layout((x_layout,), (dst_shape, src_shape)) | 342 | + _, _ = self._infer_reshape(op, x_layout, dst_shape, src_shape) |
| 301 | 343 | ||
| 302 | @patch("hyper_parallel.core.dtensor.device_mesh.platform") | 344 | @patch("hyper_parallel.core.dtensor.device_mesh.platform") |
| 303 | def test_torch_reshape_basic_split(self, mock_platform): | 345 | def test_torch_reshape_basic_split(self, mock_platform): |
| @@ -312,7 +354,7 @@ class TestParallelReshape(unittest.TestCase): | |||
| 312 | src_shape = (32, 64) | 354 | src_shape = (32, 64) |
| 313 | dst_shape = (32, 16, 4) | 355 | dst_shape = (32, 16, 4) |
| 314 | 356 | ||
| 315 | - output_layout, local_dst_shape = op_torch.infer_layout((x_layout,), (dst_shape, src_shape)) | 357 | + output_layout, local_dst_shape = self._infer_reshape(op_torch, x_layout, dst_shape, src_shape) |
| 316 | 358 | ||
| 317 | expected_map = (1, 0, -1) | 359 | expected_map = (1, 0, -1) |
| 318 | 360 | ||
| @@ -338,7 +380,7 @@ class TestParallelReshape(unittest.TestCase): | |||
| 338 | src_shape = (32, 64) | 380 | src_shape = (32, 64) |
| 339 | dst_shape = (32, -1) | 381 | dst_shape = (32, -1) |
| 340 | 382 | ||
| 341 | - output_layout, local_dst_shape = op_torch.infer_layout((x_layout,), (dst_shape, src_shape)) | 383 | + output_layout, local_dst_shape = self._infer_reshape(op_torch, x_layout, dst_shape, src_shape) |
| 342 | 384 | ||
| 343 | expected_map = (1, 0) | 385 | expected_map = (1, 0) |
| 344 | assert output_layout.tensor_map == expected_map | 386 | assert output_layout.tensor_map == expected_map |
| @@ -360,11 +402,11 @@ class TestParallelReshape(unittest.TestCase): | |||
| 360 | src_shape = (16, 4, 8) | 402 | src_shape = (16, 4, 8) |
| 361 | dst_shape = (64, 8) | 403 | dst_shape = (64, 8) |
| 362 | 404 | ||
| 363 | - output_layout, _ = op_torch.infer_layout((x_layout,), (dst_shape, src_shape)) | 405 | + output_layout, _ = self._infer_reshape(op_torch, x_layout, dst_shape, src_shape) |
| 364 | assert output_layout.tensor_map == (1, 0) | 406 | assert output_layout.tensor_map == (1, 0) |
| 365 | 407 | ||
| 366 | dst_shape_2 = (16, 2, 2, 8) | 408 | dst_shape_2 = (16, 2, 2, 8) |
| 367 | - output_layout_2, _ = op_torch.infer_layout((x_layout,), (dst_shape_2, src_shape)) | 409 | + output_layout_2, _ = self._infer_reshape(op_torch, x_layout, dst_shape_2, src_shape) |
| 368 | assert output_layout_2.tensor_map == (1, -1, -1, 0) | 410 | assert output_layout_2.tensor_map == (1, -1, -1, 0) |
| 369 | 411 | ||
| 370 | @patch("hyper_parallel.core.dtensor.device_mesh.platform") | 412 | @patch("hyper_parallel.core.dtensor.device_mesh.platform") |
| @@ -379,8 +421,8 @@ class TestParallelReshape(unittest.TestCase): | |||
| 379 | 421 | ||
| 380 | dst_shape = (32, 64) | 422 | dst_shape = (32, 64) |
| 381 | 423 | ||
| 382 | - with self.assertRaisesRegex(ValueError, "reshape requires output shape and input shape."): | 424 | + with self.assertRaisesRegex(ValueError, "cache_values length"): |
| 383 | - op_torch.infer_layout((x_layout,), (dst_shape,)) | 425 | + op_torch.infer_layout([x_layout, dst_shape]) |
| 384 | 426 | ||
| 385 | @patch("hyper_parallel.core.dtensor.device_mesh.platform") | 427 | @patch("hyper_parallel.core.dtensor.device_mesh.platform") |
| 386 | def test_torch_reshape_fail_mismatch_total_size(self, mock_platform): | 428 | def test_torch_reshape_fail_mismatch_total_size(self, mock_platform): |
| @@ -396,7 +438,7 @@ class TestParallelReshape(unittest.TestCase): | |||
| 396 | dst_shape = (20, 20) | 438 | dst_shape = (20, 20) |
| 397 | 439 | ||
| 398 | with self.assertRaises(ValueError): | 440 | with self.assertRaises(ValueError): |
| 399 | - op_torch.infer_layout((x_layout,), (dst_shape, src_shape)) | 441 | + self._infer_reshape(op_torch, x_layout, dst_shape, src_shape) |
| 400 | 442 | ||
| 401 | @patch("hyper_parallel.core.dtensor.device_mesh.platform") | 443 | @patch("hyper_parallel.core.dtensor.device_mesh.platform") |
| 402 | def test_view_layout_flatten_contiguous(self, mock_platform): | 444 | def test_view_layout_flatten_contiguous(self, mock_platform): |
| @@ -412,7 +454,7 @@ class TestParallelReshape(unittest.TestCase): | |||
| 412 | src_shape = (32, 16, 8) | 454 | src_shape = (32, 16, 8) |
| 413 | dst_shape = (32, 128) | 455 | dst_shape = (32, 128) |
| 414 | 456 | ||
| 415 | - output_layout, local_dst_shape = op_view.infer_layout((x_layout,), (dst_shape, src_shape)) | 457 | + output_layout, local_dst_shape = self._infer_reshape(op_view, x_layout, dst_shape, src_shape) |
| 416 | 458 | ||
| 417 | expected_map = (1, 0) | 459 | expected_map = (1, 0) |
| 418 | assert output_layout.tensor_map == expected_map, ( | 460 | assert output_layout.tensor_map == expected_map, ( |
| @@ -436,7 +478,7 @@ class TestParallelReshape(unittest.TestCase): | |||
| 436 | src_shape = (32, 128) | 478 | src_shape = (32, 128) |
| 437 | dst_shape = (32, 16, 8) | 479 | dst_shape = (32, 16, 8) |
| 438 | 480 | ||
| 439 | - output_layout, local_dst_shape = op_view.infer_layout((x_layout,), (dst_shape, src_shape)) | 481 | + output_layout, local_dst_shape = self._infer_reshape(op_view, x_layout, dst_shape, src_shape) |
| 440 | 482 | ||
| 441 | expected_map = (1, 0, -1) | 483 | expected_map = (1, 0, -1) |
| 442 | assert output_layout.tensor_map == expected_map, ( | 484 | assert output_layout.tensor_map == expected_map, ( |
| @@ -459,7 +501,7 @@ class TestParallelReshape(unittest.TestCase): | |||
| 459 | src_shape = (8, 4, 16) | 501 | src_shape = (8, 4, 16) |
| 460 | dst_shape = (8, -1) | 502 | dst_shape = (8, -1) |
| 461 | 503 | ||
| 462 | - output_layout, local_dst_shape = op_view.infer_layout((x_layout,), (dst_shape, src_shape)) | 504 | + output_layout, local_dst_shape = self._infer_reshape(op_view, x_layout, dst_shape, src_shape) |
| 463 | 505 | ||
| 464 | expected_map = (1, 0) | 506 | expected_map = (1, 0) |
| 465 | assert output_layout.tensor_map == expected_map | 507 | assert output_layout.tensor_map == expected_map |
| @@ -481,7 +523,7 @@ class TestParallelReshape(unittest.TestCase): | |||
| 481 | dst_shape = (10, 30) | 523 | dst_shape = (10, 30) |
| 482 | 524 | ||
| 483 | with self.assertRaisesRegex(ValueError, "total elements number"): | 525 | with self.assertRaisesRegex(ValueError, "total elements number"): |
| 484 | - op_view.infer_layout((x_layout,), (dst_shape, src_shape)) | 526 | + self._infer_reshape(op_view, x_layout, dst_shape, src_shape) |
| 485 | 527 | ||
| 486 | @patch("hyper_parallel.core.dtensor.device_mesh.platform") | 528 | @patch("hyper_parallel.core.dtensor.device_mesh.platform") |
| 487 | def test_reshape_with_partial_basic(self, mock_platform): | 529 | def test_reshape_with_partial_basic(self, mock_platform): |
| @@ -498,7 +540,7 @@ class TestParallelReshape(unittest.TestCase): | |||
| 498 | src_shape = (8, 16) | 540 | src_shape = (8, 16) |
| 499 | dst_shape = (4, 4, 8) | 541 | dst_shape = (4, 4, 8) |
| 500 | 542 | ||
| 501 | - output_layout, _ = op.infer_layout((x_layout,), (dst_shape, src_shape)) | 543 | + output_layout, _ = self._infer_reshape(op, x_layout, dst_shape, src_shape) |
| 502 | 544 | ||
| 503 | assert output_layout.is_partial(), "Output layout should preserve partial status" | 545 | assert output_layout.is_partial(), "Output layout should preserve partial status" |
| 504 | assert output_layout.get_partial_by_dev_id("dp") == "sum", ( | 546 | assert output_layout.get_partial_by_dev_id("dp") == "sum", ( |
| @@ -518,7 +560,7 @@ class TestParallelReshape(unittest.TestCase): | |||
| 518 | src_shape = (8, 16) | 560 | src_shape = (8, 16) |
| 519 | dst_shape = (1, -1, 16) | 561 | dst_shape = (1, -1, 16) |
| 520 | 562 | ||
| 521 | - output_layout, _ = op.infer_layout((x_layout,), (dst_shape, src_shape)) | 563 | + output_layout, _ = self._infer_reshape(op, x_layout, dst_shape, src_shape) |
| 522 | 564 | ||
| 523 | assert output_layout.is_partial(), "Output layout should preserve partial status" | 565 | assert output_layout.is_partial(), "Output layout should preserve partial status" |
| 524 | assert output_layout.get_partial_by_dev_id("mp") == "sum", ( | 566 | assert output_layout.get_partial_by_dev_id("mp") == "sum", ( |
| @@ -538,7 +580,7 @@ class TestParallelReshape(unittest.TestCase): | |||
| 538 | src_shape = (4, 8) | 580 | src_shape = (4, 8) |
| 539 | dst_shape = (2, 16) | 581 | dst_shape = (2, 16) |
| 540 | 582 | ||
| 541 | - output_layout, _ = op.infer_layout((x_layout,), (dst_shape, src_shape)) | 583 | + output_layout, _ = self._infer_reshape(op, x_layout, dst_shape, src_shape) |
| 542 | 584 | ||
| 543 | assert output_layout.is_partial(), "Output layout should preserve partial status" | 585 | assert output_layout.is_partial(), "Output layout should preserve partial status" |
| 544 | assert output_layout.get_partial_by_dev_id("dp") == "sum" | 586 | assert output_layout.get_partial_by_dev_id("dp") == "sum" |
| @@ -558,9 +600,7 @@ class TestParallelReshape(unittest.TestCase): | |||
| 558 | src_shape = (4, 8) | 600 | src_shape = (4, 8) |
| 559 | dst_shape = (2, 16) | 601 | dst_shape = (2, 16) |
| 560 | 602 | ||
| 561 | - output_layout, _ = op_torch.infer_layout( | 603 | + output_layout, _ = self._infer_reshape(op_torch, x_layout, dst_shape, src_shape) |
| 562 | - (x_layout,), (dst_shape, src_shape) | ||
| 563 | - ) | ||
| 564 | 604 | ||
| 565 | assert output_layout.is_partial(), "Output layout should preserve partial status" | 605 | assert output_layout.is_partial(), "Output layout should preserve partial status" |
| 566 | assert output_layout.get_partial_by_dev_id("dp") == "sum" | 606 | assert output_layout.get_partial_by_dev_id("dp") == "sum" |
| @@ -577,7 +617,7 @@ class TestParallelReshape(unittest.TestCase): | |||
| 577 | src_shape = (2,) | 617 | src_shape = (2,) |
| 578 | dst_shape = (-1,) | 618 | dst_shape = (-1,) |
| 579 | 619 | ||
| 580 | - output_layout, local_dst_shape = op_torch.infer_layout((x_layout,), (dst_shape, src_shape)) | 620 | + output_layout, local_dst_shape = self._infer_reshape(op_torch, x_layout, dst_shape, src_shape) |
| 581 | expected_map = (-1,) | 621 | expected_map = (-1,) |
| 582 | assert output_layout.tensor_map == expected_map, ( | 622 | assert output_layout.tensor_map == expected_map, ( |
| 583 | f"Replicate 1D reshape to self failed. Expected {expected_map}, " | 623 | f"Replicate 1D reshape to self failed. Expected {expected_map}, " |
| @@ -603,9 +643,7 @@ class TestParallelReshape(unittest.TestCase): | |||
| 603 | src_shape = (4, 8) | 643 | src_shape = (4, 8) |
| 604 | dst_shape = (2, 16) | 644 | dst_shape = (2, 16) |
| 605 | 645 | ||
| 606 | - output_layout, _ = op_view.infer_layout( | 646 | + output_layout, _ = self._infer_reshape(op_view, x_layout, dst_shape, src_shape) |
| 607 | - (x_layout,), (dst_shape, src_shape) | ||
| 608 | - ) | ||
| 609 | 647 | ||
| 610 | assert output_layout.is_partial(), "Output layout should preserve partial status" | 648 | assert output_layout.is_partial(), "Output layout should preserve partial status" |
| 611 | assert output_layout.get_partial_by_dev_id("dp") == "sum" | 649 | assert output_layout.get_partial_by_dev_id("dp") == "sum" |
| @@ -15,10 +15,10 @@ | |||
| 15 | """test ut for parallel slice infer_layout""" | 15 | """test ut for parallel slice infer_layout""" |
| 16 | import os | 16 | import os |
| 17 | import unittest | 17 | import unittest |
| 18 | -from unittest.mock import patch | 18 | +from unittest.mock import MagicMock, patch |
| 19 | import numpy as np | 19 | import numpy as np |
| 20 | 20 | ||
| 21 | -from hyper_parallel.core.dtensor.dtensor import _build_layout | 21 | +from hyper_parallel.core.dtensor.dtensor import _build_layout, _LAYOUT_CACHE |
| 22 | from hyper_parallel.core.dtensor.placement_types import Shard, Replicate | 22 | from hyper_parallel.core.dtensor.placement_types import Shard, Replicate |
| 23 | from hyper_parallel.core.shard.ops.parallel_slice import SliceDistributedOp | 23 | from hyper_parallel.core.shard.ops.parallel_slice import SliceDistributedOp |
| 24 | from hyper_parallel.core.dtensor.device_mesh import ( | 24 | from hyper_parallel.core.dtensor.device_mesh import ( |
| @@ -40,11 +40,13 @@ class TestParallelSlice(unittest.TestCase): | |||
| 40 | """ | 40 | """ |
| 41 | EXISTING_COMM_GROUPS.clear() | 41 | EXISTING_COMM_GROUPS.clear() |
| 42 | _DEVICE_MESH_MAP.clear() | 42 | _DEVICE_MESH_MAP.clear() |
| 43 | + _LAYOUT_CACHE.clear() | ||
| 43 | 44 | ||
| 44 | def tearDown(self): | 45 | def tearDown(self): |
| 45 | """Clean up after each test method.""" | 46 | """Clean up after each test method.""" |
| 46 | EXISTING_COMM_GROUPS.clear() | 47 | EXISTING_COMM_GROUPS.clear() |
| 47 | _DEVICE_MESH_MAP.clear() | 48 | _DEVICE_MESH_MAP.clear() |
| 49 | + _LAYOUT_CACHE.clear() | ||
| 48 | 50 | ||
| 49 | def _setup_mock_platform(self, mock_platform, platform_type=None, world_size=8): | 51 | def _setup_mock_platform(self, mock_platform, platform_type=None, world_size=8): |
| 50 | """Configure common mock-platform attributes used across tests. | 52 | """Configure common mock-platform attributes used across tests. |
| @@ -82,18 +84,22 @@ class TestParallelSlice(unittest.TestCase): | |||
| 82 | mesh = self._make_2x4_mesh(mock_platform) | 84 | mesh = self._make_2x4_mesh(mock_platform) |
| 83 | x_layout = _build_layout(mesh, (Shard(0), Replicate()), 2) | 85 | x_layout = _build_layout(mesh, (Shard(0), Replicate()), 2) |
| 84 | 86 | ||
| 85 | - output_info = op.infer_layout((x_layout,), ((0, 0), (8, 4), (8, 8))) | 87 | + cache_values = [x_layout, (0, 0), (8, 4), (8, 8)] |
| 86 | - assert x_layout == output_info[0] | 88 | + output_layouts, extra_info = op.infer_layout(cache_values) |
| 87 | - assert output_info[1] == (0, 0) | 89 | + assert output_layouts == (x_layout,) |
| 88 | - assert output_info[2] == (4, 4) | 90 | + assert extra_info[0] == (0, 0) |
| 91 | + assert extra_info[1] == (4, 4) | ||
| 89 | 92 | ||
| 90 | - # Since `get_expand_impl` is not overridden, it returns None by default. | 93 | + calls = [] |
| 91 | - # The same applies to other test classes, so it is unnecessary to test its return value. | 94 | + |
| 92 | - assert op.get_expand_impl(None, output_info, (x_layout,), | 95 | + def fake_slice(x, begin, end): |
| 93 | - ((0, 0), (8, 4), (8, 8))) is None, ( | 96 | + calls.append((x, begin, end)) |
| 94 | - f"get_expand_impl test failed. Expected None, " | 97 | + return "sliced" |
| 95 | - f"got {op.get_expand_impl(None, output_info, (x_layout,), ((0, 0), (8, 4), (8, 8)))}" | 98 | + |
| 96 | - ) | 99 | + impl = op.get_expand_impl(fake_slice, (output_layouts, extra_info), cache_values) |
| 100 | + assert impl is not None | ||
| 101 | + assert impl("local_x", (0, 0), (8, 4)) == "sliced" | ||
| 102 | + assert calls == [("local_x", (0, 0), (4, 4))] | ||
| 97 | 103 | ||
| 98 | @patch("hyper_parallel.core.dtensor.device_mesh.platform") | 104 | @patch("hyper_parallel.core.dtensor.device_mesh.platform") |
| 99 | def test_slice_layout_2(self, mock_platform): | 105 | def test_slice_layout_2(self, mock_platform): |
| @@ -106,8 +112,28 @@ class TestParallelSlice(unittest.TestCase): | |||
| 106 | x_layout = _build_layout(mesh, (Shard(0), Replicate()), 2) | 112 | x_layout = _build_layout(mesh, (Shard(0), Replicate()), 2) |
| 107 | 113 | ||
| 108 | with self.assertRaises(ValueError): | 114 | with self.assertRaises(ValueError): |
| 109 | - _ = op.infer_layout((x_layout,), ((0, 0), (4, 4), (8, 8))) | 115 | + _ = op.infer_layout([x_layout, (0, 0), (4, 4), (8, 8)]) |
| 110 | 116 | ||
| 117 | + @patch("hyper_parallel.core.dtensor.device_mesh.platform") | ||
| 118 | + def test_slice_preprocess(self, mock_platform): | ||
| 119 | + """ | ||
| 120 | + Feature: Slice preprocess | ||
| 121 | + Description: Verify local args and cache values are built in new dispatch format | ||
| 122 | + Expectation: Success | ||
| 123 | + """ | ||
| 124 | + mesh = self._make_2x4_mesh(mock_platform) | ||
| 125 | + x_layout = _build_layout(mesh, (Shard(0), Replicate()), 2) | ||
| 126 | + local_tensor = MagicMock() | ||
| 127 | + input_tensor = MagicMock() | ||
| 128 | + input_tensor.layout = x_layout | ||
| 129 | + input_tensor.shape = (8, 8) | ||
| 130 | + input_tensor.to_local.return_value = local_tensor | ||
| 131 | + | ||
| 132 | + local_args, local_kwargs, cache_values = op.preprocess((input_tensor, (0, 0), (8, 4)), {}) | ||
| 133 | + | ||
| 134 | + assert local_args == (local_tensor, (0, 0), (8, 4)) | ||
| 135 | + assert not local_kwargs | ||
| 136 | + assert cache_values == [x_layout, (0, 0), (8, 4), (8, 8)] | ||
| 111 | 137 | ||
| 112 | if __name__ == "__main__": | 138 | if __name__ == "__main__": |
| 113 | unittest.main() | 139 | unittest.main() |
| @@ -13,12 +13,11 @@ | |||
| 13 | # limitations under the License. | 13 | # limitations under the License. |
| 14 | # ============================================================================ | 14 | # ============================================================================ |
| 15 | """parallel_squeeze test""" | 15 | """parallel_squeeze test""" |
| 16 | -import os | ||
| 17 | import unittest | 16 | import unittest |
| 18 | -from unittest.mock import patch | 17 | +from unittest.mock import MagicMock, patch |
| 19 | import numpy as np | 18 | import numpy as np |
| 20 | 19 | ||
| 21 | -from hyper_parallel.core.dtensor.dtensor import _build_layout | 20 | +from hyper_parallel.core.dtensor.dtensor import _build_layout, _LAYOUT_CACHE |
| 22 | from hyper_parallel.core.dtensor.placement_types import Shard, Replicate | 21 | from hyper_parallel.core.dtensor.placement_types import Shard, Replicate |
| 23 | from hyper_parallel.core.shard.ops.parallel_squeeze import SqueezeDistributedOp | 22 | from hyper_parallel.core.shard.ops.parallel_squeeze import SqueezeDistributedOp |
| 24 | from hyper_parallel.core.dtensor.device_mesh import ( | 23 | from hyper_parallel.core.dtensor.device_mesh import ( |
| @@ -38,12 +37,14 @@ class TestParallelSqueeze(unittest.TestCase): | |||
| 38 | """ | 37 | """ |
| 39 | EXISTING_COMM_GROUPS.clear() | 38 | EXISTING_COMM_GROUPS.clear() |
| 40 | _DEVICE_MESH_MAP.clear() | 39 | _DEVICE_MESH_MAP.clear() |
| 40 | + _LAYOUT_CACHE.clear() | ||
| 41 | self.op = SqueezeDistributedOp("Squeeze") | 41 | self.op = SqueezeDistributedOp("Squeeze") |
| 42 | 42 | ||
| 43 | def tearDown(self): | 43 | def tearDown(self): |
| 44 | """Clean up after each test method.""" | 44 | """Clean up after each test method.""" |
| 45 | EXISTING_COMM_GROUPS.clear() | 45 | EXISTING_COMM_GROUPS.clear() |
| 46 | _DEVICE_MESH_MAP.clear() | 46 | _DEVICE_MESH_MAP.clear() |
| 47 | + _LAYOUT_CACHE.clear() | ||
| 47 | 48 | ||
| 48 | def _setup_mock_platform(self, mock_platform, platform_type=None, world_size=8): | 49 | def _setup_mock_platform(self, mock_platform, platform_type=None, world_size=8): |
| 49 | """Configure common mock-platform attributes used across tests. | 50 | """Configure common mock-platform attributes used across tests. |
| @@ -73,12 +74,9 @@ class TestParallelSqueeze(unittest.TestCase): | |||
| 73 | 74 | ||
| 74 | def _run_scenario(self, x_layout, expected_map, axis, input_shape): | 75 | def _run_scenario(self, x_layout, expected_map, axis, input_shape): |
| 75 | """Infer layout of Squeeze operator""" | 76 | """Infer layout of Squeeze operator""" |
| 76 | - if axis is not None: | 77 | + cache_values = [x_layout, axis, input_shape] |
| 77 | - extra_args = [axis, [input_shape]] | 78 | + output_layouts, extra_info = self.op.infer_layout(cache_values) |
| 78 | - else: | 79 | + output_layout = output_layouts[0] |
| 79 | - extra_args = [[input_shape]] | ||
| 80 | - | ||
| 81 | - output_layout = self.op.infer_layout((x_layout,), extra_args) | ||
| 82 | actual_map = output_layout.tensor_map | 80 | actual_map = output_layout.tensor_map |
| 83 | assert actual_map == expected_map, ( | 81 | assert actual_map == expected_map, ( |
| 84 | f"Squeeze failed. Expected {expected_map}, got {actual_map}" | 82 | f"Squeeze failed. Expected {expected_map}, got {actual_map}" |
| @@ -86,11 +84,35 @@ class TestParallelSqueeze(unittest.TestCase): | |||
| 86 | 84 | ||
| 87 | # Since `get_expand_impl` is not overridden, it returns None by default. | 85 | # Since `get_expand_impl` is not overridden, it returns None by default. |
| 88 | # The same applies to other test classes, so it is unnecessary to test its return value. | 86 | # The same applies to other test classes, so it is unnecessary to test its return value. |
| 89 | - assert self.op.get_expand_impl(None, output_layout, (x_layout,), extra_args) is None, ( | 87 | + infer_result = (output_layouts, extra_info) |
| 88 | + assert self.op.get_expand_impl(None, infer_result, cache_values) is None, ( | ||
| 90 | f"get_expand_impl test failed. Expected None, " | 89 | f"get_expand_impl test failed. Expected None, " |
| 91 | - f"got {self.op.get_expand_impl(None, output_layout, (x_layout,), extra_args)}" | 90 | + f"got {self.op.get_expand_impl(None, infer_result, cache_values)}" |
| 92 | ) | 91 | ) |
| 93 | 92 | ||
| 93 | + def test_squeeze_preprocess(self): | ||
| 94 | + """ | ||
| 95 | + Feature: preprocess. | ||
| 96 | + Description: convert DTensor inputs to local args and collect layout, axis, and shape cache. | ||
| 97 | + Expectation: local args are positional and kwargs are empty. | ||
| 98 | + """ | ||
| 99 | + x_layout = MagicMock() | ||
| 100 | + local_tensor = MagicMock() | ||
| 101 | + input_tensor = MagicMock() | ||
| 102 | + input_tensor.layout = x_layout | ||
| 103 | + input_tensor.shape = (8, 1, 4) | ||
| 104 | + input_tensor.to_local.return_value = local_tensor | ||
| 105 | + | ||
| 106 | + local_args, local_kwargs, cache_values = self.op.preprocess((input_tensor,), {}) | ||
| 107 | + assert local_args == (local_tensor,) | ||
| 108 | + assert not local_kwargs | ||
| 109 | + assert cache_values == [x_layout, None, (8, 1, 4)] | ||
| 110 | + | ||
| 111 | + local_args, local_kwargs, cache_values = self.op.preprocess((input_tensor, 1), {}) | ||
| 112 | + assert local_args == (local_tensor, 1) | ||
| 113 | + assert not local_kwargs | ||
| 114 | + assert cache_values == [x_layout, 1, (8, 1, 4)] | ||
| 115 | + | ||
| 94 | @patch("hyper_parallel.core.dtensor.device_mesh.platform") | 116 | @patch("hyper_parallel.core.dtensor.device_mesh.platform") |
| 95 | def test_squeeze_data_parallel_1(self, mock_platform): | 117 | def test_squeeze_data_parallel_1(self, mock_platform): |
| 96 | """ | 118 | """ |
| @@ -301,13 +323,16 @@ class TestParallelSqueeze(unittest.TestCase): | |||
| 301 | x_layout = _build_layout(mesh, (Shard(0), Shard(1), Shard(2)), 3) | 323 | x_layout = _build_layout(mesh, (Shard(0), Shard(1), Shard(2)), 3) |
| 302 | 324 | ||
| 303 | with self.assertRaises(ValueError): | 325 | with self.assertRaises(ValueError): |
| 304 | - self.op.infer_layout((x_layout,), []) | 326 | + self.op.infer_layout([]) |
| 305 | 327 | ||
| 306 | with self.assertRaises(ValueError): | 328 | with self.assertRaises(ValueError): |
| 307 | - self.op.infer_layout((x_layout,), [1]) | 329 | + self.op.infer_layout([x_layout, 1]) |
| 308 | 330 | ||
| 309 | with self.assertRaises(ValueError): | 331 | with self.assertRaises(ValueError): |
| 310 | - self.op.infer_layout((x_layout,), None) | 332 | + self.op.infer_layout([x_layout, None, None]) |
| 333 | + | ||
| 334 | + with self.assertRaises(ValueError): | ||
| 335 | + self.op.infer_layout([x_layout, None, 1]) | ||
| 311 | 336 | ||
| 312 | @patch("hyper_parallel.core.dtensor.device_mesh.platform") | 337 | @patch("hyper_parallel.core.dtensor.device_mesh.platform") |
| 313 | def test_squeeze_scalar_with_axis_17(self, mock_platform): | 338 | def test_squeeze_scalar_with_axis_17(self, mock_platform): |
| @@ -321,7 +346,7 @@ class TestParallelSqueeze(unittest.TestCase): | |||
| 321 | input_shape = [] | 346 | input_shape = [] |
| 322 | 347 | ||
| 323 | with self.assertRaises(ValueError): | 348 | with self.assertRaises(ValueError): |
| 324 | - self.op.infer_layout((x_layout,), [0, [input_shape]]) | 349 | + self.op.infer_layout([x_layout, 0, input_shape]) |
| 325 | 350 | ||
| 326 | 351 | ||
| 327 | if __name__ == "__main__": | 352 | if __name__ == "__main__": |
| @@ -13,13 +13,12 @@ | |||
| 13 | # limitations under the License. | 13 | # limitations under the License. |
| 14 | # ============================================================================ | 14 | # ============================================================================ |
| 15 | """Unit tests for TupleElementWiseDistributedOp""" | 15 | """Unit tests for TupleElementWiseDistributedOp""" |
| 16 | -import os | ||
| 17 | import unittest | 16 | import unittest |
| 18 | -from unittest.mock import patch | 17 | +from unittest.mock import MagicMock, patch |
| 19 | 18 | ||
| 20 | import numpy as np | 19 | import numpy as np |
| 21 | 20 | ||
| 22 | -from hyper_parallel.core.dtensor.dtensor import _build_layout | 21 | +from hyper_parallel.core.dtensor.dtensor import _build_layout, _LAYOUT_CACHE |
| 23 | from hyper_parallel.core.dtensor.placement_types import Shard, Replicate | 22 | from hyper_parallel.core.dtensor.placement_types import Shard, Replicate |
| 24 | from hyper_parallel.core.dtensor.device_mesh import init_device_mesh, _DEVICE_MESH_MAP | 23 | from hyper_parallel.core.dtensor.device_mesh import init_device_mesh, _DEVICE_MESH_MAP |
| 25 | from hyper_parallel.core.shard.ops.parallel_tuple_elementwise import TupleElementWiseDistributedOp | 24 | from hyper_parallel.core.shard.ops.parallel_tuple_elementwise import TupleElementWiseDistributedOp |
| @@ -42,6 +41,16 @@ def _make_mesh(mock_platform, mesh_shape, mesh_dim_names): | |||
| 42 | ) | 41 | ) |
| 43 | 42 | ||
| 44 | 43 | ||
| 44 | +def _infer_layout(op, cache_values): | ||
| 45 | + """Infer layouts with the new preprocess/infer_layout protocol.""" | ||
| 46 | + result = op.infer_layout(cache_values) | ||
| 47 | + if result is None: | ||
| 48 | + return None | ||
| 49 | + output_layouts, extra_info = result | ||
| 50 | + assert extra_info is None | ||
| 51 | + return output_layouts | ||
| 52 | + | ||
| 53 | + | ||
| 45 | class TestTupleElementWiseDistributedOpInferLayout(unittest.TestCase): | 54 | class TestTupleElementWiseDistributedOpInferLayout(unittest.TestCase): |
| 46 | """ | 55 | """ |
| 47 | Feature: TupleElementWiseDistributedOp.infer_layout | 56 | Feature: TupleElementWiseDistributedOp.infer_layout |
| @@ -52,16 +61,18 @@ class TestTupleElementWiseDistributedOpInferLayout(unittest.TestCase): | |||
| 52 | def setUp(self): | 61 | def setUp(self): |
| 53 | EXISTING_COMM_GROUPS.clear() | 62 | EXISTING_COMM_GROUPS.clear() |
| 54 | _DEVICE_MESH_MAP.clear() | 63 | _DEVICE_MESH_MAP.clear() |
| 64 | + _LAYOUT_CACHE.clear() | ||
| 55 | 65 | ||
| 56 | def tearDown(self): | 66 | def tearDown(self): |
| 57 | EXISTING_COMM_GROUPS.clear() | 67 | EXISTING_COMM_GROUPS.clear() |
| 58 | _DEVICE_MESH_MAP.clear() | 68 | _DEVICE_MESH_MAP.clear() |
| 69 | + _LAYOUT_CACHE.clear() | ||
| 59 | 70 | ||
| 60 | @patch("hyper_parallel.core.dtensor.device_mesh.platform") | 71 | @patch("hyper_parallel.core.dtensor.device_mesh.platform") |
| 61 | def test_empty_layouts_returns_none(self, mock_platform): | 72 | def test_empty_layouts_returns_none(self, mock_platform): |
| 62 | """Empty input layouts returns None.""" | 73 | """Empty input layouts returns None.""" |
| 63 | op = TupleElementWiseDistributedOp("tuple_elementwise_empty") | 74 | op = TupleElementWiseDistributedOp("tuple_elementwise_empty") |
| 64 | - result = op.infer_layout(()) | 75 | + result = op.infer_layout([]) |
| 65 | self.assertIsNone(result) | 76 | self.assertIsNone(result) |
| 66 | 77 | ||
| 67 | @patch("hyper_parallel.core.dtensor.device_mesh.platform") | 78 | @patch("hyper_parallel.core.dtensor.device_mesh.platform") |
| @@ -69,10 +80,10 @@ class TestTupleElementWiseDistributedOpInferLayout(unittest.TestCase): | |||
| 69 | """Single input layout is returned unchanged as the layouts object.""" | 80 | """Single input layout is returned unchanged as the layouts object.""" |
| 70 | mesh = _make_mesh(mock_platform, (2,), ("dp",)) | 81 | mesh = _make_mesh(mock_platform, (2,), ("dp",)) |
| 71 | layout = _build_layout(mesh, (Replicate(),), 2) | 82 | layout = _build_layout(mesh, (Replicate(),), 2) |
| 72 | - inputs = (layout,) | ||
| 73 | op = TupleElementWiseDistributedOp("tuple_elementwise_single") | 83 | op = TupleElementWiseDistributedOp("tuple_elementwise_single") |
| 74 | - result = op.infer_layout(inputs) | 84 | + result = _infer_layout(op, [layout]) |
| 75 | - self.assertIs(result, inputs) | 85 | + self.assertEqual(result, (layout,)) |
| 86 | + self.assertIsNot(result[0], layout) | ||
| 76 | 87 | ||
| 77 | @patch("hyper_parallel.core.dtensor.device_mesh.platform") | 88 | @patch("hyper_parallel.core.dtensor.device_mesh.platform") |
| 78 | def test_multiple_layouts_returned_as_is(self, mock_platform): | 89 | def test_multiple_layouts_returned_as_is(self, mock_platform): |
| @@ -80,10 +91,9 @@ class TestTupleElementWiseDistributedOpInferLayout(unittest.TestCase): | |||
| 80 | mesh = _make_mesh(mock_platform, (2, 2), ("dp", "mp")) | 91 | mesh = _make_mesh(mock_platform, (2, 2), ("dp", "mp")) |
| 81 | layout1 = _build_layout(mesh, (Replicate(), Replicate()), 2) | 92 | layout1 = _build_layout(mesh, (Replicate(), Replicate()), 2) |
| 82 | layout2 = _build_layout(mesh, (Shard(0), Replicate()), 2) | 93 | layout2 = _build_layout(mesh, (Shard(0), Replicate()), 2) |
| 83 | - inputs = (layout1, layout2) | ||
| 84 | op = TupleElementWiseDistributedOp("tuple_elementwise_multi") | 94 | op = TupleElementWiseDistributedOp("tuple_elementwise_multi") |
| 85 | - result = op.infer_layout(inputs) | 95 | + result = _infer_layout(op, [layout1, layout2]) |
| 86 | - self.assertIs(result, inputs) | 96 | + self.assertEqual(result, (layout1, layout2)) |
| 87 | 97 | ||
| 88 | @patch("hyper_parallel.core.dtensor.device_mesh.platform") | 98 | @patch("hyper_parallel.core.dtensor.device_mesh.platform") |
| 89 | def test_sharded_layout_returned(self, mock_platform): | 99 | def test_sharded_layout_returned(self, mock_platform): |
| @@ -91,17 +101,17 @@ class TestTupleElementWiseDistributedOpInferLayout(unittest.TestCase): | |||
| 91 | mesh = _make_mesh(mock_platform, (4,), ("dp",)) | 101 | mesh = _make_mesh(mock_platform, (4,), ("dp",)) |
| 92 | layout = _build_layout(mesh, (Shard(0),), 2) | 102 | layout = _build_layout(mesh, (Shard(0),), 2) |
| 93 | op = TupleElementWiseDistributedOp("tuple_elementwise_sharded") | 103 | op = TupleElementWiseDistributedOp("tuple_elementwise_sharded") |
| 94 | - result = op.infer_layout((layout,)) | 104 | + result = _infer_layout(op, [layout]) |
| 95 | self.assertEqual(result, (layout,)) | 105 | self.assertEqual(result, (layout,)) |
| 96 | 106 | ||
| 97 | @patch("hyper_parallel.core.dtensor.device_mesh.platform") | 107 | @patch("hyper_parallel.core.dtensor.device_mesh.platform") |
| 98 | - def test_extra_args_ignored(self, mock_platform): | 108 | + def test_none_layout_in_inputs(self, mock_platform): |
| 99 | - """extra_args parameter is ignored for tuple element-wise ops.""" | 109 | + """None layout in inputs is returned as part of the result.""" |
| 100 | mesh = _make_mesh(mock_platform, (2,), ("dp",)) | 110 | mesh = _make_mesh(mock_platform, (2,), ("dp",)) |
| 101 | layout = _build_layout(mesh, (Replicate(),), 2) | 111 | layout = _build_layout(mesh, (Replicate(),), 2) |
| 102 | - op = TupleElementWiseDistributedOp("tuple_elementwise_extra_args") | 112 | + op = TupleElementWiseDistributedOp("tuple_elementwise_none") |
| 103 | - result = op.infer_layout((layout,), extra_args=(42, "some_arg")) | 113 | + result = _infer_layout(op, [layout, None]) |
| 104 | - self.assertEqual(result, (layout,)) | 114 | + self.assertEqual(result, (layout, None)) |
| 105 | 115 | ||
| 106 | @patch("hyper_parallel.core.dtensor.device_mesh.platform") | 116 | @patch("hyper_parallel.core.dtensor.device_mesh.platform") |
| 107 | def test_three_layouts_returned(self, mock_platform): | 117 | def test_three_layouts_returned(self, mock_platform): |
| @@ -110,20 +120,19 @@ class TestTupleElementWiseDistributedOpInferLayout(unittest.TestCase): | |||
| 110 | layout1 = _build_layout(mesh, (Replicate(), Replicate(), Replicate()), 3) | 120 | layout1 = _build_layout(mesh, (Replicate(), Replicate(), Replicate()), 3) |
| 111 | layout2 = _build_layout(mesh, (Shard(0), Replicate(), Replicate()), 3) | 121 | layout2 = _build_layout(mesh, (Shard(0), Replicate(), Replicate()), 3) |
| 112 | layout3 = _build_layout(mesh, (Replicate(), Replicate(), Shard(2)), 3) | 122 | layout3 = _build_layout(mesh, (Replicate(), Replicate(), Shard(2)), 3) |
| 113 | - inputs = (layout1, layout2, layout3) | ||
| 114 | op = TupleElementWiseDistributedOp("tuple_elementwise_three") | 123 | op = TupleElementWiseDistributedOp("tuple_elementwise_three") |
| 115 | - result = op.infer_layout(inputs) | 124 | + result = _infer_layout(op, [layout1, layout2, layout3]) |
| 116 | - self.assertIs(result, inputs) | 125 | + self.assertEqual(result, (layout1, layout2, layout3)) |
| 117 | 126 | ||
| 118 | @patch("hyper_parallel.core.dtensor.device_mesh.platform") | 127 | @patch("hyper_parallel.core.dtensor.device_mesh.platform") |
| 119 | - def test_none_layout_in_inputs(self, mock_platform): | 128 | + def test_partial_input_raises_error(self, mock_platform): |
| 120 | - """None layout in inputs is returned as part of the result.""" | 129 | + """Partial input layout is rejected.""" |
| 121 | mesh = _make_mesh(mock_platform, (2,), ("dp",)) | 130 | mesh = _make_mesh(mock_platform, (2,), ("dp",)) |
| 122 | layout = _build_layout(mesh, (Replicate(),), 2) | 131 | layout = _build_layout(mesh, (Replicate(),), 2) |
| 123 | - inputs = (layout, None) | 132 | + layout.set_partial_by_dev_axis("dp", "sum") |
| 124 | - op = TupleElementWiseDistributedOp("tuple_elementwise_none") | 133 | + op = TupleElementWiseDistributedOp("tuple_elementwise_partial") |
| 125 | - result = op.infer_layout(inputs) | 134 | + with self.assertRaisesRegex(ValueError, "Partial status"): |
| 126 | - self.assertIs(result, inputs) | 135 | + op.infer_layout([layout]) |
| 127 | 136 | ||
| 128 | 137 | ||
| 129 | class TestTupleElementWiseInheritance(unittest.TestCase): | 138 | class TestTupleElementWiseInheritance(unittest.TestCase): |
| @@ -134,15 +143,36 @@ class TestTupleElementWiseInheritance(unittest.TestCase): | |||
| 134 | Expectation: Inherited methods behave as defined in DistributedOp. | 143 | Expectation: Inherited methods behave as defined in DistributedOp. |
| 135 | """ | 144 | """ |
| 136 | 145 | ||
| 137 | - def test_preprocess_returns_none(self): | 146 | + def test_preprocess_expands_tuple_and_unwraps_dtensors(self): |
| 138 | - """Inherited preprocess returns None (legacy fallback).""" | 147 | + """preprocess unwraps local tensors and expands tuple/list layouts for cache values.""" |
| 148 | + layout1 = object() | ||
| 149 | + layout2 = object() | ||
| 150 | + local1 = object() | ||
| 151 | + local2 = object() | ||
| 152 | + tensor1 = MagicMock() | ||
| 153 | + tensor1._layout = layout1 | ||
| 154 | + tensor1.layout = layout1 | ||
| 155 | + tensor1.to_local.return_value = local1 | ||
| 156 | + tensor2 = MagicMock() | ||
| 157 | + tensor2._layout = layout2 | ||
| 158 | + tensor2.layout = layout2 | ||
| 159 | + tensor2.to_local.return_value = local2 | ||
| 160 | + plain_slot = object() | ||
| 161 | + | ||
| 139 | op = TupleElementWiseDistributedOp("tuple_elementwise_preprocess") | 162 | op = TupleElementWiseDistributedOp("tuple_elementwise_preprocess") |
| 140 | - self.assertIsNone(op.preprocess((1, 2), {})) | 163 | + local_args, local_kwargs, cache_values = op.preprocess( |
| 164 | + ((tensor1, plain_slot, tensor2),), | ||
| 165 | + {"kw": tensor1} | ||
| 166 | + ) | ||
| 167 | + | ||
| 168 | + self.assertEqual(local_args, ((local1, plain_slot, local2),)) | ||
| 169 | + self.assertEqual(local_kwargs, {"kw": local1}) | ||
| 170 | + self.assertEqual(cache_values, [layout1, None, layout2, layout1]) | ||
| 141 | 171 | ||
| 142 | def test_get_expand_impl_returns_none(self): | 172 | def test_get_expand_impl_returns_none(self): |
| 143 | """Inherited get_expand_impl returns None.""" | 173 | """Inherited get_expand_impl returns None.""" |
| 144 | op = TupleElementWiseDistributedOp("tuple_elementwise_expand") | 174 | op = TupleElementWiseDistributedOp("tuple_elementwise_expand") |
| 145 | - self.assertIsNone(op.get_expand_impl(None, None, ())) | 175 | + self.assertIsNone(op.get_expand_impl(None, ((object(),), None), [object()])) |
| 146 | 176 | ||
| 147 | 177 | ||
| 148 | if __name__ == "__main__": | 178 | if __name__ == "__main__": |
| @@ -1067,7 +1067,7 @@ class TestDispatchNew(unittest.TestCase): | |||
| 1067 | return "py_output" | 1067 | return "py_output" |
| 1068 | 1068 | ||
| 1069 | d = object.__new__(OpDispatcher) | 1069 | d = object.__new__(OpDispatcher) |
| 1070 | - result = d._dispatch_new(mock_func, mock_dist_op, ([], {}, cache_values)) | 1070 | + result = d._dispatch_new(mock_func, mock_dist_op, None, ([], {}, cache_values)) |
| 1071 | self.assertEqual(result, "wrapped_output") | 1071 | self.assertEqual(result, "wrapped_output") |
| 1072 | mock_dist_op.infer_layout.assert_called_once() | 1072 | mock_dist_op.infer_layout.assert_called_once() |
| 1073 | mock_dist_op.wrap_output.assert_called_once() | 1073 | mock_dist_op.wrap_output.assert_called_once() |
| @@ -1101,7 +1101,7 @@ class TestDispatchNew(unittest.TestCase): | |||
| 1101 | return "should_not_be_called" | 1101 | return "should_not_be_called" |
| 1102 | 1102 | ||
| 1103 | d = object.__new__(OpDispatcher) | 1103 | d = object.__new__(OpDispatcher) |
| 1104 | - result = d._dispatch_new(mock_func, mock_dist_op, (["arg"], {}, cache_values)) | 1104 | + result = d._dispatch_new(mock_func, mock_dist_op, None, (["arg"], {}, cache_values)) |
| 1105 | mock_dist_op.infer_layout.assert_not_called() | 1105 | mock_dist_op.infer_layout.assert_not_called() |
| 1106 | mock_impl.assert_called_once() | 1106 | mock_impl.assert_called_once() |
| 1107 | 1107 | ||
| @@ -1683,10 +1683,11 @@ class TestWithLayoutInfer(unittest.TestCase): | |||
| 1683 | prim_obj = MagicMock() | 1683 | prim_obj = MagicMock() |
| 1684 | func = MagicMock(return_value=MagicMock()) | 1684 | func = MagicMock(return_value=MagicMock()) |
| 1685 | 1685 | ||
| 1686 | - # packed form: (prim_obj, "OpName", (arg1, arg2)) | 1686 | + # NOTE: aclop args are now normalized upstream by _normalize_aclop_args. |
| 1687 | - result = d._with_layout_infer(func, prim_obj, "UnpackOp", (10, 20)) | 1687 | + # Pass unpacked real args positionally and _packed_call via keyword. |
| 1688 | + result = d._with_layout_infer(func, 10, 20, _packed_call=(prim_obj, "UnpackOp")) | ||
| 1688 | 1689 | ||
| 1689 | - # op_impl should be called with the packed form | 1690 | + # op_impl should be called with the re-packed form |
| 1690 | func.assert_called_once() | 1691 | func.assert_called_once() |
| 1691 | mock_dtensor_cls.from_local.assert_called_once() | 1692 | mock_dtensor_cls.from_local.assert_called_once() |
| 1692 | 1693 | ||
仅建议：后续开发过程中如果觉得infer_layout里面的一些高频操作重复代码太多，可以抽取一些utils函数或者基类函数使用，这个慢慢总结就行。