Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / platform / mindspore / pipeline_parallel / _utils.py: 17%

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"""pipeline parallel utils""" 

16import io 

17import pickle 

18 

19from mindspore import nn, Tensor, mint, ops 

20from mindspore.common import dtype as mstype 

21from mindspore.communication import GlobalComm 

22from mindspore.mint.distributed.distributed import _object_to_tensor, send, recv 

23 

24import hyper_parallel 

25from hyper_parallel.core.shard.custom_shard import custom_shard 

26 

27 

28class _MicroBatch(nn.Cell): 

29 """ 

30 Split inputs into micro_batch in pipeline parallel. 

31 

32 Args: 

33 micro_batch_num (int): The number of micro-batch. 

34 args_batch_dim (list, optional): Specify the batch dim of the args. 

35 Default ``None``. 

36 kwargs_batch_dim(dict, optional): Specify the batch dim of the kwargs. 

37 Default ``None``. 

38 Inputs: 

39 - **args** (list) - Input args. 

40 - **kwargs** (dict) - Input kwargs. 

41 

42 Outputs: 

43 - **args_after_split** (list) - Input args after split into micro_batches. 

44 - **kwargs_after_split** (list) - Input kwargs after split into micro_batches. 

45 """ 

46 

47 def __init__(self, micro_batch_num, args_batch_dim=None, kwargs_batch_dim=None): 

48 super().__init__() 

49 self.micro_batch_num = micro_batch_num 

50 self.args_batch_dim = args_batch_dim 

51 self.kwargs_batch_dim = kwargs_batch_dim 

52 

53 def construct(self, args, kwargs): 

54 """Construct of _MicroBatch""" 

55 args_after_split = [] 

56 kwargs_after_split = [] 

57 for micro_idx in range(self.micro_batch_num): 

58 micro_args = [] 

59 micro_kwargs = {} 

60 for arg_idx, cur_arg in enumerate(args): 

61 cur_arg_batch_dim = 0 

62 if self.args_batch_dim and self.args_batch_dim[arg_idx] is not None: 

63 cur_arg_batch_dim = self.args_batch_dim[arg_idx].batch_dim 

64 if isinstance(cur_arg, hyper_parallel.DTensor): 

65 micro_arg = self.split_inputs_with_custom_shard(cur_arg, cur_arg_batch_dim, micro_idx) 

66 else: 

67 micro_arg = self.split_inputs(cur_arg, cur_arg_batch_dim, micro_idx) 

68 micro_args.append(micro_arg) 

69 args_after_split.append(micro_args) 

70 

71 for key, cur_kwarg in kwargs.items(): 

72 cur_kwarg_batch_dim = 0 

73 if self.kwargs_batch_dim is not None: 

74 cur_kwarg_batch_dim = self.kwargs_batch_dim[key].batch_dim 

75 if isinstance(cur_kwarg, hyper_parallel.DTensor): 

76 micro_kwarg = self.split_inputs_with_custom_shard(cur_kwarg, cur_kwarg_batch_dim, micro_idx) 

77 else: 

78 micro_kwarg = self.split_inputs(cur_kwarg, cur_kwarg_batch_dim, micro_idx) 

79 micro_kwargs[key] = micro_kwarg 

80 kwargs_after_split.append(micro_kwargs) 

81 return args_after_split, kwargs_after_split 

82 

83 def split_inputs_with_custom_shard(self, input_tensor, cur_arg_batch_dim, micro_idx): 

84 """Split a DTensor input along the batch dimension using its custom shard layout.""" 

85 if not isinstance(input_tensor, hyper_parallel.DTensor): 

86 raise TypeError(f"Input type {type(input_tensor)} is not DTensor.") 

87 input_layout = input_tensor.layout 

88 func_wrap = custom_shard(self.split_inputs, 

89 device_mesh=input_layout.mesh, 

90 out_placements=(input_layout.placements,), 

91 in_placements=(input_layout.placements, None, None) 

92 ) 

