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

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

2# 

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

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

5# You may obtain a copy of the License at 

6# 

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

8# 

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

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

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

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

13# limitations under the License. 

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

15"""_op_dispatch""" 

16import atexit 

17import copy 

18import glob 

19import importlib 

20import logging 

21import os 

22import sys 

23import warnings 

24from contextvars import ContextVar 

25from itertools import chain 

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

27 

28import yaml 

29 

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

31from hyper_parallel.core.dtensor.dtensor import DTensor 

32from hyper_parallel.core.dtensor.layout import RaggedShardInfo 

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

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

35from hyper_parallel.platform import get_platform 

36from hyper_parallel.platform.platform import PlatformType 

37 

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

39from hyper_parallel.core.tensor_parallel.loss_parallel import is_loss_parallel_active 

40from hyper_parallel.core.tensor_parallel.loss_parallel_ops_common import _is_shard_on_last_dim 

41 

42platform = get_platform() 

43Tensor = platform.Tensor 

44 

45logger = logging.getLogger(__name__) 

46 

47 

48def _apply_shard_offset_to_rng_args(args, offset_incr): 

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

50 

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

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

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

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

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

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

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

58 

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

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

61 device arguments). 

62 

63 Args: 

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

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

66 

67 Returns: 

68 list: Modified args with the offset tensor adjusted. 

69 """ 

70 int64_dtype = platform.tensor_dtype.int64 

71 last_int64_idx = -1 

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

73 arg = args[i] 

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

75 if last_int64_idx == i + 1: 

76 offset_idx = i + 1 

77 new_args = list(args) 

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

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

80 return new_args 

81 last_int64_idx = i 

82 return args 

83 

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

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

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

87 

88_RAGGED_ELEMENTWISE_OPS = { 

89 "abs": "unary", "absolute": "unary", "clone": "unary", "cos": "unary", 

90 "exp": "unary", "gelu": "unary", "isinf": "unary", "isnan": "unary", 

91 "log": "unary", "neg": "unary", "negative": "unary", "relu": "unary", 

92 "rsqrt": "unary", "sigmoid": "unary", "silu": "unary", "sin": "unary", 

93 "sqrt": "unary", "square": "unary", 

94 "add": "binary", "div": "binary", "mul": "binary", "pow": "binary", 

95 "real_div": "binary", "sub": "binary", "__rsub__": "binary", 

96 "__rpow__": "binary", "true_divide": "binary", 

97} 

98 

99 

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

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

102 return _no_skip_ops.get() 

103 

104 

105def get_dtensor_dispatch() -> bool: 

106 """ 

107 Get the current DTensor dispatch status. 

108 

109 Returns: 

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

111 """ 

112 return not _dtensor_dispatch_disabled.get() 

113 

114 

115class LayoutCacheKey: 

116 """Immutable layout cache key.""" 

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

118 

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

120 self._tuple = tuple(layout_ids) 

121 self._hash = hash(self._tuple) 

122 

123 @classmethod 

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

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

126 

127 Args: 

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

129 

130 Returns: 

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

132 """ 

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

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

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

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

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

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

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

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

141 for v in cache_values]) 

142 

143 def __eq__(self, other): 

144 if not isinstance(other, LayoutCacheKey): 

145 return False 

146 return self._tuple == other._tuple 

147 

148 def __hash__(self): 

149 return self._hash 

150 

151 def __repr__(self): 

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

153 

154 

155class LayoutCacheManager: 

156 """ 

157 Cache layout in infer layout. 

158 

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

160 It caches the inferred layouts and operation implementations to avoid 

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

162 """ 

163 _instance = None 

164 

165 def __init__(self): 

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

167 atexit.register(self.clear_cache) 

168 

169 @classmethod 

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

171 """ 

172 Get the singleton instance of LayoutCacheManager. 

173 

174 Returns: 

175 LayoutCacheManager: The singleton instance. 

176 """ 

177 if cls._instance is None: 

178 cls._instance = LayoutCacheManager() 

179 return cls._instance 

180 

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

182 """ 

183 Get the layout cache dictionary. 

184 

185 Returns: 

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

187 operation names to their layout caches. 

188 """ 

