Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / core / tensor_parallel / style.py: 98%

1# Copyright 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"""Parallel styles for declarative tensor-parallel module sharding. 

16 

17Provides :class:`ParallelStyle` (ABC) and concrete implementations 

18:class:`ColwiseParallel`, :class:`RowwiseParallel`, :class:`SequenceParallel`, 

19:class:`PrepareModuleInput`, :class:`PrepareModuleInputOutput`, and 

20:class:`PrepareModuleOutput` aligned with ``torch.distributed.tensor.parallel.style``. 

21""" 

22from abc import ABC, abstractmethod 

23from typing import Any, Dict, Optional, Tuple, Union 

24 

25from hyper_parallel.core.dtensor.device_mesh import DeviceMesh 

26from hyper_parallel.core.dtensor.dtensor import ( 

27 DTensor, 

28 distribute_module, 

29 distribute_tensor, 

30 _distribute_module_iter_params, 

31 _distribute_module_new_parameter, 

32 _distribute_module_param_source, 

33 _distribute_module_set_param, 

34) 

35from hyper_parallel.core.dtensor.placement_types import Partial, Placement, Replicate, Shard 

36from hyper_parallel.platform import get_platform 

37 

38platform = get_platform() 

39Module = platform.Module 

40 

41 

42def _src_data_rank_for_tensor(tensor: Any, src_data_rank: Optional[int]) -> Optional[int]: 

43 """Disable source-rank communication while sharding meta tensors.""" 

44 if getattr(tensor, "is_meta", False): 

45 return None 

46 return src_data_rank 

47 

48 

49__all__ = [ 

50 "ParallelStyle", 

51 "ColwiseParallel", 

52 "RowwiseParallel", 

53 "SequenceParallel", 

54 "PrepareModuleInput", 

55 "PrepareModuleInputOutput", 

56 "PrepareModuleOutput", 

57 "NoParallel", 

58] 

59 

60 

61class ParallelStyle(ABC): 

62 """Abstract base class for parallel styles applied to nn.Module submodules. 

63 

64 Subclasses implement ``apply`` to wrap a module with the desired 

65 parallel communication behaviour (e.g. all-to-all for context parallel). 

66 

67 ``src_data_rank`` mirrors PyTorch's tensor-parallel contract: it can be set by 

68 :func:`parallelize_module` for styles that scatter/broadcast global tensors. 

69 HyperParallel styles may ignore it until they integrate ``distribute_tensor``. 

70 """ 

71 

72 src_data_rank: Optional[int] = 0 

73 

74 @abstractmethod 

75 def apply(self, module: Module, device_mesh: DeviceMesh) -> Module: 

76 """Apply this parallel style to *module* in-place and return it. 

77 

78 Args: 

79 module: The submodule to be parallelised. 

80 device_mesh: The device mesh describing the cluster topology. 

81 

82 Returns: 

83 The (possibly wrapped) module with parallelism applied. 

84 """ 

85 

86 

87class ColwiseParallel(ParallelStyle): 

88 """Partition a compatible module in a column-wise fashion. 

89 

90 Currently supports Linear and Embedding modules (framework-agnostic via 

91 ``platform.is_linear_module`` / ``platform.is_embedding_module``). 

92 Compose with :class:`RowwiseParallel` to shard MLP or Attention blocks. 

93 

94 Keyword Args: 

95 input_layouts (Placement, optional): 

96 DTensor layout for the module input. Used to annotate the input 

97 tensor as a DTensor. Defaults to ``Replicate()``. 

98 output_layouts (Placement, optional): 

99 Desired DTensor layout of the module output. Defaults to 

100 ``Shard(-1)`` (sharded on the last dimension). 

101 use_local_output (bool, optional): 

102 If ``True`` (default), convert the output DTensor back to a local 

103 tensor via ``to_local()``. 

104 

105 Returns: 

106 A :class:`ParallelStyle` that applies column-wise sharding. 

107 

108 Example:: 

109 

110 >>> from hyper_parallel import parallelize_module, ColwiseParallel, init_device_mesh 

111 >>> m = Model(...) 

112 >>> tp_mesh = init_device_mesh("npu", (8,), mesh_dim_names=("tp",)) 

113 >>> parallelize_module(m, tp_mesh, {"linear1": ColwiseParallel()}) 