93 return func_wrap(input_tensor, cur_arg_batch_dim, micro_idx) 

94 

95 def split_inputs(self, input_tensor, cur_arg_batch_dim, micro_idx): 

96 """ 

97 Split the input along the specified batch_dim and micro_idx 

98 """ 

99 if cur_arg_batch_dim == -1: 

100 return input_tensor 

101 batch_dim_shape = input_tensor.shape[cur_arg_batch_dim] 

102 if batch_dim_shape % self.micro_batch_num != 0: 

103 raise ValueError(f"Batch dimension size {batch_dim_shape} is not divisible by \ 

104 micro_batch_num {self.micro_batch_num}") 

105 micro_batch_begin = (batch_dim_shape // self.micro_batch_num) * micro_idx 

106 micro_batch_end = (batch_dim_shape // self.micro_batch_num) * (micro_idx + 1) 

107 strided_slice_begin = [0] * input_tensor.ndim 

108 strided_slice_strides = [1] * input_tensor.ndim 

109 strided_slice_end = list(input_tensor.shape) 

110 strided_slice_begin[cur_arg_batch_dim] = micro_batch_begin 

111 strided_slice_end[cur_arg_batch_dim] = micro_batch_end 

112 micro_input = ops.strided_slice(input_tensor, strided_slice_begin, strided_slice_end, strided_slice_strides) 

113 return micro_input 

114 

115 

116def send_object_list(obj, dst=0, group=None): 

117 """ 

118 Send the input Python object to dst rank. 

119 

120 Args: 

121 obj (Any): The input tensor to be send. 

122 dst (int, optional): Specifies the global rank that send the Python object to. 

123 Default: ``0``. 

124 group (str, optional): Communication group. Default: ``None``. 

125 """ 

126 if group is None: 

127 group = GlobalComm.WORLD_COMM_GROUP 

128 if not isinstance(group, str): 

129 raise TypeError(f"For 'send_object', the argument 'group' must be type of string, \ 

130 but got 'group' type : {type(group)}.") 

131 if not isinstance(dst, int): 

132 raise TypeError("For send_object, the dst must be int.") 

133 obj_tensor, tensor_size = _object_to_tensor(obj) 

134 obj_size = Tensor([tensor_size], dtype=mstype.int32) 

135 send(obj_size, dst, group) 

136 send(obj_tensor, dst, group) 

137 

138 

139def recv_object_list(recv_obj, src=0, group=None): 

140 """ 

141 receive Python object from src rank. 

142 

143 Args: 

144 recv_obj (list): list to recv python objects. 

145 src (int, optional): Specifies the global rank that receive the Python object. 

146 Default: ``0`` . 

147 group (str, optional): Communication group. Default: ``None``. 

148 """ 

149 if group is None: 

150 group = GlobalComm.WORLD_COMM_GROUP 

151 if not isinstance(group, str): 

152 raise TypeError(f"For 'recv_object', the argument 'group' must be type of string, \ 

153 but got 'group' type : {type(group)}.") 

154 if not isinstance(src, int): 

155 raise TypeError("For recv_object, the src must be int.") 

156 obj_size = mint.zeros((1,), dtype=mstype.int32) 

157 recv(obj_size, src, group) 

158 # MindSpore PyNative ``recv`` only does a comm-stream wait; bridge to host 

159 # so the subsequent ``.item()`` reads the freshly-received value instead 

160 # of the original buffer. 

161 size_val = int(obj_size.item()) 

162 obj_tensor = mint.zeros((size_val,), dtype=mstype.int8) 

163 recv(obj_tensor, src, group) 

164 buf = obj_tensor.asnumpy().tobytes()[:size_val] 

165 recv_obj.clear() 

166 recv_obj.append(pickle.Unpickler(io.BytesIO(buf)).load()[0])