wehub-resource-sync
276 lines
11 KiB
Python
276 lines
11 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# DeepSpeed Team
|
|
"""
|
|
Benchmark: AutoSP multimodal sequence parallelism (ViT SP + fusion adapter).
|
|
|
|
Measures per-iteration latency, throughput, and peak GPU memory for the
|
|
ViT-SP + fusion-adapter pipeline at a given SP degree.
|
|
|
|
Launch (from repo root):
|
|
|
|
# SP degree 2 — two GPUs:
|
|
NCCL_P2P_DISABLE=1 torchrun --nproc_per_node=2 \\
|
|
benchmarks/autosp/bench_multimodal_sp.py [args]
|
|
|
|
# Baseline — single GPU (all-gather/scatter are no-ops):
|
|
torchrun --nproc_per_node=1 \\
|
|
benchmarks/autosp/bench_multimodal_sp.py [args]
|
|
|
|
Compare the two output tables to quantify memory savings and throughput scaling.
|
|
|
|
Arguments:
|
|
--arch {internvl, qwen2vl} architecture to simulate (default: internvl)
|
|
--batch-size N samples per batch (default: 2)
|
|
--seq-len N text sequence length (default: 512)
|
|
--visual-tokens N total visual tokens per sample (default: 256)
|
|
--hidden N hidden dimension (default: 1024)
|
|
--num-layers N ViT and LLM layers each (default: 2)
|
|
--iters N measured iterations (default: 50)
|
|
--warmup N warmup iterations (default: 10)
|
|
"""
|
|
|
|
import argparse
|
|
import logging
|
|
import statistics
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
|
|
import deepspeed
|
|
import deepspeed.comm as dist
|
|
from deepspeed.accelerator import get_accelerator
|
|
from deepspeed.sequence.auto_sp import auto_wrap_model_for_sp
|
|
from deepspeed.sequence.autosp_fusion import InternVLFusionAdapter, Qwen2VLFusionAdapter
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Token IDs
|
|
# ---------------------------------------------------------------------------
|
|
|
|
_INTERNVL_CONTEXT_ID = 92546
|
|
_QWEN2VL_START_ID = 151652
|
|
_QWEN2VL_END_ID = 151653
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Mock attention classes — names match autosp_detector registries exactly
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class InternVisionAttention(nn.Module):
|
|
|
|
def forward(self, hidden_states, **kwargs):
|
|
return hidden_states
|
|
|
|
|
|
class InternLM2Attention(nn.Module):
|
|
|
|
def forward(self, hidden_states, **kwargs):
|
|
return hidden_states
|
|
|
|
|
|
class Qwen2VLVisionAttention(nn.Module):
|
|
|
|
def forward(self, hidden_states, **kwargs):
|
|
return hidden_states
|
|
|
|
|
|
class Qwen2Attention(nn.Module):
|
|
|
|
def forward(self, hidden_states, **kwargs):
|
|
return hidden_states
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Model building blocks
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class _ViTBlock(nn.Module):
|
|
"""One ViT transformer block: attention (to be SP-wrapped) + linear FFN."""
|
|
|
|
def __init__(self, attn_cls, hidden: int) -> None:
|
|
super().__init__()
|
|
self.attn = attn_cls()
|
|
self.ffn = nn.Linear(hidden, hidden, bias=False)
|
|
|
|
def forward(self, x, **kwargs):
|
|
out = self.attn(x, **kwargs)
|
|
if isinstance(out, (tuple, list)):
|
|
out = out[0]
|
|
return self.ffn(out)
|
|
|
|
|
|
class _MinimalInternVLModel(nn.Module):
|
|
"""InternVL-like benchmark model.
|
|
|
|
Module paths detected by autosp_detector:
|
|
- ``vision_encoder.*.attn`` -> InternVisionAttention (_VIT_ATTN_CLASSNAMES)
|
|
- ``mm_projector`` -> keyword in _VISION_PROJ_KEYWORDS
|
|
|
|
``language_model`` uses plain nn.Linear layers so it is NOT wrapped by
|
|
DistributedAttention (avoids the Q/K/V interface requirement) yet still
|
|
contributes realistic compute on the scattered fused sequence.
|
|
"""
|
|
|
|
def __init__(self, hidden: int, num_layers: int) -> None:
|
|
super().__init__()
|
|
self.vision_encoder = nn.Sequential(*[_ViTBlock(InternVisionAttention, hidden) for _ in range(num_layers)])
|
|
self.mm_projector = nn.Identity()
|
|
self.language_model = nn.Sequential(*[nn.Linear(hidden, hidden, bias=False) for _ in range(num_layers)])
|
|
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)
|
|
local_fused = self.fusion(local_visual, text_embeds, input_ids)
|
|
return self.language_model(local_fused)
|
|
|
|
|
|
class _MinimalQwen2VLModel(nn.Module):
|
|
"""Qwen2VL-like benchmark model."""
|
|
|
|
def __init__(self, hidden: int, num_layers: int) -> None:
|
|
super().__init__()
|
|
self.visual = nn.Sequential(*[_ViTBlock(Qwen2VLVisionAttention, hidden) for _ in range(num_layers)])
|
|
self.multi_modal_projector = nn.Identity()
|
|
self.model = nn.Sequential(*[nn.Linear(hidden, hidden, bias=False) for _ in range(num_layers)])
|
|
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)
|
|
local_fused = self.fusion(local_visual, text_embeds, input_ids)
|
|
return self.model(local_fused)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Setup helpers
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def _build_model_and_inputs(arch: str, args, sp_group, device):
|
|
rank = dist.get_rank(sp_group)
|
|
world_size = dist.get_world_size(sp_group)
|
|
|
|
local_v = args.visual_tokens // world_size
|
|
bs, text_len, hidden = args.batch_size, args.seq_len, args.hidden
|
|
|
|
torch.manual_seed(0)
|
|
local_patches = torch.randn(bs, local_v, hidden, device=device)
|
|
text_embeds = torch.randn(bs, text_len, hidden, device=device)
|
|
input_ids = torch.zeros(bs, text_len, dtype=torch.long, device=device)
|
|
|
|
if arch == "internvl":
|
|
num_ctx = min(local_v * world_size, text_len - 2)
|
|
input_ids[:, 2:2 + num_ctx] = _INTERNVL_CONTEXT_ID
|
|
|
|
model = _MinimalInternVLModel(hidden, args.num_layers).to(device)
|
|
# Suppress the Phase 2 projection-layer warning: we wrap manually below.
|
|
_auto_sp_logger = logging.getLogger("deepspeed.sequence.auto_sp")
|
|
_prev_level = _auto_sp_logger.level
|
|
_auto_sp_logger.setLevel(logging.ERROR)
|
|
auto_wrap_model_for_sp(model, sp_group)
|
|
_auto_sp_logger.setLevel(_prev_level)
|
|
model.fusion = InternVLFusionAdapter(model.mm_projector, sp_group,
|
|
image_token_id=_INTERNVL_CONTEXT_ID).to(device)
|
|
else: # qwen2vl
|
|
num_inner = min(local_v * world_size, text_len - 3)
|
|
input_ids[:, 1] = _QWEN2VL_START_ID
|
|
input_ids[:, 2 + num_inner] = _QWEN2VL_END_ID
|
|
|
|
model = _MinimalQwen2VLModel(hidden, args.num_layers).to(device)
|
|
_auto_sp_logger = logging.getLogger("deepspeed.sequence.auto_sp")
|
|
_prev_level = _auto_sp_logger.level
|
|
_auto_sp_logger.setLevel(logging.ERROR)
|
|
auto_wrap_model_for_sp(model, sp_group)
|
|
_auto_sp_logger.setLevel(_prev_level)
|
|
model.fusion = Qwen2VLFusionAdapter(model.multi_modal_projector,
|
|
sp_group,
|
|
vision_start_token_id=_QWEN2VL_START_ID,
|
|
vision_end_token_id=_QWEN2VL_END_ID).to(device)
|
|
|
|
return model, local_patches, text_embeds, input_ids
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Benchmark runner
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def _run(arch: str, args) -> None:
|
|
deepspeed.init_distributed(dist_backend="nccl")
|
|
rank = dist.get_rank()
|
|
world_size = dist.get_world_size()
|
|
device = torch.device(get_accelerator().device_name(), rank % get_accelerator().device_count())
|
|
get_accelerator().set_device(rank % get_accelerator().device_count())
|
|
|
|
sp_group = dist.new_group(ranks=list(range(world_size)))
|
|
model, local_patches, text_embeds, input_ids = _build_model_and_inputs(arch, args, sp_group, device)
|
|
model.eval()
|
|
|
|
# Warmup
|
|
with torch.no_grad():
|
|
for _ in range(args.warmup):
|
|
model(local_patches, text_embeds, input_ids)
|
|
get_accelerator().synchronize()
|
|
get_accelerator().reset_peak_memory_stats()
|
|
|
|
# Timed iterations using CUDA events for accurate GPU-side measurement.
|
|
latencies_ms = []
|
|
with torch.no_grad():
|
|
for _ in range(args.iters):
|
|
t_start = get_accelerator().Event(enable_timing=True)
|
|
t_end = get_accelerator().Event(enable_timing=True)
|
|
t_start.record()
|
|
model(local_patches, text_embeds, input_ids)
|
|
t_end.record()
|
|
get_accelerator().synchronize()
|
|
latencies_ms.append(t_start.elapsed_time(t_end))
|
|
|
|
peak_mem_mb = get_accelerator().max_memory_allocated() / 1024**2
|
|
mean_ms = statistics.mean(latencies_ms)
|
|
std_ms = statistics.stdev(latencies_ms) if len(latencies_ms) > 1 else 0.0
|
|
# tokens/s: fused sequence length approximated by seq_len (length-preserving adapters).
|
|
throughput = (args.batch_size * args.seq_len) / (mean_ms / 1000.0)
|
|
|
|
if rank == 0:
|
|
sep = "=" * 62
|
|
print(f"\n{sep}")
|
|
print(f" AutoSP Benchmark arch={arch} sp_degree={world_size}")
|
|
print(sep)
|
|
print(f" batch_size : {args.batch_size}")
|
|
print(f" seq_len : {args.seq_len}")
|
|
print(f" visual_tokens : {args.visual_tokens} (local={args.visual_tokens // world_size}/rank)")
|
|
print(f" hidden : {args.hidden}")
|
|
print(f" num_layers : {args.num_layers}")
|
|
print(f" warmup / iters : {args.warmup} / {args.iters}")
|
|
print(f" {'─' * 58}")
|
|
print(f" Latency : {mean_ms:.2f} ± {std_ms:.2f} ms/iter")
|
|
print(f" Throughput : {throughput:,.0f} tokens/s")
|
|
print(f" Peak GPU memory : {peak_mem_mb:.1f} MB")
|
|
print(f"{sep}\n")
|
|
|
|
dist.destroy_process_group()
|
|
|
|
|
|
def main() -> None:
|
|
parser = argparse.ArgumentParser(description="AutoSP multimodal SP benchmark")
|
|
parser.add_argument("--arch",
|
|
choices=["internvl", "qwen2vl"],
|
|
default="internvl",
|
|
help="Model architecture to simulate")
|
|
parser.add_argument("--batch-size", type=int, default=2)
|
|
parser.add_argument("--seq-len", type=int, default=512)
|
|
parser.add_argument("--visual-tokens",
|
|
type=int,
|
|
default=256,
|
|
help="Total visual tokens (must be divisible by --nproc_per_node)")
|
|
parser.add_argument("--hidden", type=int, default=1024)
|
|
parser.add_argument("--num-layers", type=int, default=2, help="Number of ViT blocks and LLM linear layers each")
|
|
parser.add_argument("--iters", type=int, default=50)
|
|
parser.add_argument("--warmup", type=int, default=10)
|
|
args = parser.parse_args()
|
|
|
|
_run(args.arch, args)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|