[feat](sft): support Qwen VLM MoE fine-tuning (#2156)

kt-kernel/python/sft/__init__.py

@@ -25,6 +25,7 @@ from .arch import (
     KTAMXConfigError,
 )
 from .autograd import KTMoEFunction
+from .conv3d_compat import is_vlm_conv3d_compatible, patch_vlm_conv3d
 from .layer import KTMoELayerWrapper
 from .weights import (
     extract_moe_weights,
@@ -69,6 +70,8 @@ __all__ = [
     "KTAMXModelNotSupportedError",
     "KTAMXConfigError",
     "KTMoEFunction",
+    "patch_vlm_conv3d",
+    "is_vlm_conv3d_compatible",
     "KTMoELayerWrapper",
     "extract_moe_weights",
     "load_experts_from_checkpoint_files",

kt-kernel/python/sft/arch.py

@@ -98,7 +98,7 @@ def get_moe_arch_config(config) -> MOEArchConfig:
             has_shared_experts=getattr(config, "n_shared_experts", 0) > 0,
             router_type="deepseek_gate",
         )
-    if "Qwen2Moe" in arch or "Qwen3Moe" in arch or "Qwen3_5Moe" in arch:
+    if "Qwen2Moe" in arch or "Qwen3Moe" in arch or "Qwen3VLMoe" in arch or "Qwen3_5Moe" in arch:
         cfg = getattr(config, "text_config", config)
         return MOEArchConfig(
             moe_layer_attr="mlp",
@@ -136,7 +136,8 @@ def get_moe_arch_config(config) -> MOEArchConfig:
 
     raise KTAMXModelNotSupportedError(
         f"Model architecture {arch} not supported for KT AMX. "
-        "Supported architectures: DeepseekV2, DeepseekV3, Qwen2Moe, Qwen3Moe, Qwen3_5Moe, Glm4Moe, Mixtral"
+        "Supported architectures: DeepseekV2, DeepseekV3, Qwen2Moe, Qwen3Moe, Qwen3VLMoe, "
+        "Qwen3_5Moe, Glm4Moe, Mixtral"
     )
 
 
@@ -167,7 +168,7 @@ def detect_fused_experts(experts: nn.Module) -> bool:
 
 def _get_layers_prefix(config) -> str:
     arch = config.architectures[0] if getattr(config, "architectures", None) else ""
-    if "Qwen3_5Moe" in arch:
+    if "Qwen3_5Moe" in arch or "Qwen3VLMoe" in arch:
         return "model.language_model.layers"
     return "model.layers"
 

kt-kernel/python/sft/conv3d_compat.py

@@ -0,0 +1,150 @@
+# SPDX-License-Identifier: Apache-2.0
+
+"""Instance-scoped Qwen VLM compatibility for the KT SFT stack."""
+
+from __future__ import annotations
+
+from types import MethodType
+from typing import Any
+
+from packaging.version import Version
+
+
+_SUPPORTED_MODEL_TYPES = {
+    "qwen2_vl",
+    "qwen2_5_vl",
+    "qwen3_vl",
+    "qwen3_vl_moe",
+    "qwen3_5",
+    "qwen3_5_moe",
+}
+_COMPATIBLE_ATTR = "_kt_conv3d_compatible"
+_ORIGINAL_FORWARD_ATTR = "_kt_original_conv3d_forward"
+
+
+def _canonicalize_qwen3_vl_fused_expert_weights(
+    gate_up: Any,
+    down: Any,
+    moe_config: Any,
+) -> tuple[Any, Any, Any] | None:
+    """Convert only the transposed Qwen3-VL fused expert layout to KT layout."""
+    expert_num = moe_config.expert_num
+    intermediate_size = moe_config.intermediate_size
+    if gate_up.dim() != 3 or down.dim() != 3:
+        return None
+
+    hidden_size = gate_up.shape[1]
+    if tuple(gate_up.shape) != (expert_num, hidden_size, 2 * intermediate_size):
+        return None
+    if tuple(down.shape) != (expert_num, intermediate_size, hidden_size):
+        return None
+
+    gate_proj, up_proj = (
+        tensor.transpose(1, 2).contiguous()
+        for tensor in gate_up.split(intermediate_size, dim=2)
+    )
+    return gate_proj, up_proj, down.transpose(1, 2).contiguous()
+
+
+def _requires_conv3d_patch() -> bool:
+    import torch
+
+    version = Version(torch.__version__.split("+", 1)[0])
+    return Version("2.9.0") <= version < Version("2.10.0")
+
+
+def _is_supported_vlm(config: Any) -> bool:
+    return (
+        getattr(config, "model_type", None) in _SUPPORTED_MODEL_TYPES
+        and getattr(config, "vision_config", None) is not None
+    )
+
+
+def _contract_errors(module: Any) -> list[str]:
+    errors = []
+    if tuple(module.stride) != tuple(module.kernel_size):
+        errors.append(f"stride={module.stride} != kernel_size={module.kernel_size}")
+    if any(value != 0 for value in module.padding):
+        errors.append(f"padding={module.padding} != 0")
+    if any(value != 1 for value in module.dilation):
+        errors.append(f"dilation={module.dilation} != 1")
+    if module.groups != 1:
+        errors.append(f"groups={module.groups} != 1")
+    return errors
+
+
+def _kt_conv3d_forward(module: Any, inputs: Any) -> Any:
+    """Evaluate the non-overlapping Conv3D contract using autograd-native ops."""
+    import torch.nn.functional as functional
+
+    batch_size = inputs.shape[0]
+    kernel = tuple(module.kernel_size)
+    inputs = (
+        inputs.unfold(2, kernel[0], kernel[0])
+        .unfold(3, kernel[1], kernel[1])
+        .unfold(4, kernel[2], kernel[2])
+    )
+    depth, height, width = inputs.shape[2:5]
+    inputs = inputs.permute(0, 2, 3, 4, 1, 5, 6, 7).reshape(
+        -1, module.in_channels * kernel[0] * kernel[1] * kernel[2]
+    )
+    outputs = functional.linear(
+        inputs,
+        module.weight.reshape(module.out_channels, -1),
+        module.bias,
+    )
+    return outputs.view(batch_size, depth, height, width, module.out_channels).permute(
+        0, 4, 1, 2, 3
+    )
+
+
+def patch_vlm_conv3d(model: Any) -> list[str]:
+    """Patch supported Conv3D instances without changing ``torch.nn.Conv3d`` globally."""
+    if not _requires_conv3d_patch() or not _is_supported_vlm(
+        getattr(model, "config", None)
+    ):
+        return []
+
+    import torch
+
+    modules = [
+        (name or "<root>", module)
+        for name, module in model.named_modules()
+        if isinstance(module, torch.nn.Conv3d)
+    ]
+    unsupported = []
+    for name, module in modules:
+        errors = _contract_errors(module)
+        if errors:
+            unsupported.append(f"{name}: {', '.join(errors)}")
+
+    if unsupported:
+        raise RuntimeError(
+            "the KT Conv3D fallback does not support: " + "; ".join(unsupported)
+        )
+
+    patched_names = []
+    for name, module in modules:
+        if not getattr(module, _COMPATIBLE_ATTR, False):
+            setattr(module, _ORIGINAL_FORWARD_ATTR, module.forward)
+            module.forward = MethodType(_kt_conv3d_forward, module)
+            setattr(module, _COMPATIBLE_ATTR, True)
+        patched_names.append(name)
+
+    setattr(model, "_kt_vlm_conv3d_compatible", bool(modules))
+    return patched_names
+
+
+def is_vlm_conv3d_compatible(model: Any) -> bool:
+    """Return whether every Conv3D instance on a torch 2.9 KT VLM was patched."""
+    if not _requires_conv3d_patch():
+        return True
+
+    import torch
+
+    modules = [
+        module for module in model.modules() if isinstance(module, torch.nn.Conv3d)
+    ]
+    return bool(modules) and all(
+        getattr(module, _COMPATIBLE_ATTR, False) for module in modules
+    )

kt-kernel/python/sft/weights.py

@@ -14,6 +14,7 @@ import torch
 import torch.nn as nn
 
 from .arch import MOEArchConfig
+from .conv3d_compat import _canonicalize_qwen3_vl_fused_expert_weights
 from .dist_utils import _maybe_zero3_gathered_parameters
 
 logger = logging.getLogger(__name__)
@@ -54,6 +55,9 @@ def extract_moe_weights(
     if detect_fused_experts(experts):
         gate_up = getattr(experts, "gate_up_proj").data
         down_fused = getattr(experts, "down_proj").data
+        vlm_weights = _canonicalize_qwen3_vl_fused_expert_weights(gate_up, down_fused, moe_config)
+        if vlm_weights is not None:
+            return vlm_weights
         # gate_up_proj is [E, 2*I, H], split into gate [E, I, H] and up [E, I, H]
         intermediate = gate_up.shape[1] // 2
         gate_proj = gate_up[:, :intermediate, :].contiguous()
@@ -393,10 +397,15 @@ def load_experts_from_checkpoint_files(
         if gate_up is None or down is None:
             raise FileNotFoundError(f"Missing fused expert weights for layer {layer_idx}")
         gate_up = gate_up.cpu().to(torch.bfloat16).contiguous()
-        I = gate_up.shape[1] // 2
-        gate_proj = gate_up[:, :I, :].contiguous()
-        up_proj = gate_up[:, I:, :].contiguous()
-        down_proj = down.cpu().to(torch.bfloat16).contiguous()
+        down = down.cpu().to(torch.bfloat16)
+        vlm_weights = _canonicalize_qwen3_vl_fused_expert_weights(gate_up, down, moe_config)
+        if vlm_weights is not None:
+            gate_proj, up_proj, down_proj = vlm_weights
+        else:
+            I = gate_up.shape[1] // 2
+            gate_proj = gate_up[:, :I, :].contiguous()
+            up_proj = gate_up[:, I:, :].contiguous()
+            down_proj = down.contiguous()
         del gate_up
         print(
             f"[kt_moe] Layer {layer_idx}: fused expert format — "

kt-kernel/python/sft/wrapper.py

@@ -25,6 +25,7 @@ from .layer import KTMoELayerWrapper
 from .lora import LoRAExperts
 from .base import _supports_authoritative_optimizer_grads
 from .checkpoint import load_full_weight_layer, resolve_full_weight_checkpoint
+from .conv3d_compat import patch_vlm_conv3d
 from .dist_utils import _distributed_rank_world_size
 from .weights import (
     _clear_original_expert_weights,
@@ -160,6 +161,9 @@ def wrap_moe_layers_with_kt_wrapper(model: nn.Module, kt_plugin: Any) -> list[KT
     is_rank_0 = distributed_rank == 0
 
     moe_config = get_moe_arch_config(model.config)
+    patched_conv3d = patch_vlm_conv3d(model)
+    if patched_conv3d:
+        logger.info(f"Patched KT VLM Conv3D modules: {patched_conv3d}")
     _text_cfg = getattr(model.config, "text_config", model.config)
     hidden_size = _text_cfg.hidden_size