Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / core / shard / _op_dispatch.py: 82%

1# Copyright 2025-2026 Huawei Technologies Co., Ltd 

2# 

3# Licensed under the Apache License, Version 2.0 (the "License"); 

4# you may not use this file except in compliance with the License. 

5# You may obtain a copy of the License at 

6# 

7# http://www.apache.org/licenses/LICENSE-2.0 

8# 

9# Unless required by applicable law or agreed to in writing, software 

10# distributed under the License is distributed on an "AS IS" BASIS, 

11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 

12# See the License for the specific language governing permissions and 

13# limitations under the License. 

14# ============================================================================ 

15"""_op_dispatch""" 

16import atexit 

17import glob 

18import importlib 

19import logging 

20import os 

21import sys 

22import warnings 

23from contextvars import ContextVar 

24from itertools import chain 

25from typing import Any, Dict, FrozenSet, List, Optional 

26 

27import yaml 

28 

29from hyper_parallel.core.shard.ops.parallel_ops_register import get_distributed_op 

30from hyper_parallel.core.dtensor.dtensor import DTensor 

31from hyper_parallel.core.dtensor.random import OffsetBasedRNGTracker, is_rng_supported_mesh 

32from hyper_parallel.core.dtensor.debug._dispatch_logger import log_dispatch_enter, log_dispatch_exit 

33from hyper_parallel.platform import get_platform 

34from hyper_parallel.platform.platform import PlatformType 

35 

36from hyper_parallel.core.tensor_parallel._ce_op_registry import is_loss_parallel_op, is_decomposed_ce_op 

37from hyper_parallel.core.tensor_parallel.loss_parallel import is_loss_parallel_active 

38from hyper_parallel.core.tensor_parallel.loss_parallel_ops_common import _is_shard_on_last_dim 

39 

40platform = get_platform() 

41Tensor = platform.Tensor 

42 

43logger = logging.getLogger(__name__) 

44 

45 

46def _apply_shard_offset_to_rng_args(args, offset_incr): 

47 """Apply per-shard offset increment to seed/offset tensors in MindSpore random op args. 

48 

49 MindSpore random ops (e.g. ``randn_like_``) receive ``(seed, offset)`` as 

50 explicit int64 scalar tensors from ``default_generator._step()`` in the 

51 Python wrapper *before* the C++ dispatch triggers ``__fallback__``. By the 

52 time ``_dispatch_random_op`` is called, the kernel will use whatever 

53 ``(seed, offset)`` values are in the args—it does **not** read the 

54 generator again. This function finds the offset tensor and adds the 

55 per-rank offset increment so each shard gets a unique random stream. 

56 

57 The (seed, offset) pair is identified as the last two consecutive int64 

58 0-dim tensors in *args* (scanning from the end to skip trailing dtype / 

59 device arguments). 

60 

61 Args: 

62 args: The list of local args for the random op. 

63 offset_incr (int): Per-shard offset increment. 

64 

65 Returns: 

66 list: Modified args with the offset tensor adjusted. 

67 """ 

68 int64_dtype = platform.tensor_dtype.int64 

69 last_int64_idx = -1 

70 for i in range(len(args) - 1, -1, -1): 

71 arg = args[i] 

72 if isinstance(arg, Tensor) and arg.dtype == int64_dtype and arg.ndim == 0: 

73 if last_int64_idx == i + 1: 

74 offset_idx = i + 1 

75 new_args = list(args) 

76 new_offset = int(new_args[offset_idx].item()) + offset_incr 

77 new_args[offset_idx] = platform.tensor([new_offset], dtype=int64_dtype).reshape(()) 

78 return new_args 

79 last_int64_idx = i 

80 return args 

81 

82_dtensor_dispatch_disabled: ContextVar[bool] = ContextVar('_dtensor_dispatch_disabled', default=False) 

83_no_skip_ops: ContextVar[FrozenSet[str]] = ContextVar('_no_skip_ops', default=frozenset()) 

84_debug_mode_observer: ContextVar = ContextVar('_debug_mode_observer', default=None) 

85 

86 

87def get_no_skip_ops() -> FrozenSet[str]: 

88 """Return the set of op names that are exempt from SkipDTensorDispatch.""" 

89 return _no_skip_ops.get() 

90 

91 

92def get_dtensor_dispatch() -> bool: 

93 """ 

94 Get the current DTensor dispatch status. 

95 

96 Returns: 

97 bool: True if DTensor dispatch is enabled, False otherwise. 

98 """ 

99 return not _dtensor_dispatch_disabled.get() 

100 

101 

102class LayoutCacheKey: 

103 """Immutable layout cache key.""" 

104 __slots__ = ('_tuple', '_hash') 

105 

106 def __init__(self, layout_ids: List[str]): 

107 self._tuple = tuple(layout_ids) 

108 self._hash = hash(self._tuple) 

109 

110 @classmethod 

111 def from_cache_values(cls, cache_values: list) -> "LayoutCacheKey": 

