Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / core / tensor_parallel / api.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"""High-level API to apply parallel styles to modules (aligned with PyTorch ``parallelize_module``).""" 

16from __future__ import annotations 

17 

18import warnings 

19from contextlib import contextmanager 

20from fnmatch import fnmatch 

21from typing import Iterator, Optional, Union 

22 

23from torch import nn 

24 

25from hyper_parallel.core.dtensor.device_mesh import DeviceMesh, _mesh_resources 

26from hyper_parallel.core.tensor_parallel.style import ParallelStyle 

27 

28__all__ = ["parallelize_module"] 

29 

30 

31def _named_children(module: nn.Module): 

32 """Return the immediate named children of a PyTorch module.""" 

33 return module.named_children() 

34 

35 

36@contextmanager 

37def _tensor_parallel_mesh_context(device_mesh: DeviceMesh) -> Iterator[DeviceMesh]: 

38 """Internal: same thread-local stack as ``with device_mesh:`` for ``parallelize_module(..., None)``. 

39 

40 Prefer user code using ``with mesh:``; this exists for tests and library helpers. 

41 """ 

42 with device_mesh: 

43 yield device_mesh 

44 

45 

46def _validate_tp_mesh_dim(device_mesh: DeviceMesh) -> None: 

47 """Require a 1-D mesh, matching PyTorch tensor-parallel constraints.""" 

48 if device_mesh.ndim > 1: 

49 raise ValueError( 

50 f"Tensor Parallel only accepts a 1D DeviceMesh, but found {device_mesh.ndim}D! " 

51 f'If you have a 2-D or N-D device_mesh, consider passing in device_mesh["tp"] ' 

52 f'or another 1-D sub-mesh slice (e.g. mesh["cp"]).' 

53 ) 

54 

55 

56def parallelize_module( # type: ignore[return] 

57 module: nn.Module, 

58 device_mesh: Optional[DeviceMesh] = None, 

59 parallelize_plan: Optional[Union[ParallelStyle, dict[str, ParallelStyle]]] = None, 

60 *, 

61 src_data_rank: Optional[int] = 0, 

62) -> nn.Module: 

63 """Apply parallel styles to *module* or submodules (PyTorch-compatible interface). 

64 

65 Behaviour follows ``torch.distributed.tensor.parallel.parallelize_module``: 

66 

67 - *device_mesh* should normally be passed explicitly. Omitting it (``None``) requires an 

68 active mesh context: ``with mesh:`` (see :meth:`hyper_parallel.core.dtensor.device_mesh.DeviceMesh.__enter__`) 

69 or :func:`_tensor_parallel_mesh_context` for tests/libraries. 

70 - *parallelize_plan* may be a single :class:`ParallelStyle` (applied to *module*) 

71 or a dict mapping submodule paths to styles. Path segments support ``fnmatch`` 

72 patterns (e.g. ``\"layers.*\"``) like PyTorch FQN rules. 

73 - Only **1-D** :class:`DeviceMesh` is accepted; slice a sub-mesh from a multi-dim mesh first. 

74 - *src_data_rank* is stored on the style (``style.src_data_rank``) before ``apply``; styles 

75 that shard parameters from a logical global tensor may use it (see PyTorch TP). 

76 

77 Note: 

78 When ``parallelize_plan`` is a single :class:`ParallelStyle` (not a dict), this 

79 function modifies it in-place by setting ``parallelize_plan.src_data_rank``. 

80 The caller should be aware that the passed object will be mutated. 

81 

82 Args: 

83 module: Root module to parallelize. 

84 device_mesh: Mesh for this TP/CP slice. Use ``None`` only inside ``with mesh:`` (or 

85 :func:`_tensor_parallel_mesh_context`) so ``_mesh_resources.get_current_mesh()`` resolves 

86 (see PyTorch ``distribute_module``). 

87 parallelize_plan: A :class:`ParallelStyle` or dict ``{path: ParallelStyle}``. 

88 src_data_rank: Source rank for global tensor semantics; ``None`` means use local data only 

89 (PyTorch parity). Default ``0``. 

90 

91 Returns: 

92 *module* after in-place parallelization. 

93 """ 

94 if device_mesh is None: 

95 device_mesh = _mesh_resources.get_current_mesh() 

96 _validate_tp_mesh_dim(device_mesh) 

97 

98 if parallelize_plan is None: 

99 warnings.warn( 

100 "No parallelize_plan is provided and auto-parallel is not supported " 

101 "at the moment, so this parallelize_module call will do nothing.", 

102 stacklevel=2, 

103 ) 

104 return module 

105 

106 if isinstance(parallelize_plan, ParallelStyle): 

107 parallelize_plan.src_data_rank = src_data_rank 

108 return parallelize_plan.apply(module, device_mesh) 

109 if isinstance(parallelize_plan, dict): 

110 

111 def _apply_path( 

112 current_module: nn.Module, 

113 atoms: list[str], 

114 style: ParallelStyle, 

115 src_rank: Optional[int], 

116 ) -> bool: 

117 atom = atoms[0] 

118 matched_children = list( 

119 filter( 

120 lambda t, pattern=atom: fnmatch(t[0], pattern), 

121 _named_children(current_module), 

122 ) 

123 ) 

124 applied = False 

125 for _, submodule in matched_children: 

126 if len(atoms) == 1: 

127 parallelize_module( 

128 submodule, 

129 device_mesh, 

130 style, 

131 src_data_rank=src_rank, 

132 ) 

133 applied = True 

134 else: 

135 applied = _apply_path(submodule, atoms[1:], style, src_rank) or applied 

136 return applied 

137 

138 for module_path, parallelize_style in parallelize_plan.items(): 

139 if not isinstance(parallelize_style, ParallelStyle): 

140 raise TypeError( 

141 "Expect ParallelStyle values in parallelize_plan dict, but got " 

142 f"{type(parallelize_style)} for path '{module_path}'." 

143 ) 

144 path_splits = module_path.split(".") 

145 if module_path == "" or any(path == "" for path in path_splits): 

146 raise ValueError( 

147 f"Expect module path to be non-empty dot-separated atoms, but got '{module_path}'." 

148 ) 

149 if not _apply_path(module, path_splits, parallelize_style, src_data_rank): 

150 warnings.warn( 

151 f"parallelize_plan path '{module_path}' has no matches, so this path is skipped.", 

152 stacklevel=2, 

153 ) 

154 return module 

155 raise TypeError( 

156 "Expect Union[ParallelStyle, Dict[str, ParallelStyle]] for" 

157 f" parallelize_plan, {type(parallelize_plan)} found!" 

158 )