114 """ 

115 

116 def __init__( 

117 self, 

118 *, 

119 input_layouts: Optional[Placement] = None, 

120 output_layouts: Optional[Placement] = None, 

121 use_local_output: Optional[bool] = None, 

122 ) -> None: 

123 super().__init__() 

124 self._input_layouts_arg = input_layouts 

125 self._output_layouts_arg = output_layouts 

126 self._use_local_output_arg = use_local_output 

127 

128 self.input_layouts: Tuple[Placement, ...] = (input_layouts or Replicate(),) 

129 self.output_layouts: Tuple[Placement, ...] = (output_layouts or Shard(-1),) 

130 self.desired_input_layouts: Tuple[Placement, ...] = (Replicate(),) 

131 self.use_local_output = use_local_output if use_local_output is not None else True 

132 

133 def __repr__(self) -> str: 

134 return ( 

135 f"{self.__class__.__name__}(" 

136 f"input_layouts={self.input_layouts}, " 

137 f"output_layouts={self.output_layouts}, " 

138 f"use_local_output={self.use_local_output})" 

139 ) 

140 

141 @staticmethod 

142 def _prepare_input_fn( 

143 input_layouts: Tuple[Placement, ...], 

144 desired_input_layouts: Tuple[Placement, ...], 

145 inputs: Any, 

146 device_mesh: DeviceMesh, 

147 ) -> Any: 

148 """Annotate or redistribute the first positional input.""" 

149 input_tensor = inputs[0] 

150 if not isinstance(input_tensor, DTensor): 

151 input_tensor = DTensor.from_local( 

152 input_tensor, device_mesh, input_layouts, 

153 ) 

154 

155 if input_layouts != desired_input_layouts: 

156 input_tensor = input_tensor.redistribute( 

157 device_mesh, desired_input_layouts, 

158 ) 

159 # MindSpore requires tuple return from pre-hook 

160 return (input_tensor,) 

161 

162 def _partition_linear_fn(self, module: Any, device_mesh: DeviceMesh) -> None: 

163 """Shard Linear weight/bias along ``Shard(0)`` (column-wise).""" 

164 for key, param in _distribute_module_iter_params(module): 

165 if param is None: 

166 continue 

167 src = _distribute_module_param_source(param) 

168 requires_grad = bool(getattr(param, "requires_grad", True)) 

169 dt = distribute_tensor( 

170 src, 

171 device_mesh, 

172 [Shard(0)], 

173 src_data_rank=_src_data_rank_for_tensor(src, self.src_data_rank), 

174 ) 

175 new_param = _distribute_module_new_parameter(key, dt, requires_grad) 

176 _distribute_module_set_param(module, key, new_param) 

177 

178 def _partition_embedding_fn(self, module: Any, device_mesh: DeviceMesh) -> None: 

179 """Shard Embedding weight along ``Shard(1)`` (column-wise).""" 

180 for key, param in _distribute_module_iter_params(module): 

181 if param is None: 

182 continue 

183 src = _distribute_module_param_source(param) 

184 requires_grad = bool(getattr(param, "requires_grad", True)) 

185 dt = distribute_tensor( 

186 src, 

187 device_mesh, 

188 [Shard(1)], 

189 src_data_rank=_src_data_rank_for_tensor(src, self.src_data_rank), 

190 ) 

191 new_param = _distribute_module_new_parameter(key, dt, requires_grad) 

192 _distribute_module_set_param(module, key, new_param) 

193 

194 @staticmethod 

195 def _prepare_output_fn( 

196 output_layouts: Tuple[Placement, ...], 

197 use_local_output: bool, 

198 outputs: Any, 

199 device_mesh: DeviceMesh, 

200 ) -> Any: 

201 """Redistribute output to desired layout and optionally convert to local.""" 

202 if outputs.placements != output_layouts: 

203 outputs = outputs.redistribute(device_mesh, output_layouts) 

204 if use_local_output: 

205 return outputs.to_local() 

206 return outputs 

207 

208 def apply(self, module: Module, device_mesh: DeviceMesh) -> Module: 

209 """Apply column-wise parallelism to *module*. 

210 

211 Args: 

212 module: A Linear or Embedding module to be sharded. 

213 device_mesh: 1-D device mesh for tensor parallelism. 

214 

215 Returns: 

216 The module with distributed parameters and I/O hooks attached. 

217 

218 Raises: 

219 NotImplementedError: If *module* is not a supported type. 

220 """ 

221 if platform.is_linear_module(module): 

222 

223 def partition_fn(submodule_path, submodule, device_mesh): 

224 self._partition_linear_fn(submodule, device_mesh) 

225 

226 elif platform.is_embedding_module(module): 

227 

228 def partition_fn(submodule_path, submodule, device_mesh): 

229 self._partition_embedding_fn(submodule, device_mesh) 

230 

231 else: 

232 raise NotImplementedError( 

233 "ColwiseParallel currently only supports Linear and Embedding modules!" 

234 ) 

235 

236 def input_fn(forward_module, forward_inputs, device_mesh): 

237 return self._prepare_input_fn( 

238 self.input_layouts, 

239 self.desired_input_layouts, 

240 forward_inputs, 

241 device_mesh, 

242 ) 

243 

244 def output_fn(forward_module, forward_outputs, device_mesh): 

245 return self._prepare_output_fn( 

246 self.output_layouts, 

247 self.use_local_output, 

248 forward_outputs, 

249 device_mesh, 

250 ) 

251 

252 return distribute_module( 

253 module, 

254 device_mesh, 

255 partition_fn, 

256 input_fn, 

257 output_fn, 

258 ) 

259 

260 

261class RowwiseParallel(ParallelStyle): 

262 """Partition a compatible module in a row-wise fashion. 

