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

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__all__ = [ 

42 "ParallelStyle", 

43 "ColwiseParallel", 

44 "RowwiseParallel", 

45 "SequenceParallel", 

46 "PrepareModuleInput", 

47 "PrepareModuleInputOutput", 

48 "PrepareModuleOutput", 

49 "NoParallel", 

50] 

51 

52 

53class ParallelStyle(ABC): 

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

55 

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

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

58 

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

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

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

62 """ 

63 

64 src_data_rank: Optional[int] = 0 

65 

66 @abstractmethod 

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

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

69 

70 Args: 

71 module: The submodule to be parallelised. 

72 device_mesh: The device mesh describing the cluster topology. 

73 

74 Returns: 

75 The (possibly wrapped) module with parallelism applied. 

76 """ 

77 

78 

79class ColwiseParallel(ParallelStyle): 

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

81 

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

83 ``platform.is_linear_module`` / ``platform.is_embedding_module``). 

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

85 

86 Keyword Args: 

87 input_layouts (Placement, optional): 

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

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

90 output_layouts (Placement, optional): 

91 Desired DTensor layout of the module output. Defaults to 

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

93 use_local_output (bool, optional): 

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

95 tensor via ``to_local()``. 

96 

97 Returns: 

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

99 

100 Example:: 

101 

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

103 >>> m = Model(...) 

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

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

106 """ 

107 

108 def __init__( 

109 self, 

110 *, 

111 input_layouts: Optional[Placement] = None, 

112 output_layouts: Optional[Placement] = None, 

113 use_local_output: Optional[bool] = None, 

114 ) -> None: 

115 super().__init__() 

116 self._input_layouts_arg = input_layouts 

117 self._output_layouts_arg = output_layouts 

118 self._use_local_output_arg = use_local_output 

119 

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

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

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

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

124 

125 def __repr__(self) -> str: 

126 return ( 

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

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

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

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

131 ) 

132 

133 @staticmethod 

134 def _prepare_input_fn( 

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

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

137 inputs: Any, 

138 device_mesh: DeviceMesh, 

139 ) -> Any: 

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

141 input_tensor = inputs[0] 

142 if not isinstance(input_tensor, DTensor): 

143 input_tensor = DTensor.from_local( 

144 input_tensor, device_mesh, input_layouts, 

145 ) 

146 

147 if input_layouts != desired_input_layouts: 

148 input_tensor = input_tensor.redistribute( 

149 device_mesh, desired_input_layouts, 

150 ) 

151 # MindSpore requires tuple return from pre-hook 

152 return (input_tensor,) 

153 

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

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

156 for key, param in _distribute_module_iter_params(module): 

157 if param is None: 

158 continue 

159 src = _distribute_module_param_source(param) 

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

161 dt = distribute_tensor(src, device_mesh, [Shard(0)]) 

162 new_param = _distribute_module_new_parameter(key, dt, requires_grad) 

163 _distribute_module_set_param(module, key, new_param) 

164 

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

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

167 for key, param in _distribute_module_iter_params(module): 

168 if param is None: 

169 continue 

170 src = _distribute_module_param_source(param) 

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

172 dt = distribute_tensor(src, device_mesh, [Shard(1)]) 

173 new_param = _distribute_module_new_parameter(key, dt, requires_grad) 

174 _distribute_module_set_param(module, key, new_param) 

175 

176 @staticmethod 

177 def _prepare_output_fn( 

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

179 use_local_output: bool, 

180 outputs: Any, 

181 device_mesh: DeviceMesh, 

182 ) -> Any: 

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

184 if outputs.placements != output_layouts: 

185 outputs = outputs.redistribute(device_mesh, output_layouts) 

186 if use_local_output: 

187 return outputs.to_local() 

188 return outputs 

189 

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

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

192 

193 Args: 

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

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

196 

197 Returns: 

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

199 

200 Raises: 

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

202 """ 

203 if platform.is_linear_module(module): 

204 

205 def partition_fn(submodule_path, submodule, device_mesh): 

206 self._partition_linear_fn(submodule, device_mesh) 

207 

208 elif platform.is_embedding_module(module): 

209 

210 def partition_fn(submodule_path, submodule, device_mesh): 

211 self._partition_embedding_fn(submodule, device_mesh) 

212 

213 else: 

