| 133 | |
| 134 | |
| 135 | def _conv_split(input_, dim, kernel_size): |
| 136 | cp_world_size = get_context_parallel_world_size() |
| 137 | |
| 138 | # Bypass the function if context parallel is 1 |
| 139 | if cp_world_size == 1: |
| 140 | return input_ |
| 141 | |
| 142 | # print('in _conv_split, cp_rank:', cp_rank, 'input_size:', input_.shape) |
| 143 | |
| 144 | cp_rank = get_context_parallel_rank() |
| 145 | |
| 146 | dim_size = (input_.size()[dim] - kernel_size) // cp_world_size |
| 147 | |
| 148 | if cp_rank == 0: |
| 149 | output = input_.transpose(dim, 0)[: dim_size + kernel_size].transpose(dim, 0) |
| 150 | else: |
| 151 | # output = input_.transpose(dim, 0)[cp_rank * dim_size + 1:(cp_rank + 1) * dim_size + kernel_size].transpose(dim, 0) |
| 152 | output = input_.transpose(dim, 0)[ |
| 153 | cp_rank * dim_size + kernel_size : (cp_rank + 1) * dim_size + kernel_size |
| 154 | ].transpose(dim, 0) |
| 155 | output = output.contiguous() |
| 156 | |
| 157 | # print('out _conv_split, cp_rank:', cp_rank, 'input_size:', output.shape) |
| 158 | |
| 159 | return output |
| 160 | |
| 161 | |
| 162 | def _conv_gather(input_, dim, kernel_size): |