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 

25import torch 

26from torch import nn 

27 

28from hyper_parallel.core.dtensor.device_mesh import DeviceMesh 

29from hyper_parallel.core.dtensor.dtensor import ( 

30 DTensor, 

31 distribute_module, 

32 distribute_tensor, 

33 _distribute_module_iter_params, 

34 _distribute_module_new_parameter, 

35 _distribute_module_param_source, 

36 _distribute_module_set_param, 

37) 

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

39 

40 

41def _cast_fp_tensor(dtype: torch.dtype, tensor: torch.Tensor) -> torch.Tensor: 

42 """Cast a floating-point tensor to ``dtype`` while preserving other tensors.""" 

43 if not torch.is_floating_point(tensor) or tensor.dtype == dtype: 

44 return tensor 

45 return tensor.to(dtype) 

46 

47 

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

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

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

51 return None 

52 return src_data_rank 

53 

54 

55__all__ = [ 

56 "ParallelStyle", 

57 "ColwiseParallel", 

58 "RowwiseParallel", 

59 "SequenceParallel", 

60 "PrepareModuleInput", 

61 "PrepareModuleInputOutput", 

62 "PrepareModuleOutput", 

63 "NoParallel", 

64] 

65 

66 

67class ParallelStyle(ABC): 

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

69 

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

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

72 

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

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

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

76 """ 

77 

78 src_data_rank: Optional[int] = 0 

79 

80 @abstractmethod 

81 def apply(self, module: nn.Module, device_mesh: DeviceMesh) -> nn.Module: 

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

83 

84 Args: 

85 module: The submodule to be parallelised. 

86 device_mesh: The device mesh describing the cluster topology. 

87 

88 Returns: 

89 The (possibly wrapped) module with parallelism applied. 

90 """ 

91 

92 

93class ColwiseParallel(ParallelStyle): 

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

95 

96 Currently supports PyTorch ``Linear`` and ``Embedding`` modules. 

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

98 

99 Keyword Args: 

100 input_layouts (Placement, optional): 

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

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

103 output_layouts (Placement, optional): 

104 Desired DTensor layout of the module output. Defaults to 

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

106 use_local_output (bool, optional): 

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

108 tensor via ``to_local()``. 

109 

110 Returns: 

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

112 

113 Example:: 

114 

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

116 >>> m = Model(...) 

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

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

119 """ 

120 

121 def __init__( 

122 self, 

123 *, 

124 input_layouts: Optional[Placement] = None, 

125 output_layouts: Optional[Placement] = None, 

126 use_local_output: Optional[bool] = None, 

127 ) -> None: 

128 super().__init__() 

129 self._input_layouts_arg = input_layouts 

130 self._output_layouts_arg = output_layouts 

131 self._use_local_output_arg = use_local_output 

132 

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

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

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

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

137 

138 def __repr__(self) -> str: 

139 return ( 

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

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

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

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

144 ) 

145 

146 @staticmethod 

147 def _prepare_input_fn( 

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

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

150 inputs: Any, 

151 device_mesh: DeviceMesh, 

152 ) -> Any: 

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

154 input_tensor = inputs[0] 

155 if not isinstance(input_tensor, DTensor): 

156 input_tensor = DTensor.from_local( 

157 input_tensor, device_mesh, input_layouts, 

158 ) 

159 

160 if input_layouts != desired_input_layouts: 

161 input_tensor = input_tensor.redistribute( 

162 device_mesh, desired_input_layouts, 

163 ) 

164 # Pre-hook contract: return the possibly-redistributed input as a tuple. 

165 return (input_tensor,) 

166 

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

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

169 for key, param in _distribute_module_iter_params(module): 

170 if param is None: 

171 continue 

172 src = _distribute_module_param_source(param) 

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

174 dt = distribute_tensor( 

175 src, 

176 device_mesh, 

177 [Shard(0)], 

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

179 ) 

180 new_param = _distribute_module_new_parameter(key, dt, requires_grad) 

181 _distribute_module_set_param(module, key, new_param) 

182 

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

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

185 for key, param in _distribute_module_iter_params(module): 

186 if param is None: 

187 continue 

188 src = _distribute_module_param_source(param) 

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

190 dt = distribute_tensor( 

191 src, 

192 device_mesh, 

193 [Shard(1)], 

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

195 ) 

196 new_param = _distribute_module_new_parameter(key, dt, requires_grad) 

197 _distribute_module_set_param(module, key, new_param) 

198 

199 @staticmethod 

200 def _prepare_output_fn( 

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

202 use_local_output: bool, 

203 outputs: Any, 

204 device_mesh: DeviceMesh, 

205 ) -> Any: 

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

207 if outputs.placements != output_layouts: 

208 outputs = outputs.redistribute(device_mesh, output_layouts) 

209 if use_local_output: 

210 return outputs.to_local() 

211 return outputs 

212 

213 def apply(self, module: nn.Module, device_mesh: DeviceMesh) -> nn.Module: 

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

215 

216 Args: 

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

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

219 

220 Returns: 

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

222 

223 Raises: 

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

225 """ 

