wehub-resource-sync
278 lines
12 KiB
Python
278 lines
12 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# DeepSpeed Team
|
|
"""
|
|
End-to-end integration tests for AutoSP multimodal sequence parallelism.
|
|
|
|
Each test builds a minimal mock model whose attention-layer class names match
|
|
the autosp_detector registry, then verifies two things:
|
|
|
|
1. auto_wrap_model_for_sp correctly identifies and wraps ViT attention modules
|
|
(with the correct has_cls_token value from the registry) and emits warnings
|
|
for HF-style LLM attention without wrapping them.
|
|
2. The full pipeline (SP-wrapped ViT -> fusion adapter) produces fused output
|
|
numerically equivalent to the single-device splice reference.
|
|
|
|
These tests require 2 GPUs.
|
|
Run with:
|
|
|
|
NCCL_P2P_DISABLE=1 python -m pytest tests/unit/sequence_parallelism/test_autosp_integration.py -v
|
|
"""
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
|
|
import deepspeed.comm as dist
|
|
from deepspeed.sequence.auto_sp import auto_wrap_model_for_sp
|
|
from deepspeed.sequence.autosp_vit import UlyssesSPViTAttention
|
|
from deepspeed.sequence.autosp_fusion import InternVLFusionAdapter, Qwen2VLFusionAdapter
|
|
from deepspeed.accelerator import get_accelerator
|
|
|
|
from unit.common import DistributedTest
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Token IDs
|
|
# ---------------------------------------------------------------------------
|
|
|
|
_INTERNVL_CONTEXT_ID = 92546
|
|
_QWEN2VL_START_ID = 151652
|
|
_QWEN2VL_END_ID = 151653
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Mock attention classes
|
|
#
|
|
# Class names must match exactly the entries in autosp_detector._VIT_ATTN_CLASSNAMES
|
|
# and _LLM_ATTN_CLASSNAMES so that auto_wrap_model_for_sp detects them.
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class InternVisionAttention(nn.Module):
|
|
"""Mock ViT attention for InternVL (registered in _VIT_ATTN_CLASSNAMES)."""
|
|
|
|
def forward(self, hidden_states, **kwargs):
|
|
return hidden_states
|
|
|
|
|
|
class InternLM2Attention(nn.Module):
|
|
"""Mock LLM attention for InternVL (registered in _LLM_ATTN_CLASSNAMES)."""
|
|
|
|
def forward(self, hidden_states, **kwargs):
|
|
return hidden_states
|
|
|
|
|
|
class Qwen2VLVisionAttention(nn.Module):
|
|
"""Mock ViT attention for Qwen2-VL (registered in _VIT_ATTN_CLASSNAMES)."""
|
|
|
|
def forward(self, hidden_states, **kwargs):
|
|
return hidden_states
|
|
|
|
|
|
class Qwen2Attention(nn.Module):
|
|
"""Mock LLM attention for Qwen2-VL (registered in _LLM_ATTN_CLASSNAMES)."""
|
|
|
|
def forward(self, hidden_states, **kwargs):
|
|
return hidden_states
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Model skeleton helpers
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class _AttnLayer(nn.Module):
|
|
"""Generic transformer block that holds an attention submodule.
|
|
|
|
auto_wrap_model_for_sp scans named_modules() and replaces ``self.attn``
|
|
when its class name is in the detector's registry.
|
|
"""
|
|
|
|
def __init__(self, attn: nn.Module) -> None:
|
|
super().__init__()
|
|
self.attn = attn
|
|
|
|
def forward(self, x, **kwargs):
|
|
return self.attn(x, **kwargs)
|
|
|
|
|
|
class _MinimalInternVLModel(nn.Module):
|
|
"""Minimal InternVL-like skeleton for integration testing.
|
|
|
|
Module paths recognised by autosp_detector:
|
|
- ``vision_encoder.0.attn`` -> InternVisionAttention (_VIT_ATTN_CLASSNAMES)
|
|
- ``language_model.0.attn`` -> InternLM2Attention (_LLM_ATTN_CLASSNAMES)
|
|
- ``mm_projector`` -> keyword in _VISION_PROJ_KEYWORDS
|
|
|
|
``forward`` exercises only the ViT + fusion path; ``language_model`` is
|
|
present to verify that auto_wrap does NOT wrap HF-style LLM attention.
|
|
"""
|
|
|
|
def __init__(self) -> None:
|
|
super().__init__()
|
|
self.vision_encoder = nn.Sequential(_AttnLayer(InternVisionAttention()))
|
|
self.mm_projector = nn.Identity()
|
|
self.language_model = nn.Sequential(_AttnLayer(InternLM2Attention()))
|
|
self.fusion = None
|
|
|
|
def forward(self, local_patches: torch.Tensor, text_embeds: torch.Tensor, input_ids: torch.Tensor) -> torch.Tensor:
|
|
local_visual = self.vision_encoder(local_patches)
|
|
return self.fusion(local_visual, text_embeds, input_ids)
|
|
|
|
|
|
class _MinimalQwen2VLModel(nn.Module):
|
|
"""Minimal Qwen2-VL-like skeleton for integration testing.
|
|
|
|
Module paths recognised by autosp_detector:
|
|
- ``visual.0.attn`` -> Qwen2VLVisionAttention (_VIT_ATTN_CLASSNAMES)
|
|
- ``model.0.attn`` -> Qwen2Attention (_LLM_ATTN_CLASSNAMES)
|
|
- ``multi_modal_projector`` -> keyword in _VISION_PROJ_KEYWORDS
|
|
"""
|
|
|
|
def __init__(self) -> None:
|
|
super().__init__()
|
|
self.visual = nn.Sequential(_AttnLayer(Qwen2VLVisionAttention()))
|
|
self.multi_modal_projector = nn.Identity()
|
|
self.model = nn.Sequential(_AttnLayer(Qwen2Attention()))
|
|
self.fusion = None
|
|
|
|
def forward(self, local_patches: torch.Tensor, text_embeds: torch.Tensor, input_ids: torch.Tensor) -> torch.Tensor:
|
|
local_visual = self.visual(local_patches)
|
|
return self.fusion(local_visual, text_embeds, input_ids)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# InternVL integration tests
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestInternVLIntegration(DistributedTest):
|
|
"""Integration tests for the InternVL multimodal SP pipeline."""
|
|
|
|
world_size = 2
|
|
|
|
def test_auto_wrap_detects_and_wraps_modules(self):
|
|
"""auto_wrap_model_for_sp must replace InternVisionAttention with
|
|
UlyssesSPViTAttention (has_cls_token=False) and must NOT wrap
|
|
InternLM2Attention (HF-style, incompatible with DistributedAttention)."""
|
|
sp_group = dist.new_group(ranks=list(range(self.world_size)))
|
|
model = _MinimalInternVLModel().to(get_accelerator().device_name())
|
|
auto_wrap_model_for_sp(model, sp_group)
|
|
|
|
assert isinstance(
|
|
model.vision_encoder[0].attn,
|
|
UlyssesSPViTAttention), ("Expected vision_encoder[0].attn to be UlyssesSPViTAttention after auto_wrap")
|
|
assert not model.vision_encoder[0].attn.has_cls_token, (
|
|
"InternVisionAttention has no CLS token; has_cls_token must be False")
|
|
assert isinstance(model.language_model[0].attn,
|
|
InternLM2Attention), ("HF-style LLM attention must NOT be wrapped by auto_wrap")
|
|
|
|
def test_full_pipeline_visual_to_fused(self):
|
|
"""SP-wrapped ViT -> InternVLFusionAdapter must produce fused output
|
|
numerically equivalent to the single-device splice reference."""
|
|
sp_group = dist.new_group(ranks=list(range(self.world_size)))
|
|
rank = dist.get_rank(sp_group)
|
|
|
|
bs, local_v, text_len, hidden = 1, 4, 10, 8
|
|
num_ctx = local_v * self.world_size
|
|
|
|
torch.manual_seed(20)
|
|
full_visual = torch.randn(bs, local_v * self.world_size, hidden).to(get_accelerator().device_name())
|
|
text = torch.randn(bs, text_len, hidden).to(get_accelerator().device_name())
|
|
ids = torch.zeros(bs, text_len, dtype=torch.long).to(get_accelerator().device_name())
|
|
ids[:, 2:2 + num_ctx] = _INTERNVL_CONTEXT_ID
|
|
|
|
local_patches = full_visual[:, rank * local_v:(rank + 1) * local_v, :]
|
|
|
|
model = _MinimalInternVLModel().to(get_accelerator().device_name())
|
|
auto_wrap_model_for_sp(model, sp_group)
|
|
model.fusion = InternVLFusionAdapter(model.mm_projector, sp_group,
|
|
image_token_id=_INTERNVL_CONTEXT_ID).to(get_accelerator().device_name())
|
|
|
|
local_out = model(local_patches, text, ids)
|
|
|
|
gathered = [torch.zeros_like(local_out) for _ in range(self.world_size)]
|
|
dist.all_gather(gathered, local_out, group=sp_group)
|
|
full_sp_out = torch.cat(gathered, dim=1)
|
|
|
|
# Single-device reference: splice without SP scatter.
|
|
ref_adapter = InternVLFusionAdapter(nn.Identity(), sp_group,
|
|
image_token_id=_INTERNVL_CONTEXT_ID).to(get_accelerator().device_name())
|
|
ref_fused = ref_adapter._splice_visual_into_text(text, full_visual, ids)
|
|
pad = (self.world_size - ref_fused.shape[1] % self.world_size) % self.world_size
|
|
if pad > 0:
|
|
ref_fused = F.pad(ref_fused, (0, 0, 0, pad))
|
|
|
|
assert torch.allclose(full_sp_out, ref_fused,
|
|
atol=1e-5), (f"rank={rank} InternVL full pipeline output differs from reference: "
|
|
f"max_diff={(full_sp_out - ref_fused).abs().max().item():.2e}")
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Qwen2-VL integration tests
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestQwen2VLIntegration(DistributedTest):
|
|
"""Integration tests for the Qwen2-VL multimodal SP pipeline."""
|
|
|
|
world_size = 2
|
|
|
|
def test_auto_wrap_detects_and_wraps_modules(self):
|
|
"""auto_wrap_model_for_sp must replace Qwen2VLVisionAttention with
|
|
UlyssesSPViTAttention (has_cls_token=False) and must NOT wrap
|
|
Qwen2Attention (HF-style, incompatible with DistributedAttention)."""
|
|
sp_group = dist.new_group(ranks=list(range(self.world_size)))
|
|
model = _MinimalQwen2VLModel().to(get_accelerator().device_name())
|
|
auto_wrap_model_for_sp(model, sp_group)
|
|
|
|
assert isinstance(
|
|
model.visual[0].attn,
|
|
UlyssesSPViTAttention), ("Expected visual[0].attn to be UlyssesSPViTAttention after auto_wrap")
|
|
assert not model.visual[0].attn.has_cls_token, (
|
|
"Qwen2VLVisionAttention has no CLS token; has_cls_token must be False")
|
|
assert isinstance(model.model[0].attn,
|
|
Qwen2Attention), ("HF-style LLM attention must NOT be wrapped by auto_wrap")
|
|
|
|
def test_full_pipeline_visual_to_fused(self):
|
|
"""SP-wrapped ViT -> Qwen2VLFusionAdapter must produce fused output
|
|
numerically equivalent to the single-device splice reference."""
|
|
sp_group = dist.new_group(ranks=list(range(self.world_size)))
|
|
rank = dist.get_rank(sp_group)
|
|
|
|
bs, local_v, text_len, hidden = 1, 3, 10, 8
|
|
num_inner = local_v * self.world_size
|
|
|
|
torch.manual_seed(21)
|
|
full_visual = torch.randn(bs, local_v * self.world_size, hidden).to(get_accelerator().device_name())
|
|
text = torch.randn(bs, text_len, hidden).to(get_accelerator().device_name())
|
|
ids = torch.zeros(bs, text_len, dtype=torch.long).to(get_accelerator().device_name())
|
|
ids[:, 1] = _QWEN2VL_START_ID
|
|
ids[:, 2 + num_inner] = _QWEN2VL_END_ID
|
|
|
|
local_patches = full_visual[:, rank * local_v:(rank + 1) * local_v, :]
|
|
|
|
model = _MinimalQwen2VLModel().to(get_accelerator().device_name())
|
|
auto_wrap_model_for_sp(model, sp_group)
|
|
model.fusion = Qwen2VLFusionAdapter(model.multi_modal_projector,
|
|
sp_group,
|
|
vision_start_token_id=_QWEN2VL_START_ID,
|
|
vision_end_token_id=_QWEN2VL_END_ID).to(get_accelerator().device_name())
|
|
|
|
local_out = model(local_patches, text, ids)
|
|
|
|
gathered = [torch.zeros_like(local_out) for _ in range(self.world_size)]
|
|
dist.all_gather(gathered, local_out, group=sp_group)
|
|
full_sp_out = torch.cat(gathered, dim=1)
|
|
|
|
ref_adapter = Qwen2VLFusionAdapter(nn.Identity(),
|
|
sp_group,
|
|
vision_start_token_id=_QWEN2VL_START_ID,
|
|
vision_end_token_id=_QWEN2VL_END_ID).to(get_accelerator().device_name())
|
|
ref_fused = ref_adapter._splice_visual_into_text(text, full_visual, ids)
|
|
pad = (self.world_size - ref_fused.shape[1] % self.world_size) % self.world_size
|
|
if pad > 0:
|
|
ref_fused = F.pad(ref_fused, (0, 0, 0, pad))
|
|
|
|
assert torch.allclose(full_sp_out, ref_fused,
|
|
atol=1e-5), (f"rank={rank} Qwen2VL full pipeline output differs from reference: "
|
|
f"max_diff={(full_sp_out - ref_fused).abs().max().item():.2e}")
|