wehub-resource-sync
176 lines
7.4 KiB
Python
176 lines
7.4 KiB
Python
# Copyright (c) Microsoft Corporation.
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
# DeepSpeed Team
|
|
"""
|
|
Ulysses-style sequence-parallel wrapper for ViT encoder attention layers.
|
|
|
|
Design notes
|
|
------------
|
|
ViT self-attention is non-causal: every patch token attends to every other
|
|
patch token. This means a straightforward per-rank local attention (as used
|
|
for causal LLMs) would be *incorrect* — each rank must have access to the
|
|
full key/value context.
|
|
|
|
We therefore use a **gather-compute-scatter** pattern:
|
|
|
|
1. Input arrives already sharded along the sequence dimension (each rank owns
|
|
``local_patches = num_patches // world_size`` consecutive patches).
|
|
2. Before attention we **all-gather** patch tokens so that every rank runs the
|
|
full ViT attention over the complete patch sequence. This keeps the
|
|
computation equivalent to single-device execution.
|
|
3. The output is **scattered** back so that each rank returns only its local
|
|
slice, matching the sharded input contract expected by downstream layers.
|
|
|
|
Memory benefit: activations *outside* the attention block (e.g. feed-forward
|
|
layers, layer norms) are stored only locally, reducing per-rank memory
|
|
proportional to ``world_size``.
|
|
|
|
The ``cls`` token (if present) is replicated on every rank and is not split
|
|
across the sequence dimension. Each rank appends its local patches to the
|
|
same ``cls`` token before calling the wrapped attention.
|
|
|
|
Padding: when ``num_patches % world_size != 0``, shorter shards are
|
|
zero-padded to a uniform size for ``all_gather``. The padding is stripped
|
|
*before* the attention call by trimming each rank's contribution to its true
|
|
length, so the wrapped attention always sees exactly ``num_patches`` real
|
|
tokens — identical to single-device execution and free of softmax pollution
|
|
from dummy tokens.
|
|
"""
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
|
|
import deepspeed.comm as dist
|
|
|
|
|
|
class UlyssesSPViTAttention(nn.Module):
|
|
"""Sequence-parallel wrapper for an opaque ViT attention module.
|
|
|
|
Parameters
|
|
----------
|
|
attn:
|
|
The original ViT attention layer (any ``nn.Module`` that maps
|
|
``hidden_states`` → ``hidden_states`` or a tuple whose first element
|
|
is the attention output tensor).
|
|
process_group:
|
|
The sequence-parallel process group.
|
|
has_cls_token:
|
|
Set to ``True`` (default) when the first token in the sequence is a
|
|
``[CLS]`` token that should be replicated on every rank rather than
|
|
sharded.
|
|
"""
|
|
|
|
def __init__(self, attn: nn.Module, process_group, has_cls_token: bool = True) -> None:
|
|
super().__init__()
|
|
self.attn = attn
|
|
self.process_group = process_group
|
|
self.world_size = dist.get_world_size(process_group)
|
|
self.has_cls_token = has_cls_token
|
|
|
|
# ------------------------------------------------------------------
|
|
# forward
|
|
# ------------------------------------------------------------------
|
|
|
|
def forward(self, hidden_states: torch.Tensor, **kwargs):
|
|
"""
|
|
Parameters
|
|
----------
|
|
hidden_states:
|
|
Shape ``[bs, local_seq_len, hidden_dim]`` where
|
|
``local_seq_len = (1 + local_patches)`` if ``has_cls_token`` else
|
|
``local_patches``. Each rank holds a contiguous slice of patches.
|
|
**kwargs:
|
|
Passed through to the wrapped attention (e.g. ``attention_mask``,
|
|
``head_mask``, ``output_attentions``).
|
|
|
|
Returns
|
|
-------
|
|
Same shape as input (or a tuple whose first element matches the input
|
|
shape, preserving whatever the wrapped module returns).
|
|
"""
|
|
bs, local_seq_len, hidden_dim = hidden_states.shape
|
|
|
|
if self.has_cls_token:
|
|
# CLS token is replicated on every rank — not part of the sharded seq
|
|
cls_token = hidden_states[:, :1, :]
|
|
local_patches = hidden_states[:, 1:, :]
|
|
else:
|
|
local_patches = hidden_states
|
|
|
|
local_patch_len = local_patches.shape[1]
|
|
|
|
# -------------------------------------------------------------------
|
|
# 1. All-gather patches from all ranks to reconstruct the full sequence
|
|
# -------------------------------------------------------------------
|
|
# When num_patches % world_size != 0, ranks hold different shard sizes.
|
|
# We all-gather every rank's local_patch_len so we can:
|
|
# (a) zero-pad shorter slices to uniform size for all_gather, and
|
|
# (b) strip the padding per rank *before* calling attention, so that
|
|
# the wrapped module never sees dummy tokens (which would corrupt
|
|
# the softmax normalisation).
|
|
len_bufs = [torch.zeros(1, dtype=torch.long, device=local_patches.device) for _ in range(self.world_size)]
|
|
dist.all_gather(len_bufs,
|
|
torch.tensor([local_patch_len], dtype=torch.long, device=local_patches.device),
|
|
group=self.process_group)
|
|
all_lens = [int(t.item()) for t in len_bufs]
|
|
max_local_len = max(all_lens)
|
|
|
|
pad_len = max_local_len - local_patch_len
|
|
if pad_len > 0:
|
|
# Append zero rows so this rank's buffer matches the largest shard.
|
|
local_patches_padded = F.pad(local_patches, (0, 0, 0, pad_len))
|
|
else:
|
|
local_patches_padded = local_patches
|
|
|
|
gathered = [
|
|
torch.zeros(bs, max_local_len, hidden_dim, dtype=local_patches.dtype, device=local_patches.device)
|
|
for _ in range(self.world_size)
|
|
]
|
|
dist.all_gather(gathered, local_patches_padded.contiguous(), group=self.process_group)
|
|
|
|
# Strip per-rank padding before concatenation so attention only sees
|
|
# the true num_patches tokens, identical to single-device execution.
|
|
real_parts = [gathered[r][:, :all_lens[r], :] for r in range(self.world_size)]
|
|
full_patches = torch.cat(real_parts, dim=1) # [bs, total_real_patches, hidden_dim]
|
|
|
|
# -------------------------------------------------------------------
|
|
# 2. Build the full input (prepend CLS if needed) and call attention
|
|
# -------------------------------------------------------------------
|
|
if self.has_cls_token:
|
|
full_input = torch.cat([cls_token, full_patches], dim=1)
|
|
else:
|
|
full_input = full_patches
|
|
|
|
attn_out = self.attn(full_input, **kwargs)
|
|
|
|
# Unwrap tuple: some ViT implementations return (attn_output, attn_weights)
|
|
if isinstance(attn_out, (tuple, list)):
|
|
full_out, *extra = attn_out
|
|
else:
|
|
full_out = attn_out
|
|
extra = []
|
|
|
|
# -------------------------------------------------------------------
|
|
# 3. Scatter output: each rank keeps its local slice of the real patches.
|
|
# Because padding was stripped before attention, scatter offsets are
|
|
# the cumulative sums of all_lens, not rank * max_local_len.
|
|
# -------------------------------------------------------------------
|
|
if self.has_cls_token:
|
|
cls_out = full_out[:, :1, :]
|
|
patch_out = full_out[:, 1:, :]
|
|
else:
|
|
patch_out = full_out
|
|
|
|
rank = dist.get_rank(self.process_group)
|
|
start = sum(all_lens[:rank])
|
|
local_out = patch_out[:, start:start + local_patch_len, :].contiguous()
|
|
|
|
if self.has_cls_token:
|
|
local_out = torch.cat([cls_out, local_out], dim=1)
|
|
|
|
if extra:
|
|
return (local_out, *extra)
|
|
return local_out
|