189 return self.layout_cache 

190 

191 @staticmethod 

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

193 """ 

194 Get the distributed operation implementation by name. 

195 

196 Args: 

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

198 

199 Returns: 

200 Any: The distributed operation class or implementation. 

201 """ 

202 op = get_distributed_op(op_name) 

203 return op 

204 

205 def clear_cache(self) -> None: 

206 """ 

207 Clear all cached layouts. 

208 

209 This method is automatically registered with atexit to ensure 

210 cache is cleared when the program exits. 

211 """ 

212 self.layout_cache.clear() 

213 

214 

215class OpDispatcher: 

216 """ 

217 OpDispatcher 

218 """ 

219 

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

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

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

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

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

225 _INPLACE_BYPASS_OPS = frozenset( 

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

227 

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

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

230 _RANDOM_INPLACE_MS_OPS = frozenset({ 

231 "InplaceBernoulliScalar", 

232 "InplaceBernoulliTensor", 

233 "InplaceNormal", 

234 "InplaceRandom", 

235 "InplaceUniform", 

236 }) 

237 

238 def __init__(self): 

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

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

241 # The following attributes are initialized in _setup_yaml_dir() 

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

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

244 

245 self._setup_paths_from_env() 

246 

247 self.layout_infer_ops = self.safe_load_yaml_from_dir() 

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

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

250 "DistCommIrecv", "DistCommBroadcast", "DistCommAllReduce", "DistCommAllGather", 

251 "DistCommBatchIsendIrecv", 

252 "DistCommReduceScatter", "requires_grad_", "item", "__get__", "__set__", 

253 "register_hook", 

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

255 "_has_compatible_shallow_copy_type", "is_floating_point", "is_contiguous"}) 

256 

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

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

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

260 

261 self._random_ops = { 

262 "normal_", "uniform_", "bernoulli", "bernoulli_", 

263 "native_dropout", "rand", "rand_like", "randn", 

264 "randn_like", "randint_like", "kaiming_uniform_", 

265 "multinomial", 

266 } 

267 # Only mint random op support 

268 # MindSpore use the actual kernel name. 

269 self._random_ms_ops = { 

270 "BernoulliExt", "MultinomialExt", 

271 "InplaceBernoulliScalar", "InplaceBernoulliTensor", 

272 "InplaceNormal", "InplaceRandom", "InplaceUniform", 

273 "NormalFloatFloat", "NormalFloatTensor", "NormalTensorFloat", "NormalTensorTensor", 

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

275 "FuncDropoutExt", "UniformExt", 

276 } 

277 self._rng_tracker: Optional[OffsetBasedRNGTracker] = None 

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

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

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

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

282 self._non_loss_ops: set = set() 

283 

284 self._register_distributed_ops() 

285 

286 def _setup_paths_from_env(self): 

287 """ 

288 Setup YAML directory and Python path from environment variables. 

289 

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

291 environment variables HYPER_PARALLEL_OPS_YAML_DIR and HYPER_PARALLEL_OPS_PYTHON_PATH. 

292 """ 

293 self._setup_yaml_dir(self._env_yaml_dir) 

294 self._extend_sys_path(self._env_python_path) 

295 

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

297 """ 

298 Feature: Configure yaml_dir/work_dir for OpDispatcher 

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

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

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

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

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

304 safe_load_yaml_from_dir(); no functional behavior is changed. 

305 """ 

306 if env_yaml_dir: 

307 if os.path.isabs(env_yaml_dir): 

308 self.yaml_dir = env_yaml_dir 

309 self.work_dir = "" 

310 else: 

311 self.work_dir = os.path.normpath( 

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

313 ) 

314 self.yaml_dir = env_yaml_dir 

315 else: 

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

317 self.work_dir = os.path.normpath( 

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

319 ) 

320 

321 @staticmethod 

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

323 if not env_python_path: 

324 return 

325 python_paths = env_python_path.split(":") 

326 for path in python_paths: 

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

328 sys.path.append(path) 

329 

330 def _register_distributed_ops(self): 

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

332 self._register_single_distributed_op(op_name, config) 