112 """Build a LayoutCacheKey from a cache_values list. 

113 

114 Args: 

115 cache_values (list): Mixed list of Layout objects (with compact_str) and raw scalars. 

116 

117 Returns: 

118 LayoutCacheKey: Immutable key derived from the string representation of each value. 

119 """ 

120 # Read the cached ``_compact_str`` attribute directly instead of going through 

121 # the ``compact_str`` property getter (one fewer Python frame per Layout), and 

122 # build the key in one comprehension. The resulting tuple is byte-for-byte the 

123 # legacy string key, so eq/hash semantics are unchanged. 

124 # NOTE: the key stays a string tuple by design (cross-checked against manually 

125 # built legacy keys in the UTs); CPython caches each compact_str's hash on the 

126 # string object, so this is not the bottleneck a full integer key would target. 

127 return cls([cs if (cs := getattr(v, '_compact_str', None)) is not None else str(v) 

128 for v in cache_values]) 

129 

130 def __eq__(self, other): 

131 if not isinstance(other, LayoutCacheKey): 

132 return False 

133 return self._tuple == other._tuple 

134 

135 def __hash__(self): 

136 return self._hash 

137 

138 def __repr__(self): 

139 return f"LayoutCacheKey({self._tuple})" 

140 

141 

142class LayoutCacheManager: 

143 """ 

144 Cache layout in infer layout. 

145 

146 A singleton class that manages layout caches for distributed operations. 

147 It caches the inferred layouts and operation implementations to avoid 

148 redundant computation during repeated calls with the same input layouts. 

149 """ 

150 _instance = None 

151 

152 def __init__(self): 

153 self.layout_cache: Dict[str, Dict[LayoutCacheKey, Any]] = {} 

154 atexit.register(self.clear_cache) 

155 

156 @classmethod 

157 def get_instance(cls) -> "LayoutCacheManager": 

158 """ 

159 Get the singleton instance of LayoutCacheManager. 

160 

161 Returns: 

162 LayoutCacheManager: The singleton instance. 

163 """ 

164 if cls._instance is None: 

165 cls._instance = LayoutCacheManager() 

166 return cls._instance 

167 

168 def get_layout_cache(self) -> Dict[str, Dict[LayoutCacheKey, Any]]: 

169 """ 

170 Get the layout cache dictionary. 

171 

172 Returns: 

173 Dict[str, Dict[LayoutCacheKey, Any]]: The nested dictionary mapping 

174 operation names to their layout caches. 

175 """ 

176 return self.layout_cache 

177 

178 @staticmethod 

179 def distributed_op(op_name: str) -> Any: 

180 """ 

181 Get the distributed operation implementation by name. 

182 

183 Args: 

184 op_name (str): The name of the distributed operation. 

185 

186 Returns: 

187 Any: The distributed operation class or implementation. 

188 """ 

189 op = get_distributed_op(op_name) 

190 return op 

191 

192 def clear_cache(self) -> None: 

193 """ 

194 Clear all cached layouts. 

195 

196 This method is automatically registered with atexit to ensure 

197 cache is cleared when the program exits. 

198 """ 

199 self.layout_cache.clear() 

200 

201 

202class OpDispatcher: 

203 """ 

204 OpDispatcher 

205 """ 

206 

207 # Whitelisted ops that mutate args[0]'s storage in place. The dispatch bypass 

208 # must return the original DTensor self for these, not the unwrapped local 

209 # result, or it demotes a DTensor accumulator to a plain Tensor and breaks the 

210 # next op that adds a DTensor to it (e.g. grad-accumulation `loss += micro_loss`). 

211 # Class-level so it stays available on instances built via __new__ (e.g. tests). 

212 _INPLACE_BYPASS_OPS = frozenset( 

213 {"InplaceAddExt", "InplaceSubExt", "InplaceMul", "InplaceDiv"}) 

214 

215 # MindSpore random kernels that always mutate an existing tensor in place. 

216 # Out-of-place random kernels belong in _random_ms_ops only, not here. 

217 _RANDOM_INPLACE_MS_OPS = frozenset({ 

218 "InplaceBernoulliScalar", 

219 "InplaceBernoulliTensor", 

220 "InplaceNormal", 

221 "InplaceRandom", 

222 "InplaceUniform", 

223 }) 

224 

225 def __init__(self): 

226 self._env_yaml_dir: Optional[str] = os.environ.get("HYPER_PARALLEL_OPS_YAML_DIR") 

227 self._env_python_path: Optional[str] = os.environ.get("HYPER_PARALLEL_OPS_PYTHON_PATH") 

228 # The following attributes are initialized in _setup_yaml_dir() 

229 self.work_dir = "" # Initialized in _setup_yaml_dir() 

230 self.yaml_dir = "" # Initialized in _setup_yaml_dir() 

231 

232 self._setup_paths_from_env() 

233 

234 self.layout_infer_ops = self.safe_load_yaml_from_dir() 

235 # frozenset for O(1) membership (checked on every dispatch's bypass test). 

