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 if not isinstance(module, nn.Linear): 

38 raise NotImplementedError( 

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

40 ) 

41 module = MC2Linear.from_linear(module) 

42 module.configure_mc2( 

43 mode, 

44 device_mesh.get_group(), 

45 device_mesh.size(), 

46 sequence_dim=sequence_dim, 

47 ) 

48 return module 

49 

50 

51class MC2ColwiseParallel(ColwiseParallel): 

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

53 

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

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

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

57 """ 

58 

59 def __init__( 

60 self, 

61 *, 

62 input_layouts: Optional[Placement] = None, 

63 output_layouts: Optional[Placement] = None, 

64 use_local_output: Optional[bool] = None, 

65 ) -> None: 

66 super().__init__( 

67 input_layouts=input_layouts, 

68 output_layouts=output_layouts, 

69 use_local_output=use_local_output, 

70 ) 

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

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

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

74 self.desired_input_layouts = self.input_layouts 

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

76 

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

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

79 module = _replace_with_mc2_linear( 

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

81 ) 

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

83 

84 

85class MC2RowwiseParallel(RowwiseParallel): 

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

87 

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

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

90 """ 

91 

92 def __init__( 

93 self, 

94 *, 

95 input_layouts: Optional[Placement] = None, 

96 output_layouts: Optional[Placement] = None, 

97 reduce_dtype=None, 

98 use_local_output: bool = True, 

99 ) -> None: 

100 super().__init__( 

101 input_layouts=input_layouts, 

102 output_layouts=output_layouts, 

103 reduce_dtype=reduce_dtype, 

104 use_local_output=use_local_output, 

105 ) 

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

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

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

109 

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

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

112 module = _replace_with_mc2_linear( 

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

114 ) 

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