263 

264 Currently supports Linear and Embedding modules (framework-agnostic via 

265 ``platform.is_linear_module`` / ``platform.is_embedding_module``). 

266 Compose with :class:`ColwiseParallel` to shard MLP or Attention blocks. 

267 

268 Keyword Args: 

269 input_layouts (Placement, optional): 

270 DTensor layout for the module input. Defaults to ``Shard(-1)`` 

271 (sharded on the last dimension). 

272 output_layouts (Placement, optional): 

273 Desired DTensor layout of the module output. Defaults to 

274 ``Replicate()`` (all-reduce / reduce-scatter from partial). 

275 reduce_dtype (dtype, optional): 

276 Floating-point dtype used for reducing Partial outputs; the reduced 

277 output remains in this dtype. Defaults to ``None`` (reduce in the 

278 output dtype). 

279 use_local_output (bool, optional): 

280 If ``True`` (default), convert the output DTensor back to a local 

281 tensor via ``to_local()``. 

282 

283 Returns: 

284 A :class:`ParallelStyle` that applies row-wise sharding. 

285 

286 Example:: 

287 >>> from hyper_parallel import parallelize_module, RowwiseParallel, init_device_mesh 

288 >>> m = Model(...) 

289 >>> tp_mesh = init_device_mesh("npu", (8,), mesh_dim_names=("tp",)) 

290 >>> parallelize_module(m, tp_mesh, {"linear2": RowwiseParallel()}) 

291 """ 

292 

293 def __init__( 

294 self, 

295 *, 

296 input_layouts: Optional[Placement] = None, 

297 output_layouts: Optional[Placement] = None, 

298 reduce_dtype: Optional[Any] = None, 

299 use_local_output: bool = True, 

300 ) -> None: 

301 super().__init__() 

302 self.input_layouts: Tuple[Placement, ...] = (input_layouts or Shard(-1),) 

303 self.output_layouts: Tuple[Placement, ...] = (output_layouts or Replicate(),) 

304 self.desired_input_layouts: Tuple[Placement, ...] = (Shard(-1),) 

305 self.reduce_dtype = reduce_dtype 

306 self.use_local_output = use_local_output 

307 

308 def __repr__(self) -> str: 

309 return ( 

310 f"{self.__class__.__name__}(" 

311 f"input_layouts={self.input_layouts}, " 

312 f"output_layouts={self.output_layouts}, " 

313 f"reduce_dtype={self.reduce_dtype}, " 

314 f"use_local_output={self.use_local_output})" 

315 ) 

316 

317 @staticmethod 

318 def _prepare_input_fn( 

319 input_layouts: Tuple[Placement, ...], 

320 desired_input_layouts: Tuple[Placement, ...], 

321 inputs: Any, 

322 device_mesh: DeviceMesh, 

323 ) -> Any: 

324 """Annotate or redistribute the first positional input.""" 

325 input_tensor = inputs[0] 

326 if not isinstance(input_tensor, DTensor): 

327 input_tensor = DTensor.from_local( 

328 input_tensor, device_mesh, input_layouts, 

329 ) 

330 

331 if input_layouts != desired_input_layouts: 

332 input_tensor = input_tensor.redistribute( 

333 device_mesh, desired_input_layouts, 

334 ) 

335 # MindSpore requires tuple return from pre-hook 

336 return (input_tensor,) 

337 

338 def _partition_linear_fn(self, module: Any, device_mesh: DeviceMesh) -> None: 

339 """Shard Linear weight along ``Shard(1)`` (row-wise); bias to ``Replicate()``.""" 

340 for key, param in _distribute_module_iter_params(module): 

341 if param is None: 

342 continue 

343 src = _distribute_module_param_source(param) 

344 requires_grad = bool(getattr(param, "requires_grad", True)) 

345 placement = [Shard(1)] if key == "weight" else [Replicate()] 

346 dt = distribute_tensor( 

347 src, 

348 device_mesh, 

349 placement, 

350 src_data_rank=_src_data_rank_for_tensor(src, self.src_data_rank), 

351 ) 

352 new_param = _distribute_module_new_parameter(key, dt, requires_grad) 

353 _distribute_module_set_param(module, key, new_param) 

354 

355 def _partition_embedding_fn(self, module: Any, device_mesh: DeviceMesh) -> None: 

356 """Shard Embedding weight along ``Shard(0)`` (row-wise).""" 

357 for key, param in _distribute_module_iter_params(module): 

358 if param is None: 

359 continue 

360 src = _distribute_module_param_source(param) 

361 requires_grad = bool(getattr(param, "requires_grad", True)) 

362 dt = distribute_tensor( 

363 src, 

364 device_mesh, 

365 [Shard(0)], 

366 src_data_rank=_src_data_rank_for_tensor(src, self.src_data_rank), 

367 ) 

368 new_param = _distribute_module_new_parameter(key, dt, requires_grad) 

369 _distribute_module_set_param(module, key, new_param) 

370 

371 @staticmethod 

372 def _prepare_output_fn( 

373 output_layouts: Tuple[Placement, ...], 

374 use_local_output: bool, 

375 outputs: Any, 

376 device_mesh: DeviceMesh, 

377 module: Optional[Module] = None, 

378 reduce_dtype: Optional[Any] = None, 

379 ) -> Any: 

380 """Redistribute partial output and optionally convert to local.""" 

381 if not isinstance(outputs, DTensor): 

382 # ``nn.Embedding.forward`` returns a plain tensor even when weight is sharded; 

383 # treat the local values as partial along the TP mesh (sum) before redistributing. 

384 if module is not None and platform.is_embedding_module(module): 

385 outputs = DTensor.from_local(outputs, device_mesh, [Partial("sum")]) 

386 else: 

387 raise TypeError( 

388 "RowwiseParallel expects a DTensor from Linear outputs; " 

389 f"got {type(outputs)}. If this is an unsupported module, extend I/O hooks." 

390 ) 

391 if reduce_dtype is not None and tuple(outputs.placements) != tuple(output_layouts): 

392 local_output = platform.cast_fp_tensor(reduce_dtype, outputs.to_local()) 

393 outputs = DTensor.from_local(local_output, outputs.device_mesh, outputs.placements) 

394 if tuple(outputs.placements) != tuple(output_layouts): 

395 outputs = outputs.redistribute(device_mesh, output_layouts) 

396 if use_local_output: 

397 return outputs.to_local() 

398 return outputs 

399 

400 def apply(self, module: Module, device_mesh: DeviceMesh) -> Module: 

401 """Apply row-wise parallelism to *module*. 