236 self.whitelist = frozenset({"typeof", "DistCommIsend", 

237 "DistCommIrecv", "DistCommBroadcast", "DistCommAllReduce", "DistCommAllGather", 

238 "DistCommBatchIsendIrecv", 

239 "DistCommReduceScatter", "requires_grad_", "item", "__get__", "__set__", 

240 "register_hook", 

241 "is_complex", "chunk", "__bool__", "__len__", "__format__", "dim", 

242 "_has_compatible_shallow_copy_type", "is_floating_point", "is_contiguous"}) 

243 

244 # Ops requiring args unpacking for layout inference (packed as prim, name, real_args). 

245 # frozenset so the aclop-normalization gate in _dispatch_layout_infer is O(1). 

246 self.unpack_ops = frozenset({"ScatterUpdate", "Mod", "GatherNd", "StopGradient"}) 

247 

248 self._random_ops = { 

249 "normal_", "uniform_", "bernoulli", "bernoulli_", 

250 "native_dropout", "rand", "rand_like", "randn", 

251 "randn_like", "randint_like", "kaiming_uniform_", 

252 "multinomial", 

253 } 

254 # Only mint random op support 

255 # MindSpore use the actual kernel name. 

256 self._random_ms_ops = { 

257 "BernoulliExt", "MultinomialExt", 

258 "InplaceBernoulliScalar", "InplaceBernoulliTensor", 

259 "InplaceNormal", "InplaceRandom", "InplaceUniform", 

260 "NormalFloatFloat", "NormalFloatTensor", "NormalTensorFloat", "NormalTensorTensor", 

261 "RandpermExt", "Randn", "RandLikeExt", "RandnLike", "RandInt", "RandIntLike", "RandExt", 

262 "FuncDropoutExt", "UniformExt", 

263 } 

264 self._rng_tracker: Optional[OffsetBasedRNGTracker] = None 

265 # Op names proven to be loss/CE-irrelevant (both is_loss_parallel_op and 

266 # is_decomposed_ce_op are False). For these the loss_parallel / decomposed-CE 

267 # guards in dispatch() are always no-ops regardless of context, so we skip 

268 # them (and their is_loss_parallel_active() contextvar reads) on later calls. 

269 self._non_loss_ops: set = set() 

270 

271 self._register_distributed_ops() 

272 

273 def _setup_paths_from_env(self): 

274 """ 

275 Setup YAML directory and Python path from environment variables. 

276 

277 This method initializes the YAML directory and extends sys.path based on 

278 environment variables HYPER_PARALLEL_OPS_YAML_DIR and HYPER_PARALLEL_OPS_PYTHON_PATH. 

279 """ 

280 self._setup_yaml_dir(self._env_yaml_dir) 

281 self._extend_sys_path(self._env_python_path) 

282 

283 def _setup_yaml_dir(self, env_yaml_dir: Optional[str]): 

284 """ 

285 Feature: Configure yaml_dir/work_dir for OpDispatcher 

286 Description: Resolve the YAML directory used to load distributed op definitions. 

287 If env_yaml_dir is an absolute path, use it directly; otherwise treat it 

288 as a path relative to the project work_dir. If env_yaml_dir is not set, 

289 fall back to the default 'shard/ops/yaml' under work_dir. 

290 Expectation: self.yaml_dir and self.work_dir are set to valid values used later by 

291 safe_load_yaml_from_dir(); no functional behavior is changed. 

292 """ 

293 if env_yaml_dir: 

294 if os.path.isabs(env_yaml_dir): 

295 self.yaml_dir = env_yaml_dir 

296 self.work_dir = "" 

297 else: 

298 self.work_dir = os.path.normpath( 

299 os.path.join(os.path.dirname(os.path.realpath(__file__)), "../") 

300 ) 

301 self.yaml_dir = env_yaml_dir 

302 else: 

303 self.yaml_dir = "shard/ops/yaml" 

304 self.work_dir = os.path.normpath( 

305 os.path.join(os.path.dirname(os.path.realpath(__file__)), "../") 

306 ) 

307 

308 @staticmethod 

309 def _extend_sys_path(env_python_path: Optional[str]): 

310 if not env_python_path: 

311 return 

312 python_paths = env_python_path.split(":") 

313 for path in python_paths: 

314 if path and os.path.isdir(path) and path not in sys.path: 

315 sys.path.append(path) 

316 

317 def _register_distributed_ops(self): 

318 for op_name, config in self.layout_infer_ops.items(): 

319 self._register_single_distributed_op(op_name, config) 

320 

321 def _register_single_distributed_op(self, op_name: str, config: dict): 

322 """ 

323 Feature: Register a single distributed op implementation 

324 Description: Import the distributed op class specified by config and instantiate it 

325 with op_name to trigger registration in the distributed op registry. 

326 Prefer 'distributed_op_module' when provided; otherwise import from 

327 built-in module prefix 'hyper_parallel.core.shard.ops.' plus 

328 'distributed_op_file'. If import fails and an external python path is 

329 provided via env, fall back to importing 'distributed_op_file' directly. 

330 Expectation: The distributed op class is imported and instantiated successfully, 

331 or the original import error is raised; no functional behavior is changed. 

332 """ 

