Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / core / shard / _op_dispatch.py: 84%
377 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-07-13 05:07 +0800
1# Copyright 2025-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"""_op_dispatch"""
16import atexit
17import glob
18import importlib
19import os
20import sys
21import warnings
22from contextvars import ContextVar
23from itertools import chain
24from typing import Any, Dict, FrozenSet, List, Optional
26import yaml
28from hyper_parallel.core.shard.ops.parallel_ops_register import get_distributed_op
29from hyper_parallel.core.dtensor.dtensor import DTensor
30from hyper_parallel.core.dtensor.random import OffsetBasedRNGTracker, is_rng_supported_mesh
31from hyper_parallel.platform import get_platform
32from hyper_parallel.platform.platform import PlatformType
34from hyper_parallel.core.tensor_parallel._ce_op_registry import is_loss_parallel_op, is_decomposed_ce_op
35from hyper_parallel.core.tensor_parallel.loss_parallel import is_loss_parallel_active
36from hyper_parallel.core.tensor_parallel.loss_parallel_ops_common import _is_shard_on_last_dim
38platform = get_platform()
39Tensor = platform.Tensor
42def _apply_shard_offset_to_rng_args(args, offset_incr):
43 """Apply per-shard offset increment to seed/offset tensors in MindSpore random op args.
45 MindSpore random ops (e.g. ``randn_like_``) receive ``(seed, offset)`` as
46 explicit int64 scalar tensors from ``default_generator._step()`` in the
47 Python wrapper *before* the C++ dispatch triggers ``__fallback__``. By the
48 time ``_dispatch_random_op`` is called, the kernel will use whatever
49 ``(seed, offset)`` values are in the args—it does **not** read the
50 generator again. This function finds the offset tensor and adds the
51 per-rank offset increment so each shard gets a unique random stream.
53 The (seed, offset) pair is identified as the last two consecutive int64
54 0-dim tensors in *args* (scanning from the end to skip trailing dtype /
55 device arguments).
57 Args:
58 args: The list of local args for the random op.
59 offset_incr (int): Per-shard offset increment.
61 Returns:
62 list: Modified args with the offset tensor adjusted.
63 """
64 int64_dtype = platform.tensor_dtype.int64
65 last_int64_idx = -1
66 for i in range(len(args) - 1, -1, -1):
67 arg = args[i]
68 if isinstance(arg, Tensor) and arg.dtype == int64_dtype and arg.ndim == 0:
69 if last_int64_idx == i + 1:
70 offset_idx = i + 1
71 new_args = list(args)
72 new_offset = int(new_args[offset_idx].item()) + offset_incr
73 new_args[offset_idx] = platform.tensor([new_offset], dtype=int64_dtype).reshape(())
74 return new_args
75 last_int64_idx = i
76 return args
78_dtensor_dispatch_disabled: ContextVar[bool] = ContextVar('_dtensor_dispatch_disabled', default=False)
79_no_skip_ops: ContextVar[FrozenSet[str]] = ContextVar('_no_skip_ops', default=frozenset())
82def get_no_skip_ops() -> FrozenSet[str]:
83 """Return the set of op names that are exempt from SkipDTensorDispatch."""
84 return _no_skip_ops.get()
87def get_dtensor_dispatch() -> bool:
88 """
89 Get the current DTensor dispatch status.
91 Returns:
92 bool: True if DTensor dispatch is enabled, False otherwise.
93 """
94 return not _dtensor_dispatch_disabled.get()
97class LayoutCacheKey:
98 """Immutable layout cache key."""
99 __slots__ = ('_tuple', '_hash')
101 def __init__(self, layout_ids: List[str]):
102 self._tuple = tuple(layout_ids)
103 self._hash = hash(self._tuple)
105 @classmethod
106 def from_cache_values(cls, cache_values: list) -> "LayoutCacheKey":
107 """Build a LayoutCacheKey from a cache_values list.
109 Args:
110 cache_values (list): Mixed list of Layout objects (with compact_str) and raw scalars.
112 Returns:
113 LayoutCacheKey: Immutable key derived from the string representation of each value.
114 """
115 key_values = []
116 for v in cache_values:
117 if hasattr(v, 'compact_str'):
118 key_values.append(str(v.compact_str))
119 else:
120 key_values.append(str(v))
121 return cls(key_values)
123 def __eq__(self, other):
124 if not isinstance(other, LayoutCacheKey):
125 return False
126 return self._tuple == other._tuple
128 def __hash__(self):
129 return self._hash
131 def __repr__(self):
132 return f"LayoutCacheKey({self._tuple})"
135class LayoutCacheManager:
136 """
137 Cache layout in infer layout.
139 A singleton class that manages layout caches for distributed operations.
140 It caches the inferred layouts and operation implementations to avoid
141 redundant computation during repeated calls with the same input layouts.
142 """
143 _instance = None
145 def __init__(self):
146 self.layout_cache: Dict[str, Dict[LayoutCacheKey, Any]] = {}
147 atexit.register(self.clear_cache)
149 @classmethod
150 def get_instance(cls) -> "LayoutCacheManager":
151 """
152 Get the singleton instance of LayoutCacheManager.
154 Returns:
155 LayoutCacheManager: The singleton instance.
156 """
157 if cls._instance is None:
158 cls._instance = LayoutCacheManager()
159 return cls._instance
161 def get_layout_cache(self) -> Dict[str, Dict[LayoutCacheKey, Any]]:
162 """
163 Get the layout cache dictionary.
165 Returns:
166 Dict[str, Dict[LayoutCacheKey, Any]]: The nested dictionary mapping
167 operation names to their layout caches.
168 """
169 return self.layout_cache
171 @staticmethod
172 def distributed_op(op_name: str) -> Any:
173 """
174 Get the distributed operation implementation by name.
176 Args:
177 op_name (str): The name of the distributed operation.
179 Returns:
180 Any: The distributed operation class or implementation.
181 """
182 op = get_distributed_op(op_name)
183 return op
185 def clear_cache(self) -> None:
186 """
187 Clear all cached layouts.
189 This method is automatically registered with atexit to ensure
190 cache is cleared when the program exits.
191 """
192 self.layout_cache.clear()
195class OpDispatcher:
196 """
197 OpDispatcher
198 """
200 # Whitelisted ops that mutate args[0]'s storage in place. The dispatch bypass
201 # must return the original DTensor self for these, not the unwrapped local
202 # result, or it demotes a DTensor accumulator to a plain Tensor and breaks the
203 # next op that adds a DTensor to it (e.g. grad-accumulation `loss += micro_loss`).
204 # Class-level so it stays available on instances built via __new__ (e.g. tests).
205 _INPLACE_BYPASS_OPS = frozenset(
206 {"InplaceAddExt", "InplaceSubExt", "InplaceMul", "InplaceDiv"})
208 # MindSpore random kernels that always mutate an existing tensor in place.
209 # Out-of-place random kernels belong in _random_ms_ops only, not here.
210 _RANDOM_INPLACE_MS_OPS = frozenset({
211 "InplaceBernoulliScalar",
212 "InplaceBernoulliTensor",
213 "InplaceNormal",
214 "InplaceRandom",
215 "InplaceUniform",
216 })
218 def __init__(self):
219 self._env_yaml_dir: Optional[str] = os.environ.get("HYPER_PARALLEL_OPS_YAML_DIR")
220 self._env_python_path: Optional[str] = os.environ.get("HYPER_PARALLEL_OPS_PYTHON_PATH")
221 # The following attributes are initialized in _setup_yaml_dir()
222 self.work_dir = "" # Initialized in _setup_yaml_dir()
223 self.yaml_dir = "" # Initialized in _setup_yaml_dir()
225 self._setup_paths_from_env()
227 self.layout_infer_ops = self.safe_load_yaml_from_dir()
228 self.whitelist = ["typeof", "DistCommIsend",
229 "DistCommIrecv", "DistCommBroadcast", "DistCommAllReduce", "DistCommAllGather", "DistCommBatchIsendIrecv",
230 "DistCommReduceScatter", "requires_grad_", "item", "__get__", "__set__", "register_hook",
231 "is_complex", "chunk", "__bool__", "__len__", "__format__", "dim",
232 "_has_compatible_shallow_copy_type", "is_floating_point", "is_contiguous"]
234 # Ops requiring args unpacking for layout inference (packed as prim, name, real_args).
235 self.unpack_ops = ["ScatterUpdate", "Mod", "GatherNd", "StopGradient"]
237 self._random_ops = {
238 "normal_", "uniform_", "bernoulli", "bernoulli_",
239 "native_dropout", "rand", "rand_like", "randn",
240 "randn_like", "randint_like", "kaiming_uniform_",
241 "multinomial",
242 }
243 # Only mint random op support
244 # MindSpore use the actual kernel name.
245 self._random_ms_ops = {
246 "BernoulliExt", "MultinomialExt",
247 "InplaceBernoulliScalar", "InplaceBernoulliTensor",
248 "InplaceNormal", "InplaceRandom", "InplaceUniform",
249 "NormalFloatFloat", "NormalFloatTensor", "NormalTensorFloat", "NormalTensorTensor",
250 "RandpermExt", "Randn", "RandLikeExt", "RandnLike", "RandInt", "RandIntLike", "RandExt",
251 "FuncDropoutExt", "UniformExt",
252 }
253 self._rng_tracker: Optional[OffsetBasedRNGTracker] = None
255 self._register_distributed_ops()
257 def _setup_paths_from_env(self):
258 """
259 Setup YAML directory and Python path from environment variables.
261 This method initializes the YAML directory and extends sys.path based on
262 environment variables HYPER_PARALLEL_OPS_YAML_DIR and HYPER_PARALLEL_OPS_PYTHON_PATH.
263 """
264 self._setup_yaml_dir(self._env_yaml_dir)
265 self._extend_sys_path(self._env_python_path)
267 def _setup_yaml_dir(self, env_yaml_dir: Optional[str]):
268 """
269 Feature: Configure yaml_dir/work_dir for OpDispatcher
270 Description: Resolve the YAML directory used to load distributed op definitions.
271 If env_yaml_dir is an absolute path, use it directly; otherwise treat it
272 as a path relative to the project work_dir. If env_yaml_dir is not set,
273 fall back to the default 'shard/ops/yaml' under work_dir.
274 Expectation: self.yaml_dir and self.work_dir are set to valid values used later by
275 safe_load_yaml_from_dir(); no functional behavior is changed.
276 """
277 if env_yaml_dir:
278 if os.path.isabs(env_yaml_dir):
279 self.yaml_dir = env_yaml_dir
280 self.work_dir = ""
281 else:
282 self.work_dir = os.path.normpath(
283 os.path.join(os.path.dirname(os.path.realpath(__file__)), "../")
284 )
285 self.yaml_dir = env_yaml_dir
286 else:
287 self.yaml_dir = "shard/ops/yaml"
288 self.work_dir = os.path.normpath(
289 os.path.join(os.path.dirname(os.path.realpath(__file__)), "../")
290 )
292 @staticmethod
293 def _extend_sys_path(env_python_path: Optional[str]):
294 if not env_python_path:
295 return
296 python_paths = env_python_path.split(":")
297 for path in python_paths:
298 if path and os.path.isdir(path) and path not in sys.path:
299 sys.path.append(path)
301 def _register_distributed_ops(self):
302 for op_name, config in self.layout_infer_ops.items():
303 self._register_single_distributed_op(op_name, config)
305 def _register_single_distributed_op(self, op_name: str, config: dict):
306 """
307 Feature: Register a single distributed op implementation
308 Description: Import the distributed op class specified by config and instantiate it
309 with op_name to trigger registration in the distributed op registry.
310 Prefer 'distributed_op_module' when provided; otherwise import from
311 built-in module prefix 'hyper_parallel.core.shard.ops.' plus
312 'distributed_op_file'. If import fails and an external python path is
313 provided via env, fall back to importing 'distributed_op_file' directly.
314 Expectation: The distributed op class is imported and instantiated successfully,
315 or the original import error is raised; no functional behavior is changed.
316 """
317 class_name = config["distributed_op_class"]
319 if "distributed_op_module" in config:
320 module_name = config["distributed_op_module"]
321 module = importlib.import_module(module_name)
322 op_class = getattr(module, class_name)
323 _ = op_class(op_name)
324 return
326 module_file = config["distributed_op_file"]
327 try:
328 module_name = "hyper_parallel.core.shard.ops." + module_file
329 module = importlib.import_module(module_name)
330 op_class = getattr(module, class_name)
331 _ = op_class(op_name)
332 except (ModuleNotFoundError, ImportError):
333 if self._env_python_path:
334 module = importlib.import_module(module_file)
335 op_class = getattr(module, class_name)
336 _ = op_class(op_name)
337 else:
338 raise
340 @staticmethod
341 def _merge_default(config: dict):
342 """Apply __default__ values to all ops in this YAML file."""
343 if "__default__" not in config:
344 return config
346 default_cfg = config["__default__"]
347 merged = {}
349 for op_name, op_cfg in config.items():
350 if op_name == "__default__":
351 continue
353 new_cfg = default_cfg.copy()
354 new_cfg.update(op_cfg)
355 merged[op_name] = new_cfg
357 return merged
359 def safe_load_yaml_from_dir(self) -> dict:
360 """
361 Load yaml dictionary from directory.
363 Returns:
364 dict: Merged dictionary of all operator configurations loaded from YAML files.
365 """
366 yaml_dict = {}
367 yaml_path = os.path.join(self.work_dir, self.yaml_dir) if self.work_dir else self.yaml_dir
368 if not os.path.isdir(yaml_path):
369 raise ValueError(f"Invalid yaml directory path: {yaml_path}")
371 for yaml_file_path in glob.glob(os.path.join(yaml_path, '*.yaml')):
372 with open(yaml_file_path, 'r', encoding="utf-8") as f:
373 yaml_data = yaml.safe_load(f)
375 yaml_data = OpDispatcher._merge_default(yaml_data)
376 for name, data in yaml_data.items():
377 if name in yaml_dict:
378 raise ValueError(f"Duplicate yaml object with name '{name}'.")
379 yaml_dict[name] = data
381 return yaml_dict
383 def _dispatch_random_op(self, op_name: str, op_call: callable, args, kwargs):
384 """Handle dispatch for random ops that operate on DTensors."""
385 first_arg = next(
386 (x for x in chain(args, kwargs.values()) if isinstance(x, DTensor)),
387 None,
388 )
389 # Fall back to the default op if no DTensor is found.
390 if first_arg is None:
391 return op_call(*args, **kwargs)
393 local_args = [arg.to_local() if isinstance(arg, DTensor) else arg for arg in args]
394 local_kwargs = {k: v.to_local() if isinstance(v, DTensor) else v for k, v in kwargs.items()}
395 first_local_arg = first_arg.to_local()
397 if self._rng_tracker is None and is_rng_supported_mesh(first_arg.device_mesh):
398 self._rng_tracker = OffsetBasedRNGTracker()
400 maybe_user_generator = local_kwargs.pop("generator", None)
401 if (
402 self._rng_tracker is not None
403 and not first_local_arg.is_meta
404 and self._rng_tracker.distribute_region_enabled
405 ):
406 # pylint: disable=W0212
407 with self._rng_tracker._distribute_region(
408 device_mesh=first_arg.device_mesh,
409 placements=first_arg.placements,
410 global_shape=first_arg.shape,
411 generator=maybe_user_generator,
412 ):
413 # MindSpore random ops (e.g. mint.randn_like) extract (seed, offset)
414 # from default_generator._step() in the Python wrapper *before* the
415 # C++ dispatch triggers __fallback__. The callback reuses these
416 # pre-fetched tensor args, so set_rng_state inside _distribute_region
417 # has no effect on the kernel. Fix: apply the per-shard offset
418 # increment directly to the offset tensor in the args.
419 if platform.platform_type == PlatformType.MINDSPORE:
420 offset_incr = self._rng_tracker.compute_offset_incr(
421 first_arg.device_mesh, first_arg.placements, first_arg.shape,
422 )
423 local_args = _apply_shard_offset_to_rng_args(local_args, offset_incr)
424 local_results = op_call(*local_args, **local_kwargs)
425 else:
426 if maybe_user_generator is not None:
427 local_kwargs["generator"] = maybe_user_generator
428 local_results = op_call(*local_args, **local_kwargs)
430 return self._wrap_random_result(op_name, local_results, first_arg, args, kwargs)
432 @staticmethod
433 def _func_dropout_ext_inplace(args, kwargs) -> bool:
434 """Return True when FuncDropoutExt is invoked with inplace=True."""
435 # Kernel signature: (input, p, training, inplace, seed, offset).
436 if len(args) >= 4:
437 return bool(args[3])
438 return bool(kwargs.get("inplace", False))
440 @staticmethod
441 def _random_op_returns_self(op_name: str, args, kwargs) -> bool:
442 """Return True when a random op mutates an existing DTensor in place."""
443 if op_name in OpDispatcher._RANDOM_INPLACE_MS_OPS:
444 return True
445 if op_name == "FuncDropoutExt":
446 return OpDispatcher._func_dropout_ext_inplace(args, kwargs)
447 # Torch random inplace ops follow the ATen '_' suffix convention.
448 return op_name.endswith('_')
450 @staticmethod
451 def _wrap_random_result(op_name, local_results, first_arg, args, kwargs):
452 """Wrap a random op's local result(s) back into DTensor(s).
454 In-place ops return the input DTensor itself. Torch random inplace ops use
455 the ATen '_' suffix; MindSpore inplace random kernels are listed in
456 ``_RANDOM_INPLACE_MS_OPS``. ``FuncDropoutExt`` is handled separately
457 because the same kernel serves both modes via its ``inplace`` argument.
458 """
459 if OpDispatcher._random_op_returns_self(op_name, args, kwargs):
460 return first_arg
461 mesh = first_arg.device_mesh
462 placements = first_arg.layout.alias_placements
463 # Some ops return tuple/list, e.g. native_dropout returns (output, mask).
464 if isinstance(local_results, (tuple, list)):
465 return tuple(
466 DTensor.from_local(r, mesh, placements) if isinstance(r, Tensor) else r
467 for r in local_results
468 )
469 if isinstance(local_results, Tensor):
470 return DTensor.from_local(local_results, mesh, placements)
471 # Fallback: return as-is for non-Tensor results (currently unreachable with existing _random_ops).
472 return local_results
474 @staticmethod
475 def _unwrap_value(value: object) -> object:
476 """Replace DTensor with its local tensor; pass scalars and plain tensors through.
478 Args:
479 value (object): A single argument value from an op call.
481 Returns:
482 object: The local tensor if value is a DTensor, otherwise value unchanged.
483 """
484 if isinstance(value, DTensor):
485 return value.to_local()
486 if isinstance(value, tuple):
487 return tuple(OpDispatcher._unwrap_value(e) for e in value)
488 if isinstance(value, list):
489 return [OpDispatcher._unwrap_value(e) for e in value]
490 return value
492 @staticmethod
493 def _unwrap_args(args: tuple) -> list:
494 """Strip DTensor wrappers from args, preserving tuple/list container structure.
496 Args:
497 args: Op call positional arguments, may contain DTensor instances.
499 Returns:
500 List of args with DTensor replaced by their local tensors.
501 """
502 return [OpDispatcher._unwrap_value(arg) for arg in args]
504 @staticmethod
505 def _unwrap_kwargs(kwargs: dict) -> dict:
506 """Strip DTensor wrappers from kwargs values, preserving tuple/list container structure.
508 Args:
509 kwargs: Op call keyword arguments, values may contain DTensor instances.
511 Returns:
512 Dict of kwargs with DTensor values replaced by their local tensors.
513 """
514 return {k: OpDispatcher._unwrap_value(v) for k, v in kwargs.items()}
516 @staticmethod
517 def _gather_dtensors_to_full(args: tuple, kwargs: dict) -> tuple:
518 """Gather all DTensor arguments to full tensors for fallback execution.
520 Used when an operator has no parallel layout implementation. All DTensor
521 arguments are gathered to full tensors before calling the standard operator.
523 Args:
524 args: Op call positional arguments, may contain DTensor instances.
525 kwargs: Op call keyword arguments, may contain DTensor instances.
527 Returns:
528 Tuple of (unwrapped_args, unwrapped_kwargs) with DTensor values
529 replaced by their full tensor representations.
531 Warning:
532 This fallback performs all-gather which may consume significant memory.
533 Operators without layout implementations should be registered properly.
534 """
535 def gather(value: object) -> object:
536 if isinstance(value, DTensor):
537 return value.full_tensor()
538 if isinstance(value, tuple):
539 return tuple(gather(e) for e in value)
540 if isinstance(value, list):
541 return [gather(e) for e in value]
542 return value
544 gathered_args = [gather(arg) for arg in args]
545 gathered_kwargs = {k: gather(v) for k, v in kwargs.items()}
547 warnings.warn(
548 "Operator has no distributed layout implementation. "
549 "Falling back to all-gather which may consume significant memory. "
550 "Consider registering a proper distributed operator.",
551 UserWarning,
552 stacklevel=4
553 )
555 return gathered_args, gathered_kwargs
557 def _should_bypass_dispatch(self, op_name: str) -> bool:
558 """Return True if the op should bypass DTensor dispatch and run locally.
560 Args:
561 op_name: Canonical operator name from platform.get_op_name().
563 Returns:
564 True when the op is whitelisted or DTensor dispatch is globally disabled.
565 """
566 skip_dispatch = get_dtensor_dispatch() is False and op_name not in get_no_skip_ops()
567 return op_name in self.whitelist or op_name in self._INPLACE_BYPASS_OPS or skip_dispatch
569 @staticmethod
570 def _validate_inplace_partial_inputs(op_name: str, args: tuple, kwargs: dict) -> None:
571 """Reject local in-place add/sub when Partial contributions need gating."""
572 if op_name not in {"InplaceAddExt", "InplaceSubExt"} or not args:
573 return
574 first = args[0]
575 if len(args) >= 2:
576 second = args[1]
577 elif "other" in kwargs:
578 second = kwargs["other"]
579 else:
580 return
581 if not isinstance(first, DTensor):
582 return
583 mesh_ndim = len(first.layout.partial)
584 first_partial = tuple(first.layout.partial)
585 if isinstance(second, DTensor):
586 second_partial = tuple(second.layout.partial)
587 if len(second_partial) != mesh_ndim:
588 raise ValueError(
589 f"For {op_name}, in-place input mesh dimensions must match, "
590 f"but got {mesh_ndim} and {len(second_partial)}."
591 )
592 else:
593 second_partial = (None,) * mesh_ndim
594 if first_partial != second_partial:
595 raise ValueError(
596 f"For {op_name}, input Partial placements must be identical for "
597 f"local in-place execution, but got {first_partial} and {second_partial}."
598 )
600 def _should_dispatch_loss_parallel(self, op_name: str) -> bool:
601 """Check if should dispatch through loss_parallel path.
603 Args:
604 op_name: Canonical operator name from platform.get_op_name().
606 Returns:
607 True when in loss_parallel context and op is a CE entry point.
608 """
609 return is_loss_parallel_active() and is_loss_parallel_op(op_name)
611 def _check_decomposed_ce_op_in_loss_parallel(self, op_name: str, args: tuple, kwargs: dict):
612 """Check if decomposed CE ops are called in loss_parallel context.
614 Args:
615 op_name: Canonical operator name.
616 args: Positional arguments for op_call.
617 kwargs: Keyword arguments for op_call.
619 Raises:
620 ValueError: If decomposed CE op is called in loss_parallel context
621 with vocab-sharded DTensor input.
622 """
623 if not is_loss_parallel_active() or not is_decomposed_ce_op(op_name):
624 return
626 has_vocab_sharded_dtensor = False
627 for arg in args:
628 if isinstance(arg, DTensor) and _is_shard_on_last_dim(arg):
629 has_vocab_sharded_dtensor = True
630 break
631 if not has_vocab_sharded_dtensor:
632 for val in kwargs.values():
633 if isinstance(val, DTensor) and _is_shard_on_last_dim(val):
634 has_vocab_sharded_dtensor = True
635 break
637 if has_vocab_sharded_dtensor:
638 raise ValueError(
639 f"Operator '{op_name}' is a decomposed component of cross_entropy and should not be called "
640 f"directly within loss_parallel() context. Use F.cross_entropy(logits, targets) instead. "
641 f"For example, replace:\n"
642 f" with loss_parallel():\n"
643 f" log_probs = F.log_softmax(logits, dim=-1)\n"
644 f" loss = F.nll_loss(log_probs, targets)\n"
645 f"with:\n"
646 f" with loss_parallel():\n"
647 f" loss = F.cross_entropy(logits, targets)"
648 )
650 def _dispatch_loss_parallel(self, op_call: callable, args: tuple, kwargs: dict):
651 """Dispatch cross_entropy through the loss_parallel distributed kernel.
653 Args:
654 op_call: The raw operator callable.
655 args: Positional arguments for op_call.
656 kwargs: Keyword arguments for op_call.
658 Returns:
659 Result of the distributed cross_entropy computation.
660 """
661 if platform.platform_type == PlatformType.PYTORCH:
662 # pylint: disable=C0415
663 from hyper_parallel.platform.torch.loss_parallel_ops import distributed_cross_entropy_from_op_call
664 elif platform.platform_type == PlatformType.MINDSPORE:
665 # pylint: disable=C0415
666 from hyper_parallel.platform.mindspore.loss_parallel_ops import distributed_cross_entropy_from_op_call
667 else:
668 raise RuntimeError(f"Unsupported platform for loss_parallel: {platform.platform_type}")
669 return distributed_cross_entropy_from_op_call(op_call, args, kwargs)
671 def _check_ce_op_without_loss_parallel_context(self, op_name: str, args: tuple):
672 """Check if CE op is called with Shard(-1) DTensor outside loss_parallel context.
674 Args:
675 op_name: Canonical operator name.
676 args: Positional arguments for op_call.
678 Raises:
679 ValueError: If CE op is called with Shard(-1) logits outside loss_parallel context.
680 """
681 if is_loss_parallel_active() or not is_loss_parallel_op(op_name):
682 return
684 if len(args) == 0 or not isinstance(args[0], DTensor):
685 return
687 logits = args[0]
688 if _is_shard_on_last_dim(logits):
689 raise ValueError(
690 f"Operator '{op_name}' requires loss_parallel context when input logits are "
691 f"sharded on the vocabulary dimension (Shard(-1)). Please wrap your forward "
692 f"and backward pass with loss_parallel():\n"
693 f" with loss_parallel():\n"
694 f" loss = F.cross_entropy(logits, targets)\n"
695 f" loss.backward()\n"
696 f"If you intentionally want to gather all shards to compute cross_entropy "
697 f"(not recommended for large vocabulary), use logits.full_tensor() explicitly."
698 )
700 @staticmethod
701 def _normalize_aclop_args(op_name: str, unpack_ops: list, args: tuple) -> tuple:
702 """
703 Normalize aclop-packed arguments for MindSpore backend operators.
705 NOTE: This handles MindSpore aclop operators whose kernel signature packs
706 arguments as ``(prim, op_name_str, (real_arg0, real_arg1, ...))``. The
707 ``prim`` and ``op_name_str`` are preserved as ``packed_call`` for the
708 final kernel invocation, while the real tensor arguments are extracted
709 for layout inference and preprocessing.
711 **aclop is planned for deprecation.** Once aclop is fully removed, this
712 normalization and the associated ``unpack_ops`` list can be deleted.
714 Args:
715 op_name (str): Canonical operator name.
716 unpack_ops (list): List of op names that may use aclop packed format.
717 args (tuple): Raw positional arguments from the op call.
719 Returns:
720 tuple: ``(packed_call, normalized_args)``
721 - **packed_call**: ``(prim, op_name_str)`` tuple for kernel
722 invocation, or ``None`` if no unpacking was performed.
723 - **normalized_args**: The real tensor arguments (unpacked if
724 the packed format was detected, otherwise the original args).
725 """
726 if OpDispatcher._is_aclop_packed(op_name, unpack_ops, args):
727 return (args[0], args[1]), tuple(args[2])
728 return None, args
730 @staticmethod
731 def _is_aclop_packed(op_name: str, unpack_ops: list, args: tuple) -> bool:
732 """Check if arguments use aclop packed format."""
733 return (
734 op_name in unpack_ops
735 and len(args) == 3
736 and isinstance(args[1], str)
737 and isinstance(args[2], (tuple, list))
738 )
740 @staticmethod
741 def _call_op_impl(op_impl: callable, packed_call, args, kwargs: dict):
742 """Invoke *op_impl* with optional aclop packed-call wrapping.
744 When *packed_call* is not ``None`` the MindSpore aclop kernel expects
745 ``(prim, op_name, (arg0, arg1, ...))``. Otherwise *args* are spread
746 as positional arguments in the usual way.
748 Args:
749 op_impl: The op implementation callable.
750 packed_call: ``(prim, op_name)`` tuple or ``None``.
751 args: Local tensor arguments (list or tuple).
752 kwargs: Keyword arguments dict.
754 Returns:
755 Result of the *op_impl* invocation.
756 """
757 if packed_call is not None:
758 return op_impl(packed_call[0], packed_call[1], tuple(args), **kwargs)
759 return op_impl(*args, **kwargs)
761 def _dispatch_layout_infer(
762 self, op_name: str, op_call: callable, args: tuple, kwargs: dict
763 ):
764 """Dispatch an op through the layout-inference path.
766 Args:
767 op_name: Canonical operator name.
768 op_call: The raw operator callable.
769 args: Positional arguments for op_call.
770 kwargs: Keyword arguments for op_call.
772 Returns:
773 Result of the layout-infer dispatch.
775 Raises:
776 RuntimeError: If op_name is not registered, or preprocess returns None
777 (operator not migrated to three-phase dispatch).
778 """
779 if op_name not in self.layout_infer_ops:
780 has_dtensor = any(isinstance(arg, DTensor) for arg in args)
781 has_dtensor = has_dtensor or any(isinstance(v, DTensor) for v in kwargs.values())
782 if has_dtensor:
783 self._check_ce_op_without_loss_parallel_context(op_name, args)
785 if not is_loss_parallel_op(op_name):
786 raise RuntimeError(
787 f"Operator {op_name} does not contain parallel layout infer func. "
788 f"DTensor dispatch requires explicit layout inference registration. "
789 f"Please register a distributed operator for '{op_name}' or use local tensors."
790 )
792 gathered_args, gathered_kwargs = self._gather_dtensors_to_full(args, kwargs)
794 # Special handling for cross_entropy with 3D logits (only when NOT in loss_parallel context)
795 # PyTorch expects: logits [N, C], targets [N]
796 # But LLM forward returns: logits [batch, seq, vocab], targets [batch, seq]
797 # Note: nll_loss input is log_probs, typically already 2D, so we only reshape for cross_entropy
798 if op_name == "cross_entropy" and len(gathered_args) >= 2:
799 logits = gathered_args[0]
800 targets = gathered_args[1]
801 if isinstance(logits, Tensor) and isinstance(targets, Tensor):
802 if logits.ndim > 2 and targets.ndim > 1 and targets.ndim == logits.ndim - 1:
803 vocab_size = logits.shape[-1]
804 gathered_args[0] = logits.reshape(-1, vocab_size)
805 gathered_args[1] = targets.reshape(-1)
807 return op_call(*gathered_args, **gathered_kwargs)
808 raise RuntimeError(f"Operator {op_name} does not contain parallel layout infer func.")
810 cache_manager = LayoutCacheManager.get_instance()
811 distribute_op = cache_manager.distributed_op(op_name)
813 # Normalize aclop-packed args before any per-op processing.
814 # This allows all distributed op classes (ElementWiseDistributedOp,
815 # GatherNdDistributedOp, etc.) to receive clean unpacked args,
816 # avoiding duplicated unpack logic in each class's preprocess.
817 packed_call, args = self._normalize_aclop_args(op_name, getattr(self, 'unpack_ops', []), args)
819 result = distribute_op.preprocess(args, kwargs)
820 if result is None:
821 raise RuntimeError(
822 f"Operator '{op_name}' has not been migrated to the three-phase dispatch flow. "
823 f"Please implement preprocess() to return (local_args, local_kwargs, cache_values)."
824 )
825 output = self._dispatch_new(op_call, distribute_op, packed_call, result)
826 return self._restore_inplace_dtensor_result(op_name, args, output)
828 @staticmethod
829 def _restore_inplace_dtensor_result(op_name: str, args: tuple, output: Any) -> Any:
830 """Return the original DTensor wrapper after a local in-place operation."""
831 if op_name in {"add_", "sub_"} and args and isinstance(args[0], DTensor):
832 return args[0]
833 return output
835 @staticmethod
836 def _dispatch_new(func, distribute_op, packed_call, result) -> Tensor:
837 """New dispatch flow using preprocess result.
839 Args:
840 func: Original function.
841 distribute_op: Distributed operation instance.
842 packed_call: (prim, op_name_str) tuple for aclop kernel invocation,
843 or None for regular ops.
844 result: Preprocessed result (local_args, local_kwargs, cache_values).
846 Returns:
847 Tensor: Dispatched result as DTensor.
848 """
849 local_args, local_kwargs, cache_values = result
850 cache_key = LayoutCacheKey.from_cache_values(cache_values)
851 func_name = platform.get_op_name(func)
853 infer_result, op_impl = OpDispatcher._lookup_or_infer_layout_new(
854 func, func_name, cache_key, cache_values, distribute_op
855 )
857 op_impl = func if op_impl is None else op_impl
858 py_output = OpDispatcher._call_op_impl(op_impl, packed_call, local_args, local_kwargs)
859 return distribute_op.wrap_output(py_output, infer_result[0])
861 @staticmethod
862 def _lookup_or_infer_layout_new(func, func_name, cache_key, cache_values, distribute_op):
863 """Look up cached layout or compute via distributed op (new three-phase API).
865 Returns:
866 (infer_result, op_impl)
867 """
868 cache_manager = LayoutCacheManager.get_instance()
869 layout_cache = cache_manager.get_layout_cache()
870 if func_name not in layout_cache:
871 layout_cache[func_name] = {}
872 op_layout_cache = layout_cache[func_name]
873 if cache_key in op_layout_cache:
874 return op_layout_cache[cache_key]
875 infer_result = distribute_op.infer_layout(cache_values)
876 op_impl = distribute_op.get_expand_impl(func, infer_result, cache_values)
877 op_layout_cache[cache_key] = (infer_result, op_impl)
878 return infer_result, op_impl
880 def dispatch(self, op_call: callable, args: tuple, kwargs: dict) -> object:
881 """Route an op call through the appropriate DTensor dispatch path.
883 Args:
884 op_call: The raw operator callable.
885 args: Positional arguments for op_call.
886 kwargs: Keyword arguments for op_call.
888 Returns:
889 Result of the dispatched op call.
890 """
891 op_name = platform.get_op_name(op_call)
893 if self._should_bypass_dispatch(op_name):
894 self._validate_inplace_partial_inputs(op_name, args, kwargs)
895 result = op_call(*self._unwrap_args(args), **self._unwrap_kwargs(kwargs))
896 if op_name in self._INPLACE_BYPASS_OPS and args and isinstance(args[0], DTensor):
897 return args[0]
898 return result
900 if op_name in self._random_ops or op_name in self._random_ms_ops:
901 return self._dispatch_random_op(op_name, op_call, args, kwargs)
903 self._check_decomposed_ce_op_in_loss_parallel(op_name, args, kwargs)
905 if self._should_dispatch_loss_parallel(op_name):
906 return self._dispatch_loss_parallel(op_call, args, kwargs)
908 if op_name not in self.layout_infer_ops and get_distributed_op(op_name) is not None:
909 self.layout_infer_ops[op_name] = {}
911 return self._dispatch_layout_infer(op_name, op_call, args, kwargs)
913_OP_DISPATCHER = OpDispatcher()