wehub-resource-sync
514 lines
24 KiB
Python
514 lines
24 KiB
Python
# Copyright (c) DeepSpeed Team.
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
# DeepSpeed Team
|
|
"""AutoEP + AutoTP folding topology helpers.
|
|
|
|
The functions in this module are pure topology math unless a caller passes
|
|
runtime process-group handles into :class:`FoldingGroupHandles`.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass
|
|
from typing import Iterable
|
|
|
|
import torch
|
|
|
|
AUTOEP_FOLDING_PARAM_FAMILY_ATTR = "ds_autoep_folding_param_family"
|
|
AUTOEP_FOLDING_ROUTER_GATE_REPLICATED_PARAM = "router_gate_replicated"
|
|
AUTOEP_FOLDING_ROUTER_GATE_PARTIAL_PARAM = "router_gate_partial"
|
|
AUTOEP_FOLDING_SP_SHARDED_LAYERNORM_PARAM = "sp_sharded_layernorm"
|
|
AUTOEP_FOLDING_GRAD_CORRECTED_ATTR = "ds_autoep_folding_grad_corrected"
|
|
AUTOEP_FOLDING_GRAD_REDUCE_SKIP = "skip"
|
|
AUTOEP_FOLDING_GRAD_REDUCE_SUM = "sum"
|
|
AUTOEP_FOLDING_GRAD_REDUCE_AVERAGE = "average"
|
|
# Divide by tp_size with NO TP all_reduce. Used for routed-expert parameters: the
|
|
# folded forward all-gathers expert outputs into a replicated full view in
|
|
# ``restore_combined``, whose backward injects a ``tp_size`` factor (same factor the
|
|
# replicated router cancels via AVERAGE). Routed experts are not TP-replicated, so
|
|
# they must not be TP all_reduced; they only need that spurious ``tp_size`` factor
|
|
# divided out. The remaining data-parallel reduction is owned by the expert-data
|
|
# -parallel (EDP) path, and ``/tp_size`` is linear so it composes with that EDP
|
|
# all_reduce in either order.
|
|
AUTOEP_FOLDING_GRAD_REDUCE_EXPERT_TP_CANCEL = "expert_tp_cancel"
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class ParallelFoldingSpec:
|
|
world_size: int
|
|
pp_size: int
|
|
stage_size: int
|
|
tp_size: int
|
|
dp_size: int
|
|
ep_size: int
|
|
etp_size: int
|
|
edp_size: int
|
|
mp_mode: str = "tp"
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class FoldingGroupTables:
|
|
tp_groups: tuple[tuple[int, ...], ...]
|
|
dense_dp_groups: tuple[tuple[int, ...], ...]
|
|
ep_groups: tuple[tuple[int, ...], ...]
|
|
edp_groups: tuple[tuple[int, ...], ...]
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class FoldingGroupHandles:
|
|
spec: ParallelFoldingSpec
|
|
tp_group: object
|
|
dense_dp_group: object
|
|
ep_group: object
|
|
edp_group: object
|
|
ep_group_name: str
|
|
tp_ranks: tuple[int, ...]
|
|
dense_dp_ranks: tuple[int, ...]
|
|
ep_ranks: tuple[int, ...]
|
|
edp_ranks: tuple[int, ...]
|
|
|
|
|
|
def _divisors(value: int) -> list[int]:
|
|
return [candidate for candidate in range(1, value + 1) if value % candidate == 0]
|
|
|
|
|
|
def _require_positive(name: str, value: int) -> None:
|
|
if not isinstance(value, int) or value < 1:
|
|
raise ValueError(f"{name} must be a positive integer, got {value!r}")
|
|
|
|
|
|
def build_folding_spec(
|
|
*,
|
|
world_size: int,
|
|
pp_size: int,
|
|
tp_size: int,
|
|
ep_size: int,
|
|
etp_size: int = 1,
|
|
mp_mode: str = "tp",
|
|
) -> ParallelFoldingSpec:
|
|
"""Build the immutable per-stage folding spec from public config sizes."""
|
|
for name, value in (
|
|
("world_size", world_size),
|
|
("pp_size", pp_size),
|
|
("tensor_parallel.autotp_size", tp_size),
|
|
("expert_parallel.autoep_size", ep_size),
|
|
("expert_parallel.expert_tensor_parallel_size", etp_size),
|
|
):
|
|
_require_positive(name, value)
|
|
|
|
if world_size % pp_size != 0:
|
|
raise ValueError(f"pp_size={pp_size} must divide world_size={world_size}. "
|
|
f"Valid pp_size values: {_divisors(world_size)}")
|
|
|
|
stage_size = world_size // pp_size
|
|
if stage_size % tp_size != 0:
|
|
raise ValueError(f"tensor_parallel.autotp_size={tp_size} must divide the stage size "
|
|
f"(world_size={world_size} / pp_size={pp_size} = {stage_size}). "
|
|
f"Computed dp would be non-integral. Valid autotp_size values: {_divisors(stage_size)}")
|
|
|
|
expert_width = ep_size * etp_size
|
|
if stage_size % expert_width != 0:
|
|
raise ValueError(f"expert_parallel.autoep_size * expert_parallel.expert_tensor_parallel_size "
|
|
f"({ep_size} * {etp_size} = {expert_width}) must divide the stage size "
|
|
f"(world_size={world_size} / pp_size={pp_size} = {stage_size}). "
|
|
f"Computed edp would be non-integral. Valid expert-width values: {_divisors(stage_size)}")
|
|
|
|
return ParallelFoldingSpec(
|
|
world_size=world_size,
|
|
pp_size=pp_size,
|
|
stage_size=stage_size,
|
|
tp_size=tp_size,
|
|
dp_size=stage_size // tp_size,
|
|
ep_size=ep_size,
|
|
etp_size=etp_size,
|
|
edp_size=stage_size // expert_width,
|
|
mp_mode=mp_mode,
|
|
)
|
|
|
|
|
|
def expected_folding_group_tables(spec: ParallelFoldingSpec) -> FoldingGroupTables:
|
|
"""Derive TP, dense-DP, EP, and EDP rank tables without process groups."""
|
|
tp_groups: list[tuple[int, ...]] = []
|
|
dense_dp_groups: list[tuple[int, ...]] = []
|
|
ep_groups: list[tuple[int, ...]] = []
|
|
edp_groups: list[tuple[int, ...]] = []
|
|
|
|
for stage_start in range(0, spec.world_size, spec.stage_size):
|
|
stage_ranks = list(range(stage_start, stage_start + spec.stage_size))
|
|
|
|
for dp_idx in range(spec.dp_size):
|
|
start = dp_idx * spec.tp_size
|
|
tp_groups.append(tuple(stage_ranks[start:start + spec.tp_size]))
|
|
for tp_lane in range(spec.tp_size):
|
|
dense_dp_groups.append(tuple(stage_ranks[tp_lane::spec.tp_size]))
|
|
|
|
local_ep_groups = [
|
|
tuple(stage_ranks[start:start + spec.ep_size]) for start in range(0, len(stage_ranks), spec.ep_size)
|
|
]
|
|
ep_groups.extend(local_ep_groups)
|
|
for pos in range(spec.ep_size):
|
|
edp_groups.append(tuple(group[pos] for group in local_ep_groups))
|
|
|
|
return FoldingGroupTables(
|
|
tp_groups=tuple(tp_groups),
|
|
dense_dp_groups=tuple(dense_dp_groups),
|
|
ep_groups=tuple(ep_groups),
|
|
edp_groups=tuple(edp_groups),
|
|
)
|
|
|
|
|
|
def local_folding_ranks(global_rank: int, spec: ParallelFoldingSpec) -> dict[str, tuple[int, ...]]:
|
|
tables = expected_folding_group_tables(spec)
|
|
result = {}
|
|
for name, groups in (
|
|
("tp", tables.tp_groups),
|
|
("dense_dp", tables.dense_dp_groups),
|
|
("ep", tables.ep_groups),
|
|
("edp", tables.edp_groups),
|
|
):
|
|
result[name] = next(group for group in groups if global_rank in group)
|
|
return result
|
|
|
|
|
|
def _mpu_world_size(mpu, *names: str) -> int | None:
|
|
if mpu is None:
|
|
return None
|
|
for name in names:
|
|
getter = getattr(mpu, name, None)
|
|
if getter is not None:
|
|
return getter()
|
|
return None
|
|
|
|
|
|
def validate_folding_global(
|
|
spec: ParallelFoldingSpec,
|
|
*,
|
|
zero_stage: int = 0,
|
|
sp_size: int = 1,
|
|
deepcompile_enabled: bool = False,
|
|
use_data_before_expert_parallel: bool = False,
|
|
mpu=None,
|
|
autoep_enabled: bool = True,
|
|
tp_preset: str | None = None,
|
|
ep_preset: str | None = None,
|
|
zero_offload_optimizer: bool = False,
|
|
zero_offload_param: bool = False,
|
|
) -> None:
|
|
"""Validate global folding policy before any process group is created."""
|
|
if not autoep_enabled:
|
|
return
|
|
|
|
if deepcompile_enabled and spec.tp_size > 1:
|
|
raise ValueError("DeepCompile with AutoEP+AutoTP folding is not supported; "
|
|
"disable compile.deepcompile or use non-folded AutoEP with tensor_parallel.autotp_size=1.")
|
|
|
|
if spec.tp_size > 1 and spec.pp_size > 1:
|
|
raise ValueError("AutoEP+AutoTP folding currently supports pp_size=1 only; "
|
|
f"got pp_size={spec.pp_size}. Pipeline-parallel validation is planned separately.")
|
|
|
|
if spec.tp_size > 1 and sp_size > 1:
|
|
raise ValueError("tensor_parallel.autotp_size and Ulysses sequence parallelism are mutually exclusive "
|
|
f"for AutoEP folding (autotp_size={spec.tp_size}, sp_size={sp_size}).")
|
|
|
|
if spec.etp_size != 1:
|
|
raise ValueError(f"expert_parallel.expert_tensor_parallel_size={spec.etp_size} is reserved for "
|
|
"expert-internal tensor parallelism and is not supported yet. Use "
|
|
"expert_tensor_parallel_size=1; ETP support is planned as follow-up work.")
|
|
|
|
# Cross-lane expert parallelism (expert_width = ep * etp need NOT be a subset of
|
|
# the dense data-parallel size) is supported: ``expected_folding_group_tables``
|
|
# lays EP groups across consecutive stage ranks while dense DP remains TP-lane
|
|
# strided, so an EP group may span TP lanes and dense-DP ranks while preserving
|
|
# node-local EP groups under node-contiguous rank mappings. The only structural
|
|
# requirement is that the expert width tiles the stage cleanly, which
|
|
# ``build_folding_spec`` already enforces (``stage_size % expert_width == 0``,
|
|
# so ``edp`` is integral). The gradient convention holds across the pool
|
|
# because each family's reduction is keyed to its replication structure, not
|
|
# the EP layout: router/gate and dense/LayerNorm AVERAGE over the TP
|
|
# (token-replication) group; routed experts cancel the restore ``tp_size``
|
|
# factor (EXPERT_TP_CANCEL) and reduce data-parallel over
|
|
# the EDP group. The earlier ``expert_width <= dp`` / ``dp % expert_width == 0``
|
|
# fail-fast limitation is therefore removed; only genuinely non-tiling shapes are
|
|
# rejected above (in ``build_folding_spec``).
|
|
|
|
if tp_preset is not None and ep_preset is not None and tp_preset != ep_preset:
|
|
raise ValueError("tensor_parallel.preset_model and expert_parallel.preset_model must match when both "
|
|
f"are set (tensor_parallel.preset_model={tp_preset!r}, "
|
|
f"expert_parallel.preset_model={ep_preset!r}).")
|
|
|
|
if spec.tp_size > 1 and spec.ep_size == 1:
|
|
raise ValueError("AutoEP+AutoTP folding requires expert_parallel.autoep_size > 1. "
|
|
"The ep=1 local-computation path would duplicate routed-token gradients across TP lanes.")
|
|
|
|
if spec.tp_size > 1 and use_data_before_expert_parallel:
|
|
raise ValueError("expert_parallel with use_data_before_expert_parallel_ is not supported with "
|
|
"AutoEP+AutoTP folding. Disable use_data_before_expert_parallel_.")
|
|
|
|
if spec.tp_size > 1 and zero_stage == 3:
|
|
raise ValueError("AutoEP+AutoTP with ZeRO stage 3 is reserved for the separate ZeRO-3 composition lane. "
|
|
"Use ZeRO stage 0, 1, or 2 for this folding MVP.")
|
|
|
|
if spec.tp_size > 1 and (zero_offload_optimizer or zero_offload_param):
|
|
raise ValueError("ZeRO optimizer/parameter offload with AutoEP+AutoTP folding is not validated yet. "
|
|
"Disable offload or run a follow-up proof for per-family replica groups.")
|
|
|
|
mpu_tp = _mpu_world_size(mpu, "get_tensor_model_parallel_world_size", "get_model_parallel_world_size")
|
|
if mpu_tp not in (None, 1, spec.tp_size):
|
|
raise ValueError(f"mpu tensor/model parallel world size ({mpu_tp}) conflicts with "
|
|
f"tensor_parallel.autotp_size={spec.tp_size}.")
|
|
mpu_pp = _mpu_world_size(mpu, "get_pipeline_model_parallel_world_size", "get_pipeline_parallel_world_size")
|
|
if mpu_pp not in (None, spec.pp_size):
|
|
raise ValueError(f"mpu pipeline parallel world size ({mpu_pp}) conflicts with pp_size={spec.pp_size}.")
|
|
|
|
|
|
def mark_autoep_folding_router_parameter(param) -> None:
|
|
"""Tag a router/gate parameter as the *replicated* folded family (AVERAGE).
|
|
|
|
This is the ONLY family marker applied on the live forward path today:
|
|
``AutoEPMoELayer.__init__`` marks every ``router.*`` parameter with it. The
|
|
folded router runs redundantly on every TP peer (same tokens, same routing)
|
|
and its gradient is reconstructed into a replicated full view by the restore
|
|
all-gather (see ``deepspeed.moe.ep_tp_dispatch._AllGatherVariableRows`` and
|
|
``restore_combined``). That all-gather backward scales each peer's slice by
|
|
``tp_size``, so the extra TP reduction must AVERAGE (all_reduce then divide
|
|
by ``tp_size``); SUM would leave the ``tp_size`` factor, i.e. the 2.0x
|
|
parity regression the CPU/Gloo tests guard.
|
|
"""
|
|
setattr(param, AUTOEP_FOLDING_PARAM_FAMILY_ATTR, AUTOEP_FOLDING_ROUTER_GATE_REPLICATED_PARAM)
|
|
|
|
|
|
def mark_autoep_folding_partial_router_parameter(param) -> None:
|
|
"""Tag a router/gate parameter as a *routed-token partial* family (SUM).
|
|
|
|
Forward-looking contract; NOT used on the current forward path -- only the
|
|
unit tests in ``tests/unit/v1/moe/test_autoep_autotp_grad_parity.py`` set
|
|
it. Use it only for a future design where the router's per-token work is
|
|
genuinely partitioned across peers and the slices are NOT all-gathered back
|
|
into a replicated full view, so each peer holds a real partial gradient that
|
|
must be SUMed. Such a router is a SUM partial in any token-partitioned mode
|
|
(``mp_mode in {"tp", "sp"}``) because its partition can ride the existing
|
|
expert-dispatch all-to-all without changing the dense activation layout.
|
|
Prove the SUM with a parity test (like the existing router/gate cases)
|
|
before enabling it on a real forward path.
|
|
"""
|
|
setattr(param, AUTOEP_FOLDING_PARAM_FAMILY_ATTR, AUTOEP_FOLDING_ROUTER_GATE_PARTIAL_PARAM)
|
|
|
|
|
|
def mark_autoep_folding_sp_sharded_layernorm_parameter(param) -> None:
|
|
"""Tag a LayerNorm parameter as *SP-sequence-sharded* family (SUM under SP).
|
|
|
|
Forward-looking contract; NOT used on the current forward path -- only the
|
|
unit tests set it. Unlike the router, a LayerNorm has no adjacent dispatch
|
|
all-to-all to ride on, so the only way to token-partition it is to shard the
|
|
sequence dimension of the dense activations, which is Sequence Parallel by
|
|
definition. It therefore becomes a SUM partial only when ``mp_mode == "sp"``
|
|
and otherwise falls back to the replicated AVERAGE. Today ``tp_size > 1``
|
|
with sequence parallelism is rejected in ``validate_folding_global``; this
|
|
marker is the explicit contract for when that restriction is lifted, and
|
|
must be backed by a parity test before use.
|
|
"""
|
|
setattr(param, AUTOEP_FOLDING_PARAM_FAMILY_ATTR, AUTOEP_FOLDING_SP_SHARDED_LAYERNORM_PARAM)
|
|
|
|
|
|
def _is_moe_param_marker(param) -> bool:
|
|
return hasattr(param, "allreduce") and not param.allreduce
|
|
|
|
|
|
def _is_model_parallel_param_marker(param) -> bool:
|
|
return bool(getattr(param, "model_parallel", False) or getattr(param, "tensor_model_parallel", False))
|
|
|
|
|
|
def _autoep_folding_param_family(param, *, param_name: str | None = None) -> str | None:
|
|
"""Resolve a parameter's folded reduction family.
|
|
|
|
An explicit ``mark_autoep_folding_*`` tag always wins. The ``.router.`` name
|
|
match is only a redundant safety net: ``AutoEPMoELayer`` already tags router
|
|
params, so this fallback merely keeps the conservative *replicated* (AVERAGE)
|
|
classification if some router param ever reaches the reducer untagged. It
|
|
never returns a SUM family by name -- SUM families are opt-in via explicit
|
|
markers only, so any unrecognized replicated/dense/LayerNorm param falls
|
|
through to the AVERAGE default rather than being silently over-scaled.
|
|
"""
|
|
family = getattr(param, AUTOEP_FOLDING_PARAM_FAMILY_ATTR, None)
|
|
if family is not None:
|
|
return family
|
|
if param_name is not None and ".router." in param_name:
|
|
return AUTOEP_FOLDING_ROUTER_GATE_REPLICATED_PARAM
|
|
return None
|
|
|
|
|
|
def autoep_folding_gradient_reduction_strategy(
|
|
folding_spec: ParallelFoldingSpec | None,
|
|
param,
|
|
*,
|
|
param_name: str | None = None,
|
|
) -> str:
|
|
"""Classify one folded TP/SP gradient as ``sum``, ``average``, or ``skip``.
|
|
|
|
TP means Tensor Parallel and SP means Sequence Parallel. The parallel mode
|
|
alone is not a safe SUM-vs-AVG selector because different parameter
|
|
families see different backward semantics:
|
|
|
|
- Router/gate parameters that are explicitly marked as routed-token
|
|
partials in TP/SP token-partitioned modes receive one partial gradient per
|
|
lane, so their TP/SP reduction is a SUM. The current AutoEP folded router
|
|
gate is marked ``router_gate_replicated`` because the full-flow backward
|
|
reaches this reducer as a lane-replicated gradient; that family uses the
|
|
same AVERAGE normalization as other replicated parameters.
|
|
- Dense and LayerNorm parameters that are merely replicated by TP folding
|
|
are not routed-token partials; blindly SUMing them scales gradients by
|
|
the TP size, so their extra TP reduction is an AVERAGE.
|
|
- A true SP-sharded LayerNorm would be a partial-gradient parameter and
|
|
should SUM. The current AutoEP folding path does not mark runtime
|
|
LayerNorm parameters that way; the marker and strategy boundary exist so
|
|
future SP support has an explicit contract instead of reusing the dense
|
|
replicated default by accident.
|
|
- Model-parallel (genuinely TP-sharded) parameters are SKIP because the
|
|
TP-specific path owns their reduction.
|
|
- Routed-expert parameters are EXPERT_TP_CANCEL: their data-parallel
|
|
reduction is owned by the EP/EDP path, but the folded forward all-gathers
|
|
their outputs into a replicated full view in ``restore_combined`` (whose
|
|
backward injects a ``tp_size`` factor), so the expert-weight gradient
|
|
reaches the optimizer ``tp_size`` times too large. Experts are not
|
|
TP-replicated, so the fix is a plain ``/tp_size`` (no TP all_reduce), which
|
|
is linear and composes with the EDP all_reduce in any order. Without this,
|
|
folded expert gradients are over-scaled by ``tp_size`` -- invisible to
|
|
scale-invariant Adam but real for SGD/Lion/Muon and for gradient clipping
|
|
(it inflates the expert contribution to the global grad norm).
|
|
|
|
Underlying rule and mechanism: a folded parameter is replicated (AVERAGE)
|
|
when the forward reconstructs its partitioned work into an identical full
|
|
view inside the layer, and a genuine partial (SUM) only when the shard is
|
|
kept all the way to the loss. Today the router/gate is partitioned across
|
|
TP peers for dispatch but then all-gathered back by ``restore_combined``
|
|
(see ``deepspeed.moe.ep_tp_dispatch``), whose backward scales each peer's
|
|
gradient by ``tp_size``; the TP all_reduce then yields ``tp_size *
|
|
full_grad`` and AVERAGE divides it out. Reducing with SUM would leave that
|
|
factor -- the 2.0x router/gate parity regression the CPU/Gloo tests guard.
|
|
The router can be a SUM partial in either ``tp`` or ``sp`` mode because its
|
|
token partition can ride the existing dispatch all-to-all, whereas a
|
|
LayerNorm becomes a partial only under true ``sp`` (sequence sharding): it
|
|
has no adjacent all-to-all, so partitioning it requires changing the dense
|
|
activation layout, which is Sequence Parallel by definition.
|
|
|
|
Both the DeepSpeedEngine path and the ZeRO-2 path call this helper so the
|
|
policy cannot silently drift between optimizers.
|
|
"""
|
|
if folding_spec is None or getattr(folding_spec, "tp_size", 1) <= 1:
|
|
return AUTOEP_FOLDING_GRAD_REDUCE_SKIP
|
|
if _is_model_parallel_param_marker(param):
|
|
# Genuinely TP-sharded (column/row-parallel) params: the TP-specific path
|
|
# owns their reduction. Not produced by the folded skip-partition MVP.
|
|
return AUTOEP_FOLDING_GRAD_REDUCE_SKIP
|
|
if _is_moe_param_marker(param):
|
|
# Routed-expert params. Their EP/EDP data-parallel reduction is owned by
|
|
# the expert path, but the folded forward routes their outputs through the
|
|
# ``restore_combined`` all-gather, whose backward leaves a ``tp_size``
|
|
# factor on the expert-weight gradient (the same factor the replicated
|
|
# router cancels with AVERAGE). Experts are NOT TP-replicated, so they must
|
|
# not be TP all_reduced; the factor is cancelled with a plain ``/tp_size``.
|
|
return AUTOEP_FOLDING_GRAD_REDUCE_EXPERT_TP_CANCEL
|
|
|
|
family = _autoep_folding_param_family(param, param_name=param_name)
|
|
mp_mode = getattr(folding_spec, "mp_mode", "tp")
|
|
token_partitioned_mode = mp_mode in ("tp", "sp")
|
|
if family == AUTOEP_FOLDING_ROUTER_GATE_PARTIAL_PARAM:
|
|
return AUTOEP_FOLDING_GRAD_REDUCE_SUM if token_partitioned_mode else AUTOEP_FOLDING_GRAD_REDUCE_AVERAGE
|
|
if family == AUTOEP_FOLDING_ROUTER_GATE_REPLICATED_PARAM:
|
|
return AUTOEP_FOLDING_GRAD_REDUCE_AVERAGE
|
|
if family == AUTOEP_FOLDING_SP_SHARDED_LAYERNORM_PARAM and mp_mode == "sp":
|
|
return AUTOEP_FOLDING_GRAD_REDUCE_SUM
|
|
return AUTOEP_FOLDING_GRAD_REDUCE_AVERAGE
|
|
|
|
|
|
def reduce_autoep_folding_gradient(
|
|
folding_spec: ParallelFoldingSpec | None,
|
|
param,
|
|
grad,
|
|
*,
|
|
tp_group,
|
|
param_name: str | None = None,
|
|
) -> str:
|
|
strategy = autoep_folding_gradient_reduction_strategy(folding_spec, param, param_name=param_name)
|
|
if strategy == AUTOEP_FOLDING_GRAD_REDUCE_SKIP or grad is None or grad.data.is_sparse:
|
|
return strategy
|
|
|
|
from deepspeed import comm as dist
|
|
|
|
grad_data = grad.data
|
|
tp_world_size = dist.get_world_size(group=tp_group)
|
|
|
|
# Routed experts: cancel the ``tp_size`` factor the restore all-gather leaves,
|
|
# WITHOUT a TP all_reduce (experts are not TP-replicated; cross-TP summation of
|
|
# disjoint expert-token slices is owned by the EDP all_reduce). ``/tp_size`` is
|
|
# linear, so it composes with that EDP reduction in either order.
|
|
if strategy == AUTOEP_FOLDING_GRAD_REDUCE_EXPERT_TP_CANCEL:
|
|
if tp_world_size > 1:
|
|
grad_data.div_(tp_world_size)
|
|
return strategy
|
|
|
|
if grad_data.dtype != torch.float32:
|
|
reduced = grad_data.float()
|
|
dist.all_reduce(reduced, group=tp_group)
|
|
if strategy == AUTOEP_FOLDING_GRAD_REDUCE_AVERAGE:
|
|
reduced.div_(tp_world_size)
|
|
grad_data.copy_(reduced.to(grad_data.dtype))
|
|
return strategy
|
|
|
|
dist.all_reduce(grad_data, group=tp_group)
|
|
if strategy == AUTOEP_FOLDING_GRAD_REDUCE_AVERAGE:
|
|
grad_data.div_(tp_world_size)
|
|
return strategy
|
|
|
|
|
|
def is_autoep_folding_gradient_corrected(param) -> bool:
|
|
return bool(getattr(param, AUTOEP_FOLDING_GRAD_CORRECTED_ATTR, False))
|
|
|
|
|
|
def clear_autoep_folding_gradient_corrected(param) -> None:
|
|
if hasattr(param, AUTOEP_FOLDING_GRAD_CORRECTED_ATTR):
|
|
setattr(param, AUTOEP_FOLDING_GRAD_CORRECTED_ATTR, False)
|
|
|
|
|
|
def apply_folding_correction_to_grad_buffer(
|
|
folding_spec: ParallelFoldingSpec | None,
|
|
param,
|
|
grad,
|
|
*,
|
|
tp_group,
|
|
param_name: str | None = None,
|
|
use_correction_marker: bool = True,
|
|
) -> str:
|
|
if use_correction_marker and is_autoep_folding_gradient_corrected(param):
|
|
return AUTOEP_FOLDING_GRAD_REDUCE_SKIP
|
|
|
|
strategy = reduce_autoep_folding_gradient(folding_spec, param, grad, tp_group=tp_group, param_name=param_name)
|
|
if use_correction_marker and strategy != AUTOEP_FOLDING_GRAD_REDUCE_SKIP:
|
|
setattr(param, AUTOEP_FOLDING_GRAD_CORRECTED_ATTR, True)
|
|
return strategy
|
|
|
|
|
|
def _normalize_rank_groups(groups: Iterable[Iterable[int]]) -> set[tuple[int, ...]]:
|
|
return {tuple(int(rank) for rank in group) for group in groups}
|
|
|
|
|
|
def assert_group_matches_spec(existing_rank_lists, spec: ParallelFoldingSpec, *, group_kind: str = "ep_edp") -> None:
|
|
"""Ensure cached ``ep_size_N`` rank lists match the requested folding spec."""
|
|
tables = expected_folding_group_tables(spec)
|
|
expected_ep = _normalize_rank_groups(tables.ep_groups)
|
|
expected_edp = _normalize_rank_groups(tables.edp_groups)
|
|
|
|
if isinstance(existing_rank_lists, dict):
|
|
observed_ep = existing_rank_lists.get("ep", [])
|
|
observed_edp = existing_rank_lists.get("edp", [])
|
|
else:
|
|
observed_ep, observed_edp = existing_rank_lists
|
|
|
|
for group in _normalize_rank_groups(observed_ep):
|
|
if group not in expected_ep:
|
|
raise RuntimeError(f"Cached expert-parallel group {group} does not match folding spec {spec}.")
|
|
for group in _normalize_rank_groups(observed_edp):
|
|
if group not in expected_edp:
|
|
raise RuntimeError(f"Cached expert-data-parallel group {group} does not match folding spec {spec}.")
|