Skip to content

Commit 66fef47

Browse files
jinyouzhiCopilot
andauthored
[AffineIR] Storage maps of AutoEP experts under EDP and ZeRO (#8941)
## Summary #8800 describes AutoEP expert weights as one packed `[local_experts, ...]` tensor per EP rank. What a process actually stores is different: - under ZeRO-0/1/2, every expert-data-parallel (EDP) replica holds that tensor; - under ZeRO-3, each EDP rank holds a padded fragment of it. So a plan built from #8800's maps cannot read from training storage until the storage is first normalized to packed EP tensors. This PR describes the storage itself. It nests #8800's placement map with #8885's ZeRO storage maps, so `plan_transfer` (#8623) reads directly from runtime storage: ``` AutoEP placement map (#8800) --nest--> ZeRO storage map per EP rank (#8885) -> storage map in global ranks -> plan_transfer (#8623) ``` Related: #8252, #8800, #8885. ## Changes - `AutoEPMoELayer.expert_storage_affine_maps()` returns one storage map each for `w1`, `w2` and `w3`. - Ranks are global ranks; the logical shape is the parameter's own shape. - Each shard is the rank's flat storage, as `zero_param_storage_shard(param)` (#8885) returns it: `param.data` under ZeRO-0/1/2, or `param.ds_tensor` without its trailing padding under ZeRO-3. - Under ZeRO-0/1/2, all EDP replicas are listed as holders, so the planner's `replica_selector` can choose among them. - Under ZeRO-3, a rank whose partition is all padding, or whose EP rank holds no experts, has an empty shard. - Fails loudly instead of guessing: - mixed ZeRO-3 and non-ZeRO-3 expert params; - a storage group that does not match the EDP table; - a `ds_tensor` size that differs from even partitioning; - expert tensor parallelism, quantized ZeRO-3 storage and NVMe-offloaded ZeRO-3 storage (`NotImplementedError`). - `autoep_affine.autoep_storage_affine_map(descriptor, logical_shape, holders_by_ep_rank, partition_numel_by_ep_rank=None)` is the layer-independent builder. It is also convenient for target maps. - `groups` keeps every EDP group's ranks (`_get_expert_data_parallel_group_table`), not only the caller's. - Every rank already creates all groups, so the maps are built **locally, without a collective**, and any rank may call it alone. - The table is cleared together with the other group state. - `autotp_storage_affine_maps(model)` (#8885) takes the maps of AutoEP experts from their layer. AutoTP leaves experts whole on every TP rank, so before this it described them as replicated over TP: one EP rank's packed tensor described as the whole parameter. With `optimizer_states=True` it raises `NotImplementedError` for AutoEP experts before any collective. - Docs: `autoep.rst` gains a section on `expert_storage_affine_maps()`; `affine_ir_spec.md` §10 is updated (below). ## Authoritative and derived copies Storage maps list only **authoritative** copies: the storage the optimizer step updates and the next forward reads (`param.data`, ZeRO-3 `ds_tensor`). A write through a map therefore lasts. The runtime also keeps **derived** copies that it regenerates from authoritative ones: | Derived copy | Regenerated by | |---|---| | gathered ZeRO-3 persistent parameter | re-gather after each step | | ZeRO++ hpZ secondary partition | forward; invalidated after the step | | fused inference weights (hybrid engine) | gather + fuse for generation | | low-precision params **on a training target** | cast from the fp32 master weights at every step | A write to a derived copy alone is lost or leaves the authoritative copy stale. So maps leave derived copies out, and `zero_param_storage_shard` refuses a gathered parameter (call `engine.empty_partition_cache()` first). Whether a copy is authoritative depends on the side: the low-precision weights are authoritative for a rollout engine, which is what these maps serve, but derived on a training engine, where loading weights needs the fp32 master (not described). ## Not in scope - Executing plans (#8809), fp32 master / optimizer state of experts, and expert tensor parallelism. - Replica selection policy (#8252). ## Follow-ups 1. **Placement changes (EPLB).** The maps are computed locally because the placement, the EDP group table and the ZeRO-3 even split are known to every rank. Once the runtime can change placement at run time (expert load balancing, redundant experts), this stays correct only if every rank uses the same placement: maps and plans must be rebuilt after every rank has switched, ideally checked with a placement version. The descriptor already expresses uneven and redundant placements (redundant experts become extra holders); today the runtime packs only the contiguous, uniform placement, and anything else raises `NotImplementedError`. 2. **Expert fp32 master weights and optimizer states.** Needed when a training engine is the target or for training-side resharding. Unlike the weights, this needs an all-gather over the EDP group and an exchange over EP, then nesting as here. 3. **Derived copies as read-only sources.** A fresh derived copy (e.g. a gathered ZeRO-3 persistent parameter after the step) holds the same values as the authoritative one and could be offered to the planner as an extra source, never as a target. Its validity is time-dependent, so it would be an explicit opt-in that does not affect cached plans. ## Tests - `tests/unit/v1/moe/test_autoep_checkpoint.py::TestAutoEPExpertStorageMaps` (4 ranks, EP 2), with ZeRO-2, ZeRO-3, and ZeRO-2/3 + folded AutoTP 2: - On a real engine, `extract(full, rank)` equals each rank's `zero_param_storage_shard(param)` bit for bit. `full` comes from `GatheredParameters` and the runtime's packing order, so the reference does not depend on the affine maps. - Rank 0 builds the maps alone while the others wait at a barrier: a hidden collective would deadlock. The result equals what every rank builds. - Planning from the storage maps into a different 4-rank placement `[[3], [0], [2], [1]]` and executing it into NaN-filled targets reproduces the expected experts exactly. - With AutoTP, `autotp_storage_affine_maps(model)` gives the experts the same maps as their layer (fails without the fix), and refuses expert optimizer states. - Quantized and NVMe-offloaded ZeRO-3 expert storage is refused. - `tests/unit/checkpoint/test_autoep_affine.py`: - storage maps for replicated and ZeRO-3 layouts; - every EDP replica offered to the selector; - direct ZeRO-3 → new placement plans; - explicit rank IDs; - invalid holders refused. - `tests/unit/utils/test_groups.py`: the EDP table survives subgroup destruction. Validation on 8× Intel Arc Pro B70 (torch 2.14.0.dev20260804+xpu): `test_autoep_affine`, `test_autoep_affine_universal`, `test_zero_affine`, `test_affine_transfer_plan`, `test_affine_shard_map`, `test_autoep_zero3_metadata`, `tests/unit/utils/test_groups`, `test_autoep_checkpoint`, `test_autoep_unit`: **531 passed, 8 skipped**. After the AutoTP fix: `TestAutoEPExpertStorageMaps` 6 passed, `test_zero_affine` 27 passed. --------- Signed-off-by: Jin, Youzhi <youzhi.jin@intel.com> Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
1 parent 9148ab2 commit 66fef47

9 files changed

Lines changed: 568 additions & 21 deletions

File tree

‎deepspeed/checkpoint/affine_ir_spec.md‎

Lines changed: 21 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -613,6 +613,8 @@ expert parameters, and `AutoEPMoELayer.plan_expert_transfer(target_placement)` c
613613
them with a target descriptor through `plan_transfer` into shard-to-shard copies. Plan
614614
ranks are EP-local ranks of the respective placements; the plan is not executed, and ZeRO
615615
and EDP storage must still be normalized before its offsets apply.
616+
`AutoEPMoELayer.expert_storage_affine_maps()` describes that storage in global ranks
617+
(§10).
616618

617619
`deepspeed/checkpoint/zero_affine.py` takes the first step for live training storage. It
618620
describes, in global ranks, where ZeRO keeps each element of a parameter: the low-precision
@@ -689,8 +691,12 @@ cannot be merged across those rows and keeps one piece per row. Callers:
689691
- `autotp_storage_affine_maps(model)` nests each parameter's published AutoTP layout map
690692
with the storage of every TP shard. A parameter AutoTP leaves whole in a model built under
691693
`zero.Init` stays partitioned over the whole world, so every TP rank reports the storage of
692-
one shared copy; that storage describes the parameter on its own, without nesting.
694+
one shared copy; that storage describes the parameter on its own, without nesting. The
695+
experts of an AutoEP layer are whole on every TP rank and differ by EP rank; their maps
696+
are the layer's `expert_storage_affine_maps()`.
693697
- AutoEP expert storage nests the placement's layout map with each EP rank's EDP storage.
698+
It needs no communication: the placement, the EDP group table and the ZeRO-3 even split
699+
are known to every rank.
694700

695701
**10.4 Replica choice** stays with the planner (§2.4): every holder appears in `locations`,
696702
and the planner's default or a `replica_selector` picks one.
@@ -700,5 +706,17 @@ A caller builds them once and executes the plan for every sync, until either sid
700706
placement, parallel sizes or ZeRO partitioning changes. A parameter whose buffer is
701707
reallocated keeps its plan; only the shard views must be fetched again.
702708

703-
**10.6 Not described.** Offload tiers, MiCS / hpZ, expert tensor parallelism, quantized
704-
ZeRO-3 storage, non-elementwise transforms.
709+
**10.6 Authoritative and derived copies.** A storage map lists only *authoritative* copies:
710+
storage that the optimizer step updates and the next forward reads, such as `param.data` or
711+
a ZeRO-3 `ds_tensor`. A write through the map therefore lasts. The runtime also keeps
712+
*derived* copies that it regenerates from authoritative ones: the gathered copy of a
713+
persistent ZeRO-3 parameter, ZeRO++ hpZ secondary partitions, fused inference weights.
714+
Writing a derived copy alone is lost or leaves the authoritative copy stale, so maps leave
715+
it out and shard accessors refuse it (`zero_param_storage_shard` on a gathered parameter).
716+
Whether a copy is authoritative depends on the side: the low-precision parameters are
717+
authoritative for a rollout engine, but on a training engine they are derived from the fp32
718+
master weights, which the next step casts back over them.
719+
720+
**10.7 Not described.** Offload tiers, MiCS / hpZ, expert tensor parallelism, quantized
721+
ZeRO-3 storage, non-elementwise transforms, derived copies, and AutoEP expert optimizer
722+
states.

‎deepspeed/checkpoint/autoep_affine.py‎

Lines changed: 74 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -2,13 +2,15 @@
22
# DeepSpeed Team
33
"""AutoEP expert placement descriptors and their affine lowering.
44
5-
The descriptor records placement in EP-local rank coordinates. ZeRO and EDP
6-
fragments are deliberately outside this map: callers must first normalize
7-
storage to one logical, packed expert tensor per EP rank.
5+
The descriptor records placement in EP-local rank coordinates, and
6+
``autoep_placement_to_affine_map`` describes one logical, packed expert tensor
7+
per EP rank. ``autoep_storage_affine_map`` adds the expert-data-parallel and
8+
ZeRO layer on top: it describes what each global rank actually stores.
89
"""
910

10-
from deepspeed.checkpoint.affine import AFFINE_MAP_FORMAT_VERSION, AffinePiece, ParamAffineMap
11+
from deepspeed.checkpoint.affine import AFFINE_MAP_FORMAT_VERSION, AffinePiece, ParamAffineMap, nest_affine_maps
1112
from deepspeed.checkpoint.affine_transfer import plan_transfer
13+
from deepspeed.checkpoint.zero_affine import zero_storage_affine_map
1214
from deepspeed.checkpoint.constants import (AFFINE_MAP, AFFINE_MAP_PARAMS, AFFINE_MAP_VERSION, AUTOEP_AFFINE_MAPS,
1315
AUTOEP_EXPERT_PLACEMENT, AUTOEP_PLACEMENT_EP_SIZE,
1416
AUTOEP_PLACEMENT_EXPERTS, AUTOEP_PLACEMENT_NUM_EXPERTS,
@@ -23,6 +25,7 @@
2325
'legacy_uniform_autoep_placement_descriptor',
2426
'autoep_placements_equal',
2527
'autoep_placement_to_affine_map',
28+
'autoep_storage_affine_map',
2629
'autoep_metadata_to_affine_map',
2730
'extract_autoep_rank_tensor',
2831
'plan_autoep_placement_transfer',
@@ -169,6 +172,40 @@ def autoep_placement_to_affine_map(descriptor, logical_shape):
169172
return affine_map
170173

171174

175+
def autoep_storage_affine_map(descriptor, logical_shape, holders_by_ep_rank, partition_numel_by_ep_rank=None):
176+
"""Lower one expert parameter to the storage each global rank actually holds.
177+
178+
``holders_by_ep_rank`` maps every rank ID of ``descriptor`` to the global ranks that store that
179+
EP rank's packed ``[local_experts, ...]`` tensor, in expert-data-parallel (EDP) rank order.
180+
181+
Without ``partition_numel_by_ep_rank``, every holder stores the whole packed tensor, as EDP
182+
replicas do under ZeRO-0/1/2. With it, holder ``d`` of EP rank ``r`` stores elements
183+
``[d * n, (d + 1) * n)`` of the flattened packed tensor, where ``n`` is
184+
``partition_numel_by_ep_rank[r]``, as ZeRO-3 does. ZeRO-3's trailing padding holds no
185+
parameter elements and is not part of the shard, so a holder whose partition is all padding,
186+
like every holder of an EP rank without experts, has an empty shard.
187+
188+
The returned map describes the parameter in its own shape and is keyed by global rank; each
189+
shard is the holder's flat storage. It nests the placement's layout map with ZeRO's
190+
storage map of each packed tensor, so every holder of an element, across EDP replicas and
191+
redundant experts, is listed in the piece's ``locations``; choosing among them is left to the
192+
transfer planner's replica selection.
193+
"""
194+
layout = autoep_placement_to_affine_map(descriptor, logical_shape)
195+
entries_by_rank = {entry[AUTOEP_PLACEMENT_RANK]: entry for entry in descriptor[AUTOEP_PLACEMENT_RANKS]}
196+
holders_by_ep_rank = _validated_storage_holders(holders_by_ep_rank, entries_by_rank)
197+
198+
storage_by_ep_rank = {}
199+
for ep_rank, global_ranks in holders_by_ep_rank.items():
200+
packed_numel = _product(layout.shard_shapes[ep_rank])
201+
partition_numel = None
202+
if partition_numel_by_ep_rank is not None:
203+
partition_numel = _storage_partition_numel(partition_numel_by_ep_rank, ep_rank, len(global_ranks),
204+
packed_numel)
205+
storage_by_ep_rank[ep_rank] = zero_storage_affine_map(packed_numel, global_ranks, partition_numel)
206+
return nest_affine_maps(layout, storage_by_ep_rank)
207+
208+
172209
def autoep_metadata_to_affine_map(metadata, logical_shape, param_name=None):
173210
"""Load a persisted AutoEP affine map, falling back to legacy placement lowering."""
174211
serialized_map = metadata.get(AFFINE_MAP)
@@ -252,6 +289,39 @@ def _pieces_for_rank(experts, holders, expert_shape, expert_numel, source_stride
252289
return pieces
253290

254291

292+
def _validated_storage_holders(holders_by_ep_rank, entries_by_rank):
293+
if not isinstance(holders_by_ep_rank, dict):
294+
raise TypeError('holders_by_ep_rank must map each EP rank ID to its global holder ranks.')
295+
if set(holders_by_ep_rank) != set(entries_by_rank):
296+
raise ValueError(f'holders_by_ep_rank names EP ranks {sorted(holders_by_ep_rank)}; the placement '
297+
f'descriptor names {sorted(entries_by_rank)}.')
298+
seen = set()
299+
validated = {}
300+
for ep_rank in sorted(holders_by_ep_rank):
301+
global_ranks = list(holders_by_ep_rank[ep_rank])
302+
if not global_ranks:
303+
raise ValueError(f'EP rank {ep_rank} has no global holder rank.')
304+
for global_rank in global_ranks:
305+
if not isinstance(global_rank, int) or isinstance(global_rank, bool) or global_rank < 0:
306+
raise ValueError(f'Global holder rank {global_rank!r} of EP rank {ep_rank} is not a '
307+
'non-negative integer.')
308+
if global_rank in seen:
309+
raise ValueError(f'Global rank {global_rank} holds storage for more than one EP rank entry.')
310+
seen.add(global_rank)
311+
validated[ep_rank] = global_ranks
312+
return validated
313+
314+
315+
def _storage_partition_numel(partition_numel_by_ep_rank, ep_rank, num_holders, packed_numel):
316+
partition_numel = partition_numel_by_ep_rank.get(ep_rank)
317+
if not isinstance(partition_numel, int) or isinstance(partition_numel, bool) or partition_numel < 0:
318+
raise ValueError(f'Partition size {partition_numel!r} of EP rank {ep_rank} is not a non-negative integer.')
319+
if partition_numel * num_holders < packed_numel:
320+
raise ValueError(f'EP rank {ep_rank} packs {packed_numel} elements, but its {num_holders} holders of '
321+
f'{partition_numel} elements each cannot store them.')
322+
return partition_numel
323+
324+
255325
def _expert_holders(entries_by_rank, num_experts):
256326
holders = {expert_id: set() for expert_id in range(num_experts)}
257327
for rank, entry in entries_by_rank.items():

‎deepspeed/checkpoint/zero_affine.py‎

Lines changed: 46 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -268,7 +268,8 @@ def autotp_storage_affine_maps(model, optimizer_states=False, tp_group=None):
268268
master copy and optimizer states instead, matching ``zero_optimizer_storage_shard``, and a
269269
parameter without them maps to ``None``. A parameter AutoTP did not partition is replicated
270270
over the tensor-parallel group, unless it kept the partition over the whole world that
271-
``zero.Init`` gave it, which describes it on its own.
271+
``zero.Init`` gave it, which describes it on its own. The experts of an AutoEP layer are
272+
described by its ``expert_storage_affine_maps``; their optimizer states are not supported.
272273
273274
Collective: the tensor-parallel ranks exchange their shards' storage maps in one all-gather
274275
over ``tp_group`` (the AutoTP group by default), and every rank must call this with the same
@@ -286,15 +287,47 @@ def autotp_storage_affine_maps(model, optimizer_states=False, tp_group=None):
286287
if marker is not None:
287288
marker()
288289

289-
named_params = list(model.named_parameters())
290+
# AutoTP leaves AutoEP experts whole on every tensor-parallel rank, but each EP rank holds
291+
# different experts, so a replicated layout would describe one EP rank's packed tensor as the
292+
# whole parameter. The AutoEP layer knows its placement and describes their storage instead.
293+
from deepspeed.module_inject.auto_ep_layer import AutoEPMoELayer
294+
autoep_layers = [module for module in model.modules() if isinstance(module, AutoEPMoELayer)]
295+
if optimizer_states and autoep_layers:
296+
raise NotImplementedError('AutoEP expert optimizer states have no affine storage map.')
297+
298+
# Each EP rank holds different experts, so one rank may refuse its expert storage (for example
299+
# a quantized partition) while its tensor-parallel peers accept theirs. The refusal travels with
300+
# the all-gather so that every peer raises instead of waiting in it.
301+
local_error = None
302+
expert_maps = {}
303+
try:
304+
for module in autoep_layers:
305+
for weight_name, storage_map in module.expert_storage_affine_maps().items():
306+
expert_maps[id(getattr(module.experts, weight_name))] = storage_map
307+
except (NotImplementedError, ValueError, RuntimeError) as error:
308+
local_error = error
309+
310+
named_params = []
311+
for name, param in model.named_parameters():
312+
if id(param) not in expert_maps:
313+
named_params.append((name, param))
290314
params = [param for _, param in named_params]
291-
if optimizer_states:
292-
storage_maps = zero_optimizer_storage_affine_maps(params)
315+
if local_error is not None:
316+
local = {'error': f'{type(local_error).__name__}: {local_error}'}
293317
else:
294-
storage_maps = [zero_param_storage_affine_map(param) for param in params]
295-
local = [None if storage_map is None else storage_map.to_dict() for storage_map in storage_maps]
318+
if optimizer_states:
319+
storage_maps = zero_optimizer_storage_affine_maps(params)
320+
else:
321+
storage_maps = [zero_param_storage_affine_map(param) for param in params]
322+
local = [None if storage_map is None else storage_map.to_dict() for storage_map in storage_maps]
296323
gathered = [None] * tp_world_size
297324
dist.all_gather_object(gathered, local, group=tp_group)
325+
if local_error is not None:
326+
raise local_error
327+
for tp_rank, rank_maps in enumerate(gathered):
328+
if isinstance(rank_maps, dict):
329+
raise RuntimeError(f"Tensor-parallel rank {tp_rank} cannot describe its AutoEP expert storage: "
330+
f"{rank_maps['error']}")
298331

299332
maps = {}
300333
for position, (name, param) in enumerate(named_params):
@@ -314,4 +347,10 @@ def autotp_storage_affine_maps(model, optimizer_states=False, tp_group=None):
314347
layout = replicated_map(layout.logical_shape, 1)
315348
inner = {0: inner[0]}
316349
maps[name] = nest_affine_maps(layout, inner)
317-
return maps
350+
ordered = {}
351+
for name, param in model.named_parameters():
352+
if id(param) in expert_maps:
353+
ordered[name] = expert_maps[id(param)]
354+
else:
355+
ordered[name] = maps[name]
356+
return ordered

‎deepspeed/module_inject/auto_ep_layer.py‎

Lines changed: 78 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@
1414

1515
from __future__ import annotations
1616

17+
import math
1718
from functools import lru_cache
1819
from typing import Literal, NamedTuple
1920

@@ -22,7 +23,8 @@
2223
import deepspeed.comm as dist
2324
from deepspeed.accelerator import get_accelerator
2425
from deepspeed.checkpoint.autoep_affine import (autoep_experts_for_rank, autoep_placement_to_affine_map,
25-
autoep_placements_equal, legacy_uniform_autoep_placement_descriptor,
26+
autoep_placements_equal, autoep_storage_affine_map,
27+
legacy_uniform_autoep_placement_descriptor,
2628
plan_autoep_placement_transfer)
2729
from deepspeed.checkpoint.affine import ParamAffineMap
2830
from deepspeed.checkpoint.affine_transfer import PlanBudget, TransferPlan
@@ -792,11 +794,85 @@ def expert_affine_maps(self) -> dict[str, ParamAffineMap]:
792794
"""Return the affine map of each live expert parameter (``w1``, ``w2``, ``w3``).
793795
794796
Map ranks are EP-local ranks; translate them with ``dist.get_global_rank(self.ep_group, rank)``.
795-
Each rank's shard is the unpartitioned packed ``[local_experts, ...]`` tensor.
797+
Each rank's shard is the unpartitioned packed ``[local_experts, ...]`` tensor; use
798+
:meth:`expert_storage_affine_maps` for the expert-data-parallel and ZeRO-3 storage.
796799
"""
797800
placement = self._live_expert_placement()
798801
return {name: self._live_expert_affine_map(name, placement) for name in ("w1", "w2", "w3")}
799802

803+
def _expert_storage_group(self, params: dict[str, nn.Parameter]):
804+
partitioned = [hasattr(param, "ds_id") for param in params.values()]
805+
if all(partitioned):
806+
storage_groups = {id(param.ds_process_group): param.ds_process_group for param in params.values()}
807+
if len(storage_groups) != 1:
808+
raise ValueError("AutoEP expert parameters are ZeRO-3 partitioned over different process groups.")
809+
return next(iter(storage_groups.values()))
810+
if any(partitioned):
811+
raise ValueError("Only some AutoEP expert parameters are ZeRO-3 partitioned.")
812+
if self.folding_group_handles is not None:
813+
return self.folding_group_handles.edp_group
814+
from deepspeed.utils import groups
815+
return groups._get_expert_data_parallel_group(self.ep_group_name)
816+
817+
def expert_storage_affine_maps(self) -> dict[str, ParamAffineMap]:
818+
"""Return the affine map of the storage each global rank holds for ``w1``, ``w2`` and ``w3``.
819+
820+
This is computed locally from the expert data parallel group table and the placement, so it
821+
involves no communication and any rank may call it on its own.
822+
823+
Map ranks are global ranks and the logical shape is the parameter's own shape. Each shard is
824+
a rank's flat storage: the packed ``[local_experts, ...]`` tensor under ZeRO-0/1/2, where all
825+
expert-data-parallel replicas are listed as holders, or the rank's ZeRO-3 partition
826+
``ds_tensor`` without its trailing padding. Optimizer state is not described. Quantized and
827+
NVMe-offloaded ZeRO-3 storage is not supported.
828+
"""
829+
if self.folding_group_handles is not None and self.folding_group_handles.spec.etp_size > 1:
830+
raise NotImplementedError("Expert storage maps do not describe expert tensor-parallel shards yet.")
831+
placement = self._live_expert_placement()
832+
params = {name: getattr(self.experts, name) for name in ("w1", "w2", "w3")}
833+
for name in params:
834+
self._live_expert_affine_map(name, placement)
835+
storage_group = self._expert_storage_group(params)
836+
partitioned = hasattr(params["w1"], "ds_id")
837+
if partitioned:
838+
from deepspeed.checkpoint.zero_affine import _check_zero3_param_storage
839+
for param in params.values():
840+
_check_zero3_param_storage(param)
841+
842+
# The EDP group of each EP rank's global rank holds that EP rank's packed tensor. Every rank
843+
# built all EDP groups, so the table answers this without assuming a rank layout.
844+
from deepspeed.utils import groups
845+
group_of_rank = {}
846+
for group_ranks in groups._get_expert_data_parallel_group_table(self.ep_group_name):
847+
for global_rank in group_ranks:
848+
group_of_rank[global_rank] = list(group_ranks)
849+
own_storage_ranks = dist.get_all_ranks_from_group(storage_group)
850+
if group_of_rank.get(dist.get_rank()) != own_storage_ranks:
851+
raise RuntimeError(f"Expert storage group ranks {own_storage_ranks} do not match the expert data "
852+
f"parallel group table for {self.ep_group_name}.")
853+
holders_by_ep_rank = {}
854+
for ep_rank, global_rank in enumerate(dist.get_all_ranks_from_group(self.ep_group)):
855+
holders_by_ep_rank[ep_rank] = group_of_rank[global_rank]
856+
857+
storage_maps = {}
858+
for name, param in params.items():
859+
logical_shape = self._expert_logical_shape(param)
860+
partition_numel_by_ep_rank = None
861+
if partitioned:
862+
# ZeRO-3 pads each packed tensor to a multiple of the group size and splits it evenly.
863+
expert_numel = math.prod(logical_shape[1:])
864+
partition_numel_by_ep_rank = {}
865+
for ep_rank, holders in holders_by_ep_rank.items():
866+
packed_numel = len(autoep_experts_for_rank(placement, ep_rank)) * expert_numel
867+
partition_numel_by_ep_rank[ep_rank] = -(-packed_numel // len(holders))
868+
expected = partition_numel_by_ep_rank[self.ep_rank]
869+
if param.ds_tensor.ds_numel != expected:
870+
raise RuntimeError(f"Expert parameter {name} has a ZeRO-3 partition of {param.ds_tensor.ds_numel} "
871+
f"elements, but even partitioning gives {expected}.")
872+
storage_maps[name] = autoep_storage_affine_map(placement, logical_shape, holders_by_ep_rank,
873+
partition_numel_by_ep_rank)
874+
return storage_maps
875+
800876
def plan_expert_transfer(self,
801877
target_placement: dict,
802878
limits: PlanBudget | None = None,

0 commit comments

Comments
 (0)