402 

403 Args: 

404 module: A Linear or Embedding module to be sharded. 

405 device_mesh: 1-D device mesh for tensor parallelism. 

406 

407 Returns: 

408 The module with distributed parameters and I/O hooks attached. 

409 

410 Raises: 

411 NotImplementedError: If *module* is not a supported type. 

412 """ 

413 if platform.is_linear_module(module): 

414 

415 def partition_fn(submodule_path, submodule, device_mesh): 

416 self._partition_linear_fn(submodule, device_mesh) 

417 

418 self.desired_input_layouts = (Shard(-1),) 

419 elif platform.is_embedding_module(module): 

420 

421 def partition_fn(submodule_path, submodule, device_mesh): 

422 self._partition_embedding_fn(submodule, device_mesh) 

423 

424 self.desired_input_layouts = (Replicate(),) 

425 else: 

426 raise NotImplementedError( 

427 "RowwiseParallel currently only supports Linear and Embedding modules!" 

428 ) 

429 

430 def input_fn(forward_module, forward_inputs, device_mesh): 

431 return self._prepare_input_fn( 

432 self.input_layouts, 

433 self.desired_input_layouts, 

434 forward_inputs, 

435 device_mesh, 

436 ) 

437 

438 def output_fn(forward_module, forward_outputs, device_mesh): 

439 return self._prepare_output_fn( 

440 self.output_layouts, 

441 self.use_local_output, 

442 forward_outputs, 

443 device_mesh, 

444 forward_module, 

445 self.reduce_dtype, 

446 ) 

447 

448 return distribute_module( 

449 module, 

450 device_mesh, 

451 partition_fn, 

452 input_fn, 

453 output_fn, 

454 ) 

455 

456 

457class SequenceParallel(ParallelStyle): 

458 """Replicate module parameters and run forward with the sequence axis sharded. 

459 

460 Matches ``torch.distributed.tensor.parallel.SequenceParallel``: activations are 

461 sharded on the sequence dimension while weights stay fully replicated. Typical 

462 targets are normalization and dropout layers used after row-wise / scatter 

463 projections in tensor-parallel transformers (`Reducing Activation Recomputation 

464 in Large Transformer Models <https://arxiv.org/abs/2205.05198>`__). 

465 

466 If the first positional input is a plain tensor, it is treated as the local 

467 shard along ``sequence_dim`` and wrapped as a :class:`DTensor`. If it is already 

468 a :class:`DTensor` but not sharded on that dimension, it is redistributed. 

469 

470 Keyword Args: 

471 sequence_dim (int, optional): 

472 Tensor dimension index for the sequence axis (e.g. ``1`` for ``(B, S, H)``). 

473 Default: ``1``. 

474 use_local_output (bool, optional): 

475 If ``True``, return a local tensor via ``to_local()``; otherwise keep a 

