Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / core / distributed_checkpoint / broadcast.py: 91%
137 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-10-25 05:45 +0800
1# Copyright 2026 Huawei Technologies Co., Ltd. All rights reserved.
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7# http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14# ============================================================================
15"""Sending a shard one rank read to the ranks that hold a copy of it.
17A checkpoint holds one copy of a replicated parameter, so on load exactly one rank of each
18replica set reads it and broadcasts it to the others. This module owns that half of a load:
19the communication groups it goes through, the sends themselves, and the batching that keeps
20small shards from each costing a round trip.
21"""
22from collections import deque
23from contextlib import contextmanager
24from typing import Any, Iterable, Optional, Union
26import torch
27import torch.distributed as dist
28from torch.distributed.distributed_c10d import _get_default_group
30from hyper_parallel.core.distributed_checkpoint.metadata import MetadataIndex
31from hyper_parallel.core.distributed_checkpoint.planner import BroadcastSource, ReadItem
32from hyper_parallel.core.distributed_checkpoint.ragged import get_ragged_box_tensor
33from hyper_parallel.core.distributed_checkpoint.utils import (
34 all_gather_object,
35 get_tensor_storage_size,
36 logger,
37)
38from hyper_parallel.core.dtensor.dtensor import DTensor
39from hyper_parallel.core.utils.communication import EXISTING_COMM_GROUPS
42def copy_each(dests: list, srcs: list) -> None:
43 """
44 Copy every pair in one fused operation rather than one at a time.
46 Measured on an Ascend card, five hundred copies of 64 KiB take 11.8 ms one at a
47 time and 1.4 ms fused: nearly all of what a small copy costs is starting it.
49 Args:
50 dests (list): Tensors written into.
51 srcs (list): Tensors read from, pairing with ``dests`` by position.
52 """
53 if dests:
54 torch._foreach_copy_(dests, srcs) # pylint: disable=protected-access
57def get_created_group(rank_list: Union[list[int], tuple[int, ...]]) -> Any:
58 """
59 The process group over these ranks that the process-wide cache already holds.
61 Args:
62 rank_list (Union[list[int], tuple[int, ...]]): The ranks the group would hold.
64 Returns:
65 Any: The cached group, or None when the cache holds none over these ranks.
66 """
67 return EXISTING_COMM_GROUPS.get(str(tuple(sorted(rank_list))))
70def synchronize() -> None:
71 """
72 Block the host until the work queued on the current device stream has finished.
74 Nothing can be queued on a device stream when the process has no device, and there are
75 two ways to have none: a CPU-only install carries no ``torch.npu`` at all, and a gloo run
76 on an accelerator box never initializes the one it has.
77 """
78 device = getattr(torch, "npu", None)
79 if device is None or not device.is_initialized():
80 return
81 device.current_stream().synchronize()
84# How many broadcasts one rank keeps going at once. Starting the next without waiting on
85# the last is what lets a read overlap the send before it, but each one in flight holds
86# resources inside the communication library, so the count is capped rather than left to
87# grow with the number of shards in the checkpoint.
88_MAX_BROADCASTS_IN_FLIGHT = 8
90# Shards smaller than this travel together rather than one at a time. A broadcast of
91# 64 KiB costs about as much as one of 1 MiB -- some 145 microseconds either way over
92# four Ascend ranks -- and only reaches full speed past a few megabytes, so below that a
93# load spends its time starting broadcasts rather than moving data. This sits just above
94# where the two meet, measured at around 4.6 MiB. Note that a batch is held until its
95# broadcast lands, so a load can have this much times the in-flight limit set aside.
96DEFAULT_BROADCAST_BATCH_BYTES = 6 * 1024 * 1024
99def _existing_group(group_ranks: tuple) -> Any:
100 """
101 A communication group over these ranks that is already there, or None if none is.
103 The cache is where a group that already exists is normally found: the mesh a model is
104 sharded over puts its tp columns and dp groups there, and they stay because training
105 still needs them. Only when it holds nothing is the world worth considering - a
106 parameter every rank has a copy of needs the group of every rank, which is the one the
107 job runs on. That one has been there since initialization and is not in the cache, but
108 it is no less already there, and remaking it would raise a second communicator over
109 every rank for the sake of one read.
111 Neither belongs to the load, so neither is destroyed when the load is done.
113 Args:
114 group_ranks (tuple): The ranks the group would hold.
116 Returns:
117 Any: The group, or None if it has to be created.
118 """
119 existing = get_created_group(group_ranks)
120 if existing is not None:
121 return existing
122 if group_ranks == tuple(range(dist.get_world_size())):
123 return _get_default_group()
124 return None
127def _build_broadcast_groups(
128 items: Iterable[ReadItem],
129 supplied: dict[tuple, Any],
130) -> tuple[dict[tuple, Any], set]:
131 """
132 The communication group of every shard this load broadcasts, and which of them it owns.
134 Every rank has to call this, including one with no shard to broadcast at all: it
135 all-gathers what is needed and takes part in creating all of it, and both of those are
136 collective. A rank that stayed out because it had nothing of its own to add would leave
137 the others waiting on it.
139 A group that already exists is reused rather than made a second time - the mesh a model
140 is sharded over has usually built the tp column or the dp group a replicated parameter
141 needs. But **whether to reuse has to be decided the same way on every rank**:
142 ``new_group`` is collective over the world, so a rank that skipped it because its own
143 cache held the group, while another went ahead, would hang everybody. The cache is not
144 symmetric on its own - ``DeviceMesh.from_group`` records only the group its rank is in -
145 so the decision is taken from the gathered reports rather than from the local cache: a
146 group is made afresh if it is absent on *any* rank that needs it, and reused only when
147 it is absent on none.
149 Groups made here are made with :func:`torch.distributed.new_group`, which creates exactly
150 the ranks it is given and hands the group straight back. The project's ``create_group``
151 would do more than is wanted: it takes the rank list as a *template*, expands it into a
152 whole partition of the world, and keeps every group of it in a process-wide cache. A load
153 uses each of these groups once - and the expansion refuses rank lists a pipeline-parallel
154 model produces, such as a parameter tied across two stages that are not neighbours.
156 Args:
157 items (Iterable[ReadItem]): The read items of this rank finalized plan.
158 supplied (dict[tuple, Any]): Communication groups the caller pre-built, keyed by
159 their rank tuple. Kept as they are and never counted as owned.
161 Returns:
162 tuple[dict[tuple, Any], set]: A group for every rank tuple this rank broadcasts
163 through, and the rank tuples whose groups this load made and must destroy.
164 """
165 needed = sorted({item.source.group_ranks for item in items
166 if item.source is not None and item.source.group_ranks not in supplied})
167 absent = tuple(ranks for ranks in needed if _existing_group(ranks) is None)
168 if absent:
169 logger.warning("There are missing groups %s. Then all gather the missing groups on each rank and "
170 "create them one by one, which will increase some time consumption.", absent)
172 gathered = all_gather_object(
173 (tuple(needed), absent), dist.get_world_size(), use_collectives=True
174 )
175 wanted = sorted({ranks for reported, _ in gathered for ranks in reported})
176 absent_somewhere = {ranks for _, reported in gathered for ranks in reported}
178 groups = dict(supplied)
179 owned = set()
180 for group_ranks in wanted:
181 # Both inputs to this came from the gather, so every rank takes the same branch.
182 fresh = dist.new_group(ranks=list(group_ranks)) if group_ranks in absent_somewhere else None
183 if group_ranks not in needed:
184 continue
185 if fresh is None:
186 groups[group_ranks] = _existing_group(group_ranks)
187 continue
188 groups[group_ranks] = fresh
189 owned.add(group_ranks)
190 return groups, owned
193def _destroy_broadcast_groups(groups: dict[tuple, Any], owned: set) -> None:
194 """
195 Release the groups this load made, once it is done broadcasting through them.
197 The broadcasts have been waited on but are not necessarily finished: a collective that
198 has been waited on is only ordered against the device stream, so the transfer can still
199 be queued when the call returns. Tearing the communicator down at that point kills the
200 transfer - the receiving ranks silently keep whatever their buffer held - so drain the
201 stream first.
203 Releasing is then best effort: the tensors have really arrived by the time the loop
204 runs, so a group that refuses to go away is a leak worth a warning rather than a reason
205 to fail a load that already has its data.
207 Only what this load made is passed here, and only by a rank that belongs to it:
208 ``owned`` is built from this rank own reads, so a group it never joined never reaches
209 this.
211 Args:
212 groups (dict[tuple, Any]): Every group the load broadcast through, keyed by rank tuple.
213 owned (set): Rank tuples whose groups this load made and has to release.
214 """
215 if not owned:
216 return
217 synchronize()
218 # Destroying a group is collective over its members, so keep the same deterministic
219 # order they were created in.
220 for group_ranks in sorted(owned):
221 try:
222 dist.destroy_process_group(groups[group_ranks])
223 except Exception as e: # pylint: disable=broad-except
224 logger.warning("Failed to destroy the broadcast group %s: %s", group_ranks, e)
227@contextmanager
228def broadcast_groups_for_load(
229 items: Iterable[ReadItem],
230 groups: Optional[dict[tuple, Any]] = None,
231) -> Any:
232 """
233 The communication groups one load broadcasts through, destroyed once it is done.
235 A group made here exists to carry this load's broadcasts and nothing else, so it is torn
236 down on the way out rather than left behind: a communicator holds device memory for as
237 long as it lives, and a job that loads a checkpoint has no use for these afterwards.
239 Only what this load made is destroyed. A group the caller pre-built and passed in
240 belongs to the caller, and one that was already in the process-wide cache belongs to
241 whoever put it there - the device mesh, most often, which needs it for the rest of
242 training.
244 Destroying is done in rank-tuple order, which every rank works out the same way, and
245 only by the ranks of the group - the ones that reach the end of the read together.
247 Args:
248 items (Iterable[ReadItem]): The read items of this rank finalized plan.
249 groups (Optional[dict[tuple, Any]]): Communication groups the caller pre-built.
251 Yields:
252 dict[tuple, Any]: A group for every rank tuple this rank broadcasts through.
253 """
254 built, owned = _build_broadcast_groups(items, dict(groups or {}))
255 try:
256 yield built
257 finally:
258 _destroy_broadcast_groups(built, owned)
261def _start_broadcast(
262 in_flight: deque,
263 buffer: Any,
264 source: BroadcastSource,
265 groups: dict[tuple, Any],
266 after: Optional[Any] = None,
267) -> None:
268 """
269 Start one broadcast without waiting for it, making room for it first.
271 Args:
272 in_flight (deque): Broadcasts already going, oldest first. This one is added, after
273 waiting on the oldest should too many already be going.
274 buffer (Any): Contiguous memory to send from the source and receive into elsewhere.
275 source (BroadcastSource): Which ranks take part and which of them sends.
276 groups (dict[tuple, Any]): Communication groups, as :func:`broadcast_groups_for_load`
277 returns them.
278 after (Optional[Any]): Called once this broadcast has landed, for a send that is not
279 finished when the bytes arrive -- a batch still to be dealt out to the shards it
280 was gathered from.
281 """
282 while len(in_flight) >= _MAX_BROADCASTS_IN_FLIGHT:
283 _finish_broadcast(in_flight.popleft())
284 handle = dist.broadcast(buffer, source.src_rank, groups[source.group_ranks], async_op=True)
285 if handle is None:
286 if after is not None:
287 after()
288 return
289 in_flight.append((handle, after))
292def _finish_broadcast(pending: tuple) -> None:
293 """Wait on one broadcast and do whatever was left until it had landed."""
294 handle, after = pending
295 handle.wait()
296 if after is not None:
297 after()
300def broadcast_shard(
301 in_flight: deque,
302 state_dict: dict[str, Any],
303 item: ReadItem,
304 groups: dict[tuple, Any],
305) -> None:
306 """
307 Start sending one shard between the ranks that hold it, without waiting for it.
309 A shard travels whole rather than region by region: it is one local buffer, which is
310 contiguous where a single region of it generally is not, and every rank holding it has it
311 in the same shape. Shards are otherwise unrelated -- two of one tensor are two buffers,
312 two groups and two broadcasts.
314 Callers have to reach the shards of a group in the same order on every rank of it, since
315 ranks that enter a group collectives in different orders deadlock. The order the global
316 plan puts them in is the one thing every rank agrees on without asking.
318 Args:
319 in_flight (deque): Broadcasts already going, oldest first.
320 state_dict (dict[str, Any]): Flat state dict holding the shard to send.
321 item (ReadItem): Any item of the shard, which names it and who reads it.
322 groups (dict[tuple, Any]): Communication groups, as :func:`broadcast_groups_for_load`
323 returns them.
324 """
325 buffer = _shard_buffer(state_dict[item.dest_index.fqn], item.dest_index)
326 _start_broadcast(in_flight, buffer, item.source, groups)
329def wait_broadcasts(in_flight: deque) -> None:
330 """
331 Wait on every broadcast still going, oldest first.
333 The buffers being sent are the state dict tensors themselves, so a load that carried on
334 with one still in flight would be reading into memory a collective is still writing.
336 Args:
337 in_flight (deque): Broadcasts started so far. Emptied.
338 """
339 while in_flight:
340 _finish_broadcast(in_flight.popleft())
343class BroadcastBatcher:
344 """
345 Shards small enough that sending them one at a time would cost more than moving them.
347 A broadcast costs about the same whether it carries 64 KiB or 1 MiB -- around 145
348 microseconds either way, measured over four Ascend ranks, against 33 GiB/s once the
349 shards are large. Below that crossing point a load spends its time starting broadcasts
350 rather than moving data, and a checkpoint holds a great many small tensors: norms,
351 biases, scalars, step counters. So shards under ``batch_bytes`` are gathered into one
352 buffer, sent together, and dealt out again on the far side, while larger ones are sent
353 as they are, the fixed cost being small against what they carry.
355 Shards can only travel together when they agree on the group, the sending rank and the
356 dtype, so one batch is kept per combination. Every rank of a group meets that group
357 shards in the same order and so fills and sends the same batches at the same points,
358 which is what keeps its collectives in step -- the same thing that lets shards be sent
359 one at a time without agreeing on anything first.
360 """
362 def __init__(self, batch_bytes: int, groups: dict[tuple, Any]) -> None:
363 """
364 Args:
365 batch_bytes (int): Shards this size or larger are sent on their own, and a batch
366 is sent as soon as another shard would take it past this. Zero sends every
367 shard on its own.
368 groups (dict[tuple, Any]): Communication groups, as
369 :func:`broadcast_groups_for_load` yields them.
370 """
371 self._batch_bytes = batch_bytes
372 self._groups = groups
373 self._batches: dict[tuple, list] = {}
374 self._pending_bytes: dict[tuple, int] = {}
375 self.batched = 0
376 self.sent = 0
378 def add(self, in_flight: deque, state_dict: dict[str, Any], item: ReadItem) -> None:
379 """
380 Hand one shard over to be sent, on its own or with others.
382 Args:
383 in_flight (deque): Broadcasts already going, oldest first.
384 state_dict (dict[str, Any]): Flat state dict holding the shard to send.
385 item (ReadItem): Any item of the shard, which names it and who sends it.
386 """
387 buffer = _shard_buffer(state_dict[item.dest_index.fqn], item.dest_index)
388 nbytes = get_tensor_storage_size(buffer)
389 if nbytes >= self._batch_bytes:
390 self.sent += 1
391 _start_broadcast(in_flight, buffer, item.source, self._groups)
392 return
394 key = (item.source.group_ranks, item.source.src_rank, buffer.dtype)
395 if self._pending_bytes.get(key, 0) + nbytes > self._batch_bytes:
396 self._send(in_flight, key)
397 self._batches.setdefault(key, []).append(buffer)
398 self._pending_bytes[key] = self._pending_bytes.get(key, 0) + nbytes
400 def flush(self, in_flight: deque) -> None:
401 """
402 Send whatever is still waiting to go with something else.
404 Batches go in the order they were first added to, which follows the shards and so is
405 the same on every rank that holds them.
407 Args:
408 in_flight (deque): Broadcasts already going, oldest first.
409 """
410 for key in list(self._batches):
411 self._send(in_flight, key)
413 def _send(self, in_flight: deque, key: tuple) -> None:
414 """Send one batch, and arrange for it to be dealt out once it has landed."""
415 buffers = self._batches.pop(key, [])
416 self._pending_bytes.pop(key, None)
417 if not buffers:
418 return
420 group_ranks, src_rank = key[:2]
421 source = BroadcastSource(group_ranks=group_ranks, src_rank=src_rank)
422 self.sent += 1
423 if len(buffers) == 1:
424 # Nothing to gather it with, so gathering it would only cost a round trip.
425 _start_broadcast(in_flight, buffers[0], source, self._groups)
426 return
428 self.batched += len(buffers)
429 # Left uninitialized: the views below cover it exactly, and it is written
430 # whole either by gathering the shards into it or by the broadcast landing.
431 staging = torch.empty(
432 size=(sum(buffer.numel() for buffer in buffers),),
433 dtype=buffers[0].dtype,
434 device=getattr(buffers[0], "device", None),
435 )
436 views, offset = [], 0
437 for buffer in buffers:
438 views.append(staging[offset:offset + buffer.numel()].reshape(buffer.shape))
439 offset += buffer.numel()
441 if source.src_rank == dist.get_rank():
442 # The sender already holds every shard, so it gathers them and needs nothing back.
443 copy_each(views, buffers)
444 _start_broadcast(in_flight, staging, source, self._groups)
445 return
446 _start_broadcast(
447 in_flight, staging, source, self._groups,
448 after=lambda: copy_each(buffers, views),
449 )
452def _shard_buffer(obj: Any, index: MetadataIndex) -> Any:
453 """
454 The local buffer of one shard, which is what a single broadcast carries.
456 Mirrors what :meth:`StandardLoadPlanner.acquire_tensor` narrows into, so that the rank
457 reading a shard sends the same storage the others are waiting to have written.
459 Args:
460 obj (Any): The state dict entry the shard belongs to.
461 index (MetadataIndex): Names the shard, as a read item destination does.
463 Returns:
464 Any: A tensor view over the shard, writable in place by a collective.
465 """
466 if isinstance(obj, DTensor):
467 if obj.layout is not None and obj.layout.ragged_shard is not None:
468 return get_ragged_box_tensor(obj, index).detach()
469 return obj.to_local().detach()
470 return obj.detach()