A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 14, 23:45 UTC
[Model] Support tensor parallel for timm ViT in Deepseek_vl2 (#21494)
Signed-off-by: wzqd <1057337859@qq.com>
Jason Gu committed
Jul 25, 2025 at 13:45 UTC
b38bc652ac5111d96cfd41e3575a879e9b47efbd
1 file changed
+38
-2
vllm/model_executor/models/deepseek_vl2.py
+38
-2
@@ -14,9 +14,11 @@ from einops import rearrange, repeat
14
from transformers import BatchFeature
15
16
from vllm.config import VllmConfig
17
+from vllm.distributed import get_tensor_model_parallel_world_size
18
from vllm.model_executor import SamplingMetadata
19
from vllm.model_executor.layers.quantization import QuantizationConfig
20
from vllm.model_executor.model_loader.utils import set_default_torch_dtype
21
+from vllm.model_executor.models.transformers import replace_linear_class
22
from vllm.multimodal import MULTIMODAL_REGISTRY
23
from vllm.multimodal.inputs import (MultiModalDataDict, MultiModalFieldConfig,
24
MultiModalKwargs, NestedTensors)
@@ -379,6 +381,37 @@ class DeepseekVLV2ForCausalLM(nn.Module, SupportsMultiModal, SupportsPP):
381
self.make_empty_intermediate_tensors = (
382
self.language_model.make_empty_intermediate_tensors)
383
384
+ def _get_parent_and_attr(self, root: torch.nn.Module, dotted_name: str):
385
+ """Return (parent_module, final_attr_name) for a dotted module path."""
386
+ names = dotted_name.split('.')
387
+ parent = root
388
+ for n in names[:-1]:
389
+ parent = getattr(parent, n)
390
+ return parent, names[-1]
391
+
392
+ #patch for timm ViT instance to support tensor parallel
393
+ def patch_vit_for_tp(self, vit: torch.nn.Module,
394
+ quant_config: QuantizationConfig):
395
+ try:
396
+ import timm
397
+ except ImportError as e:
398
+ raise ImportError("Please install timm") from e
399
+
400
+ for name, module in vit.named_modules():
401
+ if isinstance(module, nn.Linear):
402
+ parent, attr_name = self._get_parent_and_attr(vit, name)
403
+ if isinstance(parent, timm.layers.Mlp) and attr_name == "fc1":
404
+ new_linear = replace_linear_class(module, "colwise",
405
+ quant_config)
406
+ setattr(parent, attr_name, new_linear)
407
+ elif isinstance(parent,
408
+ timm.layers.Mlp) and attr_name == "fc2":
409
+ new_linear = replace_linear_class(module, "rowwise",
410
+ quant_config)
411
+ setattr(parent, attr_name, new_linear)
412
+
413
+ return vit
414
+
415
def _init_vision_module(
416
self,
417
vision_config: VisionEncoderConfig,
@@ -388,8 +421,8 @@ class DeepseekVLV2ForCausalLM(nn.Module, SupportsMultiModal, SupportsPP):
421
# TODO: refactor vision model through timm wrapper from transformers
422
try:
423
import timm
391
- except ImportError:
392
- raise ImportError("Please install timm") from ImportError
424
+ except ImportError as e:
425
+ raise ImportError("Please install timm") from e
426
427
with set_default_torch_dtype(torch.float16):
428
model = timm.create_model(
@@ -400,6 +433,9 @@ class DeepseekVLV2ForCausalLM(nn.Module, SupportsMultiModal, SupportsPP):
433
dynamic_img_pad=True,
434
)
435
436
+ if get_tensor_model_parallel_world_size() > 1:
437
+ model = self.patch_vit_for_tp(model, quant_config)
438
+
439
model = model.to(dtype=torch.get_default_dtype())
440
return model
441