476 :class:`DTensor`. Default: ``False`` (PyTorch default). 

477 

478 Note: 

479 Like PyTorch, this assumes sensible defaults for norm weights (e.g. ones). 

480 Custom initializations should be broadcast so every rank agrees before or 

481 after parallelization. 

482 

483 Example:: 

484 

485 >>> from hyper_parallel import parallelize_module, SequenceParallel, init_device_mesh 

486 >>> m = Model(...) 

487 >>> tp_mesh = init_device_mesh("npu", (8,), mesh_dim_names=("tp",)) 

488 >>> parallelize_module(m, tp_mesh, {"norm": SequenceParallel()}) 

489 """ 

490 

491 def __init__(self, *, sequence_dim: int = 1, use_local_output: bool = False) -> None: 

492 super().__init__() 

493 self.sequence_sharding: Tuple[Placement, ...] = (Shard(sequence_dim),) 

494 self.use_local_output = use_local_output 

495 

496 def __repr__(self) -> str: 

497 dim = self.sequence_sharding[0].dim 

498 return ( 

499 f"{self.__class__.__name__}(" 

500 f"sequence_dim={dim}, " 

501 f"use_local_output={self.use_local_output})" 

502 ) 

503 

504 @staticmethod 

505 def _prepare_input_fn( 

506 sequence_sharding: Tuple[Placement, ...], 

507 mod: Module, 

508 inputs: Any, 

509 device_mesh: DeviceMesh, 

510 ) -> Any: 

511 """Ensure the first input is a :class:`DTensor` sharded on the sequence dim.""" 

512 input_tensor = inputs[0] 

513 if isinstance(input_tensor, DTensor): 

514 if tuple(input_tensor.placements) != tuple(sequence_sharding): 

515 input_tensor = input_tensor.redistribute(device_mesh, sequence_sharding) 

516 elif platform.is_tensor(input_tensor): 

517 input_tensor = DTensor.from_local(input_tensor, device_mesh, sequence_sharding) 

518 else: 

519 raise ValueError( 

520 f"expecting input of {mod} to be a tensor or DTensor, but got {type(input_tensor)}" 

521 ) 

522 return (input_tensor,) 

523 

524 @staticmethod 

525 def _prepare_output_fn(use_local_output: bool, outputs: Any) -> Any: 

526 if use_local_output: 

527 return outputs.to_local() 

528 return outputs 

529 

530 def apply(self, module: Module, device_mesh: DeviceMesh) -> Module: 

531 """Apply sequence-parallel hooks and replicate parameters via ``distribute_module``. 

532 

533 Args: 

534 module: Submodule to parallelize (for example ``LayerNorm`` or ``Dropout``). 

535 device_mesh: One-dimensional tensor-parallel device mesh. 

536 

537 Returns: 

538 The same ``module`` instance with forward hooks attached and parameters 

539 converted to replicated DTensors where applicable. 

540 """ 

541 

542 def partition_fn(_submodule_path, _submodule, _mesh): 

543 return None 

544 

545 def input_fn(forward_module, forward_inputs, mesh): 

546 return self._prepare_input_fn( 

547 self.sequence_sharding, 

548 forward_module, 

549 forward_inputs, 

550 mesh, 

551 ) 

552 

553 def output_fn(_forward_module, forward_outputs, _mesh): 

554 return self._prepare_output_fn(self.use_local_output, forward_outputs) 

555 

556 return distribute_module( 

557 module, 

558 device_mesh, 

559 partition_fn, 

560 input_fn, 

561 output_fn, 

562 ) 

563 

564 

565class PrepareModuleInput(ParallelStyle): 

566 """Prepare module forward *args* (and optional *kwargs*) as :class:`DTensor` layouts. 

567 

568 At forward time, converts each annotated positional (or keyword) tensor from local 

569 to :class:`DTensor` using ``input_layouts``, then redistributes to 

570 ``desired_input_layouts`` when they differ. ``None`` in a layout tuple means 

571 “leave this input unchanged”. 

572 

573 Mirrors ``torch.distributed.tensor.parallel.style.PrepareModuleInput``. 

574 

575 Keyword Args: 

576 input_layouts: Placements per positional arg, or a single :class:`Placement` 

577 wrapped as a one-tuple. ``None`` entries skip conversion for that arg. 

578 desired_input_layouts: Target placements; must match ``input_layouts`` length. 

579 input_kwarg_layouts: Optional mapping kwarg name → placement for conversion. 

580 desired_input_kwarg_layouts: Target placements for those kwargs (same keys). 

581 use_local_output: If ``True``, convert prepared inputs back to local tensors 

582 before the module runs (PyTorch names this flag ``use_local_output`` on 

583 :class:`PrepareModuleInput`). 

