Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / distributed_data / packed_balancing.py: 90%
348 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-11-15 07:18 +0800
1# Copyright 2026 Huawei Technologies Co., Ltd
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"""Rebalance raw samples from already-selected rank-local packing batches."""
17from __future__ import annotations
19import math
20from collections.abc import Callable, Iterable, Iterator, Sequence
21from contextlib import nullcontext
22from dataclasses import dataclass, replace
23from threading import Thread
24from typing import TYPE_CHECKING, Any
26import torch # pylint: disable=forbidden-backend-import
28from hyper_parallel.distributed_data.balance_logging import log_balance_stats
29from hyper_parallel.distributed_data.balancing_algorithm import BalancingAlgorithm, resolve_balancing_algorithm
30from hyper_parallel.distributed_data.cost_model import CostModel, resolve_cost_model
31from hyper_parallel.distributed_data.device_prefetch import (
32 DeviceStepPrefetcher,
33 _create_device_prefetcher,
34 _resolve_device,
35)
36from hyper_parallel.distributed_data.locality import _create_locality_groups
37from hyper_parallel.distributed_data.planner import DynamicPackingPlanner
38from hyper_parallel.distributed_data.schema import (
39 DistributedPackingPlan,
40 PackingConstraints,
41 SampleKey,
42 SampleMetadata,
43)
44from hyper_parallel.distributed_data.transport import ControlGatherWork, DataPlaneTransport, PayloadExchangeWork
46if TYPE_CHECKING:
47 from hyper_parallel.distributed_data.api import DistributedDatasetConfig
50@dataclass(frozen=True)
51class _LocalBatch:
52 data: Any
53 stats: dict[str, Any] | None = None
54 host_batch: Any = None
57@dataclass(frozen=True)
58class _PreparedSourceStep:
59 """CPU-only source data prepared before the foreground collective phase."""
61 metadata: tuple[Any, ...]
62 local_payloads: dict[SampleKey, Any]
65@dataclass(frozen=True)
66class _PendingBatch:
67 """A planned step whose payload exchange may still be in flight."""
69 plan: DistributedPackingPlan
70 stats: dict[str, Any]
71 retained_payloads: dict[SampleKey, Any]
72 payload_work: PayloadExchangeWork | None
75class _LocalBalancingIterator(Iterator[Any]):
76 """Keep one step ahead with caller-owned accelerator collective launches."""
78 # This companion iterator implements the loader pipeline; these hooks are not public APIs.
79 # pylint: disable=protected-access
81 def __init__(self, loader: LocalBalancingDataLoader) -> None:
82 """Start source workers before entering the background pipeline."""
83 self._source = iter(loader.local_dataloader)
84 self._loader = loader
85 self._step = 0
86 self.finished = False
87 self._thread: Thread | None = None
88 self._result: _LocalBatch | None = None
89 self._prepared_source_step: _PreparedSourceStep | None = None
90 self._planned: Any = None
91 self._phase = "source"
92 self._error: BaseException | None = None
94 def __next__(self) -> Any:
95 """Consume a prepared step, completing omitted prefetch hooks inline."""
96 if self.finished:
97 raise StopIteration
98 try:
99 if self._loader._uses_synchronous_collectives:
100 self.prefetch()
101 if self.finished:
102 raise StopIteration
103 elif self._thread is None:
104 self._start_prefetch()
105 self._join_worker()
106 result = self._result
107 self._result = None
108 if result is None:
109 raise RuntimeError("Local balancing prefetch completed without a batch.")
110 except BaseException:
111 self.finished = True
112 raise
113 self._step += 1
114 self._loader.last_balance_stats = result.stats
115 self._start_prefetch()
116 return self._loader._deliver_batch(result, self._step)
118 def prefetch_plan(self) -> None:
119 """Launch metadata gathering and finish decoding and planning in the worker.
121 All ranks call at the same training boundary, after enqueueing compute
122 and before host synchronization. The variable-length size exchange
123 remains synchronous; no device-wide wait is introduced.
124 """
125 if not self._loader._uses_synchronous_collectives or self.finished:
126 return
127 if self._phase != "source":
128 return
129 if self._thread is None:
130 self._start_prefetch()
131 try:
132 self._join_worker()
133 except StopIteration:
134 self.finished = True
135 return
136 with self._loader._data_stream_context():
137 work = self._loader._transport.begin_all_gather_object(self._prepared_source_step.metadata)
138 self._phase = "plan"
139 self._start_worker(self._prepare_plan, work)
141 def prefetch(self) -> None:
142 """Launch next-step payload exchange on the caller, then finish off-thread.
144 All ranks must call at the same training boundary. Repeated calls
145 before consumption are harmless; they do not advance source progress.
146 """
147 self.prefetch_plan()
148 if not self._loader._uses_synchronous_collectives or self.finished:
149 return
150 if self._phase != "plan":
151 return
152 self._join_worker()
153 with self._loader._data_stream_context():
154 planned = self._planned
155 if self._loader.config.broadcast_plan:
156 planned = self._loader._transport.broadcast_from_planner(planned)
157 plan, stats = planned
158 pending = self._loader._begin_batch(self._prepared_source_step.local_payloads, plan, stats)
159 self._prepared_source_step = None
160 self._planned = None
161 self._phase = "batch"
162 self._start_worker(self._finish_batch, pending)
164 def _prepare_plan(self, work: ControlGatherWork) -> None:
165 with self._loader._data_stream_context():
166 self._planned = self._loader._make_plan(work.wait(), self._step)
168 def _finish_batch(self, pending: _PendingBatch) -> None:
169 with self._loader._data_stream_context():
170 self._result = self._loader._finish_batch(pending)
172 def _start_prefetch(self) -> None:
173 self._phase = "source"
174 self._result = None
175 self._prepared_source_step = None
176 self._start_worker(self._run_prefetch)
178 def _start_worker(self, callback: Callable[..., None], *args: Any) -> None:
179 self._error = None
180 self._thread = Thread(
181 target=self._run_worker, args=(callback, args),
182 name="hp-local-balance-prefetch", daemon=True,
183 )
184 self._thread.start()
186 def _run_worker(self, callback: Callable[..., None], args: tuple[Any, ...]) -> None:
187 try:
188 callback(*args)
189 except BaseException as exc: # pylint: disable=broad-exception-caught
190 # Relay every worker failure, including interrupts, to the consumer.
191 self._error = exc
193 def _run_prefetch(self) -> None:
194 if self._loader._uses_synchronous_collectives:
195 raw_bins = next(self._source)
196 metadata, payloads = self._loader._read_step(raw_bins, self._step)
197 self._prepared_source_step = _PreparedSourceStep(metadata, payloads)
198 else:
199 self._result = self._loader._construct_batch(next(self._source), self._step)
201 def _join_worker(self) -> None:
202 self.wait_for_prefetch()
203 self._thread = None
204 if self._error is not None:
205 raise self._error
207 def wait_for_prefetch(self) -> None:
208 """Drain the active background phase without consuming its result."""
209 if self._thread is not None:
210 self._thread.join()
212 def close(self) -> None:
213 """Drain the producer before resetting the source or its epoch."""
214 self.wait_for_prefetch()
215 self.finished = True
216 self._thread = None
217 self._result = None
218 self._prepared_source_step = None
219 self._planned = None
220 self._error = None
223class LocalBalancingDataLoader:
224 """Balance a frozen local-packer step using only one locality domain.
226 Raw samples, not collated tensors or entire indivisible packs, are routed.
227 The source controls sampling/filtering/packing and microbatch membership;
228 this wrapper never reads future samples to improve a current step's score.
229 All ranks must yield the same number of steps. Metadata and planner output
230 are trusted; local failures propagate without a collective error handshake.
231 """
233 def __init__(
234 self,
235 local_dataloader: Iterable[Sequence[Sequence[Any]]],
236 *,
237 config: DistributedDatasetConfig,
238 metadata_fn: Callable[[Any], SampleMetadata],
239 pack_fn: Callable[[Sequence[Any], int], Any],
240 collate_fn: Callable[[Sequence[Any]], Any],
241 planner: Any,
242 transport: DataPlaneTransport,
243 global_rank: int,
244 device_prefetch: DeviceStepPrefetcher | None = None,
245 balance_stats_callback: Callable[[dict[str, Any], int], None] | None = None,
246 ) -> None:
247 """Store the source, planner, and locality-scoped data plane."""
248 self.local_dataloader = local_dataloader
249 self.config = config
250 self.group_ranks = transport.ranks
251 self.prefetches_to_device = device_prefetch is not None
252 self.last_balance_stats: dict[str, Any] | None = None
253 self.last_host_batch: Any = None
254 self._metadata_fn = metadata_fn
255 self._pack_fn = pack_fn
256 self._collate_fn = collate_fn
257 self._planner = planner
258 self._transport = transport
259 self._global_rank = global_rank
260 self._data_rank = self.group_ranks.index(global_rank)
261 self._iterator: _LocalBalancingIterator | None = None
262 self._device_prefetch = device_prefetch
263 self._data_stream = None
264 self._balance_stats_callback = balance_stats_callback
266 @property
267 def _uses_synchronous_collectives(self) -> bool:
268 """Whether data collectives must stay on the training thread."""
269 backend = getattr(self._transport, "communication_backend", "gloo").lower()
270 return "hccl" in backend or "nccl" in backend
272 def _data_stream_context(self) -> Any:
273 """Reuse the data communication stream independently of the training stream."""
274 device = getattr(self._transport, "communication_device", None)
275 if device is None:
276 return nullcontext()
277 device = torch.device(device)
278 if device.type == "cpu":
279 return nullcontext()
280 accelerator = getattr(torch, device.type)
281 accelerator.set_device(device)
282 if self._data_stream is None:
283 self._data_stream = accelerator.Stream(device=device)
284 # Both launch and completion use this stream; neither waits on the
285 # training stream just to move independent next-step data.
286 return accelerator.stream(self._data_stream)
288 def prefetch_plan(self) -> None:
289 """Gather metadata and start planning at a rank-consistent compute boundary."""
290 if self._iterator is None:
291 self._iterator = _LocalBalancingIterator(self)
292 self._iterator.prefetch_plan()
294 def prefetch(self) -> None:
295 """Launch next-step exchange before synchronizing current-step computation."""
296 if self._iterator is None:
297 self._iterator = _LocalBalancingIterator(self)
298 self._iterator.prefetch()
300 def __len__(self) -> int:
301 """Return the source length when available."""
302 return len(self.local_dataloader)
304 @property
305 def step(self) -> int:
306 """Return the number of delivered steps in the current iterator."""
307 # The owned iterator holds the delivered cursor, not speculative source progress.
308 return self._iterator._step if self._iterator is not None else 0 # pylint: disable=protected-access
310 def __next__(self) -> Any:
311 """Deliver a step without requiring a client-side iterator wrapper."""
312 if self._iterator is None:
313 self._iterator = _LocalBalancingIterator(self)
314 return next(self._iterator)
316 def set_epoch(self, epoch: int) -> None:
317 """Forward epoch selection to the local source before starting iteration.
319 Args:
320 epoch: Source epoch to start.
321 """
322 if self._iterator is not None:
323 self._iterator.close()
324 self._iterator = None
325 if self._device_prefetch is not None:
326 self._device_prefetch.close()
327 self.last_host_batch = None
328 setter = getattr(self.local_dataloader, "set_epoch", None)
329 if callable(setter):
330 setter(epoch)
332 def state_dict(self) -> dict[str, Any]:
333 """Reject checkpointing until rank-local pending packs can be restored."""
334 raise NotImplementedError("Local balancing does not yet support dataloader checkpoint/resume.")
336 def load_state_dict(self, state: dict[str, Any]) -> None:
337 """Reject restoring incomplete locality/source state.
339 Args:
340 state: Unsupported loader checkpoint.
341 """
342 del state
343 raise NotImplementedError("Local balancing does not yet support dataloader checkpoint/resume.")
345 def __iter__(self) -> Iterator[Any]:
346 """Yield constructed local microbatches after scoped metadata and sample exchange."""
347 # Re-entering an unfinished iterator must not drop its speculative batch
348 # after the stateful source has already advanced to the following step.
349 if self._iterator is None or self._iterator.finished:
350 if self._iterator is not None:
351 self._iterator.close()
352 self._iterator = _LocalBalancingIterator(self)
353 return self._iterator
355 def wait_for_prefetch(self) -> None:
356 """Drain data communication before teardown, preserving the next batch.
358 All ranks must call this at the same consumed-step boundary.
359 """
360 if self._iterator is not None:
361 self._iterator.wait_for_prefetch()
362 if self._device_prefetch is not None:
363 self._device_prefetch.close()
365 def close(self) -> None:
366 """Drain the producer and release unconsumed Host/device views."""
367 if self._iterator is not None:
368 self._iterator.close()
369 if self._device_prefetch is not None:
370 self._device_prefetch.close()
371 self.last_host_batch = None
373 def _deliver_batch(self, result: _LocalBatch, step: int) -> Any:
374 batch = result.data
375 self.last_host_batch = result.host_batch
376 if self.prefetches_to_device:
377 batch = batch.take_batch()
378 if result.stats is not None and self._balance_stats_callback is not None and self._global_rank == 0:
379 self._balance_stats_callback(result.stats, step)
380 return batch
382 def _construct_batch(
383 self,
384 raw_bins: Sequence[Sequence[Any]],
385 step: int,
386 ) -> _LocalBatch:
387 metadata, local_payloads = self._read_step(raw_bins, step)
388 gathered = self._transport.all_gather_object(metadata)
389 plan, stats = self._plan_step(gathered, step)
390 return self._finish_batch(self._begin_batch(local_payloads, plan, stats))
392 def _begin_batch(
393 self, local_payloads: dict[SampleKey, Any], plan: DistributedPackingPlan, stats: dict[str, Any],
394 ) -> _PendingBatch:
395 """Start payload exchange for moved samples while retaining local destinations."""
396 if stats["moved_samples"]:
397 outgoing: dict[int, list[tuple[SampleKey, Any]]] = {}
398 retained_payloads = {}
399 for data_rank, local_batch in enumerate(plan.local_batches):
400 target_rank = self.group_ranks[data_rank]
401 owned = [
402 (key, local_payloads[key])
403 for packing_bin in local_batch
404 for key in packing_bin.sample_keys
405 if key.reader_rank == self._global_rank
406 ]
407 if target_rank == self._global_rank:
408 retained_payloads.update(owned)
409 else:
410 outgoing[target_rank] = owned
411 prepared = self._transport.prepare_exchange(outgoing)
412 work = self._transport.begin_exchange_prepared(prepared)
413 else:
414 # Every rank sees the same plan and skips an unchanged exchange.
415 retained_payloads = local_payloads
416 work = None
417 return _PendingBatch(plan, stats, retained_payloads, work)
419 def _finish_batch(self, pending: _PendingBatch) -> _LocalBatch:
420 """Complete payload exchange, pack and collate the local batch, then stage H2D."""
421 received = pending.retained_payloads
422 if pending.payload_work is not None:
423 received.update(pending.payload_work.wait())
424 target = pending.plan.local_batch_for(self._data_rank)
425 host_batch = self._collate_fn([
426 self._pack_fn([received[key] for key in packing_bin.sample_keys], self.config.seq_len)
427 for packing_bin in target
428 ])
429 batch = self._device_prefetch(host_batch) if self._device_prefetch is not None else host_batch
430 return _LocalBatch(batch, pending.stats, host_batch)
432 def _read_step(self, raw_bins: Sequence[Sequence[Any]], step: int) -> tuple[tuple[Any, ...], dict[SampleKey, Any]]:
433 """Bind raw source bins to keyed metadata and cached payloads for one step."""
434 if len(raw_bins) != self.config.local_batch_size or any(not raw_bin for raw_bin in raw_bins):
435 raise ValueError("Each source step must contain local_batch_size non-empty raw-sample bins.")
436 key_bins = []
437 entries = []
438 payloads = {}
439 for raw_bin in raw_bins:
440 keys = []
441 for sample in raw_bin:
442 item_metadata = self._metadata_fn(sample)
443 if not isinstance(item_metadata, SampleMetadata):
444 raise ValueError("metadata_fn must return SampleMetadata.")
445 key = SampleKey(self._global_rank, step, len(entries))
446 entries.append(replace(item_metadata, key=key))
447 payloads[key] = sample
448 keys.append(key)
449 key_bins.append(tuple(keys))
450 return (tuple(key_bins), tuple(entries)), payloads
452 def _plan_step(self, gathered: Sequence[Any], step: int) -> tuple[DistributedPackingPlan, dict[str, Any]]:
453 planned = self._make_plan(gathered, step)
454 return self._transport.broadcast_from_planner(planned) if self.config.broadcast_plan else planned
456 def _make_plan(self, gathered: Sequence[Any], step: int) -> Any:
457 """Build a plan and balancing statistics on the designated planning ranks."""
458 result = None
459 if not self.config.broadcast_plan or self._global_rank == self._transport.planner_rank:
460 samples = []
461 reference_bins = []
462 for bins, entries in gathered:
463 reference_bins.extend(bins)
464 samples.extend(entries)
465 plan = self._planner.plan(samples, reference_bins=reference_bins, step=step)
466 bin_stats = self._bin_statistics(gathered, plan)
467 stats = self._statistics(plan, step, bin_stats["bins_before"])
468 if self._transport.planner_rank == 0:
469 stats.update(bin_stats)
470 result = (plan, stats)
471 return result
473 def _bin_statistics(self, gathered: Sequence[Any], plan: DistributedPackingPlan) -> dict[str, Any]:
474 """Summarize original and accepted bins using gathered metadata."""
475 # Gathered metadata predates cost-model evaluation. Reuse the accepted
476 # plan's estimates for both layouts so before/after have the same units.
477 metadata_by_key = {
478 metadata.key: metadata
479 for _, entries in gathered
480 for metadata in entries
481 }
482 cost_by_key = {key: cost.llm for key, cost in self._planner.last_sample_costs.items()}
484 def _describe(keys: Sequence[SampleKey]) -> dict[str, Any]:
485 tokens = sum(metadata_by_key[key].pack_tokens for key in keys)
486 padding = max(0, self._planner.seq_len - tokens)
487 return {
488 "samples": len(keys), "seq_len": tokens, "padding": padding,
489 "cost": math.fsum(cost_by_key[key] for key in keys) + self._planner.padding_cost(padding).llm,
490 }
492 before = tuple(tuple(_describe(keys) for keys in key_bins) for key_bins, _ in gathered)
493 after = tuple(
494 tuple(_describe(packing_bin.sample_keys) for packing_bin in local_batch)
495 for local_batch in plan.local_batches
496 )
497 return {"bins_before": before, "bins_after": after}
499 def _statistics(
500 self, plan: DistributedPackingPlan, step: int, bins_before: Sequence[Sequence[dict[str, Any]]],
501 ) -> dict[str, Any]:
502 """Count cross-rank raw-sample transfers, excluding locally retained samples."""
503 before = [math.fsum(packing_bin["cost"] for packing_bin in bins) for bins in bins_before]
504 after = [cost.llm for cost in plan.rank_costs]
505 send_samples = [0] * len(self.group_ranks)
506 recv_samples = [0] * len(self.group_ranks)
507 moved = 0
508 source_indices = {rank: index for index, rank in enumerate(self.group_ranks)}
509 for data_rank, local_batch in enumerate(plan.local_batches):
510 for packing_bin in local_batch:
511 for key in packing_bin.sample_keys:
512 source_index = source_indices[key.reader_rank]
513 if source_index != data_rank:
514 send_samples[source_index] += 1
515 recv_samples[data_rank] += 1
516 moved += 1
517 mean_before = math.fsum(before) / len(before)
518 mean_after = math.fsum(after) / len(after)
519 max_before, max_after = max(before), max(after)
520 return {
521 "step": step,
522 "scope": "node",
523 "objective": self._planner.objective,
524 "group_ranks": self.group_ranks,
525 "cost_before": tuple(before),
526 "cost_after": tuple(after),
527 "max_cost_before": max_before,
528 "max_cost_after": max_after,
529 "relative_spread_before": (max_before - min(before)) / mean_before if mean_before else 0.0,
530 "relative_spread_after": (max_after - min(after)) / mean_after if mean_after else 0.0,
531 "moved_samples": moved,
532 "send_samples": tuple(send_samples),
533 "recv_samples": tuple(recv_samples),
534 **self._planner.last_balance_decision,
535 }
538def build_local_balancing_dataloader(
539 local_dataloader: Iterable[Sequence[Sequence[Any]]],
540 mesh: Any,
541 config: DistributedDatasetConfig,
542 *,
543 metadata_fn: Callable[[Any], SampleMetadata],
544 pack_fn: Callable[[Sequence[Any], int], Any],
545 collate_fn: Callable[[Sequence[Any]], Any] = list,
546 model_config: Any = None,
547 cost_model: CostModel | None = None,
548 balancing_algorithm: BalancingAlgorithm | None = None,
549 device: Any = None,
550) -> LocalBalancingDataLoader:
551 """Build the fixed node-local balancing pipeline over complete raw steps.
553 Args:
554 local_dataloader: Each yield contains exactly local_batch_size non-empty
555 raw-sample bins. The source owns sampling and worker prefetch.
556 mesh: WORLD-covering named mesh with model-parallel dimensions of size one.
557 config: Packing limits and the minimum relative improvement required.
558 metadata_fn: CPU-only per-sample metadata and cost features.
559 pack_fn: Construct one packed sequence from its assigned raw samples.
560 collate_fn: Assemble the local step as a sequence of microbatches.
561 model_config: Effective backbone dimensions for the default FLOPs model.
562 Required when cost_model is omitted.
563 cost_model: Optional user estimate replacing the default cost model.
564 balancing_algorithm: Optional sample assignment and objective policy.
565 Costs and the improvement gate remain framework-owned.
566 device: Training device; defaults to CPU with Gloo, otherwise the current
567 accelerator. Pass an explicit NPU/CUDA device for H2D with Gloo.
569 Returns:
570 A local-step loader using the configured node-local communication
571 backend, automatic one-step buffering and final H2D prefetch. Each step
572 retains its source packs unless the candidate improves its objective by
573 more than min_balance_gain.
575 Note:
576 The external source controls iteration length and all ranks must consume
577 the same number of steps. Checkpoint/resume and
578 nontrivial model parallelism are not supported by this local-step path.
579 Iteration returns device-ready batches. last_host_batch retains the CPU
580 view for metering without a device-to-host copy.
581 """
582 if config.balance_window_size != 1:
583 raise ValueError("balance_window_size > 1 requires a native batch_sampler.")
584 if config.device_prefetch_depth != 1:
585 raise ValueError("Configurable device_prefetch_depth requires a native batch_sampler.")
586 if config.repack_sequences:
587 raise ValueError("repack_sequences requires a native batch_sampler.")
588 cost_model = resolve_cost_model(cost_model, model_config)
589 balancing_algorithm = resolve_balancing_algorithm(balancing_algorithm)
590 device = _resolve_device(device, communication_backend=config.communication_backend)
591 device_prefetch = _create_device_prefetcher(device)
592 if not all(callable(callback) for callback in (metadata_fn, pack_fn, collate_fn)):
593 raise ValueError("metadata_fn, pack_fn and collate_fn must be callable.")
594 if not hasattr(local_dataloader, "__iter__"):
595 raise ValueError("local_dataloader must be iterable.")
596 if config.dataset_reader_ranks is not None or config.planner_rank is not None:
597 raise ValueError("Node-local balancing assigns readers and planners automatically.")
598 PackingConstraints(config.seq_len, config.oversized_policy, config.packing_budgets)
599 communication_device = device if config.communication_backend == "hccl" else None
600 topology, groups = _create_locality_groups(
601 mesh,
602 dp_dim_names=config.dp_dim_names,
603 balance_group_size=config.balance_group_size,
604 communication_backend=config.communication_backend,
605 communication_device=communication_device,
606 )
607 planner = DynamicPackingPlanner(
608 data_parallel_size=len(groups.data_plane_ranks),
609 seq_len=config.seq_len,
610 local_batch_size=config.local_batch_size,
611 oversized_policy=config.oversized_policy,
612 packing_budgets=config.packing_budgets,
613 cost_model=cost_model,
614 balancing_algorithm=balancing_algorithm,
615 min_balance_gain=config.min_balance_gain,
616 validate=False,
617 )
618 return LocalBalancingDataLoader(
619 local_dataloader,
620 config=config,
621 metadata_fn=metadata_fn,
622 pack_fn=pack_fn,
623 collate_fn=collate_fn,
624 planner=planner,
625 transport=DataPlaneTransport(
626 groups,
627 topology.global_rank,
628 communication_device=communication_device,
629 ),
630 global_rank=topology.global_rank,
631 device_prefetch=device_prefetch,
632 balance_stats_callback=log_balance_stats if config.log_balance else None,
633 )
636__all__ = ["LocalBalancingDataLoader", "build_local_balancing_dataloader"]