226 if isinstance(module, nn.Linear): 

227 

228 def partition_fn(submodule_path, submodule, device_mesh): 

229 self._partition_linear_fn(submodule, device_mesh) 

230 

231 elif isinstance(module, nn.Embedding): 

232 

233 def partition_fn(submodule_path, submodule, device_mesh): 

234 self._partition_embedding_fn(submodule, device_mesh) 

235 

236 else: 

237 raise NotImplementedError( 

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

239 ) 

240 

241 def input_fn(forward_module, forward_inputs, device_mesh): 

242 return self._prepare_input_fn( 

243 self.input_layouts, 

244 self.desired_input_layouts, 

245 forward_inputs, 

246 device_mesh, 

247 ) 

248 

249 def output_fn(forward_module, forward_outputs, device_mesh): 

250 return self._prepare_output_fn( 

251 self.output_layouts, 

252 self.use_local_output, 

253 forward_outputs, 

254 device_mesh, 

255 ) 

256 

257 return distribute_module( 

258 module, 

259 device_mesh, 

260 partition_fn, 

261 input_fn, 

262 output_fn, 

263 ) 

264 

265 

266class RowwiseParallel(ParallelStyle): 

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

268 

269 Currently supports PyTorch ``Linear`` and ``Embedding`` modules. 

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

271 

272 Keyword Args: 

273 input_layouts (Placement, optional): 

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

275 (sharded on the last dimension). 

276 output_layouts (Placement, optional): 

277 Desired DTensor layout of the module output. Defaults to 

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

279 reduce_dtype (dtype, optional): 

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

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

282 output dtype). 

283 use_local_output (bool, optional): 

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

285 tensor via ``to_local()``. 

286 

287 Returns: 

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

289 

290 Example:: 

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

292 >>> m = Model(...) 

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

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

295 """ 

296 

297 def __init__( 

298 self, 

299 *, 

300 input_layouts: Optional[Placement] = None, 

301 output_layouts: Optional[Placement] = None, 

302 reduce_dtype: Optional[Any] = None, 

303 use_local_output: bool = True, 

304 ) -> None: 

305 super().__init__() 

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

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

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

309 self.reduce_dtype = reduce_dtype 

310 self.use_local_output = use_local_output 

311 

312 def __repr__(self) -> str: 

313 return ( 

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

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

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

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

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

319 ) 

320 

321 @staticmethod 

322 def _prepare_input_fn( 

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

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

325 inputs: Any, 

326 device_mesh: DeviceMesh, 

327 ) -> Any: 

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

329 input_tensor = inputs[0] 

330 if not isinstance(input_tensor, DTensor): 

331 input_tensor = DTensor.from_local( 

332 input_tensor, device_mesh, input_layouts, 

333 ) 

334 

335 if input_layouts != desired_input_layouts: 

336 input_tensor = input_tensor.redistribute( 

337 device_mesh, desired_input_layouts, 

338 ) 

339 # Pre-hook contract: return the possibly-redistributed input as a tuple. 

340 return (input_tensor,) 

341 

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

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

344 for key, param in _distribute_module_iter_params(module): 

345 if param is None: 

346 continue 

347 src = _distribute_module_param_source(param) 

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

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

350 dt = distribute_tensor( 

351 src, 

352 device_mesh, 

353 placement, 

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

355 ) 

356 new_param = _distribute_module_new_parameter(key, dt, requires_grad) 

357 _distribute_module_set_param(module, key, new_param) 

358 

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

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

361 for key, param in _distribute_module_iter_params(module): 

362 if param is None: 

363 continue 

364 src = _distribute_module_param_source(param) 

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

366 dt = distribute_tensor( 

367 src, 

368 device_mesh, 

369 [Shard(0)], 

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

371 ) 

372 new_param = _distribute_module_new_parameter(key, dt, requires_grad) 

373 _distribute_module_set_param(module, key, new_param) 

374 

375 @staticmethod 

376 def _prepare_output_fn( 

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

378 use_local_output: bool, 

379 outputs: Any, 

380 device_mesh: DeviceMesh, 

381 module: Optional[nn.Module] = None, 

382 reduce_dtype: Optional[Any] = None, 

383 ) -> Any: 

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

385 if not isinstance(outputs, DTensor): 

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

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

388 if module is not None and isinstance(module, nn.Embedding): 

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

390 else: 

391 raise TypeError( 

392 "RowwiseParallel expects a DTensor from Linear outputs; " 

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

394 ) 

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

396 local_output = _cast_fp_tensor(reduce_dtype, outputs.to_local()) 

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

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

399 outputs = outputs.redistribute(device_mesh, output_layouts) 

400 if use_local_output: 

401 return outputs.to_local() 

402 return outputs 

403 

404 def apply(self, module: nn.Module, device_mesh: DeviceMesh) -> nn.Module: 

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

406 

407 Args: 

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

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

410 

411 Returns: 

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

413 

414 Raises: 

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

416 """ 