214 raise NotImplementedError( 

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

216 ) 

217 

218 def input_fn(forward_module, forward_inputs, device_mesh): 

219 return self._prepare_input_fn( 

220 self.input_layouts, 

221 self.desired_input_layouts, 

222 forward_inputs, 

223 device_mesh, 

224 ) 

225 

226 def output_fn(forward_module, forward_outputs, device_mesh): 

227 return self._prepare_output_fn( 

228 self.output_layouts, 

229 self.use_local_output, 

230 forward_outputs, 

231 device_mesh, 

232 ) 

233 

234 return distribute_module( 

235 module, 

236 device_mesh, 

237 partition_fn, 

238 input_fn, 

239 output_fn, 

240 ) 

241 

242 

243class RowwiseParallel(ParallelStyle): 

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

245 

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

247 ``platform.is_linear_module`` / ``platform.is_embedding_module``). 

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

249 

250 Keyword Args: 

251 input_layouts (Placement, optional): 

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

253 (sharded on the last dimension). 

254 output_layouts (Placement, optional): 

255 Desired DTensor layout of the module output. Defaults to 

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

257 use_local_output (bool, optional): 

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

259 tensor via ``to_local()``. 

260 

261 Returns: 

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

263 

264 Example:: 

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

266 >>> m = Model(...) 

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

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

269 """ 

270 

271 def __init__( 

272 self, 

273 *, 

274 input_layouts: Optional[Placement] = None, 

275 output_layouts: Optional[Placement] = None, 

276 use_local_output: bool = True, 

277 ) -> None: 

278 super().__init__() 

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

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

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

282 self.use_local_output = use_local_output 

283 

284 def __repr__(self) -> str: 

285 return ( 

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

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

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

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

290 ) 

291 

292 @staticmethod 

293 def _prepare_input_fn( 

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

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

296 inputs: Any, 

297 device_mesh: DeviceMesh, 

298 ) -> Any: 

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

300 input_tensor = inputs[0] 

301 if not isinstance(input_tensor, DTensor): 

302 input_tensor = DTensor.from_local( 

303 input_tensor, device_mesh, input_layouts, 

304 ) 

305 

306 if input_layouts != desired_input_layouts: 

307 input_tensor = input_tensor.redistribute( 

308 device_mesh, desired_input_layouts, 

309 ) 

310 # MindSpore requires tuple return from pre-hook 

311 return (input_tensor,) 

312 

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

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

315 for key, param in _distribute_module_iter_params(module): 

316 if param is None: 

317 continue 

318 src = _distribute_module_param_source(param) 

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

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

321 dt = distribute_tensor(src, device_mesh, placement) 

322 new_param = _distribute_module_new_parameter(key, dt, requires_grad) 

323 _distribute_module_set_param(module, key, new_param) 

324 

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

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

327 for key, param in _distribute_module_iter_params(module): 

328 if param is None: 

329 continue 

330 src = _distribute_module_param_source(param) 

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

332 dt = distribute_tensor(src, device_mesh, [Shard(0)]) 

333 new_param = _distribute_module_new_parameter(key, dt, requires_grad) 

334 _distribute_module_set_param(module, key, new_param) 

335 

336 @staticmethod 

337 def _prepare_output_fn( 

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

339 use_local_output: bool, 

340 outputs: Any, 

341 device_mesh: DeviceMesh, 

342 module: Optional[Module] = None, 

343 ) -> Any: 

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

345 if not isinstance(outputs, DTensor): 

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

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

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

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

350 else: 

351 raise TypeError( 

352 "RowwiseParallel expects a DTensor from Linear outputs; " 

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

354 ) 

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

356 outputs = outputs.redistribute(device_mesh, output_layouts) 

357 if use_local_output: 

358 return outputs.to_local() 

359 return outputs 

360 

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

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

363 

364 Args: 

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

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

367 

368 Returns: 

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

370 

371 Raises: 

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

373 """ 

374 if platform.is_linear_module(module): 

375 

376 def partition_fn(submodule_path, submodule, device_mesh): 

377 self._partition_linear_fn(submodule, device_mesh) 

378 

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

380 elif platform.is_embedding_module(module): 

381 

382 def partition_fn(submodule_path, submodule, device_mesh): 

383 self._partition_embedding_fn(submodule, device_mesh) 

384 