333 

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

335 """ 

336 Feature: Register a single distributed op implementation 

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

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

339 Prefer 'distributed_op_module' when provided; otherwise import from 

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

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

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

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

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

345 """ 

346 class_name = config["distributed_op_class"] 

347 

348 if "distributed_op_module" in config: 

349 module_name = config["distributed_op_module"] 

350 module = importlib.import_module(module_name) 

351 op_class = getattr(module, class_name) 

352 _ = op_class(op_name) 

353 return 

354 

355 module_file = config["distributed_op_file"] 

356 try: 

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

358 module = importlib.import_module(module_name) 

359 op_class = getattr(module, class_name) 

360 _ = op_class(op_name) 

361 except (ModuleNotFoundError, ImportError): 

362 if self._env_python_path: 

363 module = importlib.import_module(module_file) 

364 op_class = getattr(module, class_name) 

365 _ = op_class(op_name) 

366 else: 

367 raise 

368 

369 @staticmethod 

370 def _merge_default(config: dict): 

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

372 if "__default__" not in config: 

373 return config 

374 

375 default_cfg = config["__default__"] 

376 merged = {} 

377 

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

379 if op_name == "__default__": 

380 continue 

381 

382 new_cfg = default_cfg.copy() 

383 new_cfg.update(op_cfg) 

384 merged[op_name] = new_cfg 

385 

386 return merged 

387 

388 def safe_load_yaml_from_dir(self) -> dict: 

389 """ 

390 Load yaml dictionary from directory. 

391 

392 Returns: 

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

394 """ 

395 yaml_dict = {} 

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

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

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

399 

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

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

402 yaml_data = yaml.safe_load(f) 

403 

404 yaml_data = OpDispatcher._merge_default(yaml_data) 

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

406 if name in yaml_dict: 

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

408 yaml_dict[name] = data 

409 

410 return yaml_dict 

411 

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

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

414 first_arg = next( 

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

416 None, 

417 ) 

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

419 if first_arg is None: 

420 return op_call(*args, **kwargs) 

421 

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

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

424 first_local_arg = first_arg.to_local() 

425 

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

427 self._rng_tracker = OffsetBasedRNGTracker() 

428 

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

430 if ( 

431 self._rng_tracker is not None 

432 and not first_local_arg.is_meta 

433 and self._rng_tracker.distribute_region_enabled 

434 ): 

435 # pylint: disable=W0212 

436 with self._rng_tracker._distribute_region( 

437 device_mesh=first_arg.device_mesh, 

438 placements=first_arg.placements, 

439 global_shape=first_arg.shape, 

440 generator=maybe_user_generator, 

441 ): 

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

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

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

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

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

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

448 if platform.platform_type == PlatformType.MINDSPORE: 

449 offset_incr = self._rng_tracker.compute_offset_incr( 

450 first_arg.device_mesh, first_arg.placements, first_arg.shape, 

451 ) 

452 local_args = _apply_shard_offset_to_rng_args(local_args, offset_incr) 

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

454 else: 

455 if maybe_user_generator is not None: 

456 local_kwargs["generator"] = maybe_user_generator 

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

458 

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

460 

461 @staticmethod 

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

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

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

465 if len(args) >= 4: 

466 return bool(args[3]) 

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

468 

469 @staticmethod 

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

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

472 if op_name in OpDispatcher._RANDOM_INPLACE_MS_OPS: 

473 return True 

474 if op_name == "FuncDropoutExt": 

475 return OpDispatcher._func_dropout_ext_inplace(args, kwargs) 

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

477 return op_name.endswith('_') 

478 

479 @staticmethod 

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

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

482 

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

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

485 ``_RANDOM_INPLACE_MS_OPS``. ``FuncDropoutExt`` is handled separately 

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

487 """ 

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

489 return first_arg 

490 mesh = first_arg.device_mesh 

491 placements = first_arg.layout.alias_placements 

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

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

494 return tuple( 

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

496 for r in local_results 

497 ) 

498 if isinstance(local_results, Tensor): 

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

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

501 return local_results 

502 

503 @staticmethod 

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

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

506 

507 Args: 

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

509 

510 Returns: 

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

512 """ 