584 """ 

585 

586 def __init__( 

587 self, 

588 *, 

589 input_layouts: Optional[Union[Placement, Tuple[Optional[Placement], ...]]] = None, 

590 desired_input_layouts: Optional[ 

591 Union[Placement, Tuple[Optional[Placement], ...]] 

592 ] = None, 

593 input_kwarg_layouts: Optional[Dict[str, Placement]] = None, 

594 desired_input_kwarg_layouts: Optional[Dict[str, Placement]] = None, 

595 use_local_output: bool = False, 

596 ) -> None: 

597 super().__init__() 

598 self.input_layouts = ( 

599 (input_layouts,) if isinstance(input_layouts, Placement) else input_layouts 

600 ) 

601 self.desired_input_layouts = ( 

602 (desired_input_layouts,) 

603 if isinstance(desired_input_layouts, Placement) 

604 else desired_input_layouts 

605 ) 

606 self.use_local_output = use_local_output 

607 if self.input_layouts is not None: 

608 if self.desired_input_layouts is None: 

609 raise AssertionError("desired module inputs should not be None!") 

610 if len(self.input_layouts) != len(self.desired_input_layouts): 

611 raise AssertionError( 

612 "input_layouts and desired_input_layouts should have same length!" 

613 ) 

614 self.with_kwargs = input_kwarg_layouts is not None 

615 self.input_kwarg_layouts = input_kwarg_layouts or {} 

616 self.desired_input_kwarg_layouts = desired_input_kwarg_layouts or {} 

617 if self.with_kwargs: 

618 if len(self.input_kwarg_layouts) != len(self.desired_input_kwarg_layouts): 

619 raise AssertionError( 

620 "input_kwarg_layouts and desired_input_kwarg_layouts should have " 

621 "same length!" 

622 ) 

623 

624 def _prepare_input_arg( 

625 self, 

626 input_obj: Any, 

627 mesh: DeviceMesh, 

628 input_layout: Optional[Placement], 

629 desired_layout: Optional[Placement], 

630 ) -> Any: 

631 """Convert one input to DTensor, redistribute if needed, optionally to_local.""" 

632 if input_layout is not None: 

633 if isinstance(input_obj, DTensor): 

634 dt_inp = input_obj 

635 else: 

636 if not platform.is_tensor(input_obj): 

637 raise AssertionError("expecting input to be a framework tensor!") 

638 dt_inp = DTensor.from_local(input_obj, mesh, (input_layout,)) 

639 

640 if desired_layout is not None and input_layout != desired_layout: 

641 dt_inp = dt_inp.redistribute(mesh, (desired_layout,)) 

642 

643 return dt_inp.to_local() if self.use_local_output else dt_inp 

644 return input_obj 

645 

646 def _prepare_input_fn(self, inputs: Any, device_mesh: DeviceMesh) -> Any: 

647 """Prepare positional ``inputs`` tuple per ``input_layouts`` / ``desired_input_layouts``.""" 

648 if self.input_layouts is None: 

649 return inputs 

650 if not isinstance(inputs, tuple): 

651 inputs = (inputs,) 

652 if len(inputs) != len(self.input_layouts): 

653 raise ValueError("module inputs and input_layouts should have same length!") 

654 if self.desired_input_layouts is None: 

655 raise AssertionError("desired module inputs should not be None!") 

656 prepared_inputs = [ 

657 self._prepare_input_arg(inp, device_mesh, il, dl) 

658 for inp, il, dl in zip(inputs, self.input_layouts, self.desired_input_layouts) 

659 ] 

660 return tuple(prepared_inputs) 

661 

662 def _prepare_input_kwarg_fn( 

663 self, 

664 inputs: Any, 

665 kwarg_inputs: Dict[str, Any], 

666 device_mesh: DeviceMesh, 

667 ) -> Tuple[Any, Dict[str, Any]]: 

668 """Prepare positional and keyword tensor inputs; returns ``(args, kwargs)`` for the hook.""" 

669 prepared_arg_inputs = self._prepare_input_fn(inputs, device_mesh) 

670 prepared_kwarg_inputs: Dict[str, Any] = {} 

671 for kwarg_key in kwarg_inputs: 

672 kwarg_val = kwarg_inputs[kwarg_key] 

673 input_layout = self.input_kwarg_layouts.get(kwarg_key) 

674 desired_input_layout = self.desired_input_kwarg_layouts.get(kwarg_key) 

675 prepared_kwarg_inputs[kwarg_key] = self._prepare_input_arg( 

676 kwarg_val, device_mesh, input_layout, desired_input_layout 

677 ) 

678 return (prepared_arg_inputs, prepared_kwarg_inputs) 

679 

680 def apply(self, module: Module, device_mesh: DeviceMesh) -> Module: 

681 if self.with_kwargs: 

682 

683 def _pre_hook(_mod, inputs, kwargs): 

684 return self._prepare_input_kwarg_fn(inputs, kwargs, device_mesh) 

685 

686 platform.register_forward_pre_hook( 

687 module, _pre_hook, prepend=False, with_kwargs=True, 

688 ) 

689 else: 

690 

691 def _pre_hook(_mod, inputs): 

692 return self._prepare_input_fn(inputs, device_mesh) 

693 

694 platform.register_forward_pre_hook(module, _pre_hook, prepend=False) 

695 return module 

696 

697 def __repr__(self) -> str: 

698 return ( 

699 f"{self.__class__.__name__}(" 

700 f"input_layouts={self.input_layouts}, " 

701 f"desired_input_layouts={self.desired_input_layouts}, " 

702 f"input_kwarg_layouts={self.input_kwarg_layouts}, " 

703 f"desired_input_kwarg_layouts={self.desired_input_kwarg_layouts}, " 

704 f"use_local_output={self.use_local_output})" 

705 ) 

706 

707 

708class PrepareModuleOutput(ParallelStyle): 

709 """Prepare module forward outputs as :class:`DTensor` and redistribute layouts. 

