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-10-25 05:45 +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"})
240 self._random_ops = {
241 "normal_", "uniform_", "bernoulli", "bernoulli_",
242 "native_dropout", "rand", "rand_like", "randn",
243 "randn_like", "randint_like", "kaiming_uniform_",
244 "multinomial",
245 }
246 self._rng_tracker: Optional[OffsetBasedRNGTracker] = None
247 # Op names proven to be loss/CE-irrelevant (both is_loss_parallel_op and
248 # is_decomposed_ce_op are False). For these the loss_parallel / decomposed-CE
249 # guards in dispatch() are always no-ops regardless of context, so we skip
250 # them (and their is_loss_parallel_active() contextvar reads) on later calls.
251 self._non_loss_ops: set = set()
253 self._register_distributed_ops()
255 def _register_distributed_ops(self):
256 for op_name, config in self.layout_infer_ops.items():
257 self._register_single_distributed_op(op_name, config)
259 @staticmethod
260 @lru_cache(maxsize=None)
261 def _resolve_distributed_op_module(module_file: str) -> str:
262 """Resolve a distributed-op module after validating its package source.
264 Args:
265 module_file: Simple module name declared by a package-owned YAML file.
267 Returns:
268 Fully qualified module name within the built-in distributed-op package.
270 Raises:
271 ValueError: If module_file is not a simple Python identifier.
272 ImportError: If the module cannot be resolved from the trusted ops directory.
273 """
274 if not isinstance(module_file, str) or not module_file.isidentifier():
275 raise ValueError(f"Invalid distributed op module name: {module_file!r}")
277 module_name = f"{_DISTRIBUTED_OPS_PACKAGE}.{module_file}"
278 spec = importlib.util.find_spec(module_name)
279 if spec is None or spec.origin is None:
280 raise ImportError(f"Distributed op module cannot be resolved: {module_name}")
282 module_origin = os.path.realpath(spec.origin)
283 try:
284 common_path = os.path.commonpath((_DISTRIBUTED_OPS_DIR, module_origin))
285 except ValueError as exc:
286 raise ImportError(f"Distributed op module has an untrusted source: {module_name}") from exc
287 if common_path != _DISTRIBUTED_OPS_DIR:
288 raise ImportError(f"Distributed op module has an untrusted source: {module_name}")
289 return module_name
291 def _register_single_distributed_op(self, op_name: str, config: dict):
292 """Import and register a package-owned distributed operator implementation."""
293 if "distributed_op_module" in config:
294 raise ValueError("distributed_op_module is not supported; use distributed_op_file")
296 class_name = config["distributed_op_class"]
297 if not isinstance(class_name, str) or not class_name.isidentifier():
298 raise ValueError(f"Invalid distributed op class name: {class_name!r}")
300 module_file = config["distributed_op_file"]
301 module_name = self._resolve_distributed_op_module(module_file)
302 module = importlib.import_module(module_name)
303 op_class = getattr(module, class_name)
304 if not isinstance(op_class, type) or not issubclass(op_class, DistributedOp):
305 raise TypeError(f"Distributed op class must inherit DistributedOp: {class_name}")
306 _ = op_class(op_name)
308 @staticmethod
309 def _merge_default(config: dict):
310 """Apply __default__ values to all ops in this YAML file."""
311 if "__default__" not in config:
312 return config
314 default_cfg = config["__default__"]
315 merged = {}
317 for op_name, op_cfg in config.items():
318 if op_name == "__default__":
319 continue
321 new_cfg = default_cfg.copy()
322 new_cfg.update(op_cfg)
323 merged[op_name] = new_cfg
325 return merged
327 def safe_load_yaml_from_dir(self) -> dict:
328 """
329 Load yaml dictionary from directory.
331 Returns:
332 dict: Merged dictionary of all operator configurations loaded from YAML files.
333 """
334 yaml_dict = {}
335 yaml_path = self.yaml_dir
336 if not os.path.isdir(yaml_path):
337 raise ValueError(f"Invalid yaml directory path: {yaml_path}")
339 for yaml_file_path in glob.glob(os.path.join(yaml_path, '*.yaml')):
340 with open(yaml_file_path, 'r', encoding="utf-8") as f:
341 yaml_data = yaml.safe_load(f)
343 yaml_data = OpDispatcher._merge_default(yaml_data)
344 for name, data in yaml_data.items():
345 if name in yaml_dict:
346 raise ValueError(f"Duplicate yaml object with name '{name}'.")
347 yaml_dict[name] = data
349 return yaml_dict
351 def _dispatch_random_op(self, op_name: str, op_call: callable, args, kwargs):
352 """Handle dispatch for random ops that operate on DTensors."""
353 first_arg = next(
354 (x for x in chain(args, kwargs.values()) if isinstance(x, DTensor)),
355 None,
356 )
357 # Fall back to the default op if no DTensor is found.
358 if first_arg is None:
359 return op_call(*args, **kwargs)
361 local_args = [arg.to_local() if isinstance(arg, DTensor) else arg for arg in args]
362 local_kwargs = {k: v.to_local() if isinstance(v, DTensor) else v for k, v in kwargs.items()}
363 first_local_arg = first_arg.to_local()
365 if self._rng_tracker is None and is_rng_supported_mesh(first_arg.device_mesh):
366 self._rng_tracker = OffsetBasedRNGTracker()
368 maybe_user_generator = local_kwargs.pop("generator", None)
369 if (
370 self._rng_tracker is not None
371 and not first_local_arg.is_meta
372 and self._rng_tracker.distribute_region_enabled
373 ):
374 # pylint: disable=W0212
375 with self._rng_tracker._distribute_region(
376 device_mesh=first_arg.device_mesh,
377 placements=first_arg.placements,
378 global_shape=first_arg.shape,
379 generator=maybe_user_generator,
380 ):
381 local_results = op_call(*local_args, **local_kwargs)
382 else:
383 if maybe_user_generator is not None:
384 local_kwargs["generator"] = maybe_user_generator
385 local_results = op_call(*local_args, **local_kwargs)
387 return self._wrap_random_result(op_name, local_results, first_arg, args, kwargs)
389 @staticmethod
390 def _func_dropout_ext_inplace(args, kwargs) -> bool:
391 """Return True when FuncDropoutExt is invoked with inplace=True."""
392 # Kernel signature: (input, p, training, inplace, seed, offset).
393 if len(args) >= 4:
394 return bool(args[3])
395 return bool(kwargs.get("inplace", False))
397 @staticmethod
398 def _random_op_returns_self(op_name: str, args, kwargs) -> bool:
399 """Return True when a random op mutates an existing DTensor in place."""
400 if op_name == "FuncDropoutExt":
401 return OpDispatcher._func_dropout_ext_inplace(args, kwargs)
402 # Torch random inplace ops follow the ATen '_' suffix convention.
403 return op_name.endswith('_')
405 @staticmethod
406 def _wrap_random_result(op_name, local_results, first_arg, args, kwargs):
407 """Wrap a random op's local result(s) back into DTensor(s).
409 In-place ops return the input DTensor itself. Torch random inplace ops use
410 the ATen '_' suffix. ``FuncDropoutExt`` is handled separately because the
411 same kernel serves both modes via its ``inplace`` argument.
412 """
413 if OpDispatcher._random_op_returns_self(op_name, args, kwargs):
414 return first_arg
415 mesh = first_arg.device_mesh
416 placements = first_arg.layout.alias_placements
417 # Some ops return tuple/list, e.g. native_dropout returns (output, mask).
418 if isinstance(local_results, (tuple, list)):
419 return tuple(
420 DTensor.from_local(r, mesh, placements) if isinstance(r, torch.Tensor) else r
421 for r in local_results
422 )
423 if isinstance(local_results, torch.Tensor):
424 return DTensor.from_local(local_results, mesh, placements)
425 # Fallback: return as-is for non-Tensor results (currently unreachable with existing _random_ops).
426 return local_results
428 @staticmethod
429 def _unwrap_value(value: object) -> object:
430 """Replace DTensor with its local tensor; pass scalars and plain tensors through.
432 Args:
433 value (object): A single argument value from an op call.
435 Returns:
436 object: The local tensor if value is a DTensor, otherwise value unchanged.
437 """
438 if isinstance(value, DTensor):
439 return value.to_local()
440 if isinstance(value, tuple):
441 return tuple(OpDispatcher._unwrap_value(e) for e in value)
442 if isinstance(value, list):
443 return [OpDispatcher._unwrap_value(e) for e in value]
444 return value
446 @staticmethod
447 def _unwrap_args(args: tuple) -> list:
448 """Strip DTensor wrappers from args, preserving tuple/list container structure.
450 Args:
451 args: Op call positional arguments, may contain DTensor instances.
453 Returns:
454 List of args with DTensor replaced by their local tensors.
455 """
456 return [OpDispatcher._unwrap_value(arg) for arg in args]
458 @staticmethod
459 def _unwrap_kwargs(kwargs: dict) -> dict:
460 """Strip DTensor wrappers from kwargs values, preserving tuple/list container structure.
462 Args:
463 kwargs: Op call keyword arguments, values may contain DTensor instances.
465 Returns:
466 Dict of kwargs with DTensor values replaced by their local tensors.
467 """
468 return {k: OpDispatcher._unwrap_value(v) for k, v in kwargs.items()}
470 @staticmethod
471 def _collect_dtensors(value: object) -> List[DTensor]:
472 """Return all DTensors nested in one dispatch argument."""
473 if isinstance(value, DTensor):
474 return [value]
475 if isinstance(value, (tuple, list)):
476 return list(chain.from_iterable(
477 OpDispatcher._collect_dtensors(item) for item in value
478 ))
479 if isinstance(value, dict):
480 return list(chain.from_iterable(
481 OpDispatcher._collect_dtensors(item) for item in value.values()
482 ))
483 return []
485 def _validate_ragged_dispatch(
486 self, op_name: str, args: tuple, kwargs: dict
487 ) -> Optional[DTensor]:
488 """Return the first Ragged input for a whitelisted elementwise op."""
489 reference = next(
490 (
491 dtensor for dtensor in self._collect_dtensors((args, kwargs))
492 if isinstance(getattr(dtensor.layout, "ragged_shard", None), RaggedShardInfo)
493 ),
494 None,
495 )
496 if reference is None:
497 return None
498 if op_name not in _RAGGED_ELEMENTWISE_OPS:
499 raise RuntimeError(
500 f"Operator {op_name!r} does not support RaggedShard in phase one"
501 )
502 return reference
504 def _dispatch_ragged_elementwise(
505 self, op_call: callable, args: tuple, kwargs: dict,
506 reference: DTensor,
507 ) -> DTensor:
508 """Execute a whitelisted op locally and inherit its Ragged Layout."""
509 local_args = tuple(self._unwrap_args(args))
510 local_kwargs = self._unwrap_kwargs(kwargs)
511 py_output = op_call(*local_args, **local_kwargs)
512 op_name = get_op_name(op_call)
513 if op_name in _RAGGED_INPLACE_ELEMENTWISE_OPS:
514 if not args or not isinstance(args[0], DTensor):
515 raise ValueError(
516 f"Ragged in-place operator {op_name!r} requires a DTensor first argument"
517 )
518 return args[0]
519 return DTensor.from_local_with_layout(
520 py_output,
521 copy.deepcopy(reference.layout),
522 shape=tuple(reference.shape),
523 )
525 @staticmethod
526 def _gather_dtensors_to_full(args: tuple, kwargs: dict) -> tuple:
527 """Gather all DTensor arguments to full tensors for fallback execution.
529 Used when an operator has no parallel layout implementation. All DTensor
530 arguments are gathered to full tensors before calling the standard operator.
532 Args:
533 args: Op call positional arguments, may contain DTensor instances.
534 kwargs: Op call keyword arguments, may contain DTensor instances.
536 Returns:
537 Tuple of (unwrapped_args, unwrapped_kwargs) with DTensor values
538 replaced by their full tensor representations.
540 Warning:
541 This fallback performs all-gather which may consume significant memory.
542 Operators without layout implementations should be registered properly.
543 """
544 def gather(value: object) -> object:
545 if isinstance(value, DTensor):
546 return value.full_tensor()
547 if isinstance(value, tuple):
548 return tuple(gather(e) for e in value)
549 if isinstance(value, list):
550 return [gather(e) for e in value]
551 return value
553 gathered_args = [gather(arg) for arg in args]
554 gathered_kwargs = {k: gather(v) for k, v in kwargs.items()}
556 warnings.warn(
557 "Operator has no distributed layout implementation. "
558 "Falling back to all-gather which may consume significant memory. "
559 "Consider registering a proper distributed operator.",
560 UserWarning,
561 stacklevel=4
562 )
564 return gathered_args, gathered_kwargs
566 def _should_bypass_dispatch(self, op_name: str) -> bool:
567 """Return True if the op should bypass DTensor dispatch and run locally.
569 Args:
570 op_name: Canonical operator name from get_op_name().
572 Returns:
573 True when the op is whitelisted or DTensor dispatch is globally disabled.
574 """
575 # Cheap O(1) frozenset checks first, short-circuit before the ContextVar
576 # read (get_dtensor_dispatch) which is the priciest part of this guard.
577 if op_name in self.whitelist or op_name in self._INPLACE_BYPASS_OPS:
578 return True
579 return get_dtensor_dispatch() is False and op_name not in get_no_skip_ops()
581 @staticmethod
582 def _validate_inplace_partial_inputs(op_name: str, args: tuple, kwargs: dict) -> None:
583 """Reject local in-place add/sub when Partial contributions need gating."""
584 if op_name not in {"InplaceAddExt", "InplaceSubExt"} or not args:
585 return
586 first = args[0]
587 if len(args) >= 2:
588 second = args[1]
589 elif "other" in kwargs:
590 second = kwargs["other"]
591 else:
592 return
593 if not isinstance(first, DTensor):
594 return
595 mesh_ndim = len(first.layout.partial)
596 first_partial = tuple(first.layout.partial)
597 if isinstance(second, DTensor):
598 second_partial = tuple(second.layout.partial)
599 if len(second_partial) != mesh_ndim:
600 raise ValueError(
601 f"For {op_name}, in-place input mesh dimensions must match, "
602 f"but got {mesh_ndim} and {len(second_partial)}."
603 )
604 else:
605 second_partial = (None,) * mesh_ndim
606 if first_partial != second_partial:
607 raise ValueError(
608 f"For {op_name}, input Partial placements must be identical for "
609 f"local in-place execution, but got {first_partial} and {second_partial}."
610 )
612 def _should_dispatch_loss_parallel(self, op_name: str) -> bool:
613 """Check if should dispatch through loss_parallel path.
615 Args:
616 op_name: Canonical operator name from get_op_name().
618 Returns:
619 True when in loss_parallel context and op is a CE entry point.
620 """
621 return is_loss_parallel_active() and is_loss_parallel_op(op_name)
623 def _check_decomposed_ce_op_in_loss_parallel(self, op_name: str, args: tuple, kwargs: dict):
624 """Check if decomposed CE ops are called in loss_parallel context.
626 Args:
627 op_name: Canonical operator name.
628 args: Positional arguments for op_call.
629 kwargs: Keyword arguments for op_call.
631 Raises:
632 ValueError: If decomposed CE op is called in loss_parallel context
633 with vocab-sharded DTensor input.
634 """
635 if not is_loss_parallel_active() or not is_decomposed_ce_op(op_name):
636 return
638 has_vocab_sharded_dtensor = False
639 for arg in args:
640 if isinstance(arg, DTensor) and _is_shard_on_last_dim(arg):
641 has_vocab_sharded_dtensor = True
642 break
643 if not has_vocab_sharded_dtensor:
644 for val in kwargs.values():
645 if isinstance(val, DTensor) and _is_shard_on_last_dim(val):
646 has_vocab_sharded_dtensor = True
647 break
649 if has_vocab_sharded_dtensor:
650 raise ValueError(
651 f"Operator '{op_name}' is a decomposed component of cross_entropy and should not be called "
652 f"directly within loss_parallel() context. Use F.cross_entropy(logits, targets) instead. "
653 f"For example, replace:\n"
654 f" with loss_parallel():\n"
655 f" log_probs = F.log_softmax(logits, dim=-1)\n"
656 f" loss = F.nll_loss(log_probs, targets)\n"
657 f"with:\n"
658 f" with loss_parallel():\n"
659 f" loss = F.cross_entropy(logits, targets)"
660 )
662 def _dispatch_loss_parallel(self, op_call: callable, args: tuple, kwargs: dict):
663 """Dispatch cross_entropy through the loss_parallel distributed kernel.
665 Args:
666 op_call: The raw operator callable.
667 args: Positional arguments for op_call.
668 kwargs: Keyword arguments for op_call.
670 Returns:
671 Result of the distributed cross_entropy computation.
672 """
673 return distributed_cross_entropy_from_op_call(op_call, args, kwargs)
675 def _check_ce_op_without_loss_parallel_context(self, op_name: str, args: tuple):
676 """Check if CE op is called with Shard(-1) DTensor outside loss_parallel context.
678 Args:
679 op_name: Canonical operator name.
680 args: Positional arguments for op_call.
682 Raises:
683 ValueError: If CE op is called with Shard(-1) logits outside loss_parallel context.
684 """
685 if is_loss_parallel_active() or not is_loss_parallel_op(op_name):
686 return
688 if len(args) == 0 or not isinstance(args[0], DTensor):
689 return
691 logits = args[0]
692 if _is_shard_on_last_dim(logits):
693 raise ValueError(
694 f"Operator '{op_name}' requires loss_parallel context when input logits are "
695 f"sharded on the vocabulary dimension (Shard(-1)). Please wrap your forward "
696 f"and backward pass with loss_parallel():\n"
697 f" with loss_parallel():\n"
698 f" loss = F.cross_entropy(logits, targets)\n"
699 f" loss.backward()\n"
700 f"If you intentionally want to gather all shards to compute cross_entropy "
701 f"(not recommended for large vocabulary), use logits.full_tensor() explicitly."
702 )
704 def _handle_unregistered_op(
705 self, op_name: str, op_call: callable, args: tuple, kwargs: dict
706 ):
707 """Handle ops that have no registered layout-inference entry.
709 This is a fallback path for ops that are not registered in
710 ``layout_infer_ops``. When arguments contain DTensors it either raises
711 (with a hint to register a distributed op) or, for loss-parallel ops,
712 gathers tensors to full and dispatches through the raw callable.
714 Args:
715 op_name: Canonical operator name.
716 op_call: The raw operator callable.
717 args: Positional arguments for op_call.
718 kwargs: Keyword arguments for op_call.
720 Returns:
721 Raw dispatch result (plain Tensor, not wrapped as DTensor).
723 Raises:
724 RuntimeError: If op_name is not registered for layout inference.
725 """
726 has_dtensor = any(isinstance(arg, DTensor) for arg in args)
727 has_dtensor = has_dtensor or any(isinstance(v, DTensor) for v in kwargs.values())
728 if has_dtensor:
729 self._check_ce_op_without_loss_parallel_context(op_name, args)
731 if not is_loss_parallel_op(op_name):
732 raise RuntimeError(
733 f"Operator {op_name} does not contain parallel layout infer func. "
734 f"DTensor dispatch requires explicit layout inference registration. "
735 f"Please register a distributed operator for '{op_name}' or use local tensors."
736 )
738 gathered_args, gathered_kwargs = self._gather_dtensors_to_full(args, kwargs)
740 # Special handling for cross_entropy with 3D logits (only when NOT in loss_parallel context)
741 # PyTorch expects: logits [N, C], targets [N]
742 # But LLM forward returns: logits [batch, seq, vocab], targets [batch, seq]
743 # Note: nll_loss input is log_probs, typically already 2D, so we only reshape for cross_entropy
744 if op_name == "cross_entropy" and len(gathered_args) >= 2:
745 logits = gathered_args[0]
746 targets = gathered_args[1]
747 if isinstance(logits, torch.Tensor) and isinstance(targets, torch.Tensor):
748 if logits.ndim > 2 and targets.ndim > 1 and targets.ndim == logits.ndim - 1:
749 vocab_size = logits.shape[-1]
750 gathered_args[0] = logits.reshape(-1, vocab_size)
751 gathered_args[1] = targets.reshape(-1)
753 return op_call(*gathered_args, **gathered_kwargs)
754 raise RuntimeError(f"Operator {op_name} does not contain parallel layout infer func.")
756 def _dispatch_layout_infer(
757 self, op_name: str, op_call: callable, args: tuple, kwargs: dict
758 ):
759 """Standard dispatch through layout-inference: preprocess → infer → execute → wrap.
761 Args:
762 op_name: Canonical operator name (already resolved by the caller).
763 op_call: The raw operator callable.
764 args: Positional arguments for op_call.
765 kwargs: Keyword arguments for op_call.
767 Returns:
768 DTensor: Dispatched result wrapped as DTensor.
770 Raises:
771 RuntimeError: If op_name is not registered, or preprocess returns None.
772 """
773 if op_name not in self.layout_infer_ops:
774 return self._handle_unregistered_op(op_name, op_call, args, kwargs)
776 cache_manager = LayoutCacheManager.get_instance()
777 distribute_op = cache_manager.distributed_op(op_name)
779 result = distribute_op.preprocess(args, kwargs)
780 if result is None:
781 raise RuntimeError(
782 f"Operator '{op_name}' has not been migrated to the three-phase dispatch flow. "
783 f"Please implement preprocess() to return (local_args, local_kwargs, cache_values)."
784 )
785 local_args, local_kwargs, cache_values = result
786 cache_key = LayoutCacheKey.from_cache_values(cache_values)
788 infer_result, op_impl = OpDispatcher._lookup_or_infer_layout(
789 op_call, op_name, cache_key, cache_values, distribute_op, cache_manager
790 )
792 op_impl = op_call if op_impl is None else op_impl
793 py_output = op_impl(*local_args, **local_kwargs)
794 output = distribute_op.wrap_output(py_output, infer_result[0])
795 return OpDispatcher._restore_inplace_dtensor_result(op_name, args, output)
797 @staticmethod
798 def _restore_inplace_dtensor_result(op_name: str, args: tuple, output: Any) -> Any:
799 """Return the original DTensor wrapper after a local in-place operation."""
800 if op_name in {"add_", "sub_"} and args and isinstance(args[0], DTensor):
801 return args[0]
802 return output
804 @staticmethod
805 def _lookup_or_infer_layout(func, func_name, cache_key, cache_values, distribute_op, cache_manager):
806 """Look up cached layout or compute via distributed op.
808 Returns:
809 (infer_result, op_impl)
810 """
811 layout_cache = cache_manager.get_layout_cache()
812 if func_name not in layout_cache:
813 layout_cache[func_name] = {}
814 op_layout_cache = layout_cache[func_name]
815 if cache_key in op_layout_cache:
816 return op_layout_cache[cache_key]
817 infer_result = distribute_op.infer_layout(cache_values)
818 op_impl = distribute_op.get_expand_impl(func, infer_result, cache_values)
819 op_layout_cache[cache_key] = (infer_result, op_impl)
820 return infer_result, op_impl
822 def dispatch(self, op_call: callable, args: tuple, kwargs: dict) -> object:
823 """Route an op call through the appropriate DTensor dispatch path.
825 Args:
826 op_call: The raw operator callable.
827 args: Positional arguments for op_call.
828 kwargs: Keyword arguments for op_call.
830 Returns:
831 Result of the dispatched op call.
832 """
833 op_name = get_op_name(op_call)
834 if logger.isEnabledFor(logging.DEBUG):
835 log_dispatch_enter(op_name, args, kwargs)
837 observer = _debug_mode_observer.get()
838 if observer is not None:
839 observer.on_op_dispatch_enter(op_name, op_call, args, kwargs)
841 result = None
842 try:
843 should_bypass = self._should_bypass_dispatch(op_name)
844 ragged_dtensor = None
845 if not should_bypass or op_name not in _RAGGED_METADATA_BYPASS_OPS:
846 ragged_dtensor = self._validate_ragged_dispatch(op_name, args, kwargs)
847 if ragged_dtensor is not None:
848 result = self._dispatch_ragged_elementwise(
849 op_call, args, kwargs, ragged_dtensor
850 )
851 return result
853 if should_bypass:
854 self._validate_inplace_partial_inputs(op_name, args, kwargs)
855 result = op_call(*self._unwrap_args(args), **self._unwrap_kwargs(kwargs))
856 if op_name in self._INPLACE_BYPASS_OPS and args and isinstance(args[0], DTensor):
857 result = args[0]
858 return result
860 if op_name in self._random_ops:
861 result = self._dispatch_random_op(op_name, op_call, args, kwargs)
862 return result
864 self._check_decomposed_ce_op_in_loss_parallel(op_name, args, kwargs)
866 if self._should_dispatch_loss_parallel(op_name):
867 result = self._dispatch_loss_parallel(op_call, args, kwargs)
868 return result
870 if op_name not in self.layout_infer_ops and get_distributed_op(op_name) is not None:
871 self.layout_infer_ops[op_name] = {}
873 result = self._dispatch_layout_infer(op_name, op_call, args, kwargs)
874 return result
875 finally:
876 if logger.isEnabledFor(logging.DEBUG):
877 log_dispatch_exit(op_name, result)
879 if observer is not None:
880 observer.on_op_dispatch_exit(op_name, result)
882_OP_DISPATCHER = OpDispatcher()