(self, input_, clear_cache=True)
| 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 | ) |
| 418 | if cp_rank == 0: |
| 419 | recv_buffer = torch.empty_like(input_parallel[:, :, -self.time_kernel_size + 1 :]).contiguous() |
| 420 | torch.distributed.recv( |
| 421 | recv_buffer, global_rank - 1 + cp_world_size, group=get_context_parallel_group() |
| 422 | ) |
| 423 | self.cache_padding = recv_buffer.contiguous().detach().clone().cpu() |
| 424 | |
| 425 | padding_2d = (self.width_pad, self.width_pad, self.height_pad, self.height_pad) |
| 426 | input_parallel = F.pad(input_parallel, padding_2d, mode="constant", value=0) |
| 427 | |
| 428 | output_parallel = self.conv(input_parallel) |
| 429 | output = output_parallel |
| 430 | return output |
| 431 | |
| 432 | |
| 433 | class ContextParallelGroupNorm(torch.nn.GroupNorm): |
nothing calls this directly
no test coverage detected