Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / core / shard / _op_dispatch.py: 88%
377 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-11-15 07:18 +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 warnings
23from contextvars import ContextVar
24from functools import lru_cache
25from itertools import chain
26from typing import Any, Dict, FrozenSet, List, Optional
28import torch
29import yaml
31from hyper_parallel.core.shard.ops.parallel_ops import DistributedOp
32from hyper_parallel.core.shard.utils import distributed_cross_entropy_from_op_call, get_op_name
33from hyper_parallel.core.shard.ops.parallel_ops_register import get_distributed_op
34from hyper_parallel.core.dtensor.dtensor import DTensor
35from hyper_parallel.core.dtensor.layout import RaggedShardInfo
36from hyper_parallel.core.dtensor.random import OffsetBasedRNGTracker, is_rng_supported_mesh
37from hyper_parallel.core.dtensor.debug._dispatch_logger import log_dispatch_enter, log_dispatch_exit
39from hyper_parallel.core.tensor_parallel._ce_op_registry import is_loss_parallel_op, is_decomposed_ce_op
40from hyper_parallel.core.tensor_parallel.loss_parallel import is_loss_parallel_active
41from hyper_parallel.core.tensor_parallel.loss_parallel_ops_common import _is_shard_on_last_dim
43logger = logging.getLogger(__name__)
45_DISTRIBUTED_OPS_PACKAGE = "hyper_parallel.core.shard.ops"
46_DISTRIBUTED_OPS_DIR = os.path.realpath(os.path.join(os.path.dirname(__file__), "ops"))
47_DISTRIBUTED_OPS_YAML_DIR = os.path.join(_DISTRIBUTED_OPS_DIR, "yaml")
50_dtensor_dispatch_disabled: ContextVar[bool] = ContextVar('_dtensor_dispatch_disabled', default=False)
51_no_skip_ops: ContextVar[FrozenSet[str]] = ContextVar('_no_skip_ops', default=frozenset())
52_debug_mode_observer: ContextVar = ContextVar('_debug_mode_observer', default=None)
54_RAGGED_ELEMENTWISE_OPS = {
55 # Torch public APIs dispatch with these lowercase names.
56 "abs": "unary", "absolute": "unary", "clamp": "unary",
57 "clamp_max": "unary", "clamp_min": "unary", "clone": "unary", "cos": "unary",
58 "conj": "unary", "empty_like": "unary", "exp": "unary", "gelu": "unary",
59 "isinf": "unary", "isnan": "unary",
60 "log": "unary", "neg": "unary", "negative": "unary", "relu": "unary",
61 "rsqrt": "unary", "sigmoid": "unary", "silu": "unary", "sin": "unary",
62 "sqrt": "unary", "square": "unary", "zeros_like": "unary",
63 "add": "binary", "add_": "binary", "addcdiv_": "binary",
64 "addcmul_": "binary", "div": "binary", "lerp_": "binary",
65 "mul": "binary", "mul_": "binary", "pow": "binary",
66 "sub": "binary", "__rsub__": "binary", "__rpow__": "binary",
67 "true_divide": "binary",
68}
70_RAGGED_INPLACE_ELEMENTWISE_OPS = frozenset({
71 "add_",
72 "addcdiv_",
73 "addcmul_",
74 "lerp_",
75 "mul_",
76})
78# Tensor subclass bookkeeping must stay available when a Ragged DTensor is
79# wrapped as a Parameter. These operations only inspect or update the local
80# autograd wrapper and do not reinterpret the logical distributed shape.
81_RAGGED_METADATA_BYPASS_OPS = frozenset({
82 "requires_grad_",
83 "__get__",
84 "__set__",
85 "register_hook",
86 "_has_compatible_shallow_copy_type",
87 "is_complex",
88 "is_floating_point",
89 "is_contiguous",
90})
93def get_no_skip_ops() -> FrozenSet[str]:
94 """Return the set of op names that are exempt from SkipDTensorDispatch."""
95 return _no_skip_ops.get()
98def get_dtensor_dispatch() -> bool:
99 """
100 Get the current DTensor dispatch status.
102 Returns:
103 bool: True if DTensor dispatch is enabled, False otherwise.
104 """
105 return not _dtensor_dispatch_disabled.get()
108def _stringify_cache_value(value) -> str:
109 """Return the cached compact string when available."""
110 compact_str = getattr(value, '_compact_str', None)
111 return compact_str if compact_str is not None else str(value)
114class LayoutCacheKey:
115 """Immutable layout cache key."""
116 __slots__ = ('_tuple', '_hash')
118 def __init__(self, layout_ids: List[str]):
119 self._tuple = tuple(layout_ids)
120 self._hash = hash(self._tuple)
122 @classmethod
123 def from_cache_values(cls, cache_values: list) -> "LayoutCacheKey":
124 """Build a LayoutCacheKey from a cache_values list.
126 Args:
127 cache_values (list): Mixed list of Layout objects (with compact_str) and raw scalars.
129 Returns:
130 LayoutCacheKey: Immutable key derived from the string representation of each value.
131 """
132 # Read the cached ``_compact_str`` attribute directly instead of going through
133 # the ``compact_str`` property getter (one fewer Python frame per Layout), and
134 # build the key in one comprehension. The resulting tuple is byte-for-byte the
135 # legacy string key, so eq/hash semantics are unchanged.
136 # NOTE: the key stays a string tuple by design (cross-checked against manually
137 # built legacy keys in the UTs); CPython caches each compact_str's hash on the
138 # string object, so this is not the bottleneck a full integer key would target.
139 return cls([_stringify_cache_value(value) for value in cache_values])
141 def __eq__(self, other):
142 if not isinstance(other, LayoutCacheKey):
143 return False
144 return self._tuple == other._tuple
146 def __hash__(self):
147 return self._hash
149 def __repr__(self):
150 return f"LayoutCacheKey({self._tuple})"
153class LayoutCacheManager:
154 """
155 Cache layout in infer layout.
157 A singleton class that manages layout caches for distributed operations.
158 It caches the inferred layouts and operation implementations to avoid
159 redundant computation during repeated calls with the same input layouts.
160 """
161 _instance = None
163 def __init__(self):
164 self.layout_cache: Dict[str, Dict[LayoutCacheKey, Any]] = {}
165 atexit.register(self.clear_cache)
167 @classmethod
168 def get_instance(cls) -> "LayoutCacheManager":
169 """
170 Get the singleton instance of LayoutCacheManager.
172 Returns:
173 LayoutCacheManager: The singleton instance.
174 """
175 if cls._instance is None:
176 cls._instance = LayoutCacheManager()
177 return cls._instance
179 def get_layout_cache(self) -> Dict[str, Dict[LayoutCacheKey, Any]]:
180 """
181 Get the layout cache dictionary.
183 Returns:
184 Dict[str, Dict[LayoutCacheKey, Any]]: The nested dictionary mapping
185 operation names to their layout caches.
186 """
187 return self.layout_cache
189 @staticmethod
190 def distributed_op(op_name: str) -> Any:
191 """
192 Get the distributed operation implementation by name.
194 Args:
195 op_name (str): The name of the distributed operation.
197 Returns:
198 Any: The distributed operation class or implementation.
199 """
200 op = get_distributed_op(op_name)
201 return op
203 def clear_cache(self) -> None:
204 """
205 Clear all cached layouts.
207 This method is automatically registered with atexit to ensure
208 cache is cleared when the program exits.
209 """
210 self.layout_cache.clear()
213class OpDispatcher:
214 """
215 OpDispatcher
216 """
218 # Whitelisted ops that mutate args[0]'s storage in place. The dispatch bypass
219 # must return the original DTensor self for these, not the unwrapped local
220 # result, or it demotes a DTensor accumulator to a plain Tensor and breaks the
221 # next op that adds a DTensor to it (e.g. grad-accumulation `loss += micro_loss`).
222 # Class-level so it stays available on instances built via __new__ (e.g. tests).
223 _INPLACE_BYPASS_OPS = frozenset(
224 {"InplaceAddExt", "InplaceSubExt", "InplaceMul", "InplaceDiv"})
226 def __init__(self):
227 self.yaml_dir = _DISTRIBUTED_OPS_YAML_DIR
229 self.layout_infer_ops = self.safe_load_yaml_from_dir()
230 # frozenset for O(1) membership (checked on every dispatch's bypass test).
231 self.whitelist = frozenset({"typeof", "DistCommIsend",
232 "DistCommIrecv", "DistCommBroadcast", "DistCommAllReduce", "DistCommAllGather",
233 "DistCommBatchIsendIrecv",
234 "DistCommReduceScatter", "requires_grad_", "item", "__get__", "__set__",
235 "register_hook",
236 "is_complex", "chunk", "__bool__", "__len__", "__format__", "dim",
237 "_has_compatible_shallow_copy_type", "is_floating_point", "is_contiguous",
238 "get_device", "clamp_", "foreach_mul", "_foreach_mul", "_foreach_mul_",
239 "_foreach_add_", "_foreach_norm", "fused_adamw", "_fused_adamw_", "is_pinned",
240 "new_tensor"})
242 self._random_ops = {
243 "normal_", "uniform_", "bernoulli", "bernoulli_",
244 "native_dropout", "rand", "rand_like", "randn",
245 "randn_like", "randint_like", "kaiming_uniform_",
246 "multinomial",
247 }
248 self._rng_tracker: Optional[OffsetBasedRNGTracker] = None
249 # Op names proven to be loss/CE-irrelevant (both is_loss_parallel_op and
250 # is_decomposed_ce_op are False). For these the loss_parallel / decomposed-CE
251 # guards in dispatch() are always no-ops regardless of context, so we skip
252 # them (and their is_loss_parallel_active() contextvar reads) on later calls.
253 self._non_loss_ops: set = set()
255 self._register_distributed_ops()
257 def _register_distributed_ops(self):
258 for op_name, config in self.layout_infer_ops.items():
259 self._register_single_distributed_op(op_name, config)
261 @staticmethod
262 @lru_cache(maxsize=None)
263 def _resolve_distributed_op_module(module_file: str) -> str:
264 """Resolve a distributed-op module after validating its package source.
266 Args:
267 module_file: Simple module name declared by a package-owned YAML file.
269 Returns:
270 Fully qualified module name within the built-in distributed-op package.
272 Raises:
273 ValueError: If module_file is not a simple Python identifier.
274 ImportError: If the module cannot be resolved from the trusted ops directory.
275 """
276 if not isinstance(module_file, str) or not module_file.isidentifier():
277 raise ValueError(f"Invalid distributed op module name: {module_file!r}")
279 module_name = f"{_DISTRIBUTED_OPS_PACKAGE}.{module_file}"
280 spec = importlib.util.find_spec(module_name)
281 if spec is None or spec.origin is None:
282 raise ImportError(f"Distributed op module cannot be resolved: {module_name}")
284 module_origin = os.path.realpath(spec.origin)
285 try:
286 common_path = os.path.commonpath((_DISTRIBUTED_OPS_DIR, module_origin))
287 except ValueError as exc:
288 raise ImportError(f"Distributed op module has an untrusted source: {module_name}") from exc
289 if common_path != _DISTRIBUTED_OPS_DIR:
290 raise ImportError(f"Distributed op module has an untrusted source: {module_name}")
291 return module_name
293 def _register_single_distributed_op(self, op_name: str, config: dict):
294 """Import and register a package-owned distributed operator implementation."""
295 if "distributed_op_module" in config:
296 raise ValueError("distributed_op_module is not supported; use distributed_op_file")
298 class_name = config["distributed_op_class"]
299 if not isinstance(class_name, str) or not class_name.isidentifier():
300 raise ValueError(f"Invalid distributed op class name: {class_name!r}")
302 module_file = config["distributed_op_file"]
303 module_name = self._resolve_distributed_op_module(module_file)
304 module = importlib.import_module(module_name)
305 op_class = getattr(module, class_name)
306 if not isinstance(op_class, type) or not issubclass(op_class, DistributedOp):
307 raise TypeError(f"Distributed op class must inherit DistributedOp: {class_name}")
308 _ = op_class(op_name)
310 @staticmethod
311 def _merge_default(config: dict):
312 """Apply __default__ values to all ops in this YAML file."""
313 if "__default__" not in config:
314 return config
316 default_cfg = config["__default__"]
317 merged = {}
319 for op_name, op_cfg in config.items():
320 if op_name == "__default__":
321 continue
323 new_cfg = default_cfg.copy()
324 new_cfg.update(op_cfg)
325 merged[op_name] = new_cfg
327 return merged
329 def safe_load_yaml_from_dir(self) -> dict:
330 """
331 Load yaml dictionary from directory.
333 Returns:
334 dict: Merged dictionary of all operator configurations loaded from YAML files.
335 """
336 yaml_dict = {}
337 yaml_path = self.yaml_dir
338 if not os.path.isdir(yaml_path):
339 raise ValueError(f"Invalid yaml directory path: {yaml_path}")
341 for yaml_file_path in glob.glob(os.path.join(yaml_path, '*.yaml')):
342 with open(yaml_file_path, 'r', encoding="utf-8") as f:
343 yaml_data = yaml.safe_load(f)
345 yaml_data = OpDispatcher._merge_default(yaml_data)
346 for name, data in yaml_data.items():
347 if name in yaml_dict:
348 raise ValueError(f"Duplicate yaml object with name '{name}'.")
349 yaml_dict[name] = data
351 return yaml_dict
353 def _dispatch_random_op(self, op_name: str, op_call: callable, args, kwargs):
354 """Handle dispatch for random ops that operate on DTensors."""
355 first_arg = next(
356 (x for x in chain(args, kwargs.values()) if isinstance(x, DTensor)),
357 None,
358 )
359 # Fall back to the default op if no DTensor is found.
360 if first_arg is None:
361 return op_call(*args, **kwargs)
363 local_args = [arg.to_local() if isinstance(arg, DTensor) else arg for arg in args]
364 local_kwargs = {k: v.to_local() if isinstance(v, DTensor) else v for k, v in kwargs.items()}
365 first_local_arg = first_arg.to_local()
367 if self._rng_tracker is None and is_rng_supported_mesh(first_arg.device_mesh):
368 self._rng_tracker = OffsetBasedRNGTracker()
370 maybe_user_generator = local_kwargs.pop("generator", None)
371 if (
372 self._rng_tracker is not None
373 and not first_local_arg.is_meta
374 and self._rng_tracker.distribute_region_enabled
375 ):
376 # pylint: disable=W0212
377 with self._rng_tracker._distribute_region(
378 device_mesh=first_arg.device_mesh,
379 placements=first_arg.placements,
380 global_shape=first_arg.shape,
381 generator=maybe_user_generator,
382 ):
383 local_results = op_call(*local_args, **local_kwargs)
384 else:
385 if maybe_user_generator is not None:
386 local_kwargs["generator"] = maybe_user_generator
387 local_results = op_call(*local_args, **local_kwargs)
389 return self._wrap_random_result(op_name, local_results, first_arg, args, kwargs)
391 @staticmethod
392 def _func_dropout_ext_inplace(args, kwargs) -> bool:
393 """Return True when FuncDropoutExt is invoked with inplace=True."""
394 # Kernel signature: (input, p, training, inplace, seed, offset).
395 if len(args) >= 4:
396 return bool(args[3])
397 return bool(kwargs.get("inplace", False))
399 @staticmethod
400 def _random_op_returns_self(op_name: str, args, kwargs) -> bool:
401 """Return True when a random op mutates an existing DTensor in place."""
402 if op_name == "FuncDropoutExt":
403 return OpDispatcher._func_dropout_ext_inplace(args, kwargs)
404 # Torch random inplace ops follow the ATen '_' suffix convention.
405 return op_name.endswith('_')
407 @staticmethod
408 def _wrap_random_result(op_name, local_results, first_arg, args, kwargs):
409 """Wrap a random op's local result(s) back into DTensor(s).
411 In-place ops return the input DTensor itself. Torch random inplace ops use
412 the ATen '_' suffix. ``FuncDropoutExt`` is handled separately because the
413 same kernel serves both modes via its ``inplace`` argument.
414 """
415 if OpDispatcher._random_op_returns_self(op_name, args, kwargs):
416 return first_arg
417 mesh = first_arg.device_mesh
418 placements = first_arg.layout.alias_placements
419 # Some ops return tuple/list, e.g. native_dropout returns (output, mask).
420 if isinstance(local_results, (tuple, list)):
421 return tuple(
422 DTensor.from_local(r, mesh, placements) if isinstance(r, torch.Tensor) else r
423 for r in local_results
424 )
425 if isinstance(local_results, torch.Tensor):
426 return DTensor.from_local(local_results, mesh, placements)
427 # Fallback: return as-is for non-Tensor results (currently unreachable with existing _random_ops).
428 return local_results
430 @staticmethod
431 def _unwrap_value(value: object) -> object:
432 """Replace DTensor with its local tensor; pass scalars and plain tensors through.
434 Args:
435 value (object): A single argument value from an op call.
437 Returns:
438 object: The local tensor if value is a DTensor, otherwise value unchanged.
439 """
440 if isinstance(value, DTensor):
441 return value.to_local()
442 if isinstance(value, tuple):
443 return tuple(OpDispatcher._unwrap_value(e) for e in value)
444 if isinstance(value, list):
445 return [OpDispatcher._unwrap_value(e) for e in value]
446 return value
448 @staticmethod
449 def _unwrap_args(args: tuple) -> list:
450 """Strip DTensor wrappers from args, preserving tuple/list container structure.
452 Args:
453 args: Op call positional arguments, may contain DTensor instances.
455 Returns:
456 List of args with DTensor replaced by their local tensors.
457 """
458 return [OpDispatcher._unwrap_value(arg) for arg in args]
460 @staticmethod
461 def _unwrap_kwargs(kwargs: dict) -> dict:
462 """Strip DTensor wrappers from kwargs values, preserving tuple/list container structure.
464 Args:
465 kwargs: Op call keyword arguments, values may contain DTensor instances.
467 Returns:
468 Dict of kwargs with DTensor values replaced by their local tensors.
469 """
470 return {k: OpDispatcher._unwrap_value(v) for k, v in kwargs.items()}
472 @staticmethod
473 def _collect_dtensors(value: object) -> List[DTensor]:
474 """Return all DTensors nested in one dispatch argument."""
475 if isinstance(value, DTensor):
476 return [value]
477 if isinstance(value, (tuple, list)):
478 return list(chain.from_iterable(
479 OpDispatcher._collect_dtensors(item) for item in value
480 ))
481 if isinstance(value, dict):
482 return list(chain.from_iterable(
483 OpDispatcher._collect_dtensors(item) for item in value.values()
484 ))
485 return []
487 def _validate_ragged_dispatch(
488 self, op_name: str, args: tuple, kwargs: dict
489 ) -> Optional[DTensor]:
490 """Return the first Ragged input for a whitelisted elementwise op."""
491 reference = next(
492 (
493 dtensor for dtensor in self._collect_dtensors((args, kwargs))
494 if isinstance(getattr(dtensor.layout, "ragged_shard", None), RaggedShardInfo)
495 ),
496 None,
497 )
498 if reference is None:
499 return None
500 if op_name not in _RAGGED_ELEMENTWISE_OPS:
501 raise RuntimeError(
502 f"Operator {op_name!r} does not support RaggedShard in phase one"
503 )
504 return reference
506 def _dispatch_ragged_elementwise(
507 self, op_call: callable, args: tuple, kwargs: dict,
508 reference: DTensor,
509 ) -> DTensor:
510 """Execute a whitelisted op locally and inherit its Ragged Layout."""
511 local_args = tuple(self._unwrap_args(args))
512 local_kwargs = self._unwrap_kwargs(kwargs)
513 py_output = op_call(*local_args, **local_kwargs)
514 op_name = get_op_name(op_call)
515 if op_name in _RAGGED_INPLACE_ELEMENTWISE_OPS:
516 if not args or not isinstance(args[0], DTensor):
517 raise ValueError(
518 f"Ragged in-place operator {op_name!r} requires a DTensor first argument"
519 )
520 return args[0]
521 return DTensor.from_local_with_layout(
522 py_output,
523 copy.deepcopy(reference.layout),
524 shape=tuple(reference.shape),
525 )
527 @staticmethod
528 def _gather_dtensors_to_full(args: tuple, kwargs: dict) -> tuple:
529 """Gather all DTensor arguments to full tensors for fallback execution.
531 Used when an operator has no parallel layout implementation. All DTensor
532 arguments are gathered to full tensors before calling the standard operator.
534 Args:
535 args: Op call positional arguments, may contain DTensor instances.
536 kwargs: Op call keyword arguments, may contain DTensor instances.
538 Returns:
539 Tuple of (unwrapped_args, unwrapped_kwargs) with DTensor values
540 replaced by their full tensor representations.
542 Warning:
543 This fallback performs all-gather which may consume significant memory.
544 Operators without layout implementations should be registered properly.
545 """
546 def gather(value: object) -> object:
547 if isinstance(value, DTensor):
548 return value.full_tensor()
549 if isinstance(value, tuple):
550 return tuple(gather(e) for e in value)
551 if isinstance(value, list):
552 return [gather(e) for e in value]
553 return value
555 gathered_args = [gather(arg) for arg in args]
556 gathered_kwargs = {k: gather(v) for k, v in kwargs.items()}
558 warnings.warn(
559 "Operator has no distributed layout implementation. "
560 "Falling back to all-gather which may consume significant memory. "
561 "Consider registering a proper distributed operator.",
562 UserWarning,
563 stacklevel=4
564 )
566 return gathered_args, gathered_kwargs
568 def _should_bypass_dispatch(self, op_name: str) -> bool:
569 """Return True if the op should bypass DTensor dispatch and run locally.
571 Args:
572 op_name: Canonical operator name from get_op_name().
574 Returns:
575 True when the op is whitelisted or DTensor dispatch is globally disabled.
576 """
577 # Cheap O(1) frozenset checks first, short-circuit before the ContextVar
578 # read (get_dtensor_dispatch) which is the priciest part of this guard.
579 if op_name in self.whitelist or op_name in self._INPLACE_BYPASS_OPS:
580 return True
581 return get_dtensor_dispatch() is False and op_name not in get_no_skip_ops()
583 @staticmethod
584 def _validate_inplace_partial_inputs(op_name: str, args: tuple, kwargs: dict) -> None:
585 """Reject local in-place add/sub when Partial contributions need gating."""
586 if op_name not in {"InplaceAddExt", "InplaceSubExt"} or not args:
587 return
588 first = args[0]
589 if len(args) >= 2:
590 second = args[1]
591 elif "other" in kwargs:
592 second = kwargs["other"]
593 else:
594 return
595 if not isinstance(first, DTensor):
596 return
597 mesh_ndim = len(first.layout.partial)
598 first_partial = tuple(first.layout.partial)
599 if isinstance(second, DTensor):
600 second_partial = tuple(second.layout.partial)
601 if len(second_partial) != mesh_ndim:
602 raise ValueError(
603 f"For {op_name}, in-place input mesh dimensions must match, "
604 f"but got {mesh_ndim} and {len(second_partial)}."
605 )
606 else:
607 second_partial = (None,) * mesh_ndim
608 if first_partial != second_partial:
609 raise ValueError(
610 f"For {op_name}, input Partial placements must be identical for "
611 f"local in-place execution, but got {first_partial} and {second_partial}."
612 )
614 def _should_dispatch_loss_parallel(self, op_name: str) -> bool:
615 """Check if should dispatch through loss_parallel path.
617 Args:
618 op_name: Canonical operator name from get_op_name().
620 Returns:
621 True when in loss_parallel context and op is a CE entry point.
622 """
623 return is_loss_parallel_active() and is_loss_parallel_op(op_name)
625 def _check_decomposed_ce_op_in_loss_parallel(self, op_name: str, args: tuple, kwargs: dict):
626 """Check if decomposed CE ops are called in loss_parallel context.
628 Args:
629 op_name: Canonical operator name.
630 args: Positional arguments for op_call.
631 kwargs: Keyword arguments for op_call.
633 Raises:
634 ValueError: If decomposed CE op is called in loss_parallel context
635 with vocab-sharded DTensor input.
636 """
637 if not is_loss_parallel_active() or not is_decomposed_ce_op(op_name):
638 return
640 has_vocab_sharded_dtensor = False
641 for arg in args:
642 if isinstance(arg, DTensor) and _is_shard_on_last_dim(arg):
643 has_vocab_sharded_dtensor = True
644 break
645 if not has_vocab_sharded_dtensor:
646 for val in kwargs.values():
647 if isinstance(val, DTensor) and _is_shard_on_last_dim(val):
648 has_vocab_sharded_dtensor = True
649 break
651 if has_vocab_sharded_dtensor:
652 raise ValueError(
653 f"Operator '{op_name}' is a decomposed component of cross_entropy and should not be called "
654 f"directly within loss_parallel() context. Use F.cross_entropy(logits, targets) instead. "
655 f"For example, replace:\n"
656 f" with loss_parallel():\n"
657 f" log_probs = F.log_softmax(logits, dim=-1)\n"
658 f" loss = F.nll_loss(log_probs, targets)\n"
659 f"with:\n"
660 f" with loss_parallel():\n"
661 f" loss = F.cross_entropy(logits, targets)"
662 )
664 def _dispatch_loss_parallel(self, op_call: callable, args: tuple, kwargs: dict):
665 """Dispatch cross_entropy through the loss_parallel distributed kernel.
667 Args:
668 op_call: The raw operator callable.
669 args: Positional arguments for op_call.
670 kwargs: Keyword arguments for op_call.
672 Returns:
673 Result of the distributed cross_entropy computation.
674 """
675 return distributed_cross_entropy_from_op_call(op_call, args, kwargs)
677 def _check_ce_op_without_loss_parallel_context(self, op_name: str, args: tuple):
678 """Check if CE op is called with Shard(-1) DTensor outside loss_parallel context.
680 Args:
681 op_name: Canonical operator name.
682 args: Positional arguments for op_call.
684 Raises:
685 ValueError: If CE op is called with Shard(-1) logits outside loss_parallel context.
686 """
687 if is_loss_parallel_active() or not is_loss_parallel_op(op_name):
688 return
690 if len(args) == 0 or not isinstance(args[0], DTensor):
691 return
693 logits = args[0]
694 if _is_shard_on_last_dim(logits):
695 raise ValueError(
696 f"Operator '{op_name}' requires loss_parallel context when input logits are "
697 f"sharded on the vocabulary dimension (Shard(-1)). Please wrap your forward "
698 f"and backward pass with loss_parallel():\n"
699 f" with loss_parallel():\n"
700 f" loss = F.cross_entropy(logits, targets)\n"
701 f" loss.backward()\n"
702 f"If you intentionally want to gather all shards to compute cross_entropy "
703 f"(not recommended for large vocabulary), use logits.full_tensor() explicitly."
704 )
706 def _handle_unregistered_op(
707 self, op_name: str, op_call: callable, args: tuple, kwargs: dict
708 ):
709 """Handle ops that have no registered layout-inference entry.
711 This is a fallback path for ops that are not registered in
712 ``layout_infer_ops``. When arguments contain DTensors it either raises
713 (with a hint to register a distributed op) or, for loss-parallel ops,
714 gathers tensors to full and dispatches through the raw callable.
716 Args:
717 op_name: Canonical operator name.
718 op_call: The raw operator callable.
719 args: Positional arguments for op_call.
720 kwargs: Keyword arguments for op_call.
722 Returns:
723 Raw dispatch result (plain Tensor, not wrapped as DTensor).
725 Raises:
726 RuntimeError: If op_name is not registered for layout inference.
727 """
728 has_dtensor = any(isinstance(arg, DTensor) for arg in args)
729 has_dtensor = has_dtensor or any(isinstance(v, DTensor) for v in kwargs.values())
730 if has_dtensor:
731 self._check_ce_op_without_loss_parallel_context(op_name, args)
733 if not is_loss_parallel_op(op_name):
734 raise RuntimeError(
735 f"Operator {op_name} does not contain parallel layout infer func. "
736 f"DTensor dispatch requires explicit layout inference registration. "
737 f"Please register a distributed operator for '{op_name}' or use local tensors."
738 )
740 gathered_args, gathered_kwargs = self._gather_dtensors_to_full(args, kwargs)
742 # Special handling for cross_entropy with 3D logits (only when NOT in loss_parallel context)
743 # PyTorch expects: logits [N, C], targets [N]
744 # But LLM forward returns: logits [batch, seq, vocab], targets [batch, seq]
745 # Note: nll_loss input is log_probs, typically already 2D, so we only reshape for cross_entropy
746 if op_name == "cross_entropy" and len(gathered_args) >= 2:
747 logits = gathered_args[0]
748 targets = gathered_args[1]
749 if isinstance(logits, torch.Tensor) and isinstance(targets, torch.Tensor):
750 if logits.ndim > 2 and targets.ndim > 1 and targets.ndim == logits.ndim - 1:
751 vocab_size = logits.shape[-1]
752 gathered_args[0] = logits.reshape(-1, vocab_size)
753 gathered_args[1] = targets.reshape(-1)
755 return op_call(*gathered_args, **gathered_kwargs)
756 raise RuntimeError(f"Operator {op_name} does not contain parallel layout infer func.")
758 def _dispatch_layout_infer(
759 self, op_name: str, op_call: callable, args: tuple, kwargs: dict
760 ):
761 """Standard dispatch through layout-inference: preprocess → infer → execute → wrap.
763 Args:
764 op_name: Canonical operator name (already resolved by the caller).
765 op_call: The raw operator callable.
766 args: Positional arguments for op_call.
767 kwargs: Keyword arguments for op_call.
769 Returns:
770 DTensor: Dispatched result wrapped as DTensor.
772 Raises:
773 RuntimeError: If op_name is not registered, or preprocess returns None.
774 """
775 if op_name not in self.layout_infer_ops:
776 return self._handle_unregistered_op(op_name, op_call, args, kwargs)
778 cache_manager = LayoutCacheManager.get_instance()
779 distribute_op = cache_manager.distributed_op(op_name)
781 result = distribute_op.preprocess(args, kwargs)
782 if result is None:
783 raise RuntimeError(
784 f"Operator '{op_name}' has not been migrated to the three-phase dispatch flow. "
785 f"Please implement preprocess() to return (local_args, local_kwargs, cache_values)."
786 )
787 local_args, local_kwargs, cache_values = result
788 cache_key = LayoutCacheKey.from_cache_values(cache_values)
790 infer_result, op_impl = OpDispatcher._lookup_or_infer_layout(
791 op_call, op_name, cache_key, cache_values, distribute_op, cache_manager
792 )
794 op_impl = op_call if op_impl is None else op_impl
795 py_output = op_impl(*local_args, **local_kwargs)
796 output = distribute_op.wrap_output(py_output, infer_result[0])
797 return OpDispatcher._restore_inplace_dtensor_result(op_name, args, output)
799 @staticmethod
800 def _restore_inplace_dtensor_result(op_name: str, args: tuple, output: Any) -> Any:
801 """Return the original DTensor wrapper after a local in-place operation."""
802 if op_name in {"add_", "sub_"} and args and isinstance(args[0], DTensor):
803 return args[0]
804 return output
806 @staticmethod
807 def _lookup_or_infer_layout(func, func_name, cache_key, cache_values, distribute_op, cache_manager):
808 """Look up cached layout or compute via distributed op.
810 Returns:
811 (infer_result, op_impl)
812 """
813 layout_cache = cache_manager.get_layout_cache()
814 if func_name not in layout_cache:
815 layout_cache[func_name] = {}
816 op_layout_cache = layout_cache[func_name]
817 if cache_key in op_layout_cache:
818 return op_layout_cache[cache_key]
819 infer_result = distribute_op.infer_layout(cache_values)
820 op_impl = distribute_op.get_expand_impl(func, infer_result, cache_values)
821 op_layout_cache[cache_key] = (infer_result, op_impl)
822 return infer_result, op_impl
824 def dispatch(self, op_call: callable, args: tuple, kwargs: dict) -> object:
825 """Route an op call through the appropriate DTensor dispatch path.
827 Args:
828 op_call: The raw operator callable.
829 args: Positional arguments for op_call.
830 kwargs: Keyword arguments for op_call.
832 Returns:
833 Result of the dispatched op call.
834 """
835 op_name = get_op_name(op_call)
836 if logger.isEnabledFor(logging.DEBUG):
837 log_dispatch_enter(op_name, args, kwargs)
839 observer = _debug_mode_observer.get()
840 if observer is not None:
841 observer.on_op_dispatch_enter(op_name, op_call, args, kwargs)
843 result = None
844 try:
845 should_bypass = self._should_bypass_dispatch(op_name)
846 ragged_dtensor = None
847 if not should_bypass or op_name not in _RAGGED_METADATA_BYPASS_OPS:
848 ragged_dtensor = self._validate_ragged_dispatch(op_name, args, kwargs)
849 if ragged_dtensor is not None:
850 result = self._dispatch_ragged_elementwise(
851 op_call, args, kwargs, ragged_dtensor
852 )
853 return result
855 if should_bypass:
856 self._validate_inplace_partial_inputs(op_name, args, kwargs)
857 result = op_call(*self._unwrap_args(args), **self._unwrap_kwargs(kwargs))
858 if op_name in self._INPLACE_BYPASS_OPS and args and isinstance(args[0], DTensor):
859 result = args[0]
860 return result
862 if op_name in self._random_ops:
863 result = self._dispatch_random_op(op_name, op_call, args, kwargs)
864 return result
866 self._check_decomposed_ce_op_in_loss_parallel(op_name, args, kwargs)
868 if self._should_dispatch_loss_parallel(op_name):
869 result = self._dispatch_loss_parallel(op_call, args, kwargs)
870 return result
872 if op_name not in self.layout_infer_ops and get_distributed_op(op_name) is not None:
873 self.layout_infer_ops[op_name] = {}
875 result = self._dispatch_layout_infer(op_name, op_call, args, kwargs)
876 return result
877 finally:
878 if logger.isEnabledFor(logging.DEBUG):
879 log_dispatch_exit(op_name, result)
881 if observer is not None:
882 observer.on_op_dispatch_exit(op_name, result)
884_OP_DISPATCHER = OpDispatcher()