chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,430 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# DeepSpeed Team
|
||||
"""
|
||||
Numerical equivalence tests for AutoSP multimodal sequence parallelism.
|
||||
|
||||
Each test verifies that running the SP-wrapped path across N ranks produces
|
||||
the same result as the equivalent single-device (non-SP) computation.
|
||||
|
||||
These tests require 2 GPUs.
|
||||
Run with:
|
||||
|
||||
NCCL_P2P_DISABLE=1 python -m pytest tests/unit/sequence_parallelism/test_autosp_equivalence.py -v
|
||||
"""
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
import pytest
|
||||
|
||||
import deepspeed.comm as dist
|
||||
from deepspeed.sequence.autosp_vit import UlyssesSPViTAttention
|
||||
from deepspeed.sequence.autosp_fusion import InternVLFusionAdapter, LlavaFusionAdapter, Qwen2VLFusionAdapter
|
||||
from deepspeed.accelerator import get_accelerator
|
||||
|
||||
from unit.common import DistributedTest
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Shared identity attention — deterministic, easy to verify
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
_IMAGE_TOKEN_ID = -200
|
||||
|
||||
|
||||
class _IdentityAttn(nn.Module):
|
||||
"""Returns hidden_states unchanged so that gather-compute-scatter is a no-op."""
|
||||
|
||||
def forward(self, hidden_states, **kwargs):
|
||||
return hidden_states
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# UlyssesSPViTAttention equivalence
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestViTSPEquivalence(DistributedTest):
|
||||
"""SP-wrapped ViT attention with an identity inner module must reproduce
|
||||
the unsharded output on every rank."""
|
||||
|
||||
world_size = 2
|
||||
|
||||
@pytest.mark.parametrize("has_cls_token", [True, False])
|
||||
@pytest.mark.parametrize("num_patches", [8, 12])
|
||||
def test_output_equals_single_device(self, has_cls_token, num_patches):
|
||||
"""Each rank's local output slice must match the corresponding slice of
|
||||
the single-device output."""
|
||||
sp_group = dist.new_group(ranks=list(range(self.world_size)))
|
||||
rank = dist.get_rank(sp_group)
|
||||
bs, hidden = 2, 32
|
||||
|
||||
# --- Single-device reference ---
|
||||
# Build the full input (all ranks see the same RNG seed so the tensor
|
||||
# is identical everywhere).
|
||||
torch.manual_seed(42)
|
||||
if has_cls_token:
|
||||
full_input = torch.randn(bs, 1 + num_patches, hidden).to(get_accelerator().device_name())
|
||||
else:
|
||||
full_input = torch.randn(bs, num_patches, hidden).to(get_accelerator().device_name())
|
||||
|
||||
identity = _IdentityAttn().to(get_accelerator().device_name())
|
||||
# Single-device path is just identity — output == input.
|
||||
ref_out = identity(full_input)
|
||||
|
||||
# --- SP path ---
|
||||
local_patches = num_patches // self.world_size
|
||||
if has_cls_token:
|
||||
cls = full_input[:, :1, :]
|
||||
patch_slice = full_input[:, 1 + rank * local_patches:1 + (rank + 1) * local_patches, :]
|
||||
local_input = torch.cat([cls, patch_slice], dim=1)
|
||||
else:
|
||||
local_input = full_input[:, rank * local_patches:(rank + 1) * local_patches, :]
|
||||
|
||||
wrapper = UlyssesSPViTAttention(_IdentityAttn().to(get_accelerator().device_name()),
|
||||
sp_group,
|
||||
has_cls_token=has_cls_token).to(get_accelerator().device_name())
|
||||
sp_out = wrapper(local_input)
|
||||
|
||||
# --- Compare ---
|
||||
# sp_out is the local slice; reconstruct what slice of ref_out it maps to.
|
||||
if has_cls_token:
|
||||
ref_slice = torch.cat(
|
||||
[ref_out[:, :1, :], ref_out[:, 1 + rank * local_patches:1 + (rank + 1) * local_patches, :]], dim=1)
|
||||
else:
|
||||
ref_slice = ref_out[:, rank * local_patches:(rank + 1) * local_patches, :]
|
||||
|
||||
assert torch.allclose(sp_out, ref_slice,
|
||||
atol=1e-5), (f"rank={rank} sp_out differs from reference: "
|
||||
f"max_diff={( sp_out - ref_slice).abs().max().item():.2e}")
|
||||
|
||||
@pytest.mark.parametrize("has_cls_token", [True, False])
|
||||
def test_noneven_patches(self, has_cls_token):
|
||||
"""When num_patches % world_size != 0, the wrapper must still produce
|
||||
correct per-rank output. With 5 patches and world_size=2, rank 0
|
||||
holds 3 patches and rank 1 holds 2 patches."""
|
||||
sp_group = dist.new_group(ranks=list(range(self.world_size)))
|
||||
rank = dist.get_rank(sp_group)
|
||||
bs, hidden = 2, 16
|
||||
num_patches = 5 # not divisible by world_size=2
|
||||
|
||||
torch.manual_seed(77)
|
||||
if has_cls_token:
|
||||
full_input = torch.randn(bs, 1 + num_patches, hidden).to(get_accelerator().device_name())
|
||||
else:
|
||||
full_input = torch.randn(bs, num_patches, hidden).to(get_accelerator().device_name())
|
||||
|
||||
# Distribute: first (num_patches % world_size) ranks carry one extra patch.
|
||||
extra = num_patches % self.world_size # = 1
|
||||
base = num_patches // self.world_size # = 2
|
||||
local_v = base + (1 if rank < extra else 0)
|
||||
patch_start = rank * base + min(rank, extra)
|
||||
|
||||
if has_cls_token:
|
||||
cls = full_input[:, :1, :]
|
||||
patch_slice = full_input[:, 1 + patch_start:1 + patch_start + local_v, :]
|
||||
local_input = torch.cat([cls, patch_slice], dim=1)
|
||||
else:
|
||||
local_input = full_input[:, patch_start:patch_start + local_v, :]
|
||||
|
||||
wrapper = UlyssesSPViTAttention(_IdentityAttn().to(get_accelerator().device_name()),
|
||||
sp_group,
|
||||
has_cls_token=has_cls_token)
|
||||
sp_out = wrapper(local_input)
|
||||
|
||||
# Reference: identity wrapper — each rank's output must equal its input slice.
|
||||
if has_cls_token:
|
||||
ref_slice = torch.cat([full_input[:, :1, :], full_input[:, 1 + patch_start:1 + patch_start + local_v, :]],
|
||||
dim=1)
|
||||
else:
|
||||
ref_slice = full_input[:, patch_start:patch_start + local_v, :]
|
||||
|
||||
assert torch.allclose(sp_out, ref_slice,
|
||||
atol=1e-5), (f"rank={rank} non-even patches: sp_out differs from reference: "
|
||||
f"max_diff={(sp_out - ref_slice).abs().max().item():.2e}")
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# LlavaFusionAdapter equivalence
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestLlavaFusionEquivalence(DistributedTest):
|
||||
"""Verifies that the SP gather/scatter in LlavaFusionAdapter is a lossless
|
||||
round-trip: concatenating all ranks' output shards reproduces the full
|
||||
fused sequence that single-device splicing would produce."""
|
||||
|
||||
world_size = 2
|
||||
|
||||
def _build_inputs(self, bs, local_v, text_len, hidden, rank):
|
||||
"""Build deterministic visual and text tensors identical on every rank."""
|
||||
torch.manual_seed(0)
|
||||
# Each rank holds a contiguous slice of the visual tokens.
|
||||
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] = _IMAGE_TOKEN_ID # one image placeholder at position 1
|
||||
local_visual = full_visual[:, rank * local_v:(rank + 1) * local_v, :]
|
||||
return full_visual, local_visual, text, ids
|
||||
|
||||
def test_shards_reassemble_to_full_fused(self):
|
||||
"""Gathering all ranks' output shards must equal the single-device
|
||||
fused sequence (modulo padding zeros)."""
|
||||
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, 6, 8
|
||||
full_visual, local_visual, text, ids = self._build_inputs(bs, local_v, text_len, hidden, rank)
|
||||
|
||||
# --- SP path: each rank gets one shard ---
|
||||
adapter = LlavaFusionAdapter(nn.Identity(), sp_group,
|
||||
image_token_id=_IMAGE_TOKEN_ID).to(get_accelerator().device_name())
|
||||
local_out = adapter(local_visual, text, ids) # [bs, local_fused, hidden]
|
||||
|
||||
# Gather all shards onto every rank so we can compare globally.
|
||||
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) # [bs, padded_fused, hidden]
|
||||
|
||||
# --- Single-device reference ---
|
||||
# Simulate what a non-SP LlavaFusionAdapter would produce: project the
|
||||
# full visual tensor (identity here) and splice once.
|
||||
ref_adapter = LlavaFusionAdapter(nn.Identity(), sp_group,
|
||||
image_token_id=_IMAGE_TOKEN_ID).to(get_accelerator().device_name())
|
||||
# Call _splice_visual_into_text directly so we bypass the SP scatter.
|
||||
ref_fused = ref_adapter._splice_visual_into_text(text, full_visual, ids)
|
||||
|
||||
# Pad reference to the same padded length.
|
||||
fused_len = ref_fused.shape[1]
|
||||
pad = (self.world_size - fused_len % 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} reassembled SP output differs from reference: "
|
||||
f"max_diff={( full_sp_out - ref_fused).abs().max().item():.2e}")
|
||||
|
||||
def test_no_image_token_passthrough(self):
|
||||
"""When there are no image placeholders the SP fused output must equal
|
||||
the sharded text after padding/scatter (all-text path)."""
|
||||
sp_group = dist.new_group(ranks=list(range(self.world_size)))
|
||||
rank = dist.get_rank(sp_group)
|
||||
|
||||
bs, local_v, text_len, hidden = 1, 2, 8, 4
|
||||
torch.manual_seed(1)
|
||||
local_visual = torch.randn(bs, local_v, 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()) # no image placeholder
|
||||
|
||||
adapter = LlavaFusionAdapter(nn.Identity(), sp_group,
|
||||
image_token_id=_IMAGE_TOKEN_ID).to(get_accelerator().device_name())
|
||||
local_out = adapter(local_visual, text, ids)
|
||||
|
||||
# Gather shards and strip the padding slice from visual gather.
|
||||
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)
|
||||
|
||||
# Expected: when there is no image token, the visual tokens are ignored.
|
||||
# So the fused output should just be the text tokens.
|
||||
ref_fused = text
|
||||
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} no-image path differs from reference: "
|
||||
f"max_diff={( full_sp_out - ref_fused).abs().max().item():.2e}")
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# InternVLFusionAdapter equivalence
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
_INTERNVL_CONTEXT_TOKEN_ID = 92546
|
||||
|
||||
|
||||
class TestInternVLFusionEquivalence(DistributedTest):
|
||||
"""Verifies that the SP gather/scatter in InternVLFusionAdapter is a lossless
|
||||
round-trip: concatenating all ranks' output shards reproduces the full fused
|
||||
sequence that single-device splicing would produce.
|
||||
|
||||
InternVL replaces IMG_CONTEXT tokens 1-to-1 with visual tokens, so the
|
||||
sequence length is preserved.
|
||||
"""
|
||||
|
||||
world_size = 2
|
||||
|
||||
def _build_inputs(self, bs, local_v, text_len, hidden, rank, num_ctx_tokens):
|
||||
"""Build deterministic inputs with a run of IMG_CONTEXT tokens in the middle."""
|
||||
torch.manual_seed(2)
|
||||
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())
|
||||
# Place IMG_CONTEXT tokens starting at position 2.
|
||||
ids[:, 2:2 + num_ctx_tokens] = _INTERNVL_CONTEXT_TOKEN_ID
|
||||
local_visual = full_visual[:, rank * local_v:(rank + 1) * local_v, :]
|
||||
return full_visual, local_visual, text, ids
|
||||
|
||||
def test_shards_reassemble_to_full_fused(self):
|
||||
"""Gathering all ranks' output shards must equal the single-device
|
||||
fused sequence (modulo padding zeros)."""
|
||||
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, 8, 4
|
||||
full_visual, local_visual, text, ids = self._build_inputs(bs,
|
||||
local_v,
|
||||
text_len,
|
||||
hidden,
|
||||
rank,
|
||||
num_ctx_tokens=local_v * self.world_size)
|
||||
|
||||
# SP path.
|
||||
adapter = InternVLFusionAdapter(nn.Identity(), sp_group,
|
||||
image_token_id=_INTERNVL_CONTEXT_TOKEN_ID).to(get_accelerator().device_name())
|
||||
local_out = adapter(local_visual, 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.
|
||||
ref_adapter = InternVLFusionAdapter(nn.Identity(), sp_group, image_token_id=_INTERNVL_CONTEXT_TOKEN_ID).to(
|
||||
get_accelerator().device_name())
|
||||
ref_fused = ref_adapter._splice_visual_into_text(text, full_visual, ids)
|
||||
|
||||
fused_len = ref_fused.shape[1]
|
||||
pad = (self.world_size - fused_len % 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 reassembled output differs from reference: "
|
||||
f"max_diff={( full_sp_out - ref_fused).abs().max().item():.2e}")
|
||||
|
||||
def test_no_context_token_passthrough(self):
|
||||
"""When there are no IMG_CONTEXT tokens the fused output must equal the text."""
|
||||
sp_group = dist.new_group(ranks=list(range(self.world_size)))
|
||||
rank = dist.get_rank(sp_group)
|
||||
|
||||
bs, local_v, text_len, hidden = 1, 2, 6, 4
|
||||
torch.manual_seed(3)
|
||||
local_visual = torch.randn(bs, local_v, 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())
|
||||
|
||||
adapter = InternVLFusionAdapter(nn.Identity(), sp_group,
|
||||
image_token_id=_INTERNVL_CONTEXT_TOKEN_ID).to(get_accelerator().device_name())
|
||||
local_out = adapter(local_visual, 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_fused = text
|
||||
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 no-context path differs from reference: "
|
||||
f"max_diff={( full_sp_out - ref_fused).abs().max().item():.2e}")
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Qwen2VLFusionAdapter equivalence
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
_QWEN2VL_START_ID = 151652
|
||||
_QWEN2VL_END_ID = 151653
|
||||
|
||||
|
||||
class TestQwen2VLFusionEquivalence(DistributedTest):
|
||||
"""Verifies that the SP gather/scatter in Qwen2VLFusionAdapter is a lossless
|
||||
round-trip: concatenating all ranks' output shards reproduces the full fused
|
||||
sequence that single-device splicing would produce.
|
||||
|
||||
Qwen2-VL replaces inner placeholder tokens (between vision_start/end pairs)
|
||||
1-to-1 with visual tokens, so the sequence length is preserved.
|
||||
"""
|
||||
|
||||
world_size = 2
|
||||
|
||||
def _build_inputs(self, bs, local_v, text_len, hidden, rank, num_inner):
|
||||
"""Build inputs with a single vision_start/end block containing num_inner placeholders."""
|
||||
torch.manual_seed(4)
|
||||
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())
|
||||
# [t0, <vis_start>, pad×num_inner, <vis_end>, ...]
|
||||
ids[:, 1] = _QWEN2VL_START_ID
|
||||
ids[:, 2 + num_inner] = _QWEN2VL_END_ID
|
||||
local_visual = full_visual[:, rank * local_v:(rank + 1) * local_v, :]
|
||||
return full_visual, local_visual, text, ids
|
||||
|
||||
def test_shards_reassemble_to_full_fused(self):
|
||||
"""Gathering all ranks' output shards must equal the single-device
|
||||
fused sequence (modulo padding zeros)."""
|
||||
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, 4
|
||||
num_inner = local_v * self.world_size # inner placeholder count equals total visual tokens
|
||||
full_visual, local_visual, text, ids = self._build_inputs(bs, local_v, text_len, hidden, rank, num_inner)
|
||||
|
||||
# SP path.
|
||||
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())
|
||||
local_out = adapter(local_visual, 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.
|
||||
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)
|
||||
|
||||
fused_len = ref_fused.shape[1]
|
||||
pad = (self.world_size - fused_len % 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 reassembled output differs from reference: "
|
||||
f"max_diff={( full_sp_out - ref_fused).abs().max().item():.2e}")
|
||||
|
||||
def test_no_vision_token_passthrough(self):
|
||||
"""When there are no vision_start/end tokens the fused output must equal the text."""
|
||||
sp_group = dist.new_group(ranks=list(range(self.world_size)))
|
||||
rank = dist.get_rank(sp_group)
|
||||
|
||||
bs, local_v, text_len, hidden = 1, 2, 8, 4
|
||||
torch.manual_seed(5)
|
||||
local_visual = torch.randn(bs, local_v, 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())
|
||||
|
||||
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())
|
||||
local_out = adapter(local_visual, 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_fused = text
|
||||
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 no-vision path differs from reference: "
|
||||
f"max_diff={( full_sp_out - ref_fused).abs().max().item():.2e}")
|
||||
@@ -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}")
|
||||
@@ -0,0 +1,312 @@
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# DeepSpeed Team
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
import deepspeed.comm as dist
|
||||
from deepspeed import initialize
|
||||
from transformers import AutoModel
|
||||
from unit.common import DistributedTest
|
||||
from deepspeed.sequence.layer import _SeqAllToAll
|
||||
from deepspeed.sequence.fpdt_layer import _FPDTGPUOffloadingAttentionImpl_, FPDT_InputConstruct
|
||||
from unit.util import skip_on_arch
|
||||
from unit.simple_model import *
|
||||
from deepspeed.utils import groups
|
||||
from deepspeed.module_inject.tp_shard import get_shard_size_list
|
||||
#Use mesh device to create data and sequence parallel group
|
||||
|
||||
|
||||
class TestUlyssesUtils(DistributedTest):
|
||||
world_size = 4
|
||||
|
||||
def test_mesh_device_creation(self) -> None:
|
||||
skip_on_arch(min_arch=8)
|
||||
model = AutoModel.from_pretrained('bert-base-uncased')
|
||||
sp_size = 2
|
||||
dp_size = 2
|
||||
ds_engine, _, _, _ = initialize(
|
||||
model=model,
|
||||
config_params={
|
||||
"train_batch_size": 8,
|
||||
"data_parallel_size": dp_size,
|
||||
"sequence_parallel_size": sp_size
|
||||
},
|
||||
)
|
||||
assert ds_engine.seq_parallel_group is not None
|
||||
assert ds_engine.data_parallel_group is not None
|
||||
assert dist.get_world_size(group=ds_engine.seq_parallel_group) == sp_size
|
||||
assert dist.get_world_size(group=ds_engine.data_parallel_group) == dp_size
|
||||
assert dist.get_world_size() == sp_size * dp_size
|
||||
|
||||
|
||||
#Sweep b,s,h,d to test all2all consistency
|
||||
@pytest.mark.parametrize("d0", [2, 4]) #batch or sequence dimension
|
||||
@pytest.mark.parametrize("d1", [4, 8]) #batch or sequence dimension
|
||||
@pytest.mark.parametrize("num_heads", [4, 8])
|
||||
@pytest.mark.parametrize("head_dim", [16, 32])
|
||||
class TestUlyssesAll2All(DistributedTest):
|
||||
world_size = 4
|
||||
|
||||
def test_alltoall_output_consistency(self, d0: int, d1: int, head_dim: int, num_heads: int) -> None:
|
||||
skip_on_arch(min_arch=8)
|
||||
model = AutoModel.from_pretrained('bert-base-uncased')
|
||||
ds_engine, _, _, _ = initialize(model=model, config_params={"train_batch_size": 8}, mesh_param=(2, 2))
|
||||
#4D tensor : b,s,h,d or s,b,h,d
|
||||
input_tensor = torch.randn(d0, d1, num_heads, head_dim, device=ds_engine.device)
|
||||
scatter_idx = 2
|
||||
batch_dim_idx = 0
|
||||
outputs = []
|
||||
seq_dims = [0] #seq first API
|
||||
#TODO: Add support for batch first (that seq_dims=[0,1]) after PR for bs>1 issue with batch first is fixed
|
||||
## See discussion in : https://github.com/deepspeedai/DeepSpeed/issues/5808
|
||||
for seq_dim in seq_dims:
|
||||
gather_idx = seq_dim
|
||||
#first all2all: sequence parallel to head parallel
|
||||
s2h_tensor = _SeqAllToAll.apply(ds_engine.seq_parallel_group, input_tensor, scatter_idx, gather_idx,
|
||||
batch_dim_idx)
|
||||
|
||||
#No op
|
||||
# second all2all: head parallel to sequence parallel
|
||||
h2s_tensor = _SeqAllToAll.apply(ds_engine.seq_parallel_group, s2h_tensor, gather_idx, scatter_idx,
|
||||
batch_dim_idx)
|
||||
print(
|
||||
f'[{dist.get_rank()}] s={seq_dim} input: {input_tensor.shape} s2h: {s2h_tensor.shape} h2s_tensor: {h2s_tensor.shape}'
|
||||
)
|
||||
outputs.append(h2s_tensor)
|
||||
|
||||
# Check outputs are the same as input
|
||||
for i in range(1, len(outputs)):
|
||||
assert torch.allclose(input_tensor, outputs[i]), f"Outputs differ for sequence dim {seq_dims[i]}"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("d0", [2, 4]) #batch or sequence dimension
|
||||
@pytest.mark.parametrize("d1", [4, 8]) #batch or sequence dimension
|
||||
@pytest.mark.parametrize("num_heads", [3, 7])
|
||||
@pytest.mark.parametrize("head_dim", [16])
|
||||
class TestUlyssesAll2All_odd(DistributedTest):
|
||||
world_size = 4
|
||||
|
||||
def test_alltoall_output_consistency(self, d0: int, d1: int, head_dim: int, num_heads: int) -> None:
|
||||
|
||||
data_parallel_size = 2
|
||||
seq_parallel_size = self.world_size // data_parallel_size
|
||||
skip_on_arch(min_arch=8)
|
||||
|
||||
def seq_batch_heads_hash(d0, d1, h, offset_d0=0, offset_d1=0, offset_h=0):
|
||||
d0 += offset_d0
|
||||
d1 += offset_d1
|
||||
h += offset_h
|
||||
return d0 * 10 + h + d1 * 0.1
|
||||
|
||||
hidden_dim = 10
|
||||
model = SimpleModel(hidden_dim)
|
||||
ds_engine, _, _, _ = initialize(model=model,
|
||||
config_params={"train_batch_size": 8},
|
||||
mesh_param=(data_parallel_size, seq_parallel_size))
|
||||
|
||||
scatter_idx = 2
|
||||
outputs = []
|
||||
inputs = []
|
||||
batch_dims = [0, 1]
|
||||
seq_dims = [1, 0]
|
||||
|
||||
for idx, seq_dim in enumerate(seq_dims):
|
||||
gather_idx = seq_dim
|
||||
batch_dim_idx = batch_dims[idx]
|
||||
|
||||
#4D tensor : b,s,h,d or s,b,h,d
|
||||
#create a hash tensor from pos_id, head_id, and batch_id
|
||||
d0_indices = torch.arange(d0).reshape(-1, 1, 1, 1)
|
||||
d1_indices = torch.arange(d1).reshape(1, -1, 1, 1)
|
||||
h_indices = torch.arange(num_heads).reshape(1, 1, -1, 1)
|
||||
input_tensor = torch.randn(d0, d1, num_heads, head_dim, device=ds_engine.device)
|
||||
if batch_dim_idx == 1: #seq_len_dim : 0(d0)
|
||||
input_tensor[:] = seq_batch_heads_hash(d0_indices, d1_indices, h_indices,
|
||||
d0 * groups._get_sequence_parallel_rank(), 0)
|
||||
elif batch_dim_idx == 0: #seq_len_dim : 1(d1)
|
||||
input_tensor[:] = seq_batch_heads_hash(d0_indices, d1_indices, h_indices, 0,
|
||||
d1 * groups._get_sequence_parallel_rank())
|
||||
inputs.append(input_tensor)
|
||||
|
||||
### first all2all: sequence parallel to head parallel
|
||||
s2h_tensor = _SeqAllToAll.apply(ds_engine.seq_parallel_group, input_tensor, scatter_idx, gather_idx,
|
||||
batch_dim_idx)
|
||||
|
||||
# s2h_tensor check for the first all2all: compare with the expected ground truth
|
||||
d0_indices = torch.arange(s2h_tensor.shape[0]).reshape(-1, 1, 1, 1)
|
||||
d1_indices = torch.arange(s2h_tensor.shape[1]).reshape(1, -1, 1, 1)
|
||||
h_indices = torch.arange(s2h_tensor.shape[2]).reshape(1, 1, -1, 1)
|
||||
shard_list = get_shard_size_list(num_heads, groups._get_sequence_parallel_world_size())
|
||||
head_offset = sum(shard_list[:groups._get_sequence_parallel_rank()])
|
||||
s2h_truth = torch.zeros_like(s2h_tensor)
|
||||
s2h_truth[:] = seq_batch_heads_hash(d0_indices, d1_indices, h_indices, 0, 0, head_offset)
|
||||
|
||||
assert torch.allclose(s2h_truth,
|
||||
s2h_tensor), f"s2h_tensor differs from the expected for sequence dim: {seq_dim}"
|
||||
#No op
|
||||
### second all2all: head parallel to sequence parallel
|
||||
h2s_tensor = _SeqAllToAll.apply(ds_engine.seq_parallel_group, s2h_tensor, gather_idx, scatter_idx,
|
||||
batch_dim_idx)
|
||||
print(
|
||||
f'[{dist.get_rank()}] s={seq_dim} input: {input_tensor.shape} s2h: {s2h_tensor.shape} h2s_tensor: {h2s_tensor.shape}'
|
||||
)
|
||||
outputs.append(h2s_tensor)
|
||||
|
||||
# Check outputs for the second all2all
|
||||
for i in range(0, len(outputs)):
|
||||
assert torch.allclose(inputs[i],
|
||||
outputs[i]), f"[{dist.get_rank()}]Outputs differ for sequence dim {seq_dims[i]}"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("d0", [4, 1]) #batch dimension
|
||||
@pytest.mark.parametrize("d1", [2048, 8192]) #sequence dimension
|
||||
@pytest.mark.parametrize("chunk_size", [128, 256]) #size of chunk
|
||||
@pytest.mark.parametrize("num_heads", [8, 4])
|
||||
@pytest.mark.parametrize("head_dim", [32])
|
||||
class TestFPDTAttention(DistributedTest):
|
||||
|
||||
def test_FPDT_attention_offloading_output_consistency(self, d0: int, d1: int, chunk_size: int, head_dim: int,
|
||||
num_heads: int) -> None:
|
||||
skip_on_arch(min_arch=8)
|
||||
world_size = 2
|
||||
|
||||
try:
|
||||
from flash_attn.flash_attn_interface import _flash_attn_forward, _flash_attn_backward
|
||||
except ImportError:
|
||||
_flash_attn_forward = None
|
||||
_flash_attn_backward = None
|
||||
|
||||
if _flash_attn_forward is None or _flash_attn_backward is None:
|
||||
pytest.skip("Flash Attention is not available.")
|
||||
|
||||
model = AutoModel.from_pretrained('bert-base-uncased')
|
||||
ds_engine, _, _, _ = initialize(
|
||||
model=model,
|
||||
config_params={
|
||||
"train_batch_size": 8,
|
||||
"data_parallel_size": 1,
|
||||
"sequence_parallel_size": world_size
|
||||
},
|
||||
)
|
||||
#3D tensor : l, b, d
|
||||
dim = head_dim * num_heads
|
||||
|
||||
seed = 42
|
||||
torch.manual_seed(seed)
|
||||
get_accelerator().manual_seed_all(seed)
|
||||
|
||||
input_tensor = torch.randn(d1, d0, dim, device=ds_engine.device, dtype=torch.half) # l, b, d
|
||||
spg = ds_engine.seq_parallel_group
|
||||
|
||||
dist.broadcast(input_tensor, src=0, group=spg)
|
||||
|
||||
class args:
|
||||
|
||||
def __init__(self):
|
||||
self.ds_sequence_parallel_fpdt_chunk_size = chunk_size
|
||||
|
||||
fpdt_input_tensor = FPDT_InputConstruct(input_tensor.permute(1, 0, 2), None, None, None, None, args(),
|
||||
world_size, dist.get_rank()).generate()[0].permute(1, 0, 2)
|
||||
|
||||
if dist.get_rank() == 0:
|
||||
qkv_linear_weight = torch.nn.Parameter(
|
||||
torch.empty(dim + 2 * dim, dim, device=dist.get_rank(), dtype=torch.half))
|
||||
torch.nn.init.normal_(qkv_linear_weight, mean=0.0, std=0.02)
|
||||
|
||||
qkv_linear_bias = torch.nn.Parameter(torch.empty(dim + 2 * dim, device=dist.get_rank(), dtype=torch.half))
|
||||
torch.nn.init.normal_(qkv_linear_bias, mean=0.0, std=0.02)
|
||||
else:
|
||||
qkv_linear_weight = torch.nn.Parameter(
|
||||
torch.empty(dim + 2 * dim, dim, device=dist.get_rank(), dtype=torch.half))
|
||||
qkv_linear_bias = torch.nn.Parameter(torch.empty(dim + 2 * dim, device=dist.get_rank(), dtype=torch.half))
|
||||
|
||||
dist.broadcast(qkv_linear_weight, src=0, group=spg)
|
||||
dist.broadcast(qkv_linear_bias, src=0, group=spg)
|
||||
|
||||
num_chunks_attn = fpdt_input_tensor.shape[0] * dist.get_world_size(spg) // chunk_size
|
||||
fpdt_output = _FPDTGPUOffloadingAttentionImpl_.apply(fpdt_input_tensor, None, None, None, spg, 2, 0, dim, dim,
|
||||
head_dim, dim, qkv_linear_weight, qkv_linear_bias, 0,
|
||||
num_chunks_attn, True)
|
||||
|
||||
# baseline
|
||||
qkv = torch.matmul(input_tensor, qkv_linear_weight.t()) + qkv_linear_bias
|
||||
q = qkv[:, :, :dim].contiguous().reshape(qkv.shape[0], qkv.shape[1], -1, head_dim).permute(1, 2, 0,
|
||||
3).contiguous()
|
||||
k = qkv[:, :, dim:dim * 2].contiguous().reshape(qkv.shape[0], qkv.shape[1], -1,
|
||||
head_dim).permute(1, 2, 0, 3).contiguous()
|
||||
v = qkv[:, :, dim * 2:dim * 3].contiguous().reshape(qkv.shape[0], qkv.shape[1], -1,
|
||||
head_dim).permute(1, 2, 0,
|
||||
3).contiguous() # b, nhead, l, d
|
||||
|
||||
scores = torch.matmul(q, k.transpose(-2, -1)) / torch.sqrt(torch.tensor(dim, dtype=torch.half))
|
||||
|
||||
causal_mask = torch.triu(torch.ones(d1, d1, device=ds_engine.device), diagonal=1).bool()
|
||||
causal_mask = causal_mask.unsqueeze(0).unsqueeze(0)
|
||||
scores = scores.masked_fill(causal_mask, float('-inf'))
|
||||
attn_weights = F.softmax(scores, dim=-1)
|
||||
output = torch.matmul(attn_weights, v).permute(0, 2, 1, 3)
|
||||
|
||||
baseline_output_shuffled = FPDT_InputConstruct(output, None, None, None, None, args(), world_size,
|
||||
dist.get_rank()).generate()[0] # b, l, n, d
|
||||
|
||||
assert torch.allclose(
|
||||
fpdt_output, baseline_output_shuffled, rtol=0.01, atol=0.1
|
||||
), f"rank {dist.get_rank()}, sp size: {dist.get_world_size(spg)}, input_tensor: {input_tensor.shape}, fpdt_input_tensor: {fpdt_input_tensor.shape}, fpdt_output: {fpdt_output.shape}, baseline_output_shuffled: {baseline_output_shuffled.shape},{torch.max(torch.abs(fpdt_output - baseline_output_shuffled))}"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("sp_size", [2])
|
||||
class TestUlyssesLossBackward(DistributedTest):
|
||||
world_size = 4
|
||||
|
||||
def test_sp_loss_backward_stability(self, sp_size: int) -> None:
|
||||
"""
|
||||
Regression test for Issue #7672.
|
||||
Verifies that using all_reduce for loss aggregation is stable
|
||||
when sequence_parallel_size < world_size, preventing IndexError.
|
||||
"""
|
||||
skip_on_arch(min_arch=8)
|
||||
|
||||
# Setup
|
||||
dp_size = self.world_size // sp_size
|
||||
model = SimpleModel(4)
|
||||
ds_engine, _, _, _ = initialize(
|
||||
model=model,
|
||||
config_params={
|
||||
"train_batch_size": 8,
|
||||
"data_parallel_size": dp_size,
|
||||
"sequence_parallel_size": sp_size
|
||||
},
|
||||
)
|
||||
|
||||
sp_group = ds_engine.seq_parallel_group
|
||||
|
||||
# Simulate Loss on each rank
|
||||
rank = dist.get_rank()
|
||||
local_loss = torch.tensor(float(rank + 1), device=ds_engine.device, requires_grad=True)
|
||||
local_weight = torch.tensor(1.0, device=ds_engine.device)
|
||||
|
||||
# Numerator: Weighted Loss summation
|
||||
weighted_loss = local_loss * local_weight
|
||||
dist.all_reduce(weighted_loss, op=dist.ReduceOp.SUM, group=sp_group)
|
||||
|
||||
# B. Denominator: Sum of total weights
|
||||
total_weight = local_weight.clone()
|
||||
dist.all_reduce(total_weight, op=dist.ReduceOp.SUM, group=sp_group)
|
||||
|
||||
# C. Calculate the final loss
|
||||
dist_loss = weighted_loss / total_weight
|
||||
|
||||
# Backward Pass verification
|
||||
try:
|
||||
dist_loss.backward()
|
||||
except IndexError as e:
|
||||
pytest.fail(f"Backward crashed with IndexError: {e}")
|
||||
|
||||
# Verify Gradients
|
||||
# Loss = (L1*1 + L2*1) / 2 = 0.5*L1 + 0.5*L2
|
||||
expected_grad = 0.5
|
||||
assert torch.allclose(local_loss.grad, torch.tensor(expected_grad, device=ds_engine.device)), \
|
||||
f"Gradient mismatch! Expected {expected_grad}, got {local_loss.grad}"
|
||||
Reference in New Issue
Block a user