333 class_name = config["distributed_op_class"] 

334 

335 if "distributed_op_module" in config: 

336 module_name = config["distributed_op_module"] 

337 module = importlib.import_module(module_name) 

338 op_class = getattr(module, class_name) 

339 _ = op_class(op_name) 

340 return 

341 

342 module_file = config["distributed_op_file"] 

343 try: 

344 module_name = "hyper_parallel.core.shard.ops." + module_file 

345 module = importlib.import_module(module_name) 

346 op_class = getattr(module, class_name) 

347 _ = op_class(op_name) 

348 except (ModuleNotFoundError, ImportError): 

349 if self._env_python_path: 

350 module = importlib.import_module(module_file) 

351 op_class = getattr(module, class_name) 

352 _ = op_class(op_name) 

353 else: 

354 raise 

355 

356 @staticmethod 

357 def _merge_default(config: dict): 

358 """Apply __default__ values to all ops in this YAML file.""" 

359 if "__default__" not in config: 

360 return config 

361 

362 default_cfg = config["__default__"] 

363 merged = {} 

364 

365 for op_name, op_cfg in config.items(): 

366 if op_name == "__default__": 

367 continue 

368 

369 new_cfg = default_cfg.copy() 

370 new_cfg.update(op_cfg) 

371 merged[op_name] = new_cfg 

372 

373 return merged 

374 

375 def safe_load_yaml_from_dir(self) -> dict: 

376 """ 

377 Load yaml dictionary from directory. 

378 

379 Returns: 

380 dict: Merged dictionary of all operator configurations loaded from YAML files. 

381 """ 

382 yaml_dict = {} 

383 yaml_path = os.path.join(self.work_dir, self.yaml_dir) if self.work_dir else self.yaml_dir 

384 if not os.path.isdir(yaml_path): 

385 raise ValueError(f"Invalid yaml directory path: {yaml_path}") 

386 

387 for yaml_file_path in glob.glob(os.path.join(yaml_path, '*.yaml')): 

388 with open(yaml_file_path, 'r', encoding="utf-8") as f: 

389 yaml_data = yaml.safe_load(f) 

390 

391 yaml_data = OpDispatcher._merge_default(yaml_data) 

392 for name, data in yaml_data.items(): 

393 if name in yaml_dict: 

394 raise ValueError(f"Duplicate yaml object with name '{name}'.") 

395 yaml_dict[name] = data 

396 

397 return yaml_dict 

398 

399 def _dispatch_random_op(self, op_name: str, op_call: callable, args, kwargs): 

400 """Handle dispatch for random ops that operate on DTensors.""" 

401 first_arg = next( 

402 (x for x in chain(args, kwargs.values()) if isinstance(x, DTensor)), 

403 None, 

404 ) 

405 # Fall back to the default op if no DTensor is found. 

406 if first_arg is None: 

407 return op_call(*args, **kwargs) 

408 

409 local_args = [arg.to_local() if isinstance(arg, DTensor) else arg for arg in args] 

410 local_kwargs = {k: v.to_local() if isinstance(v, DTensor) else v for k, v in kwargs.items()} 

411 first_local_arg = first_arg.to_local() 

412 

413 if self._rng_tracker is None and is_rng_supported_mesh(first_arg.device_mesh): 

414 self._rng_tracker = OffsetBasedRNGTracker() 

415 

416 maybe_user_generator = local_kwargs.pop("generator", None) 

417 if ( 

418 self._rng_tracker is not None 

419 and not first_local_arg.is_meta 

420 and self._rng_tracker.distribute_region_enabled 

421 ): 

422 # pylint: disable=W0212 

423 with self._rng_tracker._distribute_region( 

424 device_mesh=first_arg.device_mesh, 

425 placements=first_arg.placements, 

426 global_shape=first_arg.shape, 

427 generator=maybe_user_generator, 

428 ): 

429 # MindSpore random ops (e.g. mint.randn_like) extract (seed, offset) 

430 # from default_generator._step() in the Python wrapper *before* the 

431 # C++ dispatch triggers __fallback__. The callback reuses these 

432 # pre-fetched tensor args, so set_rng_state inside _distribute_region 

433 # has no effect on the kernel. Fix: apply the per-shard offset 

434 # increment directly to the offset tensor in the args. 

435 if platform.platform_type == PlatformType.MINDSPORE: 

436 offset_incr = self._rng_tracker.compute_offset_incr( 

437 first_arg.device_mesh, first_arg.placements, first_arg.shape, 

438 ) 

439 local_args = _apply_shard_offset_to_rng_args(local_args, offset_incr) 

440 local_results = op_call(*local_args, **local_kwargs) 

441 else: 

442 if maybe_user_generator is not None: 

443 local_kwargs["generator"] = maybe_user_generator 

444 local_results = op_call(*local_args, **local_kwargs) 

445 

446 return self._wrap_random_result(op_name, local_results, first_arg, args, kwargs) 

447 

448 @staticmethod 

449 def _func_dropout_ext_inplace(args, kwargs) -> bool: 

