| 331 | |
| 332 | |
| 333 | class ContextParallelGroupNorm(torch.nn.GroupNorm): |
| 334 | def forward(self, input_): |
| 335 | if _USE_CP: |
| 336 | input_ = conv_gather_from_context_parallel_region(input_, dim=2, kernel_size=1) |
| 337 | output = super().forward(input_) |
| 338 | if _USE_CP: |
| 339 | output = conv_scatter_to_context_parallel_region(output, dim=2, kernel_size=1) |
| 340 | return output |
| 341 | |
| 342 | |
| 343 | def Normalize(in_channels, gather=False, **kwargs): # same for 3D and 2D |