|
2 | 2 | # DeepSpeed Team |
3 | 3 | """AutoEP expert placement descriptors and their affine lowering. |
4 | 4 |
|
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. |
8 | 9 | """ |
9 | 10 |
|
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 |
11 | 12 | from deepspeed.checkpoint.affine_transfer import plan_transfer |
| 13 | +from deepspeed.checkpoint.zero_affine import zero_storage_affine_map |
12 | 14 | from deepspeed.checkpoint.constants import (AFFINE_MAP, AFFINE_MAP_PARAMS, AFFINE_MAP_VERSION, AUTOEP_AFFINE_MAPS, |
13 | 15 | AUTOEP_EXPERT_PLACEMENT, AUTOEP_PLACEMENT_EP_SIZE, |
14 | 16 | AUTOEP_PLACEMENT_EXPERTS, AUTOEP_PLACEMENT_NUM_EXPERTS, |
|
23 | 25 | 'legacy_uniform_autoep_placement_descriptor', |
24 | 26 | 'autoep_placements_equal', |
25 | 27 | 'autoep_placement_to_affine_map', |
| 28 | + 'autoep_storage_affine_map', |
26 | 29 | 'autoep_metadata_to_affine_map', |
27 | 30 | 'extract_autoep_rank_tensor', |
28 | 31 | 'plan_autoep_placement_transfer', |
@@ -169,6 +172,40 @@ def autoep_placement_to_affine_map(descriptor, logical_shape): |
169 | 172 | return affine_map |
170 | 173 |
|
171 | 174 |
|
| 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 | + |
172 | 209 | def autoep_metadata_to_affine_map(metadata, logical_shape, param_name=None): |
173 | 210 | """Load a persisted AutoEP affine map, falling back to legacy placement lowering.""" |
174 | 211 | serialized_map = metadata.get(AFFINE_MAP) |
@@ -252,6 +289,39 @@ def _pieces_for_rank(experts, holders, expert_shape, expert_numel, source_stride |
252 | 289 | return pieces |
253 | 290 |
|
254 | 291 |
|
| 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 | + |
255 | 325 | def _expert_holders(entries_by_rank, num_experts): |
256 | 326 | holders = {expert_id: set() for expert_id in range(num_experts)} |
257 | 327 | for rank, entry in entries_by_rank.items(): |
|
0 commit comments