| 431 | |
| 432 | |
| 433 | class ContextParallelGroupNorm(torch.nn.GroupNorm): |
| 434 | def forward(self, input_): |
| 435 | gather_flag = input_.shape[2] > 1 |
| 436 | if gather_flag: |
| 437 | input_ = conv_gather_from_context_parallel_region(input_, dim=2, kernel_size=1) |
| 438 | output = super().forward(input_) |
| 439 | if gather_flag: |
| 440 | output = conv_scatter_to_context_parallel_region(output, dim=2, kernel_size=1) |
| 441 | return output |
| 442 | |
| 443 | |
| 444 | def Normalize(in_channels, gather=False, **kwargs): # same for 3D and 2D |