417 if isinstance(module, nn.Linear): 

418 

419 def partition_fn(submodule_path, submodule, device_mesh): 

420 self._partition_linear_fn(submodule, device_mesh) 

421 

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

423 elif isinstance(module, nn.Embedding): 

424 

425 def partition_fn(submodule_path, submodule, device_mesh): 

426 self._partition_embedding_fn(submodule, device_mesh) 

427 

428 self.desired_input_layouts = (Replicate(),) 

429 else: 

430 raise NotImplementedError( 

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

432 ) 

433 

434 def input_fn(forward_module, forward_inputs, device_mesh): 

435 return self._prepare_input_fn( 

436 self.input_layouts, 

437 self.desired_input_layouts, 

438 forward_inputs, 

439 device_mesh, 

440 ) 

441 

442 def output_fn(forward_module, forward_outputs, device_mesh): 

443 return self._prepare_output_fn( 

444 self.output_layouts, 

445 self.use_local_output, 

446 forward_outputs, 

447 device_mesh, 

448 forward_module, 

449 self.reduce_dtype, 

450 ) 

451 

452 return distribute_module( 

453 module, 

454 device_mesh, 

455 partition_fn, 

456 input_fn, 

457 output_fn, 

458 ) 

459 

460 

461class SequenceParallel(ParallelStyle): 

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

463 

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

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

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

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

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

469 

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

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

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

473 

474 Keyword Args: 

475 sequence_dim (int, optional): 

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

477 Default: ``1``. 

478 use_local_output (bool, optional): 

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

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

481 

482 Note: 

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

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

485 after parallelization. 

486 

487 Example:: 

488 

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

490 >>> m = Model(...) 

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

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

493 """ 

494 

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

496 super().__init__() 

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

498 self.use_local_output = use_local_output 

499 

500 def __repr__(self) -> str: 

501 dim = self.sequence_sharding[0].dim 

502 return ( 

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

504 f"sequence_dim={dim}, " 

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

506 ) 

507 

508 @staticmethod 

509 def _prepare_input_fn( 

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

511 mod: nn.Module, 

512 inputs: Any, 

513 device_mesh: DeviceMesh, 

514 ) -> Any: 

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

516 input_tensor = inputs[0] 

517 if isinstance(input_tensor, DTensor): 

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

519 input_tensor = input_tensor.redistribute(device_mesh, sequence_sharding) 

520 elif torch.is_tensor(input_tensor): 

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

522 else: 

523 raise ValueError( 

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

525 ) 

526 return (input_tensor,) 

527 

528 @staticmethod 

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

530 if use_local_output: 

531 return outputs.to_local() 

532 return outputs 

533 

534 def apply(self, module: nn.Module, device_mesh: DeviceMesh) -> nn.Module: 

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

536 

537 Args: 

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

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

540 

541 Returns: 

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

543 converted to replicated DTensors where applicable. 

544 """ 

545 

546 def partition_fn(_submodule_path, _submodule, _mesh): 

547 return None 

548 

549 def input_fn(forward_module, forward_inputs, mesh): 

