| 358 | |
| 359 | |
| 360 | class ContextParallelCausalConv3d(nn.Module): |
| 361 | def __init__(self, chan_in, chan_out, kernel_size: Union[int, Tuple[int, int, int]], stride=1, **kwargs): |
| 362 | super().__init__() |
| 363 | kernel_size = cast_tuple(kernel_size, 3) |
| 364 | |
| 365 | time_kernel_size, height_kernel_size, width_kernel_size = kernel_size |
| 366 | |
| 367 | assert is_odd(height_kernel_size) and is_odd(width_kernel_size) |
| 368 | |
| 369 | time_pad = time_kernel_size - 1 |
| 370 | height_pad = height_kernel_size // 2 |
| 371 | width_pad = width_kernel_size // 2 |
| 372 | |
| 373 | self.height_pad = height_pad |
| 374 | self.width_pad = width_pad |
| 375 | self.time_pad = time_pad |
| 376 | self.time_kernel_size = time_kernel_size |
| 377 | self.temporal_dim = 2 |
| 378 | |
| 379 | stride = (stride, stride, stride) |
| 380 | dilation = (1, 1, 1) |
| 381 | self.conv = Conv3d(chan_in, chan_out, kernel_size, stride=stride, dilation=dilation, **kwargs) |
| 382 | self.cache_padding = None |
| 383 | |
| 384 | def forward(self, input_, clear_cache=True): |
| 385 | # if input_.shape[2] == 1: # handle image |
| 386 | # # first frame padding |
| 387 | # input_parallel = torch.cat([input_] * self.time_kernel_size, dim=2) |
| 388 | # else: |
| 389 | # input_parallel = conv_pass_from_last_rank(input_, self.temporal_dim, self.time_kernel_size) |
| 390 | |
| 391 | # padding_2d = (self.width_pad, self.width_pad, self.height_pad, self.height_pad) |
| 392 | # input_parallel = F.pad(input_parallel, padding_2d, mode = 'constant', value = 0) |
| 393 | |
| 394 | # output_parallel = self.conv(input_parallel) |
| 395 | # output = output_parallel |
| 396 | # return output |
| 397 | |
| 398 | input_parallel = fake_cp_pass_from_previous_rank( |
| 399 | input_, self.temporal_dim, self.time_kernel_size, self.cache_padding |
| 400 | ) |
| 401 | |
| 402 | del self.cache_padding |
| 403 | self.cache_padding = None |
| 404 | if not clear_cache: |
| 405 | cp_rank, cp_world_size = get_context_parallel_rank(), get_context_parallel_world_size() |
| 406 | global_rank = torch.distributed.get_rank() |
| 407 | if cp_world_size == 1: |
| 408 | self.cache_padding = ( |
| 409 | input_parallel[:, :, -self.time_kernel_size + 1 :].contiguous().detach().clone().cpu() |
| 410 | ) |
| 411 | else: |
| 412 | if cp_rank == cp_world_size - 1: |
| 413 | torch.distributed.isend( |
| 414 | input_parallel[:, :, -self.time_kernel_size + 1 :].contiguous(), |
| 415 | global_rank + 1 - cp_world_size, |
| 416 | group=get_context_parallel_group(), |
| 417 | ) |