File size: 4,487 Bytes
5d9b308 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 | # SPDX-License-Identifier: Apache-2.0
# Spatial-parallel 3D convolution for the MiniMax H3 visual VAE.
import torch
import torch.nn as nn
import torch.nn.functional as F
from .parallel import get_parallel_state, exchange_borders
class BaseConv3d(nn.Conv3d):
def __init__(
self,
in_channels,
out_channels,
kernel_size,
stride=1,
padding=0,
bias=True,
padding_mode="zeros",
padding_mode_t=None,
causal=True,
):
super().__init__(
in_channels,
out_channels,
kernel_size=kernel_size,
stride=stride,
padding=padding,
bias=bias,
padding_mode=padding_mode,
)
padding_mode = "constant" if padding_mode == "zeros" else padding_mode
padding_mode_t = "constant" if padding_mode_t == "zeros" else padding_mode_t
self.pad_mode = padding_mode
self.pad_mode_t = padding_mode_t or ("constant" if causal else "replicate")
self.causal = causal
def _apply_temporal_padding(self, x):
B, C, D, H, W = x.shape
if D > 1:
pad_size = (
0,
0,
0,
0,
self.padding[0] * 2 if self.causal else self.padding[0],
0 if self.causal else self.padding[0],
)
return F.pad(x, pad_size, mode=self.pad_mode_t)
else:
if self.pad_mode_t == "constant":
assert self.causal, "Zeros padding is only supported for causal mode"
zeros = torch.zeros_like(x[:, :, :1, :, :]).expand(
-1, -1, self.kernel_size[0] - 1, -1, -1
)
return torch.cat([zeros, x], dim=2)
else:
return x.expand(-1, -1, self.kernel_size[0], -1, -1)
def _apply_padding(self, x):
if sum(self.padding) == 0:
return x
x = F.pad(
x,
(self.padding[2], self.padding[2], self.padding[1], self.padding[1], 0, 0),
mode=self.pad_mode,
)
x = self._apply_temporal_padding(x)
return x
def forward(self, x):
if sum(self.padding) == 0:
return super().forward(x)
x = self._apply_padding(x)
return F.conv3d(
x,
self.weight,
self.bias,
stride=self.stride,
padding=0,
dilation=self.dilation,
)
class SpatialParallelConv3d(BaseConv3d):
def __init__(
self,
in_channels,
out_channels,
kernel_size,
stride=1,
padding=0,
bias=True,
padding_mode="zeros",
padding_mode_t=None,
causal=True,
):
super().__init__(
in_channels,
out_channels,
kernel_size=kernel_size,
stride=stride,
padding=padding,
bias=bias,
padding_mode=padding_mode,
padding_mode_t=padding_mode_t,
causal=causal,
)
self.spatial_parallel = False
self.chunk_dim = -1
def _exchange_borders(self, x, sp_rank, sp_size):
if self.chunk_dim == -1:
pad = self.padding[2]
elif self.chunk_dim == -2:
pad = self.padding[1]
else:
raise ValueError(f"Invalid chunk dimension: {self.chunk_dim}")
if pad == 0:
return x
local_process_group = get_parallel_state()["sp_process_group"]
return exchange_borders(
x,
pad,
self.pad_mode,
sp_rank,
sp_size,
local_process_group,
dim=self.chunk_dim,
)
def _apply_padding(self, x):
if not self.spatial_parallel:
return super()._apply_padding(x)
state = get_parallel_state()
x = self._exchange_borders(x, state["sp_rank"], state["sp_size"])
if self.chunk_dim == -1:
x = F.pad(
x, (0, 0, self.padding[1], self.padding[1], 0, 0), mode=self.pad_mode
)
elif self.chunk_dim == -2:
x = F.pad(
x, (self.padding[2], self.padding[2], 0, 0, 0, 0), mode=self.pad_mode
)
else:
raise ValueError(f"Invalid chunk dimension: {self.chunk_dim}")
x = self._apply_temporal_padding(x)
return x
|