710 

711 Registers a forward hook that treats each return value like 

712 ``torch.distributed.tensor.parallel.style.PrepareModuleOutput``: optional 

713 ``None`` slots in ``output_layouts`` pass that output through unchanged. 

714 

715 Keyword Args: 

716 output_layouts: Current or assumed placement per output tensor. 

717 desired_output_layouts: Target placements; length must match ``output_layouts``. 

718 reduce_dtype: Floating-point dtype used before reducing Partial outputs; 

719 the reduced output remains in this dtype. 

720 use_local_output: If ``True`` (default), return local shards after redistribution. 

721 """ 

722 

723 def __init__( 

724 self, 

725 *, 

726 output_layouts: Union[Placement, Tuple[Optional[Placement], ...]], 

727 desired_output_layouts: Union[Placement, Tuple[Optional[Placement], ...]], 

728 reduce_dtype: Optional[Any] = None, 

729 use_local_output: bool = True, 

730 ) -> None: 

731 super().__init__() 

732 self.output_layouts = ( 

733 (output_layouts,) if isinstance(output_layouts, Placement) else output_layouts 

734 ) 

735 self.desired_output_layouts = ( 

736 (desired_output_layouts,) 

737 if isinstance(desired_output_layouts, Placement) 

738 else desired_output_layouts 

739 ) 

740 self.reduce_dtype = reduce_dtype 

741 self.use_local_output = use_local_output 

742 if len(self.output_layouts) != len(self.desired_output_layouts): 

743 raise AssertionError( 

744 "output_layouts and desired_output_layouts should have same length!" 

745 ) 

746 

747 def _prepare_out_fn(self, outputs: Any, device_mesh: DeviceMesh) -> Any: 

748 """Redistribute each output tensor per ``output_layouts`` / ``desired_output_layouts``.""" 

749 prepared_outputs: list = [] 

750 if not isinstance(outputs, tuple): 

751 outputs = (outputs,) 

752 if len(outputs) != len(self.output_layouts): 

753 raise ValueError("module outputs and output_layouts should have same length!") 

754 for out, out_layout, desired_out_layout in zip( 

755 outputs, self.output_layouts, self.desired_output_layouts, 

756 ): 

757 if out_layout is not None: 

758 if isinstance(out, DTensor): 

759 dt_out = out 

760 else: 

761 dt_out = DTensor.from_local(out, device_mesh, (out_layout,)) 

762 has_partial_output = any( 

763 placement.is_partial() for placement in dt_out.placements 

764 ) 

765 if ( 

766 self.reduce_dtype is not None 

767 and has_partial_output 

768 and out_layout != desired_out_layout 

769 ): 

770 local_out = platform.cast_fp_tensor(self.reduce_dtype, dt_out.to_local()) 

771 dt_out = DTensor.from_local(local_out, dt_out.device_mesh, dt_out.placements) 

772 if out_layout != desired_out_layout: 

773 dt_out = dt_out.redistribute(device_mesh, (desired_out_layout,)) 

774 prepared_outputs.append( 

775 dt_out.to_local() if self.use_local_output else dt_out 

776 ) 

777 else: 

778 prepared_outputs.append(out) 

779 if len(prepared_outputs) == 1: 

780 return prepared_outputs[0] 

781 return tuple(prepared_outputs) 

782 

783 def apply(self, module: Module, device_mesh: DeviceMesh) -> Module: 

784 

785 def _hook(_mod, _inputs, outputs): 

786 return self._prepare_out_fn(outputs, device_mesh) 

787 

788 module.register_forward_hook(_hook) 

789 return module 

790 

791 def __repr__(self) -> str: 

792 return ( 

793 f"{self.__class__.__name__}(" 

794 f"output_layouts={self.output_layouts}, " 

795 f"desired_output_layouts={self.desired_output_layouts}, " 

796 f"reduce_dtype={self.reduce_dtype}, " 

797 f"use_local_output={self.use_local_output})" 

798 ) 

799 

800 

801class PrepareModuleInputOutput(ParallelStyle): 

802 """Combine :class:`PrepareModuleInput` and :class:`PrepareModuleOutput` on one module. 

