chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,277 @@
|
||||
# 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}")
|
||||
Reference in New Issue
Block a user