fix: fix Kimi-VL encoder parallelism (#30869)
This commit is contained in:
Mick
@@ -0,0 +1,159 @@
|
||||
"""CPU-only coverage for Kimi-VL encoder parallelism wiring."""
|
||||
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
from sglang.srt.layers.linear import (
|
||||
ColumnParallelLinear,
|
||||
QKVParallelLinear,
|
||||
RowParallelLinear,
|
||||
)
|
||||
from sglang.srt.managers.schedule_batch import Modality, MultimodalDataItem
|
||||
from sglang.srt.models.kimi_vl import KimiVLForConditionalGeneration
|
||||
from sglang.srt.models.kimi_vl_moonvit import MoonVitEncoderLayer, multihead_attention
|
||||
from sglang.srt.multimodal.mm_utils import run_dp_sharded_mrope_vision_model
|
||||
from sglang.srt.runtime_context import get_parallel
|
||||
from sglang.test.ci.ci_register import register_cpu_ci
|
||||
from sglang.test.test_utils import CustomTestCase
|
||||
|
||||
register_cpu_ci(est_time=5, suite="base-a-test-cpu")
|
||||
|
||||
|
||||
class _VisionTower:
|
||||
dtype = torch.float32
|
||||
device = torch.device("cpu")
|
||||
|
||||
def __init__(self):
|
||||
self.calls = []
|
||||
|
||||
def __call__(
|
||||
self, pixel_values, image_grid_hws=None, max_seqlen=None, grid_hw=None
|
||||
):
|
||||
image_grid_hws = grid_hw if image_grid_hws is None else image_grid_hws
|
||||
self.calls.append((pixel_values, image_grid_hws))
|
||||
return [
|
||||
torch.full((1, 4, 2), index + 1.0)
|
||||
for index in range(image_grid_hws.shape[0])
|
||||
]
|
||||
|
||||
|
||||
class _Projector:
|
||||
def __init__(self):
|
||||
self.input = None
|
||||
|
||||
def __call__(self, image_features):
|
||||
self.input = image_features
|
||||
return image_features
|
||||
|
||||
|
||||
class _GridRecordingVisionTower:
|
||||
def __call__(self, pixel_values, grid_hw, max_seqlen=None):
|
||||
self.grid_thw = grid_hw
|
||||
self.max_seqlen = max_seqlen
|
||||
return pixel_values
|
||||
|
||||
|
||||
def _bare_model(*, use_data_parallel: bool):
|
||||
model = KimiVLForConditionalGeneration.__new__(KimiVLForConditionalGeneration)
|
||||
nn.Module.__init__(model)
|
||||
model.config = SimpleNamespace(text_config=SimpleNamespace(hidden_size=16))
|
||||
model.use_data_parallel = use_data_parallel
|
||||
model.vision_tower = _VisionTower()
|
||||
model.multi_modal_projector = _Projector()
|
||||
return model
|
||||
|
||||
|
||||
def _image_item(feature, grid_hws):
|
||||
return MultimodalDataItem(
|
||||
modality=Modality.IMAGE,
|
||||
offsets=[(0, 1)],
|
||||
feature=feature,
|
||||
model_specific_data={"image_grid_hws": torch.tensor(grid_hws)},
|
||||
)
|
||||
|
||||
|
||||
class TestKimiVLEncoderParallelism(CustomTestCase):
|
||||
def test_moonvit_uses_tensor_parallel_layers(self):
|
||||
with get_parallel().override(
|
||||
tp_size=1, tp_rank=0, attn_tp_size=1, attn_tp_rank=0
|
||||
):
|
||||
layer = MoonVitEncoderLayer(
|
||||
num_heads=2,
|
||||
hidden_dim=8,
|
||||
mlp_dim=16,
|
||||
prefix="vision_tower.encoder.blocks.0",
|
||||
use_tensor_parallel=True,
|
||||
)
|
||||
|
||||
self.assertIsInstance(layer.wqkv, QKVParallelLinear)
|
||||
self.assertIsInstance(layer.wo, RowParallelLinear)
|
||||
self.assertIsInstance(layer.mlp.fc0, ColumnParallelLinear)
|
||||
self.assertIsInstance(layer.mlp.fc1, RowParallelLinear)
|
||||
|
||||
def test_encoder_dp_uses_existing_mrope_sharding_helper(self):
|
||||
model = _bare_model(use_data_parallel=True)
|
||||
items = [
|
||||
_image_item(torch.randn(4, 2), [[2, 2]]),
|
||||
_image_item(torch.randn(8, 2), [[4, 2]]),
|
||||
]
|
||||
sharded_features = torch.randn(3, 4, 2)
|
||||
|
||||
with patch(
|
||||
"sglang.srt.models.kimi_vl.run_dp_sharded_mrope_vision_model",
|
||||
return_value=sharded_features,
|
||||
) as run_dp:
|
||||
output = model.get_image_feature(items)
|
||||
|
||||
run_dp.assert_called_once()
|
||||
_, pixel_values, grid_hws = run_dp.call_args.args
|
||||
self.assertEqual(pixel_values.shape, (12, 2))
|
||||
self.assertEqual(grid_hws, [[2, 2], [4, 2]])
|
||||
self.assertEqual(run_dp.call_args.kwargs, {"rope_type": "rope_2d"})
|
||||
self.assertIs(output, sharded_features)
|
||||
|
||||
def test_encoder_dp_keeps_moonvit_grid_metadata_on_vision_device(self):
|
||||
vision_tower = _GridRecordingVisionTower()
|
||||
pixel_values = torch.randn(4, 2)
|
||||
|
||||
with get_parallel().override(
|
||||
tp_size=1, tp_rank=0, attn_tp_size=1, attn_tp_rank=0
|
||||
):
|
||||
output = run_dp_sharded_mrope_vision_model(
|
||||
vision_tower, pixel_values, [[2, 2]], rope_type="rope_2d"
|
||||
)
|
||||
|
||||
self.assertIs(output, pixel_values)
|
||||
self.assertEqual(vision_tower.grid_thw.device, pixel_values.device)
|
||||
self.assertEqual(vision_tower.max_seqlen, 4)
|
||||
|
||||
def test_encoder_dp_tp1_concatenates_moonvit_image_outputs(self):
|
||||
vision_tower = _VisionTower()
|
||||
pixel_values = torch.randn(4, 2)
|
||||
with get_parallel().override(
|
||||
tp_size=1, tp_rank=0, attn_tp_size=1, attn_tp_rank=0
|
||||
):
|
||||
output = run_dp_sharded_mrope_vision_model(
|
||||
vision_tower, pixel_values, [[2, 2]], rope_type="rope_2d"
|
||||
)
|
||||
self.assertIsInstance(output, torch.Tensor)
|
||||
self.assertEqual(output.shape, (1, 4, 2))
|
||||
|
||||
def test_moonvit_attention_accepts_precomputed_max_seqlen(self):
|
||||
q = torch.randn(4, 2, 4, dtype=torch.bfloat16)
|
||||
cu_seqlens = torch.tensor([0, 4], dtype=torch.int32)
|
||||
fake_output = torch.randn_like(q)
|
||||
with patch(
|
||||
"sglang.srt.models.kimi_vl_moonvit.flash_attn_varlen_func",
|
||||
return_value=fake_output,
|
||||
) as flash_attn:
|
||||
output = multihead_attention(q, q, q, cu_seqlens, cu_seqlens, max_seqlen=4)
|
||||
self.assertTrue(torch.equal(output, fake_output.flatten(start_dim=-2)))
|
||||
self.assertEqual(flash_attn.call_args.args[5:7], (4, 4))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(pytest.main([__file__, "-v"]))
|
||||
@@ -0,0 +1,64 @@
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from sglang.test.ci.ci_register import register_cpu_ci
|
||||
|
||||
register_cpu_ci(est_time=3, suite="base-a-test-cpu")
|
||||
|
||||
from sglang.srt.models import kimi_vl_moonvit
|
||||
from sglang.srt.models.kimi_vl_moonvit import Learnable2DInterpPosEmb
|
||||
|
||||
|
||||
def test_learnable_2d_pos_emb_caches_inference_interpolation(monkeypatch):
|
||||
module = Learnable2DInterpPosEmb(height=2, width=2, dim=4).eval()
|
||||
inputs = torch.zeros(6, 4)
|
||||
grid_hw = torch.tensor([[2, 3]])
|
||||
calls = 0
|
||||
original_interpolate = torch.nn.functional.interpolate
|
||||
|
||||
def counting_interpolate(*args, **kwargs):
|
||||
nonlocal calls
|
||||
calls += 1
|
||||
return original_interpolate(*args, **kwargs)
|
||||
|
||||
monkeypatch.setattr(torch.nn.functional, "interpolate", counting_interpolate)
|
||||
first = module(inputs, grid_hw)
|
||||
second = module(inputs, grid_hw)
|
||||
torch.testing.assert_close(first, second)
|
||||
assert calls == 1
|
||||
|
||||
|
||||
def test_learnable_2d_pos_emb_does_not_cache_training_interpolation(monkeypatch):
|
||||
module = Learnable2DInterpPosEmb(height=2, width=2, dim=4).train()
|
||||
inputs = torch.zeros(6, 4)
|
||||
grid_hw = torch.tensor([[2, 3]])
|
||||
calls = 0
|
||||
original_interpolate = torch.nn.functional.interpolate
|
||||
|
||||
def counting_interpolate(*args, **kwargs):
|
||||
nonlocal calls
|
||||
calls += 1
|
||||
return original_interpolate(*args, **kwargs)
|
||||
|
||||
monkeypatch.setattr(torch.nn.functional, "interpolate", counting_interpolate)
|
||||
module(inputs, grid_hw)
|
||||
module(inputs, grid_hw)
|
||||
assert calls == 2
|
||||
|
||||
|
||||
def test_learnable_2d_pos_emb_evicts_oldest_inference_cache_entry(monkeypatch):
|
||||
monkeypatch.setattr(kimi_vl_moonvit, "_MAX_INFERENCE_POS_EMB_CACHE_ENTRIES", 1)
|
||||
module = Learnable2DInterpPosEmb(height=2, width=2, dim=4).eval()
|
||||
inputs = torch.zeros(6, 4)
|
||||
|
||||
module(inputs, torch.tensor([[2, 3]]))
|
||||
module(inputs, torch.tensor([[3, 2]]))
|
||||
|
||||
assert len(module._interpolated_pos_emb_cache) == 1
|
||||
assert ((3, 2), module.weight.dtype, module.weight.device) in (
|
||||
module._interpolated_pos_emb_cache
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(pytest.main([__file__, "-v"]))
|
||||
@@ -0,0 +1,59 @@
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
|
||||
from sglang.test.ci.ci_register import register_cpu_ci
|
||||
|
||||
register_cpu_ci(est_time=3, suite="base-a-test-cpu")
|
||||
|
||||
from sglang.srt.multimodal.vit_cuda_graph_runner import ViTCudaGraphRunner
|
||||
|
||||
|
||||
class _Block:
|
||||
def forward(self, x):
|
||||
return x
|
||||
|
||||
|
||||
def _runner(*, use_data_parallel: bool) -> ViTCudaGraphRunner:
|
||||
vit = SimpleNamespace(
|
||||
blocks=[_Block()],
|
||||
deepstack_visual_indexes=[],
|
||||
deepstack_merger_list=None,
|
||||
use_data_parallel=use_data_parallel,
|
||||
)
|
||||
return ViTCudaGraphRunner(vit)
|
||||
|
||||
|
||||
def test_dp_vit_graph_capture_does_not_enter_tp_communication_capture():
|
||||
runner = _runner(use_data_parallel=True)
|
||||
with patch(
|
||||
"sglang.srt.multimodal.vit_cuda_graph_runner.get_tp_group",
|
||||
side_effect=AssertionError("DP capture must be rank-local"),
|
||||
):
|
||||
with runner._capture_context():
|
||||
pass
|
||||
|
||||
|
||||
def test_non_dp_vit_graph_capture_uses_tp_communication_capture():
|
||||
entered = []
|
||||
|
||||
class Capture:
|
||||
def __enter__(self):
|
||||
entered.append(True)
|
||||
|
||||
def __exit__(self, *args):
|
||||
return False
|
||||
|
||||
group = SimpleNamespace(ca_comm=SimpleNamespace(capture=lambda: Capture()))
|
||||
runner = _runner(use_data_parallel=False)
|
||||
with patch(
|
||||
"sglang.srt.multimodal.vit_cuda_graph_runner.get_tp_group", return_value=group
|
||||
):
|
||||
with runner._capture_context():
|
||||
pass
|
||||
assert entered == [True]
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(pytest.main([__file__, "-v"]))
|
||||
@@ -4,12 +4,16 @@ Tests cover the pure utility functions (compat patches, config helpers,
|
||||
context length, GGUF detection, etc.) that don't require actual model files.
|
||||
"""
|
||||
|
||||
import inspect
|
||||
import tempfile
|
||||
import unittest
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import patch
|
||||
|
||||
from transformers import PretrainedConfig
|
||||
from transformers.image_processing_utils import BaseImageProcessor
|
||||
|
||||
from sglang.srt.utils import hf_transformers_patches
|
||||
from sglang.srt.utils.hf_transformers.common import (
|
||||
_is_deepseek_ocr2_model,
|
||||
_is_deepseek_ocr_model,
|
||||
@@ -27,6 +31,33 @@ from sglang.test.ci.ci_register import register_cpu_ci
|
||||
register_cpu_ci(est_time=6, suite="base-a-test-cpu")
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# _patch_image_processor_kwargs
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestImageProcessorKwargsPatch(unittest.TestCase):
|
||||
def test_filters_unsupported_kwargs_and_caches_signature(self):
|
||||
class StrictImageProcessor(BaseImageProcessor):
|
||||
model_input_names = ["pixel_values"]
|
||||
|
||||
def preprocess(self, images, accepted=None):
|
||||
return {"images": images, "accepted": accepted}
|
||||
|
||||
processor = StrictImageProcessor()
|
||||
with patch.object(
|
||||
hf_transformers_patches.inspect,
|
||||
"signature",
|
||||
wraps=inspect.signature,
|
||||
) as signature:
|
||||
first = processor("first", accepted=True, device="cuda")
|
||||
second = processor("second", accepted=False, device="cuda")
|
||||
|
||||
self.assertEqual(first, {"images": "first", "accepted": True})
|
||||
self.assertEqual(second, {"images": "second", "accepted": False})
|
||||
self.assertEqual(signature.call_count, 1)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# normalize_rope_scaling_compat
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
Reference in New Issue
Block a user