803 

804 Same keyword arguments as the two styles, with ``use_local_input`` mapping to 

805 ``PrepareModuleInput(..., use_local_output=use_local_input)`` for PyTorch parity. 

806 """ 

807 

808 def __init__( 

809 self, 

810 *, 

811 input_layouts: Optional[Union[Placement, Tuple[Optional[Placement], ...]]] = None, 

812 desired_input_layouts: Optional[ 

813 Union[Placement, Tuple[Optional[Placement], ...]] 

814 ] = None, 

815 input_kwarg_layouts: Optional[Dict[str, Placement]] = None, 

816 desired_input_kwarg_layouts: Optional[Dict[str, Placement]] = None, 

817 use_local_input: bool = False, 

818 output_layouts: Union[Placement, Tuple[Optional[Placement], ...]], 

819 desired_output_layouts: Union[Placement, Tuple[Optional[Placement], ...]], 

820 reduce_dtype: Optional[Any] = None, 

821 use_local_output: bool = True, 

822 ) -> None: 

823 super().__init__() 

824 self.prepare_module_input = PrepareModuleInput( 

825 input_layouts=input_layouts, 

826 desired_input_layouts=desired_input_layouts, 

827 input_kwarg_layouts=input_kwarg_layouts, 

828 desired_input_kwarg_layouts=desired_input_kwarg_layouts, 

829 use_local_output=use_local_input, 

830 ) 

831 self.prepare_module_output = PrepareModuleOutput( 

832 output_layouts=output_layouts, 

833 desired_output_layouts=desired_output_layouts, 

834 reduce_dtype=reduce_dtype, 

835 use_local_output=use_local_output, 

836 ) 

837 

838 def apply(self, module: Module, device_mesh: DeviceMesh) -> Module: 

839 self.prepare_module_input.apply(module, device_mesh) 

840 self.prepare_module_output.apply(module, device_mesh) 

841 return module 

842 

843 def __repr__(self) -> str: 

844 p_in = self.prepare_module_input 

845 p_out = self.prepare_module_output 

846 return ( 

847 f"{self.__class__.__name__}(" 

848 f"input_layouts={p_in.input_layouts}, " 

849 f"desired_input_layouts={p_in.desired_input_layouts}, " 

850 f"input_kwarg_layouts={p_in.input_kwarg_layouts}, " 

851 f"desired_input_kwarg_layouts={p_in.desired_input_kwarg_layouts}, " 

852 f"use_local_input={p_in.use_local_output}, " 

853 f"output_layouts={p_out.output_layouts}, " 

854 f"desired_output_layouts={p_out.desired_output_layouts}, " 

855 f"reduce_dtype={p_out.reduce_dtype}, " 

856 f"use_local_output={p_out.use_local_output})" 

857 ) 

858 

859 

860class NoParallel(ParallelStyle): 

861 """Replicate module parameters without sharding, while maintaining DTensor semantics. 

862 

863 Parameters and buffers are converted to fully replicated :class:`DTensor`, and I/O 

864 hooks ensure tensor conversion and layout alignment at module boundaries. 

865 

866 Use this style for modules that must perform identical computations across TP ranks, 

867 such as: 

868 

869 - MoE Router/Gate modules 

870 - Normalization layers when sequence parallel is disabled 

871 - Any module that should not shard weights but needs DTensor compatibility 

872 

873 Keyword Args: 

874 input_layout (Placement, optional): 

875 Layout used to annotate the first positional input if it is a plain tensor. 

876 Defaults to ``Replicate()``. 

877 output_layout (Placement, optional): 

878 Target layout for the module output. If the output :class:`DTensor` has a 

879 different placement, it will be redistributed. Defaults to ``Replicate()``. 

880 desired_input_layout (Placement, optional): 

881 Final layout for the first input after annotation/redistribution. 

882 Defaults to ``Replicate()``. If different from ``input_layout``, a 

883 redistribution is performed. 

884 use_local_output (bool, optional): 

885 If ``True``, convert the output :class:`DTensor` back to a local 

886 tensor via ``to_local()``. Defaults to ``True``. 

887 

888 Example:: 

889 

890 >>> from hyper_parallel import parallelize_module, NoParallel, init_device_mesh 

891 >>> model = Transformer(...) 

892 >>> tp_mesh = init_device_mesh("npu", (8,), mesh_dim_names=("tp",)) 

893 >>> parallelize_module(model, tp_mesh, { 

894 ... "router": NoParallel(), 

895 ... "norm": SequenceParallel() if use_sp else NoParallel(), 

896 ... }) 

897 """ 

898 

899 def __init__( 

900 self, 

901 *, 

902 input_layout: Optional[Placement] = None,