450 """Return True when FuncDropoutExt is invoked with inplace=True.""" 

451 # Kernel signature: (input, p, training, inplace, seed, offset). 

452 if len(args) >= 4: 

453 return bool(args[3]) 

454 return bool(kwargs.get("inplace", False)) 

455 

456 @staticmethod 

457 def _random_op_returns_self(op_name: str, args, kwargs) -> bool: 

458 """Return True when a random op mutates an existing DTensor in place.""" 

459 if op_name in OpDispatcher._RANDOM_INPLACE_MS_OPS: 

460 return True 

461 if op_name == "FuncDropoutExt": 

462 return OpDispatcher._func_dropout_ext_inplace(args, kwargs) 

463 # Torch random inplace ops follow the ATen '_' suffix convention. 

464 return op_name.endswith('_') 

465 

466 @staticmethod 

467 def _wrap_random_result(op_name, local_results, first_arg, args, kwargs): 

468 """Wrap a random op's local result(s) back into DTensor(s). 

469 

470 In-place ops return the input DTensor itself. Torch random inplace ops use 

471 the ATen '_' suffix; MindSpore inplace random kernels are listed in 

472 ``_RANDOM_INPLACE_MS_OPS``. ``FuncDropoutExt`` is handled separately 

473 because the same kernel serves both modes via its ``inplace`` argument. 

474 """ 

475 if OpDispatcher._random_op_returns_self(op_name, args, kwargs): 

476 return first_arg 

477 mesh = first_arg.device_mesh 

478 placements = first_arg.layout.alias_placements 

479 # Some ops return tuple/list, e.g. native_dropout returns (output, mask). 

480 if isinstance(local_results, (tuple, list)): 

481 return tuple( 

482 DTensor.from_local(r, mesh, placements) if isinstance(r, Tensor) else r 

483 for r in local_results 

484 ) 

485 if isinstance(local_results, Tensor): 

486 return DTensor.from_local(local_results, mesh, placements) 

487 # Fallback: return as-is for non-Tensor results (currently unreachable with existing _random_ops). 

488 return local_results 

489 

490 @staticmethod 

491 def _unwrap_value(value: object) -> object: 

492 """Replace DTensor with its local tensor; pass scalars and plain tensors through. 

493 

494 Args: 

495 value (object): A single argument value from an op call. 

496 

497 Returns: 

498 object: The local tensor if value is a DTensor, otherwise value unchanged. 

499 """ 

500 if isinstance(value, DTensor): 

501 return value.to_local() 

502 if isinstance(value, tuple): 

503 return tuple(OpDispatcher._unwrap_value(e) for e in value) 

504 if isinstance(value, list): 

505 return [OpDispatcher._unwrap_value(e) for e in value] 

506 return value 

507 

508 @staticmethod 

509 def _unwrap_args(args: tuple) -> list: 

510 """Strip DTensor wrappers from args, preserving tuple/list container structure. 

511 

512 Args: 

513 args: Op call positional arguments, may contain DTensor instances. 

514 

515 Returns: 

516 List of args with DTensor replaced by their local tensors. 

517 """ 

518 return [OpDispatcher._unwrap_value(arg) for arg in args] 

519 

520 @staticmethod 

521 def _unwrap_kwargs(kwargs: dict) -> dict: 

522 """Strip DTensor wrappers from kwargs values, preserving tuple/list container structure. 

523 

524 Args: 

525 kwargs: Op call keyword arguments, values may contain DTensor instances. 

526 

527 Returns: 

528 Dict of kwargs with DTensor values replaced by their local tensors. 

529 """ 

530 return {k: OpDispatcher._unwrap_value(v) for k, v in kwargs.items()} 

531 

532 @staticmethod 

533 def _gather_dtensors_to_full(args: tuple, kwargs: dict) -> tuple: 

534 """Gather all DTensor arguments to full tensors for fallback execution. 

535 

536 Used when an operator has no parallel layout implementation. All DTensor 

537 arguments are gathered to full tensors before calling the standard operator. 

538 

539 Args: 

540 args: Op call positional arguments, may contain DTensor instances. 

541 kwargs: Op call keyword arguments, may contain DTensor instances. 

542 

543 Returns: 

544 Tuple of (unwrapped_args, unwrapped_kwargs) with DTensor values 

545 replaced by their full tensor representations. 

546 

547 Warning: 

548 This fallback performs all-gather which may consume significant memory. 

549 Operators without layout implementations should be registered properly. 

550 """ 

551 def gather(value: object) -> object: 

552 if isinstance(value, DTensor): 

553 return value.full_tensor() 

554 if isinstance(value, tuple): 

555 return tuple(gather(e) for e in value) 

556 if isinstance(value, list): 

557 return [gather(e) for e in value] 

558 return value 

559 

560 gathered_args = [gather(arg) for arg in args] 

561 gathered_kwargs = {k: gather(v) for k, v in kwargs.items()} 

562 

