Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / core / pipeline_parallel / _microbatch.py: 91%

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""" 

16from typing import Optional 

17 

18import torch 

19from torch import nn 

20 

21import hyper_parallel 

22 

23 

24class _MicroBatch(nn.Module): 

25 """ 

26 Split inputs into micro_batch in pipeline parallel. 

27 

28 Args: 

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

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

31 Default ``None``. 

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

33 Default ``None``. 

34 Inputs: 

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

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

37 

38 Outputs: 

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

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

41 """ 

42 

43 def __init__(self, micro_batch_num: int, args_batch_dim: Optional[tuple] = None, 

44 kwargs_batch_dim: Optional[dict] = None) -> None: 

45 """Store the number of micro-batches and their batch dimensions.""" 

46 super().__init__() 

47 self.micro_batch_num = micro_batch_num 

48 self.args_batch_dim = args_batch_dim 

49 self.kwargs_batch_dim = kwargs_batch_dim 

50 

51 def _split_args_for_micro_batch(self, args: tuple, micro_idx: int) -> list: 

52 """Split positional arguments for one micro-batch.""" 

53 micro_args = [] 

54 for arg_idx, cur_arg in enumerate(args): 

55 cur_arg_batch_dim = 0 

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

57 cur_arg_batch_dim = self.args_batch_dim[arg_idx].batch_dim 

58 if isinstance(cur_arg, hyper_parallel.DTensor): 

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

60 else: 

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

62 micro_args.append(micro_arg) 

63 return micro_args 

64 

65 def _split_kwargs_for_micro_batch(self, kwargs: dict, micro_idx: int) -> dict: 

66 """Split keyword arguments for one micro-batch.""" 

67 micro_kwargs = {} 

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

69 cur_kwarg_batch_dim = 0 

70 if self.kwargs_batch_dim is not None: 

71 cur_kwarg_batch_dim = self.kwargs_batch_dim[key].batch_dim 

72 if isinstance(cur_kwarg, hyper_parallel.DTensor): 

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

74 else: 

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

76 micro_kwargs[key] = micro_kwarg 

77 return micro_kwargs 

78 

79 def forward(self, args: tuple, kwargs: dict) -> tuple[list, list]: 

80 """forward of _MicroBatch""" 

81 args_after_split = [] 

82 kwargs_after_split = [] 

83 for micro_idx in range(self.micro_batch_num): 

84 args_after_split.append(self._split_args_for_micro_batch(args, micro_idx)) 

85 kwargs_after_split.append(self._split_kwargs_for_micro_batch(kwargs, micro_idx)) 

86 return args_after_split, kwargs_after_split 

87 

88 def split_inputs_with_custom_shard( 

89 self, input_tensor: hyper_parallel.DTensor, 

90 cur_arg_batch_dim: int, micro_idx: int) -> hyper_parallel.DTensor: 

91 """Split a DTensor input along the batch dimension while preserving its distributed layout.""" 

92 input_layout = input_tensor.layout 

93 func_wrap = hyper_parallel.custom_shard(self.split_inputs, 

94 device_mesh=input_layout.mesh, 

95 out_placements=(input_layout.placements,), 

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

97 ) 

98 return func_wrap(input_tensor, cur_arg_batch_dim, micro_idx) 

99 

100 def split_inputs(self, input_tensor: torch.Tensor, cur_arg_batch_dim: int, micro_idx: int) -> torch.Tensor: 

101 """ 

102 Split the input along the specified batch_dim and micro_idx 

103 """ 

104 if cur_arg_batch_dim == -1: 

105 return input_tensor 

106 batch_dim_shape = input_tensor.shape[cur_arg_batch_dim] 

107 if batch_dim_shape % self.micro_batch_num != 0: 

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

109 micro_batch_num {self.micro_batch_num}") 

110 micro_batch_size = batch_dim_shape // self.micro_batch_num 

111 

112 # Calculate start and end idx 

113 start = micro_batch_size * micro_idx 

114 end = micro_batch_size * (micro_idx + 1) 

115 

116 # Create slicing tuple 

117 slices = [slice(None)] * input_tensor.ndim 

118 slices[cur_arg_batch_dim] = slice(start, end) 

119 return input_tensor[tuple(slices)]