| 293 | |
| 294 | |
| 295 | class ContextParallelCausalConv3d(nn.Module): |
| 296 | def __init__(self, chan_in, chan_out, kernel_size: Union[int, Tuple[int, int, int]], stride=1, **kwargs): |
| 297 | super().__init__() |
| 298 | kernel_size = cast_tuple(kernel_size, 3) |
| 299 | |
| 300 | time_kernel_size, height_kernel_size, width_kernel_size = kernel_size |
| 301 | |
| 302 | assert is_odd(height_kernel_size) and is_odd(width_kernel_size) |
| 303 | |
| 304 | time_pad = time_kernel_size - 1 |
| 305 | height_pad = height_kernel_size // 2 |
| 306 | width_pad = width_kernel_size // 2 |
| 307 | |
| 308 | self.height_pad = height_pad |
| 309 | self.width_pad = width_pad |
| 310 | self.time_pad = time_pad |
| 311 | self.time_kernel_size = time_kernel_size |
| 312 | self.temporal_dim = 2 |
| 313 | |
| 314 | stride = (stride, stride, stride) |
| 315 | dilation = (1, 1, 1) |
| 316 | self.conv = Conv3d(chan_in, chan_out, kernel_size, stride=stride, dilation=dilation, **kwargs) |
| 317 | |
| 318 | def forward(self, input_): |
| 319 | # temporal padding inside |
| 320 | if _USE_CP: |
| 321 | input_parallel = conv_pass_from_last_rank(input_, self.temporal_dim, self.time_kernel_size) |
| 322 | else: |
| 323 | input_ = input_.transpose(0, self.temporal_dim) |
| 324 | input_parallel = torch.cat([input_[:1]] * (self.time_kernel_size - 1) + [input_], dim=0) |
| 325 | input_parallel = input_parallel.transpose(0, self.temporal_dim) |
| 326 | padding_2d = (self.width_pad, self.width_pad, self.height_pad, self.height_pad) |
| 327 | input_parallel = F.pad(input_parallel, padding_2d, mode="constant", value=0) |
| 328 | output_parallel = self.conv(input_parallel) |
| 329 | output = output_parallel |
| 330 | return output |
| 331 | |
| 332 | |
| 333 | class ContextParallelGroupNorm(torch.nn.GroupNorm): |