513 if isinstance(value, DTensor): 

514 return value.to_local() 

515 if isinstance(value, tuple): 

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

517 if isinstance(value, list): 

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

519 return value 

520 

521 @staticmethod 

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

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

524 

525 Args: 

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

527 

528 Returns: 

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

530 """ 

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

532 

533 @staticmethod 

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

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

536 

537 Args: 

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

539 

540 Returns: 

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

542 """ 

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

544 

545 @staticmethod 

546 def _collect_dtensors(value: object) -> List[DTensor]: 

547 """Return all DTensors nested in one dispatch argument.""" 

548 if isinstance(value, DTensor): 

549 return [value] 

550 if isinstance(value, (tuple, list)): 

551 return list(chain.from_iterable( 

552 OpDispatcher._collect_dtensors(item) for item in value 

553 )) 

554 if isinstance(value, dict): 

555 return list(chain.from_iterable( 

556 OpDispatcher._collect_dtensors(item) for item in value.values() 

557 )) 

558 return [] 

559 

560 def _validate_ragged_dispatch( 

561 self, op_name: str, args: tuple, kwargs: dict 

562 ) -> Optional[DTensor]: 

563 """Return the first Ragged input for a whitelisted elementwise op.""" 

564 reference = next( 

565 ( 

566 dtensor for dtensor in self._collect_dtensors((args, kwargs)) 

567 if isinstance(getattr(dtensor.layout, "ragged_shard", None), RaggedShardInfo) 

568 ), 

569 None, 

570 ) 

571 if reference is None: 

572 return None 

573 if op_name not in _RAGGED_ELEMENTWISE_OPS: 

574 raise RuntimeError( 

575 f"Operator {op_name!r} does not support RaggedShard in phase one" 

576 ) 

577 return reference 

578 

579 def _dispatch_ragged_elementwise( 

580 self, op_call: callable, args: tuple, kwargs: dict, 

581 reference: DTensor, 

582 ) -> DTensor: 

583 """Execute a whitelisted op locally and inherit its Ragged Layout.""" 

584 local_args = tuple(self._unwrap_args(args)) 

585 local_kwargs = self._unwrap_kwargs(kwargs) 

586 py_output = op_call(*local_args, **local_kwargs) 

587 return DTensor.from_local_with_layout( 

588 py_output, 

589 copy.deepcopy(reference.layout), 

590 shape=tuple(reference.shape), 

591 ) 

592 

593 @staticmethod 

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

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

596 

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

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

599 

600 Args: 

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

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

603 

604 Returns: 

605 Tuple of (unwrapped_args, unwrapped_kwargs) with DTensor values 

606 replaced by their full tensor representations. 

607 

608 Warning: 

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

610 Operators without layout implementations should be registered properly. 

611 """ 

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

613 if isinstance(value, DTensor): 

614 return value.full_tensor() 

615 if isinstance(value, tuple): 

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

617 if isinstance(value, list): 

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

619 return value 

620 

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

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

623 

624 warnings.warn( 

625 "Operator has no distributed layout implementation. " 

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

627 "Consider registering a proper distributed operator.", 

628 UserWarning, 

629 stacklevel=4 

630 ) 

631 

632 return gathered_args, gathered_kwargs 

633 

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

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

636 

637 Args: 

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

639 

640 Returns: 

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

642 """ 

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

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

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

646 return True 

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

648 

649 @staticmethod 

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

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

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

653 return 

654 first = args[0] 

655 if len(args) >= 2: 

656 second = args[1] 

657 elif "other" in kwargs: 

658 second = kwargs["other"] 

659 else: 

660 return 

661 if not isinstance(first, DTensor): 

662 return 

663 mesh_ndim = len(first.layout.partial) 

664 first_partial = tuple(first.layout.partial) 

665 if isinstance(second, DTensor): 

666 second_partial = tuple(second.layout.partial) 

667 if len(second_partial) != mesh_ndim: 

668 raise ValueError( 

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

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

671 ) 

672 else: 

673 second_partial = (None,) * mesh_ndim 

674 if first_partial != second_partial: 

