325 lines
12 KiB
Python
325 lines
12 KiB
Python
|
"""Ragged Ulysses sequence-parallel transport for H3 inference."""
|
||
|
|
|
||
|
|
from __future__ import annotations
|
||
|
|
|
||
|
|
from dataclasses import dataclass
|
||
|
|
|
||
|
|
import torch
|
||
|
|
import torch.distributed as dist
|
||
|
|
|
||
|
|
|
||
|
|
def balanced_ranges(total: int, parts: int) -> tuple[tuple[int, int], ...]:
|
||
|
|
"""Split ``total`` ordered items into balanced contiguous non-empty ranges."""
|
||
|
|
if parts < 1:
|
||
|
|
raise ValueError("parts must be positive")
|
||
|
|
if total < parts:
|
||
|
|
raise ValueError(f"cannot split {total} items into {parts} non-empty ranges")
|
||
|
|
base, extra = divmod(total, parts)
|
||
|
|
lengths = [base + (rank < extra) for rank in range(parts)]
|
||
|
|
ranges = []
|
||
|
|
start = 0
|
||
|
|
for length in lengths:
|
||
|
|
stop = start + int(length)
|
||
|
|
ranges.append((start, stop))
|
||
|
|
start = stop
|
||
|
|
return tuple(ranges)
|
||
|
|
|
||
|
|
|
||
|
|
def range_lengths(ranges: tuple[tuple[int, int], ...]) -> tuple[int, ...]:
|
||
|
|
return tuple(stop - start for start, stop in ranges)
|
||
|
|
|
||
|
|
|
||
|
|
def localize_segments(
|
||
|
|
segments: list[tuple[int, int, int]],
|
||
|
|
shard_start: int,
|
||
|
|
shard_stop: int,
|
||
|
|
) -> list[tuple[int, int, int]]:
|
||
|
|
"""Clip global H3 AdaLN segments to one contiguous token shard."""
|
||
|
|
localized = []
|
||
|
|
for start, stop, row in segments:
|
||
|
|
local_start = max(start, shard_start)
|
||
|
|
local_stop = min(stop, shard_stop)
|
||
|
|
if local_start < local_stop:
|
||
|
|
localized.append((local_start - shard_start, local_stop - shard_start, row))
|
||
|
|
return localized
|
||
|
|
|
||
|
|
|
||
|
|
@dataclass(frozen=True)
|
||
|
|
class SequenceParallelContext:
|
||
|
|
"""One rank's ragged token and attention-head ownership."""
|
||
|
|
|
||
|
|
group: dist.ProcessGroup | None
|
||
|
|
rank: int
|
||
|
|
world_size: int
|
||
|
|
token_ranges: tuple[tuple[int, int], ...]
|
||
|
|
head_ranges: tuple[tuple[int, int], ...]
|
||
|
|
head_dim: int
|
||
|
|
|
||
|
|
@classmethod
|
||
|
|
def create(
|
||
|
|
cls,
|
||
|
|
sequence_length: int,
|
||
|
|
heads: int,
|
||
|
|
head_dim: int,
|
||
|
|
*,
|
||
|
|
group: dist.ProcessGroup | None = None,
|
||
|
|
) -> "SequenceParallelContext":
|
||
|
|
if not dist.is_initialized():
|
||
|
|
raise RuntimeError("torch.distributed process group is not initialized")
|
||
|
|
world_size = dist.get_world_size(group)
|
||
|
|
rank = dist.get_rank(group)
|
||
|
|
return cls(
|
||
|
|
group=group,
|
||
|
|
rank=rank,
|
||
|
|
world_size=world_size,
|
||
|
|
token_ranges=balanced_ranges(sequence_length, world_size),
|
||
|
|
head_ranges=balanced_ranges(heads, world_size),
|
||
|
|
head_dim=head_dim,
|
||
|
|
)
|
||
|
|
|
||
|
|
@property
|
||
|
|
def sequence_length(self) -> int:
|
||
|
|
return self.token_ranges[-1][1]
|
||
|
|
|
||
|
|
@property
|
||
|
|
def heads(self) -> int:
|
||
|
|
return self.head_ranges[-1][1]
|
||
|
|
|
||
|
|
@property
|
||
|
|
def token_lengths(self) -> tuple[int, ...]:
|
||
|
|
return range_lengths(self.token_ranges)
|
||
|
|
|
||
|
|
@property
|
||
|
|
def head_lengths(self) -> tuple[int, ...]:
|
||
|
|
return range_lengths(self.head_ranges)
|
||
|
|
|
||
|
|
@property
|
||
|
|
def local_token_range(self) -> tuple[int, int]:
|
||
|
|
return self.token_ranges[self.rank]
|
||
|
|
|
||
|
|
@property
|
||
|
|
def local_head_range(self) -> tuple[int, int]:
|
||
|
|
return self.head_ranges[self.rank]
|
||
|
|
|
||
|
|
@property
|
||
|
|
def local_token_length(self) -> int:
|
||
|
|
start, stop = self.local_token_range
|
||
|
|
return stop - start
|
||
|
|
|
||
|
|
@property
|
||
|
|
def local_head_count(self) -> int:
|
||
|
|
start, stop = self.local_head_range
|
||
|
|
return stop - start
|
||
|
|
|
||
|
|
def localize_segments(self, segments: list[tuple[int, int, int]]) -> list[tuple[int, int, int]]:
|
||
|
|
return localize_segments(segments, *self.local_token_range)
|
||
|
|
|
||
|
|
def seq_to_heads(
|
||
|
|
self,
|
||
|
|
q: torch.Tensor,
|
||
|
|
k: torch.Tensor,
|
||
|
|
v: torch.Tensor,
|
||
|
|
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
|
||
|
|
"""Exchange local tokens for full-sequence Q/K/V over locally owned heads.
|
||
|
|
|
||
|
|
Inputs use BSHD layout ``[1, local_tokens, all_heads, head_dim]``. Outputs
|
||
|
|
use ``[1, all_tokens, local_heads, head_dim]``.
|
||
|
|
"""
|
||
|
|
expected = (1, self.local_token_length, self.heads, self.head_dim)
|
||
|
|
if tuple(q.shape) != expected or tuple(k.shape) != expected or tuple(v.shape) != expected:
|
||
|
|
raise ValueError(
|
||
|
|
f"sequence-parallel Q/K/V must each have shape {expected}; "
|
||
|
|
f"got {tuple(q.shape)}, {tuple(k.shape)}, {tuple(v.shape)}"
|
||
|
|
)
|
||
|
|
if q.dtype != k.dtype or q.dtype != v.dtype or q.device != k.device or q.device != v.device:
|
||
|
|
raise ValueError("sequence-parallel Q/K/V must share dtype and device")
|
||
|
|
if self.world_size == 1:
|
||
|
|
return q, k, v
|
||
|
|
|
||
|
|
local_qkv = torch.stack((q[0], k[0], v[0]), dim=1)
|
||
|
|
send_chunks = [
|
||
|
|
local_qkv[:, :, start:stop, :].contiguous().view(-1)
|
||
|
|
for start, stop in self.head_ranges
|
||
|
|
]
|
||
|
|
input_splits = [chunk.numel() for chunk in send_chunks]
|
||
|
|
send = torch.cat(send_chunks)
|
||
|
|
|
||
|
|
output_splits = [
|
||
|
|
token_length * 3 * self.local_head_count * self.head_dim
|
||
|
|
for token_length in self.token_lengths
|
||
|
|
]
|
||
|
|
receive = torch.empty(sum(output_splits), dtype=q.dtype, device=q.device)
|
||
|
|
dist.all_to_all_single(
|
||
|
|
receive,
|
||
|
|
send,
|
||
|
|
output_split_sizes=output_splits,
|
||
|
|
input_split_sizes=input_splits,
|
||
|
|
group=self.group,
|
||
|
|
)
|
||
|
|
|
||
|
|
source_chunks = []
|
||
|
|
offset = 0
|
||
|
|
for token_length, count in zip(self.token_lengths, output_splits, strict=True):
|
||
|
|
source_chunks.append(
|
||
|
|
receive[offset : offset + count].view(
|
||
|
|
token_length, 3, self.local_head_count, self.head_dim,
|
||
|
|
)
|
||
|
|
)
|
||
|
|
offset += count
|
||
|
|
full_qkv = torch.cat(source_chunks, dim=0)
|
||
|
|
full_q, full_k, full_v = full_qkv.unbind(dim=1)
|
||
|
|
return full_q.unsqueeze(0), full_k.unsqueeze(0), full_v.unsqueeze(0)
|
||
|
|
|
||
|
|
def heads_to_seq(self, output: torch.Tensor) -> torch.Tensor:
|
||
|
|
"""Exchange full-sequence local-head output back to local-token all-head output.
|
||
|
|
|
||
|
|
Input is BSHD ``[1, all_tokens, local_heads, head_dim]``. The return value
|
||
|
|
is ``[local_tokens, all_heads, head_dim]``.
|
||
|
|
"""
|
||
|
|
expected = (1, self.sequence_length, self.local_head_count, self.head_dim)
|
||
|
|
if tuple(output.shape) != expected:
|
||
|
|
raise ValueError(f"sequence-parallel output must have shape {expected}, got {tuple(output.shape)}")
|
||
|
|
if self.world_size == 1:
|
||
|
|
return output[0]
|
||
|
|
|
||
|
|
output = output[0]
|
||
|
|
send_chunks = []
|
||
|
|
input_splits = []
|
||
|
|
token_offset = 0
|
||
|
|
for token_length in self.token_lengths:
|
||
|
|
chunk = output[token_offset : token_offset + token_length].contiguous().view(-1)
|
||
|
|
send_chunks.append(chunk)
|
||
|
|
input_splits.append(chunk.numel())
|
||
|
|
token_offset += token_length
|
||
|
|
send = torch.cat(send_chunks)
|
||
|
|
|
||
|
|
output_splits = [
|
||
|
|
self.local_token_length * head_length * self.head_dim
|
||
|
|
for head_length in self.head_lengths
|
||
|
|
]
|
||
|
|
receive = torch.empty(sum(output_splits), dtype=output.dtype, device=output.device)
|
||
|
|
dist.all_to_all_single(
|
||
|
|
receive,
|
||
|
|
send,
|
||
|
|
output_split_sizes=output_splits,
|
||
|
|
input_split_sizes=input_splits,
|
||
|
|
group=self.group,
|
||
|
|
)
|
||
|
|
|
||
|
|
head_chunks = []
|
||
|
|
offset = 0
|
||
|
|
for head_length, count in zip(self.head_lengths, output_splits, strict=True):
|
||
|
|
head_chunks.append(
|
||
|
|
receive[offset : offset + count].view(
|
||
|
|
self.local_token_length, head_length, self.head_dim,
|
||
|
|
)
|
||
|
|
)
|
||
|
|
offset += count
|
||
|
|
return torch.cat(head_chunks, dim=1).contiguous()
|
||
|
|
|
||
|
|
def target_intersection(self, target_start: int, target_stop: int) -> tuple[int, int]:
|
||
|
|
"""Return one global target span's bounds relative to this token shard."""
|
||
|
|
shard_start, shard_stop = self.local_token_range
|
||
|
|
start = max(target_start, shard_start)
|
||
|
|
stop = min(target_stop, shard_stop)
|
||
|
|
if start >= stop:
|
||
|
|
return (0, 0)
|
||
|
|
return (start - shard_start, stop - shard_start)
|
||
|
|
|
||
|
|
def target_counts(self, target_start: int, target_stop: int) -> tuple[int, ...]:
|
||
|
|
"""Return ordered target-row counts contributed by every token rank."""
|
||
|
|
counts = []
|
||
|
|
for shard_start, shard_stop in self.token_ranges:
|
||
|
|
counts.append(max(0, min(target_stop, shard_stop) - max(target_start, shard_start)))
|
||
|
|
return tuple(counts)
|
||
|
|
|
||
|
|
def all_gather_target_rows(
|
||
|
|
self,
|
||
|
|
local_rows: torch.Tensor,
|
||
|
|
target_start: int,
|
||
|
|
target_stop: int,
|
||
|
|
) -> torch.Tensor:
|
||
|
|
"""Gather a global target span's projected rows onto every rank.
|
||
|
|
|
||
|
|
Padding is transport-only and is removed before concatenation; it is never
|
||
|
|
exposed to attention or model semantics.
|
||
|
|
"""
|
||
|
|
counts = self.target_counts(target_start, target_stop)
|
||
|
|
if local_rows.ndim != 2:
|
||
|
|
raise ValueError("target rows must be rank-2 [rows, features]")
|
||
|
|
if local_rows.shape[0] != counts[self.rank]:
|
||
|
|
raise ValueError(
|
||
|
|
f"rank {self.rank} must contribute {counts[self.rank]} target rows, "
|
||
|
|
f"got {local_rows.shape[0]}"
|
||
|
|
)
|
||
|
|
if self.world_size == 1:
|
||
|
|
return local_rows
|
||
|
|
max_rows = max(counts)
|
||
|
|
padded = torch.zeros(
|
||
|
|
max_rows, local_rows.shape[1], dtype=local_rows.dtype, device=local_rows.device,
|
||
|
|
)
|
||
|
|
if local_rows.shape[0]:
|
||
|
|
padded[: local_rows.shape[0]].copy_(local_rows)
|
||
|
|
gathered = [torch.empty_like(padded) for _ in range(self.world_size)]
|
||
|
|
dist.all_gather(gathered, padded, group=self.group)
|
||
|
|
return torch.cat(
|
||
|
|
[rows[:count] for rows, count in zip(gathered, counts, strict=True) if count],
|
||
|
|
dim=0,
|
||
|
|
)
|
||
|
|
|
||
|
|
def all_gather_rows(self, local_rows: torch.Tensor) -> torch.Tensor:
|
||
|
|
"""Gather ragged token rows on every rank without exposing padding to the model."""
|
||
|
|
if local_rows.shape[0] != self.local_token_length:
|
||
|
|
raise ValueError(
|
||
|
|
f"rank {self.rank} must contribute {self.local_token_length} rows, "
|
||
|
|
f"got {local_rows.shape[0]}"
|
||
|
|
)
|
||
|
|
if self.world_size == 1:
|
||
|
|
return local_rows
|
||
|
|
max_rows = max(self.token_lengths)
|
||
|
|
padded = torch.zeros(
|
||
|
|
(max_rows, *local_rows.shape[1:]),
|
||
|
|
dtype=local_rows.dtype,
|
||
|
|
device=local_rows.device,
|
||
|
|
)
|
||
|
|
padded[: local_rows.shape[0]].copy_(local_rows)
|
||
|
|
gathered = [torch.empty_like(padded) for _ in range(self.world_size)]
|
||
|
|
dist.all_gather(gathered, padded, group=self.group)
|
||
|
|
return torch.cat(
|
||
|
|
[rows[:count] for rows, count in zip(gathered, self.token_lengths, strict=True)],
|
||
|
|
dim=0,
|
||
|
|
)
|
||
|
|
|
||
|
|
def reduce_scatter_rows(self, partial_full_rows: torch.Tensor) -> torch.Tensor:
|
||
|
|
"""Sum tensor-parallel partials and return this rank's ragged token rows."""
|
||
|
|
if partial_full_rows.shape[0] != self.sequence_length:
|
||
|
|
raise ValueError(
|
||
|
|
f"partial rows must cover sequence length {self.sequence_length}, "
|
||
|
|
f"got {partial_full_rows.shape[0]}"
|
||
|
|
)
|
||
|
|
if self.world_size == 1:
|
||
|
|
return partial_full_rows
|
||
|
|
trailing_shape = partial_full_rows.shape[1:]
|
||
|
|
row_width = partial_full_rows[0].numel()
|
||
|
|
send_chunks = []
|
||
|
|
input_splits = []
|
||
|
|
offset = 0
|
||
|
|
for token_length in self.token_lengths:
|
||
|
|
chunk = partial_full_rows[offset : offset + token_length].contiguous().view(-1)
|
||
|
|
send_chunks.append(chunk)
|
||
|
|
input_splits.append(chunk.numel())
|
||
|
|
offset += token_length
|
||
|
|
send = torch.cat(send_chunks)
|
||
|
|
|
||
|
|
output_splits = [self.local_token_length * row_width] * self.world_size
|
||
|
|
receive = torch.empty(sum(output_splits), dtype=send.dtype, device=send.device)
|
||
|
|
dist.all_to_all_single(
|
||
|
|
receive,
|
||
|
|
send,
|
||
|
|
output_split_sizes=output_splits,
|
||
|
|
input_split_sizes=input_splits,
|
||
|
|
group=self.group,
|
||
|
|
)
|
||
|
|
contributions = receive.view(self.world_size, self.local_token_length, *trailing_shape)
|
||
|
|
return contributions.sum(dim=0)
|