All-gather errors/timesteps across the appropriate group and return a list of ready-to-add ``(err_block, t_idx, pos_or_None)`` items. ★ This is a COLLECTIVE — every rank MUST reach this call together. The caller is responsible for invoking it unconditionally during the
(
self, buffer, error, index, batch_size, num_frame
)
| 476 | return result |
| 477 | |
| 478 | def _gather_errors_for_buffer( |
| 479 | self, buffer, error, index, batch_size, num_frame |
| 480 | ): |
| 481 | """All-gather errors/timesteps across the appropriate group and return |
| 482 | a list of ready-to-add ``(err_block, t_idx, pos_or_None)`` items. |
| 483 | |
| 484 | ★ This is a COLLECTIVE — every rank MUST reach this call together. |
| 485 | The caller is responsible for invoking it unconditionally during the |
| 486 | warmup window (just like SVI's ``all_gather`` outside the random |
| 487 | ``if`` blocks). Random decisions about whether to actually consume |
| 488 | the returned items belong to ``_apply_gathered_items`` instead. |
| 489 | |
| 490 | Group selection mirrors SVI's intent: |
| 491 | * **2D (num_blocks > 0)** — DP group only. Other SP ranks' samples |
| 492 | map to positions unreachable by this rank, so cross-SP gather |
| 493 | wastes bandwidth. |
| 494 | * **1D (num_blocks == 0)** — WORLD group (SVI default). Buckets |
| 495 | are pos-agnostic so every rank's errors are valid samples. |
| 496 | """ |
| 497 | import torch.distributed as dist |
| 498 | if not dist.is_initialized() or dist.get_world_size() <= 1: |
| 499 | return self._collect_local_items(buffer, error, index, batch_size, num_frame) |
| 500 | |
| 501 | if buffer.num_blocks > 0: |
| 502 | from wan_5b.distributed.sp_training import get_data_parallel_group |
| 503 | comm_group = get_data_parallel_group() |
| 504 | if comm_group is None: |
| 505 | return self._collect_local_items(buffer, error, index, batch_size, num_frame) |
| 506 | comm_size = dist.get_world_size(comm_group) |
| 507 | else: |
| 508 | comm_group = None |
| 509 | comm_size = dist.get_world_size() |
| 510 | |
| 511 | if comm_size <= 1: |
| 512 | return self._collect_local_items(buffer, error, index, batch_size, num_frame) |
| 513 | |
| 514 | err_local = error.detach().contiguous() |
| 515 | idx_local = index.detach().contiguous() |
| 516 | err_list = [torch.empty_like(err_local) for _ in range(comm_size)] |
| 517 | idx_list = [torch.empty_like(idx_local) for _ in range(comm_size)] |
| 518 | if comm_group is None: |
| 519 | dist.all_gather(err_list, err_local) |
| 520 | dist.all_gather(idx_list, idx_local) |
| 521 | else: |
| 522 | dist.all_gather(err_list, err_local, group=comm_group) |
| 523 | dist.all_gather(idx_list, idx_local, group=comm_group) |
| 524 | |
| 525 | block_size = self.num_frame_per_block |
| 526 | num_blocks = num_frame // block_size |
| 527 | items = [] |
| 528 | for err_r, idx_r in zip(err_list, idx_list): |
| 529 | idx_per_block = idx_r[:, ::block_size] |
| 530 | err_blocks = err_r.reshape( |
| 531 | batch_size, num_blocks, block_size, *err_r.shape[2:] |
| 532 | ) |
| 533 | for b in range(batch_size): |
| 534 | for blk in range(num_blocks): |
| 535 | pos = blk if buffer.num_blocks > 0 else None |
no test coverage detected