675 raise ValueError( 

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

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

678 ) 

679 

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

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

682 

683 Args: 

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

685 

686 Returns: 

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

688 """ 

689 return is_loss_parallel_active() and is_loss_parallel_op(op_name) 

690 

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

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

693 

694 Args: 

695 op_name: Canonical operator name. 

696 args: Positional arguments for op_call. 

697 kwargs: Keyword arguments for op_call. 

698 

699 Raises: 

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

701 with vocab-sharded DTensor input. 

702 """ 

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

704 return 

705 

706 has_vocab_sharded_dtensor = False 

707 for arg in args: 

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

709 has_vocab_sharded_dtensor = True 

710 break 

711 if not has_vocab_sharded_dtensor: 

712 for val in kwargs.values(): 

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

714 has_vocab_sharded_dtensor = True 

715 break 

716 

717 if has_vocab_sharded_dtensor: 

718 raise ValueError( 

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

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

721 f"For example, replace:\n" 

722 f" with loss_parallel():\n" 

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

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

725 f"with:\n" 

726 f" with loss_parallel():\n" 

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

728 ) 

729 

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

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

732 

733 Args: 

734 op_call: The raw operator callable. 

735 args: Positional arguments for op_call. 

736 kwargs: Keyword arguments for op_call. 

737 

738 Returns: 

739 Result of the distributed cross_entropy computation. 

740 """ 

741 if platform.platform_type == PlatformType.PYTORCH: 

742 # pylint: disable=C0415 

743 from hyper_parallel.platform.torch.loss_parallel_ops import distributed_cross_entropy_from_op_call 

744 elif platform.platform_type == PlatformType.MINDSPORE: 

745 # pylint: disable=C0415 

746 from hyper_parallel.platform.mindspore.loss_parallel_ops import distributed_cross_entropy_from_op_call 

747 else: 

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

749 return distributed_cross_entropy_from_op_call(op_call, args, kwargs) 

750 

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

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

753 

754 Args: 

755 op_name: Canonical operator name. 

756 args: Positional arguments for op_call. 

757 

758 Raises: 

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

760 """ 

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

762 return 

763 

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

765 return 

766 

767 logits = args[0] 

768 if _is_shard_on_last_dim(logits): 

769 raise ValueError( 

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

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

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

773 f" with loss_parallel():\n" 

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

775 f" loss.backward()\n" 

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

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

778 ) 

779 

780 @staticmethod 

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

782 """ 

783 Normalize aclop-packed arguments for MindSpore backend operators. 

784 

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

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

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

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

789 for layout inference and preprocessing. 

790 

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

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

793 

794 Args: 

795 op_name (str): Canonical operator name. 

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

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

798 

799 Returns: 

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

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

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

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

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

805 """ 

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

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

808 return None, args 

809 

810 @staticmethod 

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

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

813 return ( 

814 op_name in unpack_ops 

815 and len(args) == 3 

816 and isinstance(args[1], str) 

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

818 ) 

819 

820 @staticmethod 

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

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

823 

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

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

826 as positional arguments in the usual way. 

827 

828 Args: 

829 op_impl: The op implementation callable. 

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

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

832 kwargs: Keyword arguments dict. 

833 

834 Returns: 

835 Result of the *op_impl* invocation. 

836 """ 

837 if packed_call is not None: 

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

839 return op_impl(*args, **kwargs) 

840 

841 def _handle_unregistered_op( 

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

843 ): 

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

845 

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

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

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

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

850 

851 Args: 

852 op_name: Canonical operator name. 

853 op_call: The raw operator callable. 

854 args: Positional arguments for op_call. 

855 kwargs: Keyword arguments for op_call. 

856 

857 Returns: 

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

859 

860 Raises: 

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

862 """ 

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

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

865 if has_dtensor: 

866 self._check_ce_op_without_loss_parallel_context(op_name, args) 

867 

868 if not is_loss_parallel_op(op_name): 

869 raise RuntimeError( 

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

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

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

873 ) 

874 

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

876 

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

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

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

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

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

882 logits = gathered_args[0] 

883 targets = gathered_args[1] 

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