563 warnings.warn( 

564 "Operator has no distributed layout implementation. " 

565 "Falling back to all-gather which may consume significant memory. " 

566 "Consider registering a proper distributed operator.", 

567 UserWarning, 

568 stacklevel=4 

569 ) 

570 

571 return gathered_args, gathered_kwargs 

572 

573 def _should_bypass_dispatch(self, op_name: str) -> bool: 

574 """Return True if the op should bypass DTensor dispatch and run locally. 

575 

576 Args: 

577 op_name: Canonical operator name from platform.get_op_name(). 

578 

579 Returns: 

580 True when the op is whitelisted or DTensor dispatch is globally disabled. 

581 """ 

582 # Cheap O(1) frozenset checks first, short-circuit before the ContextVar 

583 # read (get_dtensor_dispatch) which is the priciest part of this guard. 

584 if op_name in self.whitelist or op_name in self._INPLACE_BYPASS_OPS: 

585 return True 

586 return get_dtensor_dispatch() is False and op_name not in get_no_skip_ops() 

587 

588 @staticmethod 

589 def _validate_inplace_partial_inputs(op_name: str, args: tuple, kwargs: dict) -> None: 

590 """Reject local in-place add/sub when Partial contributions need gating.""" 

591 if op_name not in {"InplaceAddExt", "InplaceSubExt"} or not args: 

592 return 

593 first = args[0] 

594 if len(args) >= 2: 

595 second = args[1] 

596 elif "other" in kwargs: 

597 second = kwargs["other"] 

598 else: 

599 return 

600 if not isinstance(first, DTensor): 

601 return 

602 mesh_ndim = len(first.layout.partial) 

603 first_partial = tuple(first.layout.partial) 

604 if isinstance(second, DTensor): 

605 second_partial = tuple(second.layout.partial) 

606 if len(second_partial) != mesh_ndim: 

607 raise ValueError( 

608 f"For {op_name}, in-place input mesh dimensions must match, " 

609 f"but got {mesh_ndim} and {len(second_partial)}." 

610 ) 

611 else: 

612 second_partial = (None,) * mesh_ndim 

613 if first_partial != second_partial: 

614 raise ValueError( 

615 f"For {op_name}, input Partial placements must be identical for " 

616 f"local in-place execution, but got {first_partial} and {second_partial}." 

617 ) 

618 

619 def _should_dispatch_loss_parallel(self, op_name: str) -> bool: 

620 """Check if should dispatch through loss_parallel path. 

621 

622 Args: 

623 op_name: Canonical operator name from platform.get_op_name(). 

624 

625 Returns: 

626 True when in loss_parallel context and op is a CE entry point. 

627 """ 

628 return is_loss_parallel_active() and is_loss_parallel_op(op_name) 

629 

630 def _check_decomposed_ce_op_in_loss_parallel(self, op_name: str, args: tuple, kwargs: dict): 

631 """Check if decomposed CE ops are called in loss_parallel context. 

632 

633 Args: 

634 op_name: Canonical operator name. 

635 args: Positional arguments for op_call. 

636 kwargs: Keyword arguments for op_call. 

637 

638 Raises: 

639 ValueError: If decomposed CE op is called in loss_parallel context 

640 with vocab-sharded DTensor input. 

641 """ 

642 if not is_loss_parallel_active() or not is_decomposed_ce_op(op_name): 

643 return 

644 

645 has_vocab_sharded_dtensor = False 

646 for arg in args: 

647 if isinstance(arg, DTensor) and _is_shard_on_last_dim(arg): 

648 has_vocab_sharded_dtensor = True 

649 break 

650 if not has_vocab_sharded_dtensor: 

651 for val in kwargs.values(): 

652 if isinstance(val, DTensor) and _is_shard_on_last_dim(val): 

653 has_vocab_sharded_dtensor = True 

654 break 

655 

656 if has_vocab_sharded_dtensor: 

657 raise ValueError( 

658 f"Operator '{op_name}' is a decomposed component of cross_entropy and should not be called " 

659 f"directly within loss_parallel() context. Use F.cross_entropy(logits, targets) instead. " 

660 f"For example, replace:\n" 

661 f" with loss_parallel():\n" 

662 f" log_probs = F.log_softmax(logits, dim=-1)\n" 

663 f" loss = F.nll_loss(log_probs, targets)\n" 

664 f"with:\n" 

665 f" with loss_parallel():\n" 

666 f" loss = F.cross_entropy(logits, targets)" 

667 ) 

668 

669 def _dispatch_loss_parallel(self, op_call: callable, args: tuple, kwargs: dict): 

670 """Dispatch cross_entropy through the loss_parallel distributed kernel. 

671 

672 Args: 

673 op_call: The raw operator callable. 

674 args: Positional arguments for op_call. 

675 kwargs: Keyword arguments for op_call. 

676 

677 Returns: 

678 Result of the distributed cross_entropy computation. 

679 """ 

680 if platform.platform_type == PlatformType.PYTORCH: 

681 # pylint: disable=C0415 

682 from hyper_parallel.platform.torch.loss_parallel_ops import distributed_cross_entropy_from_op_call 

