Need to support conv3d conversion

[seockho-kim]

What

Some VL models(ex, qwen-vl) use conv3d in their vision encoder,
so tico needs a pass to convert conv3d to other operations, like conv2d

[seockho-kim]

class Conv3dToConv2d(torch.nn.Module):
    def __init__(self, conv3d_layer):
        super().__init__()
        self.kT, self.kH, self.kW = conv3d_layer.kernel_size
        self.stride = conv3d_layer.stride
        self.padding = conv3d_layer.padding
        self.dilation = conv3d_layer.dilation
        self.bias = conv3d_layer.bias
        
        self.conv2d_layers = torch.nn.ModuleList()
        for t in range(self.kT):
            conv2d = torch.nn.Conv2d(
                in_channels=conv3d_layer.in_channels,
                out_channels=conv3d_layer.out_channels,
                kernel_size=(self.kH, self.kW),
                stride=self.stride[1:],
                padding=self.padding[1:],
                dilation=self.dilation[1:],
                bias=False
            )
            conv2d.weight.data.copy_(conv3d_layer.weight.data[:, :, t])
            self.conv2d_layers.append(conv2d)
    
    def forward(self, x):
        """
        x: [N, C, T, H, W]
        """
        N, C, T, H, W = x.shape
        
        T_out = (T + 2 * self.padding[0] - self.dilation[0] * (self.kT - 1) - 1) // self.stride[0] + 1
        
        if self.padding[0] > 0:
            zero_padding = torch.zeros(N, C, self.padding[0], H, W, device=x.device, dtype=x.dtype)
            # (N, C, padding, H, W)
            
            x_padded = torch.cat([zero_padding, x, zero_padding], dim=2)  # (N, C, T+2*padding, H, W)
            T_padded = T + 2 * self.padding[0]
        else:
            x_padded = x
            T_padded = T
        
        outputs = []
        for t_out in range(T_out):
            t_in = t_out * self.stride[0]
            
            y_t = 0
            for i, conv2d in enumerate(self.conv2d_layers):
                t_idx = t_in + i * self.dilation[0]
                
                if t_idx < T_padded:
                    x_2d = x_padded[:, :, t_idx]  # (N, C, H, W)
                    y_t = y_t + conv2d(x_2d)
            
            if self.bias is not None:
                y_t = y_t + self.bias.view(1, -1, 1, 1)
            
            outputs.append(y_t)
        
        return torch.stack(outputs, dim=2)  # [N, C_out, T_out, H_out, W_out]

[mhs4670go]

Seems that the current implementation correctly decomposes Conv3d into a sum of kT independent Conv2d ops along the temporal axis. This is mathematically correct in the general case and works for arbitrary stride/dilation/padding.

However, for Qwen2.5-VL’s patch embedding Conv3d, I think we can use a simpler and strictly equivalent transformation that is both smaller in graph size and more efficient, under specific (but common) conditions.

We can concatenate the kT temporal frames into the channel dimension and reshape the Conv3d weight accordingly. Then, we can run a single Conv2d.

This single-Conv2d replacement is exact only if:

• kernel_T == stride_T
• padding_T == 0
• dilation_T == 1
• no temporal overlap

These conditions hold for Qwen2.5-VL patch embedding, but not for general Conv3d layers. So, we can keep the current your implementation as the general Conv3d fallback.

I'll implement this logic soon. Please feel free to give your opinions.

import torch
import torch.nn.functional as F

def demo(device="cpu", dtype=torch.float32):
    torch.manual_seed(0)

    N, C_in, C_out = 1, 3, 1280
    T, H, W = 6, 28, 28
    kT, kH, kW = 2, 14, 14
    sT, sH, sW = 2, 14, 14

    x = torch.randn(N, C_in, T, H, W, device=device, dtype=dtype)
    w3 = torch.randn(C_out, C_in, kT, kH, kW, device=device, dtype=dtype)

    y3 = F.conv3d(x, w3, bias=None, stride=(sT, sH, sW), padding=0, dilation=1, groups=1)

    T_out = (T - kT) // sT + 1
    x_grp = x.reshape(N, C_in, T_out, kT, H, W)                 # (N,C,T_out,2,H,W)
    x2 = x_grp.permute(0, 2, 1, 3, 4, 5).contiguous()           # (N,T_out,C,2,H,W)
    x2 = x2.view(N * T_out, C_in * kT, H, W)                    # (N*T_out, C*2, H, W)

    w2 = w3.contiguous().view(C_out, C_in * kT, kH, kW)         # (C_out, C*2, 14, 14)

    y2 = F.conv2d(x2, w2, bias=None, stride=(sH, sW), padding=0, dilation=1, groups=1)

    H_out = (H - kH) // sH + 1
    W_out = (W - kW) // sW + 1
    y2 = y2.view(N, T_out, C_out, H_out, W_out).permute(0, 2, 1, 3, 4).contiguous()

    max_abs = (y3 - y2).abs().max().item()
    print("max |y3-y2| =", max_abs)

demo()

[mhs4670go]

4-rank only version

def demo_option_a_rank4_only(
    T=6, H=28, W=28,
    device="cpu",
    dtype=torch.float32,
):
    torch.manual_seed(0)

    N = 1
    C_in = 3
    C_out = 1280
    kT, kH, kW = 2, 14, 14
    sT, sH, sW = 2, 14, 14

    x = torch.randn(N, C_in, T, H, W, device=device, dtype=dtype)
    w3 = torch.randn(C_out, C_in, kT, kH, kW, device=device, dtype=dtype)

    # Reference conv3d (5D output) - for verification only
    y3 = F.conv3d(x, w3, bias=None, stride=(sT, sH, sW), padding=0, dilation=1, groups=1)

    T_out = (T - kT) // sT + 1
    H_out = (H - kH) // sH + 1
    W_out = (W - kW) // sW + 1

    # Weight folding matches interleaved order: [c0_t0,c0_t1,c1_t0,c1_t1,...]
    w2 = w3.contiguous().view(C_out, C_in * kT, kH, kW)

    def pack_two_frames_interleaved_rank4(x0, x1):
        # x0,x1: (1,3,H,W) -> (1,6,H,W)
        chunks = []
        for c in range(C_in):
            chunks.append(x0[:, c:c+1, :, :])
            chunks.append(x1[:, c:c+1, :, :])
        return torch.cat(chunks, dim=1)

    # HW-friendly rank<=4 path
    outs = []
    for t_out in range(T_out):
        t0 = t_out * sT
        x0 = x[:, :, t0, :, :]       # (1,3,H,W)  rank4
        x1 = x[:, :, t0 + 1, :, :]   # (1,3,H,W)  rank4

        x_packed = pack_two_frames_interleaved_rank4(x0, x1)  # (1,6,H,W) rank4

        y_t = F.conv2d(
            x_packed, w2, bias=None,
            stride=(sH, sW), padding=0, dilation=1, groups=1
        )  # (1,1280,H_out,W_out) rank4
        outs.append(y_t)

    # Pack time into channels (rank4)
    y_packed = torch.cat(outs, dim=1)  # (1, 1280*T_out, H_out, W_out)

    # Verification reshape back to 5D (not HW-constrained)
    y5 = y_packed.view(1, T_out, C_out, H_out, W_out).permute(0, 2, 1, 3, 4).contiguous()

    max_abs = (y3 - y5).abs().max().item()
    print("y3 shape:", tuple(y3.shape))
    print("y_packed shape (rank4):", tuple(y_packed.shape))
    print("max |y3 - y5| =", max_abs)
    return max_abs