385 self.desired_input_layouts = (Replicate(),) 

386 else: 

387 raise NotImplementedError( 

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

389 ) 

390 

391 def input_fn(forward_module, forward_inputs, device_mesh): 

392 return self._prepare_input_fn( 

393 self.input_layouts, 

394 self.desired_input_layouts, 

395 forward_inputs, 

396 device_mesh, 

397 ) 

398 

399 def output_fn(forward_module, forward_outputs, device_mesh): 

400 return self._prepare_output_fn( 

401 self.output_layouts, 

402 self.use_local_output, 

403 forward_outputs, 

404 device_mesh, 

405 forward_module, 

406 ) 

407 

408 return distribute_module( 

409 module, 

410 device_mesh, 

411 partition_fn, 

412 input_fn, 

413 output_fn, 

414 ) 

415 

416 

417class SequenceParallel(ParallelStyle): 

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

419 

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

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

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

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

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

425 

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

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

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

429 

430 Keyword Args: 

431 sequence_dim (int, optional): 

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

433 Default: ``1``. 

434 use_local_output (bool, optional): 

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

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

437 

438 Note: 

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

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

441 after parallelization. 

442 

443 Example:: 

444 

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

446 >>> m = Model(...) 

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

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

449 """ 

450 

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

452 super().__init__() 

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

454 self.use_local_output = use_local_output 

455 

456 def __repr__(self) -> str: 

457 dim = self.sequence_sharding[0].dim 

458 return ( 

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

460 f"sequence_dim={dim}, " 

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

462 ) 

463 

464 @staticmethod 

465 def _prepare_input_fn( 

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

467 mod: Module, 

468 inputs: Any, 

469 device_mesh: DeviceMesh, 

470 ) -> Any: 

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

472 input_tensor = inputs[0] 

473 if isinstance(input_tensor, DTensor): 

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

475 input_tensor = input_tensor.redistribute(device_mesh, sequence_sharding) 

476 elif platform.is_tensor(input_tensor): 

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

478 else: 

479 raise ValueError( 

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

481 ) 

482 return (input_tensor,) 

483 

484 @staticmethod 

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

486 if use_local_output: 

487 return outputs.to_local() 

488 return outputs 

489 

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

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

492 

493 Args: 

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

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

496 

497 Returns: 

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

499 converted to replicated DTensors where applicable. 

500 """ 

501 

502 def partition_fn(_submodule_path, _submodule, _mesh): 

503 return None 

504 

505 def input_fn(forward_module, forward_inputs, mesh): 

506 return self._prepare_input_fn( 

507 self.sequence_sharding, 

508 forward_module, 

509 forward_inputs, 

510 mesh, 

511 ) 

512 

513 def output_fn(_forward_module, forward_outputs, _mesh): 

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

515 

516 return distribute_module( 

517 module, 

518 device_mesh, 

519 partition_fn, 

520 input_fn, 

521 output_fn, 

522 ) 

523 

524 

525class PrepareModuleInput(ParallelStyle): 

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

527 

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

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

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

531 “leave this input unchanged”. 

532 

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

534 

535 Keyword Args: 

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

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

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

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

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

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

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

543 :class:`PrepareModuleInput`). 

