[VLM] avoid CUDA placement on non-CUDA platforms (#38750)
This commit is contained in:
@@ -23,6 +23,7 @@ import torch
|
||||
from PIL import Image
|
||||
from transformers import BaseImageProcessor
|
||||
|
||||
from sglang.srt import platforms
|
||||
from sglang.srt.managers.schedule_batch import (
|
||||
Modality,
|
||||
MultimodalDataItem,
|
||||
@@ -721,19 +722,23 @@ class BaseMultimodalProcessor(ABC):
|
||||
return processor, _tokenizer_of(processor)
|
||||
|
||||
def _preprocessing_competes_with_the_scheduler(self) -> bool:
|
||||
"""Whether image preprocessing submits its work to the serving GPU.
|
||||
"""Whether image preprocessing contends with the serving accelerator.
|
||||
|
||||
The fast image processor runs inside the tokenizer process but on
|
||||
``cuda:{base_gpu_id}`` -- the device the scheduler serves from. A second
|
||||
preprocessing worker there is one more competitor for that device rather
|
||||
than added parallelism.
|
||||
The fast image processor runs inside the tokenizer process but may run on
|
||||
the same accelerator as the scheduler. A second preprocessing worker there
|
||||
adds device contention rather than CPU preprocessing parallelism.
|
||||
"""
|
||||
if _is_cpu or get_exec().deterministic.rl_on_policy_target is not None:
|
||||
return False
|
||||
if self.disable_fast_image_processor:
|
||||
return False
|
||||
image_processor = getattr(self._processor, "image_processor", None)
|
||||
return isinstance(image_processor, BaseImageProcessor)
|
||||
if not isinstance(image_processor, BaseImageProcessor):
|
||||
return False
|
||||
if _is_xpu or _is_npu:
|
||||
return True
|
||||
platform = platforms.current_platform
|
||||
return platform.is_cuda_alike()
|
||||
|
||||
def _resolve_auto_mm_processor_worker_num(self) -> int:
|
||||
"""The worker count to use when the user did not ask for one.
|
||||
@@ -769,9 +774,12 @@ class BaseMultimodalProcessor(ABC):
|
||||
if _is_xpu:
|
||||
return "xpu"
|
||||
if not _is_npu:
|
||||
platform = platforms.current_platform
|
||||
if not platform.is_cuda_alike():
|
||||
return None
|
||||
# Per-worker placement travels as a constructor argument, and
|
||||
# this record is that argument.
|
||||
return f"cuda:{server_args.base_gpu_id}"
|
||||
return f"{platform.device_type}:{server_args.base_gpu_id}"
|
||||
if processor.__class__.__name__ == "MiniMaxVLProcessor":
|
||||
# MiniMax's image/video processors create 10-dim tensors during
|
||||
# patch extraction, exceeding the Ascend 8-dim limit; patch them
|
||||
|
||||
@@ -75,6 +75,21 @@ class TestMmProcessConfigValidation(CustomTestCase):
|
||||
class TestBaseProcessorConfigExtraction(CustomTestCase):
|
||||
"""Verify BaseMultimodalProcessor.__init__ extracts configs from server_args."""
|
||||
|
||||
def _patch_platform(self, cuda_alike, device_type):
|
||||
platforms = SimpleNamespace(
|
||||
current_platform=SimpleNamespace(
|
||||
is_cuda_alike=lambda: cuda_alike,
|
||||
device_type=device_type,
|
||||
)
|
||||
)
|
||||
return patch.multiple(
|
||||
"sglang.srt.multimodal.processors.base_processor",
|
||||
_is_cpu=False,
|
||||
_is_xpu=False,
|
||||
_is_npu=False,
|
||||
platforms=platforms,
|
||||
)
|
||||
|
||||
def _make_processor(
|
||||
self,
|
||||
mm_process_config,
|
||||
@@ -210,12 +225,23 @@ class TestBaseProcessorConfigExtraction(CustomTestCase):
|
||||
9.30 -> 4.02 req/s on GB300 for full-page images."""
|
||||
from transformers import BaseImageProcessor
|
||||
|
||||
with self._patch_platform(cuda_alike=True, device_type="cuda"):
|
||||
proc = self._make_processor(
|
||||
{}, image_processor=MagicMock(spec=BaseImageProcessor)
|
||||
)
|
||||
self.assertEqual(proc.mm_processor_worker_num, 1)
|
||||
self.assertIsNone(proc.mm_processor_executor)
|
||||
|
||||
def test_non_accelerator_fast_processor_gets_two_workers(self):
|
||||
from transformers import BaseImageProcessor
|
||||
|
||||
with self._patch_platform(cuda_alike=False, device_type="custom"):
|
||||
proc = self._make_processor(
|
||||
{}, image_processor=MagicMock(spec=BaseImageProcessor)
|
||||
)
|
||||
self.assertEqual(proc.mm_processor_worker_num, 2)
|
||||
self.assertIsNotNone(proc.mm_processor_executor)
|
||||
|
||||
def test_explicit_request_overrides_the_path_decision(self):
|
||||
"""The server argument wins: an operator who measured their own workload
|
||||
can still ask for concurrency on the GPU path."""
|
||||
@@ -240,7 +266,10 @@ class TestBaseProcessorConfigExtraction(CustomTestCase):
|
||||
BaseMultimodalProcessor,
|
||||
)
|
||||
|
||||
with patch.object(BaseMultimodalProcessor, "auto_mm_processor_worker_num", 3):
|
||||
with (
|
||||
patch.object(BaseMultimodalProcessor, "auto_mm_processor_worker_num", 3),
|
||||
self._patch_platform(cuda_alike=True, device_type="cuda"),
|
||||
):
|
||||
proc = self._make_processor(
|
||||
{}, image_processor=MagicMock(spec=BaseImageProcessor)
|
||||
)
|
||||
|
||||
@@ -51,10 +51,16 @@ class TestFastImageProcessorDevice(CustomTestCase):
|
||||
reset_context()
|
||||
self.addCleanup(reset_context)
|
||||
|
||||
def _device(self, processor, **platform):
|
||||
def _device(self, processor, cuda_alike=True, device_type="cuda", **platform):
|
||||
flags = {"_is_cpu": False, "_is_xpu": False, "_is_npu": False}
|
||||
flags.update(platform)
|
||||
with patch.multiple(BASE, **flags):
|
||||
platforms = SimpleNamespace(
|
||||
current_platform=SimpleNamespace(
|
||||
is_cuda_alike=lambda: cuda_alike,
|
||||
device_type=device_type,
|
||||
)
|
||||
)
|
||||
with patch.multiple(BASE, platforms=platforms, **flags):
|
||||
return processor._fast_image_processor_device(_Processor())
|
||||
|
||||
def test_device_follows_the_instance_base_gpu_id(self):
|
||||
@@ -92,6 +98,23 @@ class TestFastImageProcessorDevice(CustomTestCase):
|
||||
device = processor._fast_image_processor_device(Glm4vProcessor())
|
||||
self.assertIsNone(device)
|
||||
|
||||
def test_platform_respects_cuda_compatibility(self):
|
||||
processor = _make(base_gpu_id=3)
|
||||
for cuda_alike, device_type, expected in (
|
||||
(False, "other", None),
|
||||
(True, "cuda", "cuda:3"),
|
||||
(True, "musa", "musa:3"),
|
||||
):
|
||||
with self.subTest(cuda_alike=cuda_alike, device_type=device_type):
|
||||
self.assertEqual(
|
||||
self._device(
|
||||
processor,
|
||||
cuda_alike=cuda_alike,
|
||||
device_type=device_type,
|
||||
),
|
||||
expected,
|
||||
)
|
||||
|
||||
|
||||
class TestFastImageProcessorMemoryPool(CustomTestCase):
|
||||
def setUp(self):
|
||||
@@ -161,6 +184,18 @@ class TestFastImageProcessorMemoryPool(CustomTestCase):
|
||||
events.append("exit")
|
||||
|
||||
with (
|
||||
patch.multiple(
|
||||
BASE,
|
||||
_is_cpu=False,
|
||||
_is_xpu=False,
|
||||
_is_npu=False,
|
||||
platforms=SimpleNamespace(
|
||||
current_platform=SimpleNamespace(
|
||||
is_cuda_alike=lambda: True,
|
||||
device_type="cuda",
|
||||
)
|
||||
),
|
||||
),
|
||||
patch(f"{BASE}.BaseImageProcessor", ImageProcessor),
|
||||
patch(f"{BASE}.torch.cuda.device", return_value=nullcontext()),
|
||||
patch(f"{BASE}.torch.cuda.MemPool", return_value="pool"),
|
||||
|
||||
Reference in New Issue
Block a user