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