wehub-resource-sync
116 lines
3.8 KiB
Python
116 lines
3.8 KiB
Python
# Copyright (c) Microsoft Corporation.
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
# DeepSpeed Team
|
|
"""
|
|
Automatically detect ViT encoder and LLM decoder attention modules in
|
|
multimodal models to guide AutoSP injection.
|
|
|
|
Extend _VIT_ATTN_CLASSNAMES / _LLM_ATTN_CLASSNAMES when adding support for
|
|
new model architectures.
|
|
"""
|
|
|
|
import torch.nn as nn
|
|
from dataclasses import dataclass, field
|
|
from typing import List, Optional, Tuple
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Architecture registry
|
|
# ---------------------------------------------------------------------------
|
|
|
|
# Known ViT attention class names (HuggingFace transformers naming)
|
|
_VIT_ATTN_CLASSNAMES = {
|
|
"ViTAttention",
|
|
"CLIPAttention",
|
|
"SiglipAttention",
|
|
"InternVisionAttention",
|
|
"Qwen2VLVisionAttention",
|
|
"Idefics2VisionAttention",
|
|
"PaliGemmaVisionAttention",
|
|
}
|
|
|
|
# Whether each known ViT class uses a prepended CLS token.
|
|
# CLS is replicated on every rank and is NOT sharded across the sequence.
|
|
# Defaults to True for unknown classes (safe fallback).
|
|
_VIT_HAS_CLS_TOKEN = {
|
|
"ViTAttention": True,
|
|
"CLIPAttention": True,
|
|
"SiglipAttention": False,
|
|
"InternVisionAttention": False,
|
|
"Qwen2VLVisionAttention": False,
|
|
"Idefics2VisionAttention": False,
|
|
"PaliGemmaVisionAttention": False,
|
|
}
|
|
|
|
# Known LLM decoder attention class names
|
|
_LLM_ATTN_CLASSNAMES = {
|
|
"LlamaAttention",
|
|
"MistralAttention",
|
|
"Qwen2Attention",
|
|
"InternLM2Attention",
|
|
"GemmaAttention",
|
|
"Phi3Attention",
|
|
"GPTNeoXAttention",
|
|
"FalconAttention",
|
|
"MptAttention",
|
|
}
|
|
|
|
# Common attribute names that hold the vision-language projection layer
|
|
_VISION_PROJ_KEYWORDS = (
|
|
"visual_projection",
|
|
"mm_projector",
|
|
"vision_proj",
|
|
"multi_modal_projector",
|
|
"img_projection",
|
|
)
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Data structures
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@dataclass
|
|
class SPModelInfo:
|
|
"""Holds the detection results for a multimodal model."""
|
|
|
|
# (dotted_name, module) pairs for ViT attention layers
|
|
vit_attn_modules: List[Tuple[str, nn.Module]] = field(default_factory=list)
|
|
# (dotted_name, module) pairs for LLM decoder attention layers
|
|
llm_attn_modules: List[Tuple[str, nn.Module]] = field(default_factory=list)
|
|
# (dotted_name, module) for the outermost vision-language projection layer
|
|
vision_projection_module: Optional[Tuple[str, nn.Module]] = None
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Detection logic
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def detect_model_sp_info(model: nn.Module) -> SPModelInfo:
|
|
"""Recursively scan *model* and return an :class:`SPModelInfo`.
|
|
|
|
The function identifies:
|
|
* ViT encoder attention layers → wrapped with :class:`UlyssesSPViTAttention`
|
|
* LLM decoder attention layers → wrapped with :class:`DistributedAttention`
|
|
* The vision-language projection layer → wrapped with
|
|
:class:`ModalityFusionSPAdapter` (Phase 2)
|
|
|
|
To add support for a new architecture, simply register its attention class
|
|
names in ``_VIT_ATTN_CLASSNAMES`` or ``_LLM_ATTN_CLASSNAMES``.
|
|
"""
|
|
info = SPModelInfo()
|
|
for name, module in model.named_modules():
|
|
cls_name = type(module).__name__
|
|
if cls_name in _VIT_ATTN_CLASSNAMES:
|
|
info.vit_attn_modules.append((name, module))
|
|
elif cls_name in _LLM_ATTN_CLASSNAMES:
|
|
info.llm_attn_modules.append((name, module))
|
|
|
|
# Record only the first (outermost) match to avoid double-wrapping
|
|
# nested projection modules.
|
|
if info.vision_projection_module is None:
|
|
if any(kw in name for kw in _VISION_PROJ_KEYWORDS):
|
|
info.vision_projection_module = (name, module)
|
|
|
|
return info
|