544 """ 

545 

546 def __init__( 

547 self, 

548 *, 

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

550 desired_input_layouts: Optional[ 

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

552 ] = None, 

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

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

555 use_local_output: bool = False, 

556 ) -> None: 

557 super().__init__() 

558 self.input_layouts = ( 

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

560 ) 

561 self.desired_input_layouts = ( 

562 (desired_input_layouts,) 

563 if isinstance(desired_input_layouts, Placement) 

564 else desired_input_layouts 

565 ) 

566 self.use_local_output = use_local_output 

567 if self.input_layouts is not None: 

568 if self.desired_input_layouts is None: 

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

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

571 raise AssertionError( 

572 "input_layouts and desired_input_layouts should have same length!" 

573 ) 

574 self.with_kwargs = input_kwarg_layouts is not None 

575 self.input_kwarg_layouts = input_kwarg_layouts or {} 

576 self.desired_input_kwarg_layouts = desired_input_kwarg_layouts or {} 

577 if self.with_kwargs: 

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

579 raise AssertionError( 

580 "input_kwarg_layouts and desired_input_kwarg_layouts should have " 

581 "same length!" 

582 ) 

583 

584 def _prepare_input_arg( 

585 self, 

586 input_obj: Any, 

587 mesh: DeviceMesh, 

588 input_layout: Optional[Placement], 

589 desired_layout: Optional[Placement], 

590 ) -> Any: 

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

592 if input_layout is not None: 

593 if isinstance(input_obj, DTensor): 

594 dt_inp = input_obj 

595 else: 

596 if not platform.is_tensor(input_obj): 

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

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

599 

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

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

602 

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

604 return input_obj 

605 

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

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

608 if self.input_layouts is None: 

609 return inputs 

610 if not isinstance(inputs, tuple): 

611 inputs = (inputs,) 

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

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

614 if self.desired_input_layouts is None: 

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

616 prepared_inputs = [ 

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

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

619 ] 

620 return tuple(prepared_inputs) 

621 

622 def _prepare_input_kwarg_fn( 

623 self, 

624 inputs: Any, 

625 kwarg_inputs: Dict[str, Any], 

626 device_mesh: DeviceMesh, 

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

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

629 prepared_arg_inputs = self._prepare_input_fn(inputs, device_mesh) 

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

631 for kwarg_key in kwarg_inputs: 

632 kwarg_val = kwarg_inputs[kwarg_key] 

633 input_layout = self.input_kwarg_layouts.get(kwarg_key) 

634 desired_input_layout = self.desired_input_kwarg_layouts.get(kwarg_key) 

635 prepared_kwarg_inputs[kwarg_key] = self._prepare_input_arg( 

636 kwarg_val, device_mesh, input_layout, desired_input_layout 

637 ) 

638 return (prepared_arg_inputs, prepared_kwarg_inputs) 

639 

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

641 if self.with_kwargs: 

642 

643 def _pre_hook(_mod, inputs, kwargs): 

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

645 

646 platform.register_forward_pre_hook( 

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

648 ) 

649 else: 

650 

651 def _pre_hook(_mod, inputs): 

652 return self._prepare_input_fn(inputs, device_mesh) 

653 

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

655 return module 

656 

657 def __repr__(self) -> str: 

658 return ( 

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

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

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

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

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

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

665 ) 

666 

667 

668class PrepareModuleOutput(ParallelStyle): 

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

670 

671 Registers a forward hook that treats each return value like 

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

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

674 

675 Keyword Args: 

676 output_layouts: Current or assumed placement per output tensor. 

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

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

679 """ 

680 

681 def __init__( 

682 self, 

683 *, 

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

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

686 use_local_output: bool = True, 

687 ) -> None: 

688 super().__init__() 

689 self.output_layouts = ( 

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

691 ) 

692 self.desired_output_layouts = ( 

693 (desired_output_layouts,) 

694 if isinstance(desired_output_layouts, Placement) 

695 else desired_output_layouts 

696 ) 

697 self.use_local_output = use_local_output 

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

699 raise AssertionError( 

700 "output_layouts and desired_output_layouts should have same length!" 

701 ) 

702 

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

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

705 prepared_outputs: list = [] 

706 if not isinstance(outputs, tuple): 

707 outputs = (outputs,) 

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

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

710 for out, out_layout, desired_out_layout in zip( 

711 outputs, self.output_layouts, self.desired_output_layouts, 

712 ): 

713 if out_layout is not None: 

714 if isinstance(out, DTensor): 

715 dt_out = out 

716 else: 

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

718 if out_layout != desired_out_layout: 

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

720 prepared_outputs.append( 

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

722 ) 

723 else: 

724 prepared_outputs.append(out) 

725 if len(prepared_outputs) == 1: 

726 return prepared_outputs[0] 

727 return tuple(prepared_outputs) 

728 

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

730 

731 def _hook(_mod, _inputs, outputs): 

732 return self._prepare_out_fn(outputs, device_mesh) 

733 

734 module.register_forward_hook(_hook) 

735 return module 

736 

737 def __repr__(self) -> str: 

738 return ( 

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

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

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

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

743 ) 

744 

745 

746class PrepareModuleInputOutput(ParallelStyle): 

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

748 

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

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

