Optimize VLM PatchEmbed

[dayo09]

What?

Let's optimize the below VLM PatchEmbed module.

Q. Doesn't one-optimize resolve this?

• I checked that both

• (1) Remove redundant SelectV2 [passes] Remove redundant SelectV2 from Conv3D Legalize #474

  • Remove Reshape + Concat + Reshape at the bottom of the graph too

• (2) Lower 4-dim permute to 3-dim permute

[dayo09]

(1) Remove redundant SelectV2

Why had SelectV2 inserted?

• Conv3d to Conv2d, pseudo code: https://github.sec.samsung.net/seockho-kim/notes/issues/44#issuecomment-4022675

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
        
        # 여러 Conv2d 레이어를 내부에 포함
        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
        
        # 수동 time padding - concatenate 사용
        if self.padding[0] > 0:
            # zero padding 생성
            zero_padding = torch.zeros(N, C, self.padding[0], H, W, device=x.device, dtype=x.dtype)
            # (N, C, padding, H, W)
            
            # 앞뒤로 padding 추가
            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):
                # dilation 적용된 시간 위치
                t_idx = t_in + i * self.dilation[0]
                
                if t_idx < T_padded:
                    # 5D -> 4D로 변환하여 Conv2d에 입력
                    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]

---

Model

    @check_model_inputs
    def forward(
        self, hidden_states: torch.Tensor, grid_thw: torch.Tensor, **kwargs: Unpack[TransformersKwargs]
    ) -> tuple | BaseModelOutputWithDeepstackFeatures:
        """
        Args:
            hidden_states (`torch.Tensor` of shape `(seq_len, hidden_size)`):
                The final hidden states of the model.
            grid_thw (`torch.Tensor` of shape `(num_images_or_videos, 3)`):
                The temporal, height and width of feature shape of each image in LLM.

        Returns:
            `torch.Tensor`: hidden_states.
        """
        hidden_states = self.patch_embed(hidden_states)

        pos_embeds = self.fast_pos_embed_interpolate(grid_thw)
        hidden_states = hidden_states + pos_embeds

        rotary_pos_emb = self.rot_pos_emb(grid_thw)

class Qwen3VLVisionPatchEmbed(nn.Module):
    def __init__(self, config) -> None:
        super().__init__()
        self.patch_size = config.patch_size
        self.temporal_patch_size = config.temporal_patch_size
        self.in_channels = config.in_channels
        self.embed_dim = config.hidden_size

        kernel_size = [self.temporal_patch_size, self.patch_size, self.patch_size]
        self.proj = nn.Conv3d(self.in_channels, self.embed_dim, kernel_size=kernel_size, stride=kernel_size, bias=True)

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        target_dtype = self.proj.weight.dtype

        hidden_states = hidden_states.view(
            -1, self.in_channels, self.temporal_patch_size, self.patch_size, self.patch_size
        )
        hidden_states = self.proj(hidden_states.to(dtype=target_dtype)).view(-1, self.embed_dim)
        return hidden_states

• 'Qwen/Qwen3-VL-2B-Instruct'
  • "temporal_patch_size": 2

Qwen3VLVisionPatchEmbed(
  (proj): Conv3d(3, 1024, kernel_size=(2, 16, 16), stride=(2, 16, 16))
)

Here, SelectV2 node is utilized only when T-axis dilation goes off T_MAX limit. It creates zero paddings for that case. However, when dilation[0] (T-axis dilation) is 1, the case never happen.

[llFreetimell]

Idea with using Conv2D

If kernel_size == (D, H, W), input * kernel is exactly same with element-wise multiplication of flattened elements.
Although we can optimize with Mul operator, let's use Conv2D :)

PyTorch Codes

import torch

# Original Input
N = 123
input = torch.randint(0, 3, (N, 1536)).float()

# Original Qwen3VLVisionPatchEmbed.forward()
I1 = input.view(-1, 3, 2, 16, 16)
conv3d = torch.nn.Conv3d(in_channels=3, out_channels=1024, kernel_size=(2, 16, 16), stride=(2, 16, 16), padding=0)
R1 = conv3d(I1).view(-1, 1024)

# Suggestion : Reshape + Conv2D + Permute + Reshape
I2=input.view(1, 1, -1, 1536)
conv2d = torch.nn.Conv2d(in_channels=1, out_channels=1024, kernel_size=(1, 1536))
conv2d.weight.data = conv3d.weight.data.view(1024, 1, 1, 1536) # constant folding / no reshape
conv2d.bias.data = conv3d.bias.data
conv2d_result = conv2d(I2)
R2 = torch.permute(conv2d_result, (2, 1, 0, 3)).view(-1, 1024)

# You can check R1 == R2
# print(R1)
# print(R2)

Needed passes

• Add above special pattern to ConvertConv3dToConv2d pass
• Implement ConvertPermuteToReshape pass
  • e.g. when x.shape == [1, a, 1, b], torch.permute(x, [1,3,0,2]) can be changed to torch.reshape(x, [a, b, 1, 1])

[mhs4670go]

FYI, conv3d logic is also related with #430.

[dayo09]

To update the context:

As @mhs4670go demonstrated, Qwen3-Vl's patch embedding layer assumes that same stride == same kernel, thus can be lowered to conv2d. Not only that, it's actually Linear projection.

After figuring out that fact, we also assessed the impacts of both Conv3d-to-Conv2d conversion and Conv3d-to-linear conversion with @llFreetimell and @parjong.

• Converting it into linear may harm performance as some reshape operations could be added.
• Thus, we shortly decided to perform conv2d with (1, 1, -1, D * H * W) as the input shape comes as (Batch, D * H * W) - @llFreetimell is still inspecting the previous layer to ensure this.

this proposal is based on above condition. I will update this to the issue.

From: #430 (comment)