550 return self._prepare_input_fn( 

551 self.sequence_sharding, 

552 forward_module, 

553 forward_inputs, 

554 mesh, 

555 ) 

556 

557 def output_fn(_forward_module, forward_outputs, _mesh): 

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

559 

560 return distribute_module( 

561 module, 

562 device_mesh, 

563 partition_fn, 

564 input_fn, 

565 output_fn, 

566 ) 

567 

568 

569class PrepareModuleInput(ParallelStyle): 

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

571 

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

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

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

575 “leave this input unchanged”. 

576 

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

578 

579 Keyword Args: 

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

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

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

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

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

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

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

587 :class:`PrepareModuleInput`). 

588 """ 

589 

590 def __init__( 

591 self, 

592 *, 

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

594 desired_input_layouts: Optional[ 

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

596 ] = None, 

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

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

599 use_local_output: bool = False, 

600 ) -> None: 

601 super().__init__() 

602 self.input_layouts = ( 

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

604 ) 

605 self.desired_input_layouts = ( 

606 (desired_input_layouts,) 

607 if isinstance(desired_input_layouts, Placement) 

608 else desired_input_layouts 

609 ) 

610 self.use_local_output = use_local_output 

611 if self.input_layouts is not None: 

612 if self.desired_input_layouts is None: 

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

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

615 raise AssertionError( 

616 "input_layouts and desired_input_layouts should have same length!" 

617 ) 

618 self.with_kwargs = input_kwarg_layouts is not None 

619 self.input_kwarg_layouts = input_kwarg_layouts or {} 

620 self.desired_input_kwarg_layouts = desired_input_kwarg_layouts or {} 

621 if self.with_kwargs: 

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

623 raise AssertionError( 

624 "input_kwarg_layouts and desired_input_kwarg_layouts should have " 

625 "same length!" 

626 ) 

627 

628 def _prepare_input_arg( 

629 self, 

630 input_obj: Any, 

631 mesh: DeviceMesh, 

632 input_layout: Optional[Placement], 

633 desired_layout: Optional[Placement], 

634 ) -> Any: 

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

636 if input_layout is not None: 

637 if isinstance(input_obj, DTensor): 

638 dt_inp = input_obj 

639 else: 

640 if not torch.is_tensor(input_obj): 

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

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

643 

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

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

646 

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

648 return input_obj 

649 

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

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

652 if self.input_layouts is None: 

653 return inputs 

654 if not isinstance(inputs, tuple): 

655 inputs = (inputs,) 

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

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

658 if self.desired_input_layouts is None: 

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

660 prepared_inputs = [ 

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

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

663 ] 

664 return tuple(prepared_inputs) 

665 

666 def _prepare_input_kwarg_fn( 

667 self, 

668 inputs: Any, 

669 kwarg_inputs: Dict[str, Any], 

670 device_mesh: DeviceMesh, 

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

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

673 prepared_arg_inputs = self._prepare_input_fn(inputs, device_mesh) 

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

675 for kwarg_key in kwarg_inputs: 

676 kwarg_val = kwarg_inputs[kwarg_key] 

677 input_layout = self.input_kwarg_layouts.get(kwarg_key) 

678 desired_input_layout = self.desired_input_kwarg_layouts.get(kwarg_key) 

679 prepared_kwarg_inputs[kwarg_key] = self._prepare_input_arg( 

680 kwarg_val, device_mesh, input_layout, desired_input_layout 

681 ) 

682 return (prepared_arg_inputs, prepared_kwarg_inputs) 

683 

684 def apply(self, module: nn.Module, device_mesh: DeviceMesh) -> nn.Module: 

685 if self.with_kwargs: 

686 

687 def _pre_hook(_mod, inputs, kwargs): 

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

689 

690 module.register_forward_pre_hook(_pre_hook, prepend=False, with_kwargs=True) 

691 else: 

692 

693 def _pre_hook(_mod, inputs): 

694 return self._prepare_input_fn(inputs, device_mesh) 

695 

696 module.register_forward_pre_hook(_pre_hook, prepend=False) 

697 return module 

698 

699 def __repr__(self) -> str: 

700 return ( 

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

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

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

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

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

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

707 ) 

708 

709 

710class PrepareModuleOutput(ParallelStyle): 

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

712 

713 Registers a forward hook that treats each return value like 

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

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

716 

717 Keyword Args: 

718 output_layouts: Current or assumed placement per output tensor. 

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

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

721 the reduced output remains in this dtype. 

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

723 """ 

