Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / core / tensor_parallel / mc2_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 MC2 fused linear layers (PyTorch / Ascend).""" 

16from __future__ import annotations 

17 

18from typing import Optional 

19 

20from torch import nn 

21 

22from hyper_parallel.core.dtensor.device_mesh import DeviceMesh 

23from hyper_parallel.core.dtensor.placement_types import Placement, Shard 

24from hyper_parallel.core.tensor_parallel.mc2 import MC2Linear 

25from hyper_parallel.core.tensor_parallel.style import ColwiseParallel, RowwiseParallel 

26 

27__all__ = ["MC2ColwiseParallel", "MC2RowwiseParallel"] 

28 

29 

30def _replace_with_mc2_linear( 

31 module: nn.Module, 

32 mode: str, 

33 device_mesh: DeviceMesh, 

34 sequence_dim: int, 

35) -> MC2Linear: 

36 """Replace ``nn.Linear`` with configured ``MC2Linear`` in place.""" 

37 # Torch-only check: get_platform() is a process singleton and mixed UT may 

38 # already have cached MindSpore, whose is_linear_module rejects nn.Linear. 

39 if not isinstance(module, nn.Linear): 

40 raise NotImplementedError( 

41 f"MC2 parallel style only supports Linear modules, but got {type(module).__name__}." 

42 ) 

43 module = MC2Linear.from_linear(module) 

44 module.configure_mc2( 

45 mode, 

46 device_mesh.get_group(), 

47 device_mesh.size(), 

48 sequence_dim=sequence_dim, 

49 ) 

50 return module 

51 

52 

53class MC2ColwiseParallel(ColwiseParallel): 

54 """Column parallelism using fused all-gather and matmul (MC2). 

55 

56 Requires a sequence-sharded input layout so AllGather can be folded into the 

57 matmul. Unlike :class:`ColwiseParallel`, this style does **not** redistribute 

58 the input to ``Replicate()`` before the Linear. 

59 """ 

60 

61 def __init__( 

62 self, 

63 *, 

64 input_layouts: Optional[Placement] = None, 

65 output_layouts: Optional[Placement] = None, 

66 use_local_output: Optional[bool] = None, 

67 ) -> None: 

68 super().__init__( 

69 input_layouts=input_layouts, 

70 output_layouts=output_layouts, 

71 use_local_output=use_local_output, 

72 ) 

73 if not isinstance(self.input_layouts[0], Shard): 

74 raise ValueError("MC2ColwiseParallel requires a sharded input layout.") 

75 # Keep sequence sharding; fused kernel performs the all-gather. 

76 self.desired_input_layouts = self.input_layouts 

77 self._sequence_dim = self.input_layouts[0].dim 

78 

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

80 """Replace Linear with MC2Linear then apply column-wise sharding hooks.""" 

81 module = _replace_with_mc2_linear( 

82 module, "all_gather", device_mesh, self._sequence_dim 

83 ) 

84 return super().apply(module, device_mesh) 

85 

86 

87class MC2RowwiseParallel(RowwiseParallel): 

88 """Row parallelism using fused matmul and reduce-scatter (MC2). 

89 

90 Requires a sequence-sharded output layout so ReduceScatter replaces AllReduce 

91 and restores sequence parallelism after the row-parallel Linear. 

92 """ 

93 

94 def __init__( 

95 self, 

96 *, 

97 input_layouts: Optional[Placement] = None, 

98 output_layouts: Optional[Placement] = None, 

99 reduce_dtype=None, 

100 use_local_output: bool = True, 

101 ) -> None: 

102 super().__init__( 

103 input_layouts=input_layouts, 

104 output_layouts=output_layouts, 

105 reduce_dtype=reduce_dtype, 

106 use_local_output=use_local_output, 

107 ) 

108 if not isinstance(self.output_layouts[0], Shard): 

109 raise ValueError("MC2RowwiseParallel requires a sharded output layout.") 

110 self._sequence_dim = self.output_layouts[0].dim 

111 

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

113 """Replace Linear with MC2Linear then apply row-wise sharding hooks.""" 

114 module = _replace_with_mc2_linear( 

115 module, "reduce_scatter", device_mesh, self._sequence_dim 

116 ) 

117 return super().apply(module, device_mesh)