espnet2.speechlm.dataloader.batch.synchronize_batches

source

espnet2.speechlm.dataloader.batch.synchronize_batches(batches: List[List[T]]) → List[List[T]]

Synchronize batches across all GPU ranks in distributed training.

Ensures all GPU ranks have the same number of batches by duplicating the last few batches on ranks with fewer batches. This is useful for distributed training where each rank may have different numbers of batches due to data sharding.

• Parameters:batches – List of batches to synchronize.
• Returns: Synchronized list of batches with duplicates added if necessary.
• Raises:RuntimeError – If some ranks have batches while others have none. A rank with no batches cannot be padded, so the ranks would keep different numbers of batches.

Notes

• If torch.distributed is not initialized, returns unchanged
• If torch.distributed is initialized but no accelerator is available, raises RuntimeError (synchronization cannot be performed)
• Duplicates are taken from the end of the batch list