683 elif platform.platform_type == PlatformType.MINDSPORE: 

684 # pylint: disable=C0415 

685 from hyper_parallel.platform.mindspore.loss_parallel_ops import distributed_cross_entropy_from_op_call 

686 else: 

687 raise RuntimeError(f"Unsupported platform for loss_parallel: {platform.platform_type}") 

688 return distributed_cross_entropy_from_op_call(op_call, args, kwargs) 

689 

690 def _check_ce_op_without_loss_parallel_context(self, op_name: str, args: tuple): 

691 """Check if CE op is called with Shard(-1) DTensor outside loss_parallel context. 

692 

693 Args: 

694 op_name: Canonical operator name. 

695 args: Positional arguments for op_call. 

696 

697 Raises: 

698 ValueError: If CE op is called with Shard(-1) logits outside loss_parallel context. 

699 """ 

700 if is_loss_parallel_active() or not is_loss_parallel_op(op_name): 

701 return 

702 

703 if len(args) == 0 or not isinstance(args[0], DTensor): 

704 return 

705 

706 logits = args[0] 

707 if _is_shard_on_last_dim(logits): 

708 raise ValueError( 

709 f"Operator '{op_name}' requires loss_parallel context when input logits are " 

710 f"sharded on the vocabulary dimension (Shard(-1)). Please wrap your forward " 

711 f"and backward pass with loss_parallel():\n" 

712 f" with loss_parallel():\n" 

713 f" loss = F.cross_entropy(logits, targets)\n" 

714 f" loss.backward()\n" 

715 f"If you intentionally want to gather all shards to compute cross_entropy " 

716 f"(not recommended for large vocabulary), use logits.full_tensor() explicitly." 

717 ) 

718 

719 @staticmethod 

720 def _normalize_aclop_args(op_name: str, unpack_ops: list, args: tuple) -> tuple: 

721 """ 

722 Normalize aclop-packed arguments for MindSpore backend operators. 

723 

724 NOTE: This handles MindSpore aclop operators whose kernel signature packs 

725 arguments as ``(prim, op_name_str, (real_arg0, real_arg1, ...))``. The 

726 ``prim`` and ``op_name_str`` are preserved as ``packed_call`` for the 

727 final kernel invocation, while the real tensor arguments are extracted 

728 for layout inference and preprocessing. 

729 

730 **aclop is planned for deprecation.** Once aclop is fully removed, this 

731 normalization and the associated ``unpack_ops`` list can be deleted. 

732 

733 Args: 

734 op_name (str): Canonical operator name. 

735 unpack_ops (list): List of op names that may use aclop packed format. 

736 args (tuple): Raw positional arguments from the op call. 

737 

738 Returns: 

739 tuple: ``(packed_call, normalized_args)`` 

740 - **packed_call**: ``(prim, op_name_str)`` tuple for kernel 

741 invocation, or ``None`` if no unpacking was performed. 

742 - **normalized_args**: The real tensor arguments (unpacked if 

743 the packed format was detected, otherwise the original args). 

744 """ 

745 if OpDispatcher._is_aclop_packed(op_name, unpack_ops, args): 

746 return (args[0], args[1]), tuple(args[2]) 

747 return None, args 

748 

749 @staticmethod 

750 def _is_aclop_packed(op_name: str, unpack_ops: list, args: tuple) -> bool: 

751 """Check if arguments use aclop packed format.""" 

752 return ( 

753 op_name in unpack_ops 

754 and len(args) == 3 

755 and isinstance(args[1], str) 

756 and isinstance(args[2], (tuple, list)) 

757 ) 

758 

759 @staticmethod 

760 def _call_op_impl(op_impl: callable, packed_call, args, kwargs: dict): 

761 """Invoke *op_impl* with optional aclop packed-call wrapping. 

762 

763 When *packed_call* is not ``None`` the MindSpore aclop kernel expects 

764 ``(prim, op_name, (arg0, arg1, ...))``. Otherwise *args* are spread 

765 as positional arguments in the usual way. 

766 

767 Args: 

768 op_impl: The op implementation callable. 

769 packed_call: ``(prim, op_name)`` tuple or ``None``. 

770 args: Local tensor arguments (list or tuple). 

771 kwargs: Keyword arguments dict. 

772 

773 Returns: 

774 Result of the *op_impl* invocation. 

775 """ 

776 if packed_call is not None: 

777 return op_impl(packed_call[0], packed_call[1], tuple(args), **kwargs) 

778 return op_impl(*args, **kwargs) 

779 

780 def _handle_unregistered_op( 

781 self, op_name: str, op_call: callable, args: tuple, kwargs: dict 

782 ): 

783 """Handle ops that have no registered layout-inference entry. 

784 

785 This is a fallback path for ops that are not registered in 

786 ``layout_infer_ops``. When arguments contain DTensors it either raises 

787 (with a hint to register a distributed op) or, for loss-parallel ops, 

788 gathers tensors to full and dispatches through the raw callable. 

789 

790 Args: 

791 op_name: Canonical operator name. 

792 op_call: The raw operator callable. 

793 args: Positional arguments for op_call. 

794 kwargs: Keyword arguments for op_call. 

795 

796 Returns: 

797 Raw dispatch result (plain Tensor, not wrapped as DTensor). 

798 

799 Raises: 

800 RuntimeError: If op_name is not registered for layout inference. 

801 """ 

