Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / core / shard / _op_dispatch.py: 83%
429 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-08-25 04:27 +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 copy
18import glob
19import importlib
20import logging
21import os
22import sys
23import warnings
24from contextvars import ContextVar
25from itertools import chain
26from typing import Any, Dict, FrozenSet, List, Optional
28import yaml
30from hyper_parallel.core.shard.ops.parallel_ops_register import get_distributed_op
31from hyper_parallel.core.dtensor.dtensor import DTensor
32from hyper_parallel.core.dtensor.layout import RaggedShardInfo
33from hyper_parallel.core.dtensor.random import OffsetBasedRNGTracker, is_rng_supported_mesh
34from hyper_parallel.core.dtensor.debug._dispatch_logger import log_dispatch_enter, log_dispatch_exit
35from hyper_parallel.platform import get_platform
36from hyper_parallel.platform.platform import PlatformType
38from hyper_parallel.core.tensor_parallel._ce_op_registry import is_loss_parallel_op, is_decomposed_ce_op
39from hyper_parallel.core.tensor_parallel.loss_parallel import is_loss_parallel_active
40from hyper_parallel.core.tensor_parallel.loss_parallel_ops_common import _is_shard_on_last_dim
42platform = get_platform()
43Tensor = platform.Tensor
45logger = logging.getLogger(__name__)
48def _apply_shard_offset_to_rng_args(args, offset_incr):
49 """Apply per-shard offset increment to seed/offset tensors in MindSpore random op args.
51 MindSpore random ops (e.g. ``randn_like_``) receive ``(seed, offset)`` as
52 explicit int64 scalar tensors from ``default_generator._step()`` in the
53 Python wrapper *before* the C++ dispatch triggers ``__fallback__``. By the
54 time ``_dispatch_random_op`` is called, the kernel will use whatever
55 ``(seed, offset)`` values are in the args—it does **not** read the
56 generator again. This function finds the offset tensor and adds the
57 per-rank offset increment so each shard gets a unique random stream.
59 The (seed, offset) pair is identified as the last two consecutive int64
60 0-dim tensors in *args* (scanning from the end to skip trailing dtype /
61 device arguments).
63 Args:
64 args: The list of local args for the random op.
65 offset_incr (int): Per-shard offset increment.
67 Returns:
68 list: Modified args with the offset tensor adjusted.
69 """
70 int64_dtype = platform.tensor_dtype.int64
71 last_int64_idx = -1
72 for i in range(len(args) - 1, -1, -1):
73 arg = args[i]
74 if isinstance(arg, Tensor) and arg.dtype == int64_dtype and arg.ndim == 0:
75 if last_int64_idx == i + 1:
76 offset_idx = i + 1
77 new_args = list(args)
78 new_offset = int(new_args[offset_idx].item()) + offset_incr
79 new_args[offset_idx] = platform.tensor([new_offset], dtype=int64_dtype).reshape(())
80 return new_args
81 last_int64_idx = i
82 return args
84_dtensor_dispatch_disabled: ContextVar[bool] = ContextVar('_dtensor_dispatch_disabled', default=False)
85_no_skip_ops: ContextVar[FrozenSet[str]] = ContextVar('_no_skip_ops', default=frozenset())
86_debug_mode_observer: ContextVar = ContextVar('_debug_mode_observer', default=None)
88_RAGGED_ELEMENTWISE_OPS = {
89 "abs": "unary", "absolute": "unary", "clone": "unary", "cos": "unary",
90 "conj": "unary", "empty_like": "unary", "exp": "unary", "gelu": "unary",
91 "isinf": "unary", "isnan": "unary",
92 "log": "unary", "neg": "unary", "negative": "unary", "relu": "unary",
93 "rsqrt": "unary", "sigmoid": "unary", "silu": "unary", "sin": "unary",
94 "sqrt": "unary", "square": "unary", "zeros_like": "unary",
95 "add": "binary", "add_": "binary", "addcdiv_": "binary",
96 "addcmul_": "binary", "div": "binary", "lerp_": "binary",
97 "mul": "binary", "mul_": "binary", "pow": "binary", "real_div": "binary",
98 "sub": "binary", "__rsub__": "binary", "__rpow__": "binary",
99 "true_divide": "binary",
100}
102_RAGGED_INPLACE_ELEMENTWISE_OPS = frozenset({
103 "add_",
104 "addcdiv_",
105 "addcmul_",
106 "lerp_",
107 "mul_",
108})
110# Tensor subclass bookkeeping must stay available when a Ragged DTensor is
111# wrapped as a Parameter. These operations only inspect or update the local
112# autograd wrapper and do not reinterpret the logical distributed shape.
113_RAGGED_METADATA_BYPASS_OPS = frozenset({
114 "requires_grad_",
115 "__get__",
116 "__set__",
117 "register_hook",
118 "_has_compatible_shallow_copy_type",
119 "is_complex",
120 "is_floating_point",
121 "is_contiguous",
122})
125def get_no_skip_ops() -> FrozenSet[str]:
126 """Return the set of op names that are exempt from SkipDTensorDispatch."""
127 return _no_skip_ops.get()
130def get_dtensor_dispatch() -> bool:
131 """
132 Get the current DTensor dispatch status.
134 Returns:
135 bool: True if DTensor dispatch is enabled, False otherwise.
136 """
137 return not _dtensor_dispatch_disabled.get()
140class LayoutCacheKey:
141 """Immutable layout cache key."""
142 __slots__ = ('_tuple', '_hash')
144 def __init__(self, layout_ids: List[str]):
145 self._tuple = tuple(layout_ids)
146 self._hash = hash(self._tuple)
148 @classmethod
149 def from_cache_values(cls, cache_values: list) -> "LayoutCacheKey":
150 """Build a LayoutCacheKey from a cache_values list.
152 Args:
153 cache_values (list): Mixed list of Layout objects (with compact_str) and raw scalars.
155 Returns:
156 LayoutCacheKey: Immutable key derived from the string representation of each value.
157 """
158 # Read the cached ``_compact_str`` attribute directly instead of going through
159 # the ``compact_str`` property getter (one fewer Python frame per Layout), and
160 # build the key in one comprehension. The resulting tuple is byte-for-byte the
161 # legacy string key, so eq/hash semantics are unchanged.
162 # NOTE: the key stays a string tuple by design (cross-checked against manually
163 # built legacy keys in the UTs); CPython caches each compact_str's hash on the
164 # string object, so this is not the bottleneck a full integer key would target.
165 return cls([cs if (cs := getattr(v, '_compact_str', None)) is not None else str(v)
166 for v in cache_values])
168 def __eq__(self, other):
169 if not isinstance(other, LayoutCacheKey):
170 return False
171 return self._tuple == other._tuple
173 def __hash__(self):
174 return self._hash
176 def __repr__(self):
177 return f"LayoutCacheKey({self._tuple})"
180class LayoutCacheManager:
181 """
182 Cache layout in infer layout.
184 A singleton class that manages layout caches for distributed operations.
185 It caches the inferred layouts and operation implementations to avoid
186 redundant computation during repeated calls with the same input layouts.
187 """
188 _instance = None
190 def __init__(self):
191 self.layout_cache: Dict[str, Dict[LayoutCacheKey, Any]] = {}
192 atexit.register(self.clear_cache)
194 @classmethod
195 def get_instance(cls) -> "LayoutCacheManager":
196 """
197 Get the singleton instance of LayoutCacheManager.
199 Returns:
200 LayoutCacheManager: The singleton instance.
201 """
202 if cls._instance is None:
203 cls._instance = LayoutCacheManager()
204 return cls._instance
206 def get_layout_cache(self) -> Dict[str, Dict[LayoutCacheKey, Any]]:
207 """
208 Get the layout cache dictionary.
210 Returns:
211 Dict[str, Dict[LayoutCacheKey, Any]]: The nested dictionary mapping
212 operation names to their layout caches.
213 """
214 return self.layout_cache
216 @staticmethod
217 def distributed_op(op_name: str) -> Any:
218 """
219 Get the distributed operation implementation by name.
221 Args:
222 op_name (str): The name of the distributed operation.
224 Returns:
225 Any: The distributed operation class or implementation.
226 """
227 op = get_distributed_op(op_name)
228 return op
230 def clear_cache(self) -> None:
231 """
232 Clear all cached layouts.
234 This method is automatically registered with atexit to ensure
235 cache is cleared when the program exits.
236 """
237 self.layout_cache.clear()
240class OpDispatcher:
241 """
242 OpDispatcher
243 """
245 # Whitelisted ops that mutate args[0]'s storage in place. The dispatch bypass
246 # must return the original DTensor self for these, not the unwrapped local
247 # result, or it demotes a DTensor accumulator to a plain Tensor and breaks the
248 # next op that adds a DTensor to it (e.g. grad-accumulation `loss += micro_loss`).
249 # Class-level so it stays available on instances built via __new__ (e.g. tests).
250 _INPLACE_BYPASS_OPS = frozenset(
251 {"InplaceAddExt", "InplaceSubExt", "InplaceMul", "InplaceDiv"})
253 # MindSpore random kernels that always mutate an existing tensor in place.
254 # Out-of-place random kernels belong in _random_ms_ops only, not here.
255 _RANDOM_INPLACE_MS_OPS = frozenset({
256 "InplaceBernoulliScalar",
257 "InplaceBernoulliTensor",
258 "InplaceNormal",
259 "InplaceRandom",
260 "InplaceUniform",
261 })
263 def __init__(self):
264 self._env_yaml_dir: Optional[str] = os.environ.get("HYPER_PARALLEL_OPS_YAML_DIR")
265 self._env_python_path: Optional[str] = os.environ.get("HYPER_PARALLEL_OPS_PYTHON_PATH")
266 # The following attributes are initialized in _setup_yaml_dir()
267 self.work_dir = "" # Initialized in _setup_yaml_dir()
268 self.yaml_dir = "" # Initialized in _setup_yaml_dir()
270 self._setup_paths_from_env()
272 self.layout_infer_ops = self.safe_load_yaml_from_dir()
273 # frozenset for O(1) membership (checked on every dispatch's bypass test).
274 self.whitelist = frozenset({"typeof", "DistCommIsend",
275 "DistCommIrecv", "DistCommBroadcast", "DistCommAllReduce", "DistCommAllGather",
276 "DistCommBatchIsendIrecv",
277 "DistCommReduceScatter", "requires_grad_", "item", "__get__", "__set__",
278 "register_hook",
279 "is_complex", "chunk", "__bool__", "__len__", "__format__", "dim",
280 "_has_compatible_shallow_copy_type", "is_floating_point", "is_contiguous"})
282 # Ops requiring args unpacking for layout inference (packed as prim, name, real_args).
283 # frozenset so the aclop-normalization gate in _dispatch_layout_infer is O(1).
284 self.unpack_ops = frozenset({"ScatterUpdate", "Mod", "GatherNd", "StopGradient"})
286 self._random_ops = {
287 "normal_", "uniform_", "bernoulli", "bernoulli_",
288 "native_dropout", "rand", "rand_like", "randn",
289 "randn_like", "randint_like", "kaiming_uniform_",
290 "multinomial",
291 }
292 # Only mint random op support
293 # MindSpore use the actual kernel name.
294 self._random_ms_ops = {
295 "BernoulliExt", "MultinomialExt",
296 "InplaceBernoulliScalar", "InplaceBernoulliTensor",
297 "InplaceNormal", "InplaceRandom", "InplaceUniform",
298 "NormalFloatFloat", "NormalFloatTensor", "NormalTensorFloat", "NormalTensorTensor",
299 "RandpermExt", "Randn", "RandLikeExt", "RandnLike", "RandInt", "RandIntLike", "RandExt",
300 "FuncDropoutExt", "UniformExt",
301 }
302 self._rng_tracker: Optional[OffsetBasedRNGTracker] = None
303 # Op names proven to be loss/CE-irrelevant (both is_loss_parallel_op and
304 # is_decomposed_ce_op are False). For these the loss_parallel / decomposed-CE
305 # guards in dispatch() are always no-ops regardless of context, so we skip
306 # them (and their is_loss_parallel_active() contextvar reads) on later calls.
307 self._non_loss_ops: set = set()
309 self._register_distributed_ops()
311 def _setup_paths_from_env(self):
312 """
313 Setup YAML directory and Python path from environment variables.
315 This method initializes the YAML directory and extends sys.path based on
316 environment variables HYPER_PARALLEL_OPS_YAML_DIR and HYPER_PARALLEL_OPS_PYTHON_PATH.
317 """
318 self._setup_yaml_dir(self._env_yaml_dir)
319 self._extend_sys_path(self._env_python_path)
321 def _setup_yaml_dir(self, env_yaml_dir: Optional[str]):
322 """
323 Feature: Configure yaml_dir/work_dir for OpDispatcher
324 Description: Resolve the YAML directory used to load distributed op definitions.
325 If env_yaml_dir is an absolute path, use it directly; otherwise treat it
326 as a path relative to the project work_dir. If env_yaml_dir is not set,
327 fall back to the default 'shard/ops/yaml' under work_dir.
328 Expectation: self.yaml_dir and self.work_dir are set to valid values used later by
329 safe_load_yaml_from_dir(); no functional behavior is changed.
330 """
331 if env_yaml_dir:
332 if os.path.isabs(env_yaml_dir):
333 self.yaml_dir = env_yaml_dir
334 self.work_dir = ""
335 else:
336 self.work_dir = os.path.normpath(
337 os.path.join(os.path.dirname(os.path.realpath(__file__)), "../")
338 )
339 self.yaml_dir = env_yaml_dir
340 else:
341 self.yaml_dir = "shard/ops/yaml"
342 self.work_dir = os.path.normpath(
343 os.path.join(os.path.dirname(os.path.realpath(__file__)), "../")
344 )
346 @staticmethod
347 def _extend_sys_path(env_python_path: Optional[str]):
348 if not env_python_path:
349 return
350 python_paths = env_python_path.split(":")
351 for path in python_paths:
352 if path and os.path.isdir(path) and path not in sys.path:
353 sys.path.append(path)
355 def _register_distributed_ops(self):
356 for op_name, config in self.layout_infer_ops.items():
357 self._register_single_distributed_op(op_name, config)
359 def _register_single_distributed_op(self, op_name: str, config: dict):
360 """
361 Feature: Register a single distributed op implementation
362 Description: Import the distributed op class specified by config and instantiate it
363 with op_name to trigger registration in the distributed op registry.
364 Prefer 'distributed_op_module' when provided; otherwise import from
365 built-in module prefix 'hyper_parallel.core.shard.ops.' plus
366 'distributed_op_file'. If import fails and an external python path is
367 provided via env, fall back to importing 'distributed_op_file' directly.
368 Expectation: The distributed op class is imported and instantiated successfully,
369 or the original import error is raised; no functional behavior is changed.
370 """
371 class_name = config["distributed_op_class"]
373 if "distributed_op_module" in config:
374 module_name = config["distributed_op_module"]
375 module = importlib.import_module(module_name)
376 op_class = getattr(module, class_name)
377 _ = op_class(op_name)
378 return
380 module_file = config["distributed_op_file"]
381 try:
382 module_name = "hyper_parallel.core.shard.ops." + module_file
383 module = importlib.import_module(module_name)
384 op_class = getattr(module, class_name)
385 _ = op_class(op_name)
386 except (ModuleNotFoundError, ImportError):
387 if self._env_python_path:
388 module = importlib.import_module(module_file)
389 op_class = getattr(module, class_name)
390 _ = op_class(op_name)
391 else:
392 raise
394 @staticmethod
395 def _merge_default(config: dict):
396 """Apply __default__ values to all ops in this YAML file."""
397 if "__default__" not in config:
398 return config
400 default_cfg = config["__default__"]
401 merged = {}
403 for op_name, op_cfg in config.items():
404 if op_name == "__default__":
405 continue
407 new_cfg = default_cfg.copy()
408 new_cfg.update(op_cfg)
409 merged[op_name] = new_cfg
411 return merged
413 def safe_load_yaml_from_dir(self) -> dict:
414 """
415 Load yaml dictionary from directory.
417 Returns:
418 dict: Merged dictionary of all operator configurations loaded from YAML files.
419 """
420 yaml_dict = {}
421 yaml_path = os.path.join(self.work_dir, self.yaml_dir) if self.work_dir else self.yaml_dir
422 if not os.path.isdir(yaml_path):
423 raise ValueError(f"Invalid yaml directory path: {yaml_path}")
425 for yaml_file_path in glob.glob(os.path.join(yaml_path, '*.yaml')):
426 with open(yaml_file_path, 'r', encoding="utf-8") as f:
427 yaml_data = yaml.safe_load(f)
429 yaml_data = OpDispatcher._merge_default(yaml_data)
430 for name, data in yaml_data.items():
431 if name in yaml_dict:
432 raise ValueError(f"Duplicate yaml object with name '{name}'.")
433 yaml_dict[name] = data
435 return yaml_dict
437 def _dispatch_random_op(self, op_name: str, op_call: callable, args, kwargs):
438 """Handle dispatch for random ops that operate on DTensors."""
439 first_arg = next(
440 (x for x in chain(args, kwargs.values()) if isinstance(x, DTensor)),
441 None,
442 )
443 # Fall back to the default op if no DTensor is found.
444 if first_arg is None:
445 return op_call(*args, **kwargs)
447 local_args = [arg.to_local() if isinstance(arg, DTensor) else arg for arg in args]
448 local_kwargs = {k: v.to_local() if isinstance(v, DTensor) else v for k, v in kwargs.items()}
449 first_local_arg = first_arg.to_local()
451 if self._rng_tracker is None and is_rng_supported_mesh(first_arg.device_mesh):
452 self._rng_tracker = OffsetBasedRNGTracker()
454 maybe_user_generator = local_kwargs.pop("generator", None)
455 if (
456 self._rng_tracker is not None
457 and not first_local_arg.is_meta
458 and self._rng_tracker.distribute_region_enabled
459 ):
460 # pylint: disable=W0212
461 with self._rng_tracker._distribute_region(
462 device_mesh=first_arg.device_mesh,
463 placements=first_arg.placements,
464 global_shape=first_arg.shape,
465 generator=maybe_user_generator,
466 ):
467 # MindSpore random ops (e.g. mint.randn_like) extract (seed, offset)
468 # from default_generator._step() in the Python wrapper *before* the
469 # C++ dispatch triggers __fallback__. The callback reuses these
470 # pre-fetched tensor args, so set_rng_state inside _distribute_region
471 # has no effect on the kernel. Fix: apply the per-shard offset
472 # increment directly to the offset tensor in the args.
473 if platform.platform_type == PlatformType.MINDSPORE:
474 offset_incr = self._rng_tracker.compute_offset_incr(
475 first_arg.device_mesh, first_arg.placements, first_arg.shape,
476 )
477 local_args = _apply_shard_offset_to_rng_args(local_args, offset_incr)
478 local_results = op_call(*local_args, **local_kwargs)
479 else:
480 if maybe_user_generator is not None:
481 local_kwargs["generator"] = maybe_user_generator
482 local_results = op_call(*local_args, **local_kwargs)
484 return self._wrap_random_result(op_name, local_results, first_arg, args, kwargs)
486 @staticmethod
487 def _func_dropout_ext_inplace(args, kwargs) -> bool:
488 """Return True when FuncDropoutExt is invoked with inplace=True."""
489 # Kernel signature: (input, p, training, inplace, seed, offset).
490 if len(args) >= 4:
491 return bool(args[3])
492 return bool(kwargs.get("inplace", False))
494 @staticmethod
495 def _random_op_returns_self(op_name: str, args, kwargs) -> bool:
496 """Return True when a random op mutates an existing DTensor in place."""
497 if op_name in OpDispatcher._RANDOM_INPLACE_MS_OPS:
498 return True
499 if op_name == "FuncDropoutExt":
500 return OpDispatcher._func_dropout_ext_inplace(args, kwargs)
501 # Torch random inplace ops follow the ATen '_' suffix convention.
502 return op_name.endswith('_')
504 @staticmethod
505 def _wrap_random_result(op_name, local_results, first_arg, args, kwargs):
506 """Wrap a random op's local result(s) back into DTensor(s).
508 In-place ops return the input DTensor itself. Torch random inplace ops use
509 the ATen '_' suffix; MindSpore inplace random kernels are listed in
510 ``_RANDOM_INPLACE_MS_OPS``. ``FuncDropoutExt`` is handled separately
511 because the same kernel serves both modes via its ``inplace`` argument.
512 """
513 if OpDispatcher._random_op_returns_self(op_name, args, kwargs):
514 return first_arg
515 mesh = first_arg.device_mesh
516 placements = first_arg.layout.alias_placements
517 # Some ops return tuple/list, e.g. native_dropout returns (output, mask).
518 if isinstance(local_results, (tuple, list)):
519 return tuple(
520 DTensor.from_local(r, mesh, placements) if isinstance(r, Tensor) else r
521 for r in local_results
522 )
523 if isinstance(local_results, Tensor):
524 return DTensor.from_local(local_results, mesh, placements)
525 # Fallback: return as-is for non-Tensor results (currently unreachable with existing _random_ops).
526 return local_results
528 @staticmethod
529 def _unwrap_value(value: object) -> object:
530 """Replace DTensor with its local tensor; pass scalars and plain tensors through.
532 Args:
533 value (object): A single argument value from an op call.
535 Returns:
536 object: The local tensor if value is a DTensor, otherwise value unchanged.
537 """
538 if isinstance(value, DTensor):
539 return value.to_local()
540 if isinstance(value, tuple):
541 return tuple(OpDispatcher._unwrap_value(e) for e in value)
542 if isinstance(value, list):
543 return [OpDispatcher._unwrap_value(e) for e in value]
544 return value
546 @staticmethod
547 def _unwrap_args(args: tuple) -> list:
548 """Strip DTensor wrappers from args, preserving tuple/list container structure.
550 Args:
551 args: Op call positional arguments, may contain DTensor instances.
553 Returns:
554 List of args with DTensor replaced by their local tensors.
555 """
556 return [OpDispatcher._unwrap_value(arg) for arg in args]
558 @staticmethod
559 def _unwrap_kwargs(kwargs: dict) -> dict:
560 """Strip DTensor wrappers from kwargs values, preserving tuple/list container structure.
562 Args:
563 kwargs: Op call keyword arguments, values may contain DTensor instances.
565 Returns:
566 Dict of kwargs with DTensor values replaced by their local tensors.
567 """
568 return {k: OpDispatcher._unwrap_value(v) for k, v in kwargs.items()}
570 @staticmethod
571 def _collect_dtensors(value: object) -> List[DTensor]:
572 """Return all DTensors nested in one dispatch argument."""
573 if isinstance(value, DTensor):
574 return [value]
575 if isinstance(value, (tuple, list)):
576 return list(chain.from_iterable(
577 OpDispatcher._collect_dtensors(item) for item in value
578 ))
579 if isinstance(value, dict):
580 return list(chain.from_iterable(
581 OpDispatcher._collect_dtensors(item) for item in value.values()
582 ))
583 return []
585 def _validate_ragged_dispatch(
586 self, op_name: str, args: tuple, kwargs: dict
587 ) -> Optional[DTensor]:
588 """Return the first Ragged input for a whitelisted elementwise op."""
589 reference = next(
590 (
591 dtensor for dtensor in self._collect_dtensors((args, kwargs))
592 if isinstance(getattr(dtensor.layout, "ragged_shard", None), RaggedShardInfo)
593 ),
594 None,
595 )
596 if reference is None:
597 return None
598 if op_name not in _RAGGED_ELEMENTWISE_OPS:
599 raise RuntimeError(
600 f"Operator {op_name!r} does not support RaggedShard in phase one"
601 )
602 return reference
604 def _dispatch_ragged_elementwise(
605 self, op_call: callable, args: tuple, kwargs: dict,
606 reference: DTensor,
607 ) -> DTensor:
608 """Execute a whitelisted op locally and inherit its Ragged Layout."""
609 local_args = tuple(self._unwrap_args(args))
610 local_kwargs = self._unwrap_kwargs(kwargs)
611 py_output = op_call(*local_args, **local_kwargs)
612 op_name = platform.get_op_name(op_call)
613 if op_name in _RAGGED_INPLACE_ELEMENTWISE_OPS:
614 if not args or not isinstance(args[0], DTensor):
615 raise ValueError(
616 f"Ragged in-place operator {op_name!r} requires a DTensor first argument"
617 )
618 return args[0]
619 return DTensor.from_local_with_layout(
620 py_output,
621 copy.deepcopy(reference.layout),
622 shape=tuple(reference.shape),
623 )
625 @staticmethod
626 def _gather_dtensors_to_full(args: tuple, kwargs: dict) -> tuple:
627 """Gather all DTensor arguments to full tensors for fallback execution.
629 Used when an operator has no parallel layout implementation. All DTensor
630 arguments are gathered to full tensors before calling the standard operator.
632 Args:
633 args: Op call positional arguments, may contain DTensor instances.
634 kwargs: Op call keyword arguments, may contain DTensor instances.
636 Returns:
637 Tuple of (unwrapped_args, unwrapped_kwargs) with DTensor values
638 replaced by their full tensor representations.
640 Warning:
641 This fallback performs all-gather which may consume significant memory.
642 Operators without layout implementations should be registered properly.
643 """
644 def gather(value: object) -> object:
645 if isinstance(value, DTensor):
646 return value.full_tensor()
647 if isinstance(value, tuple):
648 return tuple(gather(e) for e in value)
649 if isinstance(value, list):
650 return [gather(e) for e in value]
651 return value
653 gathered_args = [gather(arg) for arg in args]
654 gathered_kwargs = {k: gather(v) for k, v in kwargs.items()}
656 warnings.warn(
657 "Operator has no distributed layout implementation. "
658 "Falling back to all-gather which may consume significant memory. "
659 "Consider registering a proper distributed operator.",
660 UserWarning,
661 stacklevel=4
662 )
664 return gathered_args, gathered_kwargs
666 def _should_bypass_dispatch(self, op_name: str) -> bool:
667 """Return True if the op should bypass DTensor dispatch and run locally.
669 Args:
670 op_name: Canonical operator name from platform.get_op_name().
672 Returns:
673 True when the op is whitelisted or DTensor dispatch is globally disabled.
674 """
675 # Cheap O(1) frozenset checks first, short-circuit before the ContextVar
676 # read (get_dtensor_dispatch) which is the priciest part of this guard.
677 if op_name in self.whitelist or op_name in self._INPLACE_BYPASS_OPS:
678 return True
679 return get_dtensor_dispatch() is False and op_name not in get_no_skip_ops()
681 @staticmethod
682 def _validate_inplace_partial_inputs(op_name: str, args: tuple, kwargs: dict) -> None:
683 """Reject local in-place add/sub when Partial contributions need gating."""
684 if op_name not in {"InplaceAddExt", "InplaceSubExt"} or not args:
685 return
686 first = args[0]
687 if len(args) >= 2:
688 second = args[1]
689 elif "other" in kwargs:
690 second = kwargs["other"]
691 else:
692 return
693 if not isinstance(first, DTensor):
694 return
695 mesh_ndim = len(first.layout.partial)
696 first_partial = tuple(first.layout.partial)
697 if isinstance(second, DTensor):
698 second_partial = tuple(second.layout.partial)
699 if len(second_partial) != mesh_ndim:
700 raise ValueError(
701 f"For {op_name}, in-place input mesh dimensions must match, "
702 f"but got {mesh_ndim} and {len(second_partial)}."
703 )
704 else:
705 second_partial = (None,) * mesh_ndim
706 if first_partial != second_partial:
707 raise ValueError(
708 f"For {op_name}, input Partial placements must be identical for "
709 f"local in-place execution, but got {first_partial} and {second_partial}."
710 )
712 def _should_dispatch_loss_parallel(self, op_name: str) -> bool:
713 """Check if should dispatch through loss_parallel path.
715 Args:
716 op_name: Canonical operator name from platform.get_op_name().
718 Returns:
719 True when in loss_parallel context and op is a CE entry point.
720 """
721 return is_loss_parallel_active() and is_loss_parallel_op(op_name)
723 def _check_decomposed_ce_op_in_loss_parallel(self, op_name: str, args: tuple, kwargs: dict):
724 """Check if decomposed CE ops are called in loss_parallel context.
726 Args:
727 op_name: Canonical operator name.
728 args: Positional arguments for op_call.
729 kwargs: Keyword arguments for op_call.
731 Raises:
732 ValueError: If decomposed CE op is called in loss_parallel context
733 with vocab-sharded DTensor input.
734 """
735 if not is_loss_parallel_active() or not is_decomposed_ce_op(op_name):
736 return
738 has_vocab_sharded_dtensor = False
739 for arg in args:
740 if isinstance(arg, DTensor) and _is_shard_on_last_dim(arg):
741 has_vocab_sharded_dtensor = True
742 break
743 if not has_vocab_sharded_dtensor:
744 for val in kwargs.values():
745 if isinstance(val, DTensor) and _is_shard_on_last_dim(val):
746 has_vocab_sharded_dtensor = True
747 break
749 if has_vocab_sharded_dtensor:
750 raise ValueError(
751 f"Operator '{op_name}' is a decomposed component of cross_entropy and should not be called "
752 f"directly within loss_parallel() context. Use F.cross_entropy(logits, targets) instead. "
753 f"For example, replace:\n"
754 f" with loss_parallel():\n"
755 f" log_probs = F.log_softmax(logits, dim=-1)\n"
756 f" loss = F.nll_loss(log_probs, targets)\n"
757 f"with:\n"
758 f" with loss_parallel():\n"
759 f" loss = F.cross_entropy(logits, targets)"
760 )
762 def _dispatch_loss_parallel(self, op_call: callable, args: tuple, kwargs: dict):
763 """Dispatch cross_entropy through the loss_parallel distributed kernel.
765 Args:
766 op_call: The raw operator callable.
767 args: Positional arguments for op_call.
768 kwargs: Keyword arguments for op_call.
770 Returns:
771 Result of the distributed cross_entropy computation.
772 """
773 if platform.platform_type == PlatformType.PYTORCH:
774 # pylint: disable=C0415
775 from hyper_parallel.platform.torch.loss_parallel_ops import distributed_cross_entropy_from_op_call
776 elif platform.platform_type == PlatformType.MINDSPORE:
777 # pylint: disable=C0415
778 from hyper_parallel.platform.mindspore.loss_parallel_ops import distributed_cross_entropy_from_op_call
779 else:
780 raise RuntimeError(f"Unsupported platform for loss_parallel: {platform.platform_type}")
781 return distributed_cross_entropy_from_op_call(op_call, args, kwargs)
783 def _check_ce_op_without_loss_parallel_context(self, op_name: str, args: tuple):
784 """Check if CE op is called with Shard(-1) DTensor outside loss_parallel context.
786 Args:
787 op_name: Canonical operator name.
788 args: Positional arguments for op_call.
790 Raises:
791 ValueError: If CE op is called with Shard(-1) logits outside loss_parallel context.
792 """
793 if is_loss_parallel_active() or not is_loss_parallel_op(op_name):
794 return
796 if len(args) == 0 or not isinstance(args[0], DTensor):
797 return
799 logits = args[0]
800 if _is_shard_on_last_dim(logits):
801 raise ValueError(
802 f"Operator '{op_name}' requires loss_parallel context when input logits are "
803 f"sharded on the vocabulary dimension (Shard(-1)). Please wrap your forward "
804 f"and backward pass with loss_parallel():\n"
805 f" with loss_parallel():\n"
806 f" loss = F.cross_entropy(logits, targets)\n"
807 f" loss.backward()\n"
808 f"If you intentionally want to gather all shards to compute cross_entropy "
809 f"(not recommended for large vocabulary), use logits.full_tensor() explicitly."
810 )
812 @staticmethod
813 def _normalize_aclop_args(op_name: str, unpack_ops: list, args: tuple) -> tuple:
814 """
815 Normalize aclop-packed arguments for MindSpore backend operators.
817 NOTE: This handles MindSpore aclop operators whose kernel signature packs
818 arguments as ``(prim, op_name_str, (real_arg0, real_arg1, ...))``. The
819 ``prim`` and ``op_name_str`` are preserved as ``packed_call`` for the
820 final kernel invocation, while the real tensor arguments are extracted
821 for layout inference and preprocessing.
823 **aclop is planned for deprecation.** Once aclop is fully removed, this
824 normalization and the associated ``unpack_ops`` list can be deleted.
826 Args:
827 op_name (str): Canonical operator name.
828 unpack_ops (list): List of op names that may use aclop packed format.
829 args (tuple): Raw positional arguments from the op call.
831 Returns:
832 tuple: ``(packed_call, normalized_args)``
833 - **packed_call**: ``(prim, op_name_str)`` tuple for kernel
834 invocation, or ``None`` if no unpacking was performed.
835 - **normalized_args**: The real tensor arguments (unpacked if
836 the packed format was detected, otherwise the original args).
837 """
838 if OpDispatcher._is_aclop_packed(op_name, unpack_ops, args):
839 return (args[0], args[1]), tuple(args[2])
840 return None, args
842 @staticmethod
843 def _is_aclop_packed(op_name: str, unpack_ops: list, args: tuple) -> bool:
844 """Check if arguments use aclop packed format."""
845 return (
846 op_name in unpack_ops
847 and len(args) == 3
848 and isinstance(args[1], str)
849 and isinstance(args[2], (tuple, list))
850 )
852 @staticmethod
853 def _call_op_impl(op_impl: callable, packed_call, args, kwargs: dict):
854 """Invoke *op_impl* with optional aclop packed-call wrapping.
856 When *packed_call* is not ``None`` the MindSpore aclop kernel expects
857 ``(prim, op_name, (arg0, arg1, ...))``. Otherwise *args* are spread
858 as positional arguments in the usual way.
860 Args:
861 op_impl: The op implementation callable.
862 packed_call: ``(prim, op_name)`` tuple or ``None``.
863 args: Local tensor arguments (list or tuple).
864 kwargs: Keyword arguments dict.
866 Returns:
867 Result of the *op_impl* invocation.
868 """
869 if packed_call is not None:
870 return op_impl(packed_call[0], packed_call[1], tuple(args), **kwargs)
871 return op_impl(*args, **kwargs)
873 def _handle_unregistered_op(
874 self, op_name: str, op_call: callable, args: tuple, kwargs: dict
875 ):
876 """Handle ops that have no registered layout-inference entry.
878 This is a fallback path for ops that are not registered in
879 ``layout_infer_ops``. When arguments contain DTensors it either raises
880 (with a hint to register a distributed op) or, for loss-parallel ops,
881 gathers tensors to full and dispatches through the raw callable.
883 Args:
884 op_name: Canonical operator name.
885 op_call: The raw operator callable.
886 args: Positional arguments for op_call.
887 kwargs: Keyword arguments for op_call.
889 Returns:
890 Raw dispatch result (plain Tensor, not wrapped as DTensor).
892 Raises:
893 RuntimeError: If op_name is not registered for layout inference.
894 """
895 has_dtensor = any(isinstance(arg, DTensor) for arg in args)
896 has_dtensor = has_dtensor or any(isinstance(v, DTensor) for v in kwargs.values())
897 if has_dtensor:
898 self._check_ce_op_without_loss_parallel_context(op_name, args)
900 if not is_loss_parallel_op(op_name):
901 raise RuntimeError(
902 f"Operator {op_name} does not contain parallel layout infer func. "
903 f"DTensor dispatch requires explicit layout inference registration. "
904 f"Please register a distributed operator for '{op_name}' or use local tensors."
905 )
907 gathered_args, gathered_kwargs = self._gather_dtensors_to_full(args, kwargs)
909 # Special handling for cross_entropy with 3D logits (only when NOT in loss_parallel context)
910 # PyTorch expects: logits [N, C], targets [N]
911 # But LLM forward returns: logits [batch, seq, vocab], targets [batch, seq]
912 # Note: nll_loss input is log_probs, typically already 2D, so we only reshape for cross_entropy
913 if op_name == "cross_entropy" and len(gathered_args) >= 2:
914 logits = gathered_args[0]
915 targets = gathered_args[1]
916 if isinstance(logits, Tensor) and isinstance(targets, Tensor):
917 if logits.ndim > 2 and targets.ndim > 1 and targets.ndim == logits.ndim - 1:
918 vocab_size = logits.shape[-1]
919 gathered_args[0] = logits.reshape(-1, vocab_size)
920 gathered_args[1] = targets.reshape(-1)
922 return op_call(*gathered_args, **gathered_kwargs)
923 raise RuntimeError(f"Operator {op_name} does not contain parallel layout infer func.")
925 def _dispatch_layout_infer(
926 self, op_name: str, op_call: callable, args: tuple, kwargs: dict
927 ):
928 """Standard dispatch through layout-inference: preprocess → infer → execute → wrap.
930 Args:
931 op_name: Canonical operator name (already resolved by the caller).
932 op_call: The raw operator callable.
933 args: Positional arguments for op_call.
934 kwargs: Keyword arguments for op_call.
936 Returns:
937 DTensor: Dispatched result wrapped as DTensor.
939 Raises:
940 RuntimeError: If op_name is not registered, or preprocess returns None.
941 """
942 if op_name not in self.layout_infer_ops:
943 return self._handle_unregistered_op(op_name, op_call, args, kwargs)
945 cache_manager = LayoutCacheManager.get_instance()
946 distribute_op = cache_manager.distributed_op(op_name)
948 # Normalize aclop-packed args before any per-op processing. Only the handful
949 # of (deprecation-bound) unpack_ops ever use the packed format, so gate the
950 # whole normalization behind an O(1) membership test instead of paying two
951 # function frames (_normalize_aclop_args + _is_aclop_packed) on every op.
952 if op_name in getattr(self, 'unpack_ops', ()):
953 packed_call, args = self._normalize_aclop_args(op_name, self.unpack_ops, args)
954 else:
955 packed_call = None
957 result = distribute_op.preprocess(args, kwargs)
958 if result is None:
959 raise RuntimeError(
960 f"Operator '{op_name}' has not been migrated to the three-phase dispatch flow. "
961 f"Please implement preprocess() to return (local_args, local_kwargs, cache_values)."
962 )
963 local_args, local_kwargs, cache_values = result
964 cache_key = LayoutCacheKey.from_cache_values(cache_values)
966 infer_result, op_impl = OpDispatcher._lookup_or_infer_layout(
967 op_call, op_name, cache_key, cache_values, distribute_op, cache_manager
968 )
970 op_impl = op_call if op_impl is None else op_impl
971 py_output = OpDispatcher._call_op_impl(op_impl, packed_call, local_args, local_kwargs)
972 output = distribute_op.wrap_output(py_output, infer_result[0])
973 return OpDispatcher._restore_inplace_dtensor_result(op_name, args, output)
975 @staticmethod
976 def _restore_inplace_dtensor_result(op_name: str, args: tuple, output: Any) -> Any:
977 """Return the original DTensor wrapper after a local in-place operation."""
978 if op_name in {"add_", "sub_"} and args and isinstance(args[0], DTensor):
979 return args[0]
980 return output
982 @staticmethod
983 def _lookup_or_infer_layout(func, func_name, cache_key, cache_values, distribute_op, cache_manager):
984 """Look up cached layout or compute via distributed op.
986 Returns:
987 (infer_result, op_impl)
988 """
989 layout_cache = cache_manager.get_layout_cache()
990 if func_name not in layout_cache:
991 layout_cache[func_name] = {}
992 op_layout_cache = layout_cache[func_name]
993 if cache_key in op_layout_cache:
994 return op_layout_cache[cache_key]
995 infer_result = distribute_op.infer_layout(cache_values)
996 op_impl = distribute_op.get_expand_impl(func, infer_result, cache_values)
997 op_layout_cache[cache_key] = (infer_result, op_impl)
998 return infer_result, op_impl
1000 def dispatch(self, op_call: callable, args: tuple, kwargs: dict) -> object:
1001 """Route an op call through the appropriate DTensor dispatch path.
1003 Args:
1004 op_call: The raw operator callable.
1005 args: Positional arguments for op_call.
1006 kwargs: Keyword arguments for op_call.
1008 Returns:
1009 Result of the dispatched op call.
1010 """
1011 op_name = platform.get_op_name(op_call)
1012 if logger.isEnabledFor(logging.DEBUG):
1013 log_dispatch_enter(op_name, args, kwargs)
1015 observer = _debug_mode_observer.get()
1016 if observer is not None:
1017 observer.on_op_dispatch_enter(op_name, op_call, args, kwargs)
1019 result = None
1020 try:
1021 should_bypass = self._should_bypass_dispatch(op_name)
1022 ragged_dtensor = None
1023 if not should_bypass or op_name not in _RAGGED_METADATA_BYPASS_OPS:
1024 ragged_dtensor = self._validate_ragged_dispatch(op_name, args, kwargs)
1025 if ragged_dtensor is not None:
1026 result = self._dispatch_ragged_elementwise(
1027 op_call, args, kwargs, ragged_dtensor
1028 )
1029 return result
1031 if should_bypass:
1032 self._validate_inplace_partial_inputs(op_name, args, kwargs)
1033 result = op_call(*self._unwrap_args(args), **self._unwrap_kwargs(kwargs))
1034 if op_name in self._INPLACE_BYPASS_OPS and args and isinstance(args[0], DTensor):
1035 result = args[0]
1036 return result
1038 if op_name in self._random_ops or op_name in self._random_ms_ops:
1039 result = self._dispatch_random_op(op_name, op_call, args, kwargs)
1040 return result
1042 self._check_decomposed_ce_op_in_loss_parallel(op_name, args, kwargs)
1044 if self._should_dispatch_loss_parallel(op_name):
1045 result = self._dispatch_loss_parallel(op_call, args, kwargs)
1046 return result
1048 if op_name not in self.layout_infer_ops and get_distributed_op(op_name) is not None:
1049 self.layout_infer_ops[op_name] = {}
1051 result = self._dispatch_layout_infer(op_name, op_call, args, kwargs)
1052 return result
1053 finally:
1054 if logger.isEnabledFor(logging.DEBUG):
1055 log_dispatch_exit(op_name, result)
1057 if observer is not None:
1058 observer.on_op_dispatch_exit(op_name, result)
1060_OP_DISPATCHER = OpDispatcher()