724 

725 def __init__( 

726 self, 

727 *, 

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

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

730 reduce_dtype: Optional[Any] = None, 

731 use_local_output: bool = True, 

732 ) -> None: 

733 super().__init__() 

734 self.output_layouts = ( 

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

736 ) 

737 self.desired_output_layouts = ( 

738 (desired_output_layouts,) 

739 if isinstance(desired_output_layouts, Placement) 

740 else desired_output_layouts 

741 ) 

742 self.reduce_dtype = reduce_dtype 

743 self.use_local_output = use_local_output 

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

745 raise AssertionError( 

746 "output_layouts and desired_output_layouts should have same length!" 

747 ) 

748 

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

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

751 prepared_outputs: list = [] 

752 if not isinstance(outputs, tuple): 

753 outputs = (outputs,) 

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

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

756 for out, out_layout, desired_out_layout in zip( 

757 outputs, self.output_layouts, self.desired_output_layouts, 

758 ): 

759 if out_layout is not None: 

760 if isinstance(out, DTensor): 

761 dt_out = out 

762 else: 

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

764 has_partial_output = any( 

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

766 ) 

767 if ( 

768 self.reduce_dtype is not None 

769 and has_partial_output 

770 and out_layout != desired_out_layout 

771 ): 

772 local_out = _cast_fp_tensor(self.reduce_dtype, dt_out.to_local()) 

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

774 if out_layout != desired_out_layout: 

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

776 prepared_outputs.append( 

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

778 ) 

779 else: 

780 prepared_outputs.append(out) 

781 if len(prepared_outputs) == 1: 

782 return prepared_outputs[0] 

783 return tuple(prepared_outputs) 

784 

785 def apply(self, module: nn.Module, device_mesh: DeviceMesh) -> nn.Module: 

786 

787 def _hook(_mod, _inputs, outputs): 

788 return self._prepare_out_fn(outputs, device_mesh) 

789 

790 module.register_forward_hook(_hook) 

791 return module 

792 

793 def __repr__(self) -> str: 

794 return ( 

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

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

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

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

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

800 ) 

801 

802 

803class PrepareModuleInputOutput(ParallelStyle): 

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

805 

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

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

808 """ 

809 

810 def __init__( 

811 self, 

812 *, 

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

814 desired_input_layouts: Optional[ 

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

816 ] = None, 

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

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

819 use_local_input: bool = False, 

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

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

822 reduce_dtype: Optional[Any] = None, 

823 use_local_output: bool = True, 

824 ) -> None: 

825 super().__init__() 

826 self.prepare_module_input = PrepareModuleInput( 

827 input_layouts=input_layouts, 

828 desired_input_layouts=desired_input_layouts, 

829 input_kwarg_layouts=input_kwarg_layouts, 

830 desired_input_kwarg_layouts=desired_input_kwarg_layouts, 

831 use_local_output=use_local_input, 

832 ) 

833 self.prepare_module_output = PrepareModuleOutput( 

834 output_layouts=output_layouts, 

835 desired_output_layouts=desired_output_layouts, 

836 reduce_dtype=reduce_dtype, 

837 use_local_output=use_local_output, 

838 ) 

839 

840 def apply(self, module: nn.Module, device_mesh: DeviceMesh) -> nn.Module: 

841 self.prepare_module_input.apply(module, device_mesh) 

842 self.prepare_module_output.apply(module, device_mesh) 

843 return module 

844 

845 def __repr__(self) -> str: 

846 p_in = self.prepare_module_input 

847 p_out = self.prepare_module_output 

848 return ( 

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

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

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

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

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

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

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

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

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

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

859 ) 

860 

861 

862class NoParallel(ParallelStyle): 

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

864 

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

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

867 

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

869 such as: 

870 

871 - MoE Router/Gate modules 

872 - Normalization layers when sequence parallel is disabled 

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

874 

875 Keyword Args: 

876 input_layout (Placement, optional): 

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

878 Defaults to ``Replicate()``. 

879 output_layout (Placement, optional): 

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

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

882 desired_input_layout (Placement, optional): 

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

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

885 redistribution is performed. 

886 use_local_output (bool, optional): 

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

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

889 

890 Example:: 

891 

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