wehub-resource-sync
453 lines
21 KiB
Python
453 lines
21 KiB
Python
# Copyright (c) DeepSpeed Team.
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
# DeepSpeed Team
|
|
"""Route-full / partition-dispatch helpers for AutoEP + AutoTP folding."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass
|
|
import os
|
|
|
|
import torch
|
|
import deepspeed.comm as dist
|
|
|
|
_FOLDING_DIGEST_MOD_A = 2147483647
|
|
_FOLDING_DIGEST_MOD_B = 2147483629
|
|
|
|
|
|
@dataclass
|
|
class RoutedAssignmentPayload:
|
|
token_indices: torch.Tensor
|
|
expert_indices: torch.Tensor
|
|
assignment_indices: torch.Tensor
|
|
capacity_slots: torch.Tensor
|
|
combine_weights: torch.Tensor
|
|
drop_mask: torch.Tensor
|
|
pad_mask: torch.Tensor
|
|
input_splits: list[int]
|
|
output_splits: list[int]
|
|
extra: dict[str, torch.Tensor]
|
|
|
|
|
|
@dataclass
|
|
class RestoreContext:
|
|
original_payload: RoutedAssignmentPayload
|
|
local_indices: torch.Tensor
|
|
tp_rank: int
|
|
tp_size: int
|
|
num_tokens: int
|
|
counters: dict[str, int]
|
|
|
|
|
|
def assignment_ordinals_by_expert(expert_indices: torch.Tensor) -> torch.Tensor:
|
|
"""Return stable ordinals within each contiguous expert segment."""
|
|
if expert_indices.numel() == 0:
|
|
return expert_indices.to(torch.long)
|
|
positions = torch.arange(expert_indices.numel(), device=expert_indices.device, dtype=torch.long)
|
|
starts = torch.zeros_like(positions)
|
|
starts[0] = 0
|
|
segment_start = torch.zeros(expert_indices.numel(), device=expert_indices.device, dtype=torch.bool)
|
|
segment_start[0] = True
|
|
segment_start[1:] = expert_indices[1:] != expert_indices[:-1]
|
|
starts = torch.where(segment_start, positions, starts)
|
|
starts = torch.cummax(starts, dim=0).values
|
|
return positions - starts
|
|
|
|
|
|
def _take(payload: RoutedAssignmentPayload, indices: torch.Tensor) -> RoutedAssignmentPayload:
|
|
extra = {
|
|
key:
|
|
value.index_select(0, indices)
|
|
if torch.is_tensor(value) and value.shape[:1] == payload.token_indices.shape[:1] else value
|
|
for key, value in payload.extra.items()
|
|
}
|
|
return RoutedAssignmentPayload(
|
|
token_indices=payload.token_indices.index_select(0, indices),
|
|
expert_indices=payload.expert_indices.index_select(0, indices),
|
|
assignment_indices=payload.assignment_indices.index_select(0, indices),
|
|
capacity_slots=payload.capacity_slots.index_select(0, indices),
|
|
combine_weights=payload.combine_weights.index_select(0, indices),
|
|
drop_mask=payload.drop_mask.index_select(0, indices),
|
|
pad_mask=payload.pad_mask.index_select(0, indices),
|
|
input_splits=list(payload.input_splits),
|
|
output_splits=list(payload.output_splits),
|
|
extra=extra,
|
|
)
|
|
|
|
|
|
def _recompute_input_splits(payload: RoutedAssignmentPayload) -> list[int]:
|
|
destinations = payload.extra.get("destination_ranks")
|
|
if destinations is None:
|
|
return list(payload.input_splits)
|
|
if len(payload.input_splits) == 0:
|
|
return []
|
|
counts = torch.bincount(destinations.to(torch.long), minlength=len(payload.input_splits))
|
|
return [int(value) for value in counts[:len(payload.input_splits)].cpu().tolist()]
|
|
|
|
|
|
def _tensor_digest_words(tensor: torch.Tensor) -> torch.Tensor:
|
|
tensor = tensor.detach()
|
|
if tensor.is_floating_point():
|
|
words = torch.nan_to_num(tensor.float(), nan=0.0, posinf=3.4028235e38,
|
|
neginf=-3.4028235e38).mul(1000003.0).round().to(torch.long)
|
|
else:
|
|
words = tensor.to(torch.long)
|
|
return words.reshape(-1)
|
|
|
|
|
|
def _digest_words(words: torch.Tensor, *, salt: int, modulus: int) -> torch.Tensor:
|
|
if words.numel() == 0:
|
|
return torch.tensor(salt, device=words.device, dtype=torch.long)
|
|
positions = torch.arange(1, words.numel() + 1, device=words.device, dtype=torch.long)
|
|
positions = positions.add_(salt).remainder_(modulus)
|
|
values = words.remainder(modulus)
|
|
return (values.mul(positions).remainder_(modulus).sum().add_(words.numel() * salt).remainder_(modulus))
|
|
|
|
|
|
def _payload_digest(payload: RoutedAssignmentPayload) -> torch.Tensor:
|
|
device = payload.token_indices.device
|
|
active = (~payload.drop_mask & ~payload.pad_mask).to(torch.long)
|
|
digest = torch.tensor(
|
|
[payload.token_indices.numel(),
|
|
int(sum(payload.input_splits)),
|
|
int(sum(payload.output_splits)), 0, 0],
|
|
device=device,
|
|
dtype=torch.long)
|
|
fields = (
|
|
payload.token_indices,
|
|
payload.expert_indices,
|
|
payload.assignment_indices,
|
|
payload.capacity_slots,
|
|
payload.combine_weights,
|
|
payload.drop_mask,
|
|
payload.pad_mask,
|
|
active,
|
|
payload.extra.get("destination_ranks", torch.empty(0, device=device, dtype=torch.long)),
|
|
)
|
|
for index, field in enumerate(fields, start=1):
|
|
if not torch.is_tensor(field):
|
|
continue
|
|
words = _tensor_digest_words(field)
|
|
digest[3] = digest[3].add(_digest_words(words, salt=17 * index,
|
|
modulus=_FOLDING_DIGEST_MOD_A)).remainder_(_FOLDING_DIGEST_MOD_A)
|
|
digest[4] = digest[4].add(_digest_words(words, salt=31 * index,
|
|
modulus=_FOLDING_DIGEST_MOD_B)).remainder_(_FOLDING_DIGEST_MOD_B)
|
|
return digest
|
|
|
|
|
|
def _payload_digest_components(payload: RoutedAssignmentPayload) -> dict[str, torch.Tensor]:
|
|
device = payload.token_indices.device
|
|
active = (~payload.drop_mask & ~payload.pad_mask).to(torch.long)
|
|
fields = {
|
|
"token_indices": payload.token_indices,
|
|
"expert_indices": payload.expert_indices,
|
|
"assignment_indices": payload.assignment_indices,
|
|
"capacity_slots": payload.capacity_slots,
|
|
"combine_weights": payload.combine_weights,
|
|
"drop_mask": payload.drop_mask,
|
|
"pad_mask": payload.pad_mask,
|
|
"active": active,
|
|
"destination_ranks": payload.extra.get("destination_ranks", torch.empty(0, device=device, dtype=torch.long)),
|
|
}
|
|
components: dict[str, torch.Tensor] = {}
|
|
for index, (name, field) in enumerate(fields.items(), start=1):
|
|
if not torch.is_tensor(field):
|
|
continue
|
|
words = _tensor_digest_words(field)
|
|
components[name] = torch.stack((
|
|
torch.tensor(words.numel(), device=device, dtype=torch.long),
|
|
_digest_words(words, salt=17 * index, modulus=_FOLDING_DIGEST_MOD_A),
|
|
_digest_words(words, salt=31 * index, modulus=_FOLDING_DIGEST_MOD_B),
|
|
))
|
|
return components
|
|
|
|
|
|
def _format_payload_debug(payload: RoutedAssignmentPayload, *, digest: torch.Tensor, max_digest: torch.Tensor,
|
|
min_digest: torch.Tensor, tp_group) -> str:
|
|
if os.environ.get("AUTOEP_FOLDING_DEBUG_PAYLOAD", "0") not in {"1", "true", "TRUE", "yes"}:
|
|
return ""
|
|
|
|
differing_fields = []
|
|
for name, component in _payload_digest_components(payload).items():
|
|
component_max = component.clone()
|
|
component_min = component.clone()
|
|
dist.all_reduce(component_max, op=dist.ReduceOp.MAX, group=tp_group)
|
|
dist.all_reduce(component_min, op=dist.ReduceOp.MIN, group=tp_group)
|
|
if not torch.equal(component_max, component_min):
|
|
differing_fields.append({
|
|
"field": name,
|
|
"local": [int(value) for value in component.detach().cpu().tolist()],
|
|
"min": [int(value) for value in component_min.detach().cpu().tolist()],
|
|
"max": [int(value) for value in component_max.detach().cpu().tolist()],
|
|
})
|
|
|
|
sample_limit = int(os.environ.get("AUTOEP_FOLDING_DEBUG_SAMPLE_LIMIT", "12"))
|
|
samples = {
|
|
"token_indices": payload.token_indices[:sample_limit].detach().cpu().tolist(),
|
|
"expert_indices": payload.expert_indices[:sample_limit].detach().cpu().tolist(),
|
|
"assignment_indices": payload.assignment_indices[:sample_limit].detach().cpu().tolist(),
|
|
"capacity_slots": payload.capacity_slots[:sample_limit].detach().cpu().tolist(),
|
|
"combine_weights": payload.combine_weights[:sample_limit].detach().float().cpu().tolist(),
|
|
}
|
|
try:
|
|
tp_group_ranks = dist.get_all_ranks_from_group(tp_group)
|
|
except Exception:
|
|
tp_group_ranks = []
|
|
details = {
|
|
"rank": dist.get_rank(),
|
|
"tp_rank": dist.get_rank(group=tp_group),
|
|
"tp_group_ranks": tp_group_ranks,
|
|
"digest": [int(value) for value in digest.detach().cpu().tolist()],
|
|
"digest_min": [int(value) for value in min_digest.detach().cpu().tolist()],
|
|
"digest_max": [int(value) for value in max_digest.detach().cpu().tolist()],
|
|
"differing_fields": differing_fields,
|
|
"samples": samples,
|
|
}
|
|
return f" Debug details: {details}"
|
|
|
|
|
|
def assert_tp_payload_consistent(payload: RoutedAssignmentPayload, *, tp_group, tp_size: int) -> None:
|
|
if tp_size <= 1 or not dist.is_initialized():
|
|
return
|
|
|
|
digest = _payload_digest(payload)
|
|
max_digest = digest.clone()
|
|
min_digest = digest.clone()
|
|
dist.all_reduce(max_digest, op=dist.ReduceOp.MAX, group=tp_group)
|
|
dist.all_reduce(min_digest, op=dist.ReduceOp.MIN, group=tp_group)
|
|
if not torch.equal(max_digest, min_digest):
|
|
debug_details = _format_payload_debug(payload,
|
|
digest=digest,
|
|
max_digest=max_digest,
|
|
min_digest=min_digest,
|
|
tp_group=tp_group)
|
|
raise RuntimeError("AutoEP+AutoTP routing decisions differ across tensor-parallel lanes. "
|
|
"Folded dispatch requires identical routed-token payloads before TP partitioning."
|
|
f"{debug_details}")
|
|
|
|
|
|
def partition_assignments(
|
|
payload: RoutedAssignmentPayload,
|
|
*,
|
|
tp_group,
|
|
tp_rank: int,
|
|
tp_size: int,
|
|
) -> tuple[RoutedAssignmentPayload, RestoreContext]:
|
|
"""Partition routed assignments across TP peers by stable per-expert ordinal.
|
|
|
|
Each peer keeps only ``assignment_index % tp_size == tp_rank`` of the
|
|
(token, expert) assignments and drops the rest *before* the EP dispatch
|
|
all-to-all, so the dispatch carries the full token set exactly once (split
|
|
across peers) instead of ``tp_size`` redundant copies. The dropped work is
|
|
reconstructed afterwards by ``restore_combined``'s all-gather; that
|
|
reconstruction is what makes the folded router/gate gradient replicated
|
|
(AVERAGE) rather than a true SUM partial -- see ``_AllGatherVariableRows``.
|
|
"""
|
|
active = ~payload.drop_mask & ~payload.pad_mask
|
|
if tp_size <= 1:
|
|
keep = active
|
|
else:
|
|
keep = (payload.assignment_indices.remainder(tp_size) == tp_rank) & active
|
|
local_indices = torch.nonzero(keep, as_tuple=False).flatten()
|
|
|
|
local = _take(payload, local_indices)
|
|
local.input_splits = _recompute_input_splits(local)
|
|
local.output_splits = list(local.input_splits)
|
|
ctx = RestoreContext(
|
|
original_payload=payload,
|
|
local_indices=local_indices,
|
|
tp_rank=tp_rank,
|
|
tp_size=tp_size,
|
|
num_tokens=int(payload.extra.get("num_tokens", torch.tensor(0)).item()) if torch.is_tensor(
|
|
payload.extra.get("num_tokens")) else int(payload.extra.get("num_tokens", 0)),
|
|
counters={
|
|
"assignments_total": int((~payload.drop_mask & ~payload.pad_mask).sum().item()),
|
|
"assignments_local": int(local_indices.numel()),
|
|
"padded": int(payload.pad_mask.sum().item()),
|
|
"dropped": int(payload.drop_mask.sum().item()),
|
|
"split_sum_in": int(sum(local.input_splits)),
|
|
"split_sum_out": int(sum(local.output_splits)),
|
|
},
|
|
)
|
|
return local, ctx
|
|
|
|
|
|
def _pad_rows(tensor: torch.Tensor, rows: int) -> torch.Tensor:
|
|
if tensor.shape[0] == rows:
|
|
return tensor
|
|
pad_shape = (rows - tensor.shape[0], *tensor.shape[1:])
|
|
return torch.cat((tensor, tensor.new_zeros(pad_shape)), dim=0)
|
|
|
|
|
|
class _AllGatherVariableRows(torch.autograd.Function):
|
|
"""Differentiable all-gather of row-variable tensors across the TP folding group.
|
|
|
|
Forward concatenates every TP peer's local rows into one tensor that is
|
|
identical on every peer: a replicated full view of the rows that
|
|
``partition_assignments`` had split across peers before the EP dispatch.
|
|
|
|
Backward is the matching reduce-scatter. Because the forward output is
|
|
consumed identically on every peer, each peer holds the same ``grad_output``;
|
|
summing those replicas with ``all_reduce`` and keeping this peer's own
|
|
row-slice is the correct vector-Jacobian product.
|
|
|
|
Gradient-reduction consequence (important -- this is why the folded
|
|
router/gate uses AVERAGE, not SUM): the ``all_reduce`` in backward scales
|
|
each peer's slice gradient by ``tp_size``. A parameter whose gradient flows
|
|
through this restore all-gather -- the folded router/gate scores, see
|
|
``restore_combined`` -- therefore reaches the optimizer's TP reducer carrying
|
|
``tp_size`` times its own routed-token slice. The TP reducer all_reduce then
|
|
produces ``tp_size * full_grad``, and the AVERAGE strategy in
|
|
``auto_ep_folding.autoep_folding_gradient_reduction_strategy`` divides by
|
|
``tp_size`` to recover the true gradient. Reducing with SUM instead leaves
|
|
the uncancelled ``tp_size`` factor -- exactly the 2.0x router/gate gradient
|
|
regression the CPU/Gloo parity tests guard against. The partition is
|
|
reconstructed into a replicated full view here, so it is not a genuine SUM
|
|
partial; a future true-SP path that kept the shard to the loss would be.
|
|
"""
|
|
|
|
@staticmethod
|
|
def forward(ctx, tensor, group, counts, max_rows):
|
|
ctx.group = group
|
|
ctx.counts = tuple(counts)
|
|
ctx.max_rows = max_rows
|
|
ctx.group_rank = dist.get_rank(group=group)
|
|
if max_rows == 0:
|
|
return tensor.new_empty((0, *tensor.shape[1:]))
|
|
padded = _pad_rows(tensor, max_rows)
|
|
gathered = [torch.zeros_like(padded) for _ in counts]
|
|
dist.all_gather(gathered, padded, group=group)
|
|
return torch.cat([chunk[:count] for chunk, count in zip(gathered, counts)], dim=0)
|
|
|
|
@staticmethod
|
|
def backward(ctx, grad_output):
|
|
local_count = ctx.counts[ctx.group_rank]
|
|
if ctx.max_rows == 0:
|
|
return grad_output.new_empty((0, *grad_output.shape[1:])), None, None, None
|
|
reduced_chunks = []
|
|
for chunk, count in zip(torch.split(grad_output, ctx.counts, dim=0), ctx.counts):
|
|
grad_padded = grad_output.new_zeros((ctx.max_rows, *grad_output.shape[1:]))
|
|
if count:
|
|
grad_padded[:count].copy_(chunk)
|
|
# grad_output is replicated across TP peers (the gathered full view
|
|
# is consumed identically), so this all_reduce sums tp_size copies
|
|
# and injects the tp_size factor documented in the class docstring.
|
|
dist.all_reduce(grad_padded, group=ctx.group)
|
|
reduced_chunks.append(grad_padded)
|
|
grad_padded = reduced_chunks[ctx.group_rank]
|
|
return grad_padded[:local_count].contiguous(), None, None, None
|
|
|
|
|
|
def _all_gather_variable_rows(tensor: torch.Tensor,
|
|
group,
|
|
tp_size: int,
|
|
*,
|
|
preserve_grad: bool = False) -> torch.Tensor:
|
|
if tp_size <= 1 or not dist.is_initialized():
|
|
return tensor
|
|
|
|
local_rows = torch.tensor([tensor.shape[0]], dtype=torch.long, device=tensor.device)
|
|
row_counts = [torch.zeros_like(local_rows) for _ in range(tp_size)]
|
|
dist.all_gather(row_counts, local_rows, group=group)
|
|
counts = [int(item.item()) for item in row_counts]
|
|
max_rows = max(counts) if counts else tensor.shape[0]
|
|
if preserve_grad:
|
|
return _AllGatherVariableRows.apply(tensor, group, tuple(counts), max_rows)
|
|
else:
|
|
padded = _pad_rows(tensor, max_rows)
|
|
gathered = [torch.zeros_like(padded) for _ in range(tp_size)]
|
|
dist.all_gather(gathered, padded, group=group)
|
|
return torch.cat([chunk[:count] for chunk, count in zip(gathered, counts)], dim=0)
|
|
|
|
|
|
def _debug_validate_restore_coverage(payload: RoutedAssignmentPayload, ctx: RestoreContext,
|
|
all_token_indices: torch.Tensor, all_expert_indices: torch.Tensor,
|
|
all_assignment_indices: torch.Tensor, all_capacity_slots: torch.Tensor) -> None:
|
|
active = ~payload.drop_mask & ~payload.pad_mask
|
|
expected_rows = torch.stack((
|
|
payload.token_indices[active].to(torch.long),
|
|
payload.expert_indices[active].to(torch.long),
|
|
payload.assignment_indices[active].to(torch.long),
|
|
payload.capacity_slots[active].to(torch.long),
|
|
),
|
|
dim=1)
|
|
observed_rows = torch.stack((
|
|
all_token_indices.to(torch.long),
|
|
all_expert_indices.to(torch.long),
|
|
all_assignment_indices.to(torch.long),
|
|
all_capacity_slots.to(torch.long),
|
|
),
|
|
dim=1)
|
|
if expected_rows.numel() == 0 and observed_rows.numel() == 0:
|
|
return
|
|
if observed_rows.shape[0] != expected_rows.shape[0]:
|
|
raise RuntimeError("AutoEP+AutoTP restore coverage mismatch: gathered assignment count "
|
|
f"{observed_rows.shape[0]} != expected {expected_rows.shape[0]}")
|
|
if observed_rows.shape[0] <= 4096:
|
|
expected = {tuple(row) for row in expected_rows.detach().cpu().tolist()}
|
|
observed = {tuple(row) for row in observed_rows.detach().cpu().tolist()}
|
|
if observed != expected:
|
|
missing = sorted(expected - observed)[:5]
|
|
duplicate_or_stale = sorted(observed - expected)[:5]
|
|
raise RuntimeError("AutoEP+AutoTP restore coverage mismatch: "
|
|
f"missing={missing} unexpected={duplicate_or_stale}")
|
|
|
|
|
|
def restore_combined(local_combined: torch.Tensor,
|
|
ctx: RestoreContext,
|
|
*,
|
|
tp_group,
|
|
validate_coverage: bool = False) -> torch.Tensor:
|
|
"""Gather TP-partitioned assignment outputs and combine back by token index.
|
|
|
|
The all-gather rebuilds an identical full output on every TP peer, so all
|
|
downstream compute (and the router/gate score gradient) is replicated across
|
|
the folding group. Its differentiable backward injects a ``tp_size`` factor
|
|
(see ``_AllGatherVariableRows``) that the optimizer's TP gradient reducer
|
|
cancels with the AVERAGE strategy. A future true-SP path that kept
|
|
activations sequence-sharded instead of gathering them here would make those
|
|
parameters genuine SUM partials -- the reason the SUM family markers exist
|
|
in ``deepspeed.module_inject.auto_ep_folding``.
|
|
"""
|
|
payload = ctx.original_payload
|
|
local_token_indices = payload.token_indices.index_select(0, ctx.local_indices)
|
|
local_capacity_slots = payload.capacity_slots.index_select(0, ctx.local_indices)
|
|
local_weights = payload.combine_weights.index_select(0, ctx.local_indices).to(local_combined.dtype)
|
|
|
|
all_outputs = _all_gather_variable_rows(local_combined,
|
|
tp_group,
|
|
ctx.tp_size,
|
|
preserve_grad=local_combined.requires_grad)
|
|
all_token_indices = _all_gather_variable_rows(local_token_indices, tp_group, ctx.tp_size).to(torch.long)
|
|
all_capacity_slots = _all_gather_variable_rows(local_capacity_slots, tp_group, ctx.tp_size).to(torch.long)
|
|
all_weights = _all_gather_variable_rows(local_weights,
|
|
tp_group,
|
|
ctx.tp_size,
|
|
preserve_grad=local_weights.requires_grad).to(local_combined.dtype)
|
|
if validate_coverage:
|
|
local_expert_indices = payload.expert_indices.index_select(0, ctx.local_indices)
|
|
local_assignment_indices = payload.assignment_indices.index_select(0, ctx.local_indices)
|
|
all_expert_indices = _all_gather_variable_rows(local_expert_indices, tp_group, ctx.tp_size).to(torch.long)
|
|
all_assignment_indices = _all_gather_variable_rows(local_assignment_indices, tp_group,
|
|
ctx.tp_size).to(torch.long)
|
|
_debug_validate_restore_coverage(payload, ctx, all_token_indices, all_expert_indices, all_assignment_indices,
|
|
all_capacity_slots)
|
|
|
|
if ctx.num_tokens <= 0:
|
|
ctx.num_tokens = int(payload.token_indices.max().item()) + 1 if payload.token_indices.numel() else 0
|
|
output = local_combined.new_zeros((ctx.num_tokens, local_combined.shape[-1]))
|
|
if all_outputs.numel() > 0:
|
|
weight_shape = (-1, ) + (1, ) * (all_outputs.dim() - 1)
|
|
weighted_outputs = all_outputs * all_weights.reshape(weight_shape)
|
|
# Add one top-k slot at a time so token accumulation order stays stable
|
|
# without materializing a [tokens, top_k, hidden] buffer.
|
|
for slot in torch.unique(all_capacity_slots, sorted=True).tolist():
|
|
rows = all_capacity_slots == int(slot)
|
|
output.index_add_(0, all_token_indices[rows], weighted_outputs[rows])
|
|
return output
|
|
|
|
|
|
def dispatch_counters(ctx: RestoreContext) -> dict[str, int]:
|
|
return dict(ctx.counters)
|