802 has_dtensor = any(isinstance(arg, DTensor) for arg in args) 

803 has_dtensor = has_dtensor or any(isinstance(v, DTensor) for v in kwargs.values()) 

804 if has_dtensor: 

805 self._check_ce_op_without_loss_parallel_context(op_name, args) 

806 

807 if not is_loss_parallel_op(op_name): 

808 raise RuntimeError( 

809 f"Operator {op_name} does not contain parallel layout infer func. " 

810 f"DTensor dispatch requires explicit layout inference registration. " 

811 f"Please register a distributed operator for '{op_name}' or use local tensors." 

812 ) 

813 

814 gathered_args, gathered_kwargs = self._gather_dtensors_to_full(args, kwargs) 

815 

816 # Special handling for cross_entropy with 3D logits (only when NOT in loss_parallel context) 

817 # PyTorch expects: logits [N, C], targets [N] 

818 # But LLM forward returns: logits [batch, seq, vocab], targets [batch, seq] 

819 # Note: nll_loss input is log_probs, typically already 2D, so we only reshape for cross_entropy 

820 if op_name == "cross_entropy" and len(gathered_args) >= 2: 

821 logits = gathered_args[0] 

822 targets = gathered_args[1] 

823 if isinstance(logits, Tensor) and isinstance(targets, Tensor): 

824 if logits.ndim > 2 and targets.ndim > 1 and targets.ndim == logits.ndim - 1: 

825 vocab_size = logits.shape[-1] 

826 gathered_args[0] = logits.reshape(-1, vocab_size) 

827 gathered_args[1] = targets.reshape(-1) 

828 

829 return op_call(*gathered_args, **gathered_kwargs) 

830 raise RuntimeError(f"Operator {op_name} does not contain parallel layout infer func.") 

831 

832 def _dispatch_layout_infer( 

833 self, op_name: str, op_call: callable, args: tuple, kwargs: dict 

834 ): 

835 """Standard dispatch through layout-inference: preprocess → infer → execute → wrap. 

836 

837 Args: 

838 op_name: Canonical operator name (already resolved by the caller). 

839 op_call: The raw operator callable. 

840 args: Positional arguments for op_call. 

841 kwargs: Keyword arguments for op_call. 

842 

843 Returns: 

844 DTensor: Dispatched result wrapped as DTensor. 

845 

846 Raises: 

847 RuntimeError: If op_name is not registered, or preprocess returns None. 

848 """ 

849 if op_name not in self.layout_infer_ops: 

850 return self._handle_unregistered_op(op_name, op_call, args, kwargs) 

851 

852 cache_manager = LayoutCacheManager.get_instance() 

853 distribute_op = cache_manager.distributed_op(op_name) 

854 

855 # Normalize aclop-packed args before any per-op processing. Only the handful 

856 # of (deprecation-bound) unpack_ops ever use the packed format, so gate the 

857 # whole normalization behind an O(1) membership test instead of paying two 

858 # function frames (_normalize_aclop_args + _is_aclop_packed) on every op. 

859 if op_name in getattr(self, 'unpack_ops', ()): 

860 packed_call, args = self._normalize_aclop_args(op_name, self.unpack_ops, args) 

861 else: 

862 packed_call = None 

863 

864 result = distribute_op.preprocess(args, kwargs) 

865 if result is None: 

866 raise RuntimeError( 

867 f"Operator '{op_name}' has not been migrated to the three-phase dispatch flow. " 

868 f"Please implement preprocess() to return (local_args, local_kwargs, cache_values)." 

869 ) 

870 local_args, local_kwargs, cache_values = result 

871 cache_key = LayoutCacheKey.from_cache_values(cache_values) 

872 

873 infer_result, op_impl = OpDispatcher._lookup_or_infer_layout( 

874 op_call, op_name, cache_key, cache_values, distribute_op, cache_manager 

875 ) 

876 

877 op_impl = op_call if op_impl is None else op_impl 

878 py_output = OpDispatcher._call_op_impl(op_impl, packed_call, local_args, local_kwargs) 

879 output = distribute_op.wrap_output(py_output, infer_result[0]) 

880 return OpDispatcher._restore_inplace_dtensor_result(op_name, args, output) 

881 

882 @staticmethod 

883 def _restore_inplace_dtensor_result(op_name: str, args: tuple, output: Any) -> Any: 

884 """Return the original DTensor wrapper after a local in-place operation.""" 

885 if op_name in {"add_", "sub_"} and args and isinstance(args[0], DTensor): 

886 return args[0] 

887 return output 

888 

889 @staticmethod 

890 def _lookup_or_infer_layout(func, func_name, cache_key, cache_values, distribute_op, cache_manager): 

891 """Look up cached layout or compute via distributed op.