mirror of
https://github.com/kvcache-ai/ktransformers.git
synced 2026-10-29 06:04:23 +08:00
69 lines
2.3 KiB
Python
69 lines
2.3 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
import copy
|
|
import importlib.util
|
|
import sys
|
|
from pathlib import Path
|
|
from types import SimpleNamespace
|
|
|
|
import pytest
|
|
import torch
|
|
|
|
|
|
MODULE_PATH = (
|
|
Path(__file__).resolve().parents[2] / "python" / "sft" / "conv3d_compat.py"
|
|
)
|
|
SPEC = importlib.util.spec_from_file_location(
|
|
"kt_sft_conv3d_compat_under_test", MODULE_PATH
|
|
)
|
|
assert SPEC is not None and SPEC.loader is not None
|
|
conv3d_compat = importlib.util.module_from_spec(SPEC)
|
|
sys.modules[SPEC.name] = conv3d_compat
|
|
SPEC.loader.exec_module(conv3d_compat)
|
|
|
|
|
|
def _vlm(conv3d):
|
|
model = torch.nn.Sequential(conv3d)
|
|
model.config = SimpleNamespace(model_type="qwen3_vl_moe", vision_config=object())
|
|
return model
|
|
|
|
|
|
def test_patch_vlm_conv3d_matches_native_forward_and_backward(monkeypatch):
|
|
monkeypatch.setattr(conv3d_compat, "_requires_conv3d_patch", lambda: True)
|
|
native_forward = torch.nn.Conv3d.forward
|
|
reference = _vlm(
|
|
torch.nn.Conv3d(
|
|
3, 4, kernel_size=(2, 4, 4), stride=(2, 4, 4), bias=True
|
|
).double()
|
|
)
|
|
patched = copy.deepcopy(reference)
|
|
reference_input = torch.randn(
|
|
2, 3, 4, 8, 8, dtype=torch.float64, requires_grad=True
|
|
)
|
|
patched_input = reference_input.detach().clone().requires_grad_(True)
|
|
|
|
assert conv3d_compat.patch_vlm_conv3d(patched) == ["0"]
|
|
assert conv3d_compat.patch_vlm_conv3d(patched) == ["0"]
|
|
assert torch.nn.Conv3d.forward is native_forward
|
|
assert conv3d_compat.is_vlm_conv3d_compatible(patched)
|
|
|
|
expected = reference(reference_input)
|
|
actual = patched(patched_input)
|
|
expected.square().sum().backward()
|
|
actual.square().sum().backward()
|
|
|
|
torch.testing.assert_close(actual, expected)
|
|
torch.testing.assert_close(patched_input.grad, reference_input.grad)
|
|
torch.testing.assert_close(patched[0].weight.grad, reference[0].weight.grad)
|
|
torch.testing.assert_close(patched[0].bias.grad, reference[0].bias.grad)
|
|
|
|
|
|
def test_patch_vlm_conv3d_rejects_unsupported_contract_atomically(monkeypatch):
|
|
monkeypatch.setattr(conv3d_compat, "_requires_conv3d_patch", lambda: True)
|
|
model = _vlm(torch.nn.Conv3d(3, 4, kernel_size=(2, 4, 4), stride=(1, 4, 4)))
|
|
|
|
with pytest.raises(RuntimeError, match="stride="):
|
|
conv3d_compat.patch_vlm_conv3d(model)
|
|
|
|
assert not hasattr(model[0], "_kt_conv3d_compatible")
|