751 """ 

752 

753 def __init__( 

754 self, 

755 *, 

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

757 desired_input_layouts: Optional[ 

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

759 ] = None, 

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

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

762 use_local_input: bool = False, 

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

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

765 use_local_output: bool = True, 

766 ) -> None: 

767 super().__init__() 

768 self.prepare_module_input = PrepareModuleInput( 

769 input_layouts=input_layouts, 

770 desired_input_layouts=desired_input_layouts, 

771 input_kwarg_layouts=input_kwarg_layouts, 

772 desired_input_kwarg_layouts=desired_input_kwarg_layouts, 

773 use_local_output=use_local_input, 

774 ) 

775 self.prepare_module_output = PrepareModuleOutput( 

776 output_layouts=output_layouts, 

777 desired_output_layouts=desired_output_layouts, 

778 use_local_output=use_local_output, 

779 ) 

780 

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

782 self.prepare_module_input.apply(module, device_mesh) 

783 self.prepare_module_output.apply(module, device_mesh) 

784 return module 

785 

786 def __repr__(self) -> str: 

787 p_in = self.prepare_module_input 

788 p_out = self.prepare_module_output 

789 return ( 

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

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

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

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

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

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

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

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

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

799 ) 

800 

801 

802class NoParallel(ParallelStyle): 

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

804 

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

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

807 

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

809 such as: 

810 

811 - MoE Router/Gate modules 

812 - Normalization layers when sequence parallel is disabled 

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

814 

815 Keyword Args: 

816 input_layout (Placement, optional): 

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

818 Defaults to ``Replicate()``. 

819 output_layout (Placement, optional): 

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

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

822 desired_input_layout (Placement, optional): 

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

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

825 redistribution is performed. 

826 use_local_output (bool, optional): 

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

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

829 

830 Example:: 

831 

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

833 >>> model = Transformer(...) 

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

835 >>> parallelize_module(model, tp_mesh, { 

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

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

838 ... }) 

839 """ 

840 

841 def __init__( 

842 self, 

843 *, 

844 input_layout: Optional[Placement] = None, 

845 output_layout: Optional[Placement] = None, 

846 desired_input_layout: Optional[Placement] = None, 

847 use_local_output: bool = True, 

848 ) -> None: 

849 super().__init__() 

850 self.input_layout = input_layout or Replicate() 

851 self.output_layout = output_layout or Replicate() 

852 self.desired_input_layout = desired_input_layout or Replicate() 

853 self.use_local_output = use_local_output 

854 

855 def __repr__(self) -> str: 

856 return ( 

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

858 f"input_layout={self.input_layout}, " 

859 f"output_layout={self.output_layout}, " 

860 f"desired_input_layout={self.desired_input_layout}, " 

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

862 ) 

863 

864 @staticmethod 

865 def _prepare_input_fn( 

866 input_layout: Placement, 

867 desired_input_layout: Placement, 

868 inputs: Tuple[Any, ...], 

869 device_mesh: DeviceMesh, 

870 ) -> Any: 

871 """Annotate and redistribute the first positional input. 

872 

873 If the first input is a plain tensor, wrap it as a :class:`DTensor` with 

874 ``input_layout``. If the resulting :class:`DTensor` placements differ from 

875 ``(desired_input_layout,)``, redistribute to the desired layout. 

876 

877 Args: 

878 input_layout: Layout for :meth:`DTensor.from_local` annotation. 

879 desired_input_layout: Target layout after redistribution. 

880 inputs: Tuple of module forward inputs. 

881 device_mesh: Device mesh for tensor distribution. 

882 

883 Returns: 

884 Tuple of the first input (possibly converted/redistributed) followed by 

885 any remaining positional inputs. 

886 """ 

887 input_tensor = inputs[0] 

888 if not isinstance(input_tensor, DTensor): 

889 input_tensor = DTensor.from_local( 

890 input_tensor, device_mesh, (input_layout,) 

891 ) 

892 

893 if tuple(input_tensor.placements) != (desired_input_layout,): 

894 input_tensor = input_tensor.redistribute( 

895 device_mesh, (desired_input_layout,) 

896 ) 

897 

898 return (input_tensor, *inputs[1:]) 

899 

900 @staticmethod 

901 def _prepare_output_fn( 

902 output_layout: Placement, 

903 use_local_output: bool, 

904 outputs: DTensor, 

905 device_mesh: DeviceMesh, 

906 ) -> Any: 

907 """Redistribute output and optionally convert to local tensor. 

908