source/src/speculators/train/distributed_batch_sampler.py

"""
MIT License

Copyright (c) 2023 One

Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.

Adapted from https://github.com/imoneoi/multipack_sampler.
"""

# Standard
import warnings
from heapq import heapreplace
from typing import NamedTuple

import numpy as np

# Third Party
from numpy.typing import ArrayLike, NDArray
from torch.utils.data import Sampler


## Multipack Distributed Batch Sampler
class _Bin(NamedTuple):
    """Helper named tuple for `lpt_packed_batch`"""

    fill: int  # sum of items in _Bin
    slot: int  # heap slot id (0..num_replicas-1)


def _lpt_packed_batch(
    lengths: np.ndarray,
    max_len: int,
    num_replicas: int,
    start_index: int,
    rank: int,
    rotation: int,
) -> None | list:
    """
    Check if lengths can be distributed into `num_replicas` machines with at most
    `max_len` tokens per machine and return this rank's batch.

    Uses the LPT (Longest processing time first scheduling) algorithm
    Time: O(|lengths| log |lengths| + |lengths| log replicas)

    Bins are packed to a token budget, which leaves sample counts skewed: with every
    bin empty the heap tie-breaks on slot id, so slot 0 always takes the largest sample
    and the highest slot absorbs the small ones. `rotation` re-maps slots to ranks so
    that role cycles across batches.

    Returns:
    `None` if unable to find a valid packing. Otherwise, return the batch indices that
    correspond to `rank`.
    """

    # Greedily assign lengths (in decreasing order) to the least full slot until they
    # are all assigned or we run out of space.
    local_batch = []
    heap = [_Bin(0, i) for i in range(num_replicas)]

    target_slot = (rank - rotation) % num_replicas

    # sort in descending order
    indices = np.argsort(lengths)[::-1]

    for idx, size in zip(indices, lengths[indices], strict=True):
        new_fill = heap[0].fill + size
        if new_fill > max_len:
            # Size doesn't fit in least full batch (or any others), report failure.
            return None

        if heap[0].slot == target_slot:
            # minimum bucket corresponds to this rank -> add idx to local batch
            local_batch.append(start_index + idx)

        _ = heapreplace(heap, _Bin(new_fill, heap[0].slot))

    return local_batch


def _assign_to_packed_batches(
    lengths: np.ndarray, max_len: int, rank: int, replicas: int
) -> list[NDArray]:
    """Distribute lengths to batches across all ranks, while respecting max_length.
    Uses a binary search + LPT algorithm.

    Args:
        lengths (np.ndarray): array of dataset sample lengths
        max_len (int): maximum allowed sum of lengths in batch
        rank (int): global rank to collect batches for
        replicas (int): world size to distribute batches to

    Returns:
        tuple[list, int, int]:
            - list of np.arrays containing the indices for each batch on this rank
            - sum of dataset lengths included (total sum of lengths in dataset minus any
              that were dropped at end of dataset)
            - total token capacity if each batch maxed out max_length
    """

    lengths_so_far = 0
    ind = 0
    result: list = []
    lengths_cumsum = np.cumsum(lengths)

    # binary search for max integer x such that the next x elements in shuffled lengths
    # array can be packed into `replicas` batches.
    # Add this rank's batch to `result` and repeat until end of dataset
    while True:
        if len(lengths) - ind < replicas:
            # Not enough lengths left to pack into `num_replicas` batches
            # Break and drop whatever lengths we have left
            break

        # binary search in [1, 1 + upper bound for x)
        left = 1
        right = 1 + np.searchsorted(
            lengths_cumsum[ind:], lengths_so_far + max_len * replicas, "right"
        )

        # Cycle the slot->rank mapping so no rank is permanently the many-samples one.
        rotation = len(result)

        batch = None
        while right - left > 1 and right > replicas:
            mid = (left + right) // 2
            batch = _lpt_packed_batch(
                lengths[ind : ind + mid], max_len, replicas, ind, rank, rotation
            )
            if batch is None:
                right = mid
            else:
                left = mid

        if batch is None:
            batch = _lpt_packed_batch(
                lengths[ind : ind + left], max_len, replicas, ind, rank, rotation
            )

        ind += left
        lengths_so_far = lengths_cumsum[ind - 1]

        # append only result for this rank (already filtered in lpt_packed_batch)
        result.append(batch)

    return result


class MultipackDistributedBatchSamplerV2(Sampler):
    def __init__(
        self,
        batch_max_length: int,
        lengths: ArrayLike,
        num_replicas: int,
        rank: int,
        truncate_long_samples: bool = True,
        seed: int = 0,
        max_batches: int | None = None,
    ):
        """Efficient distributed packing sampler for linear attention style models

        Args:
            batch_max_length (int): max number of tokens in a single batch per device
            lengths (ArrayLike[int]): the lengths of each sample in the dataset
            num_replicas (int): The number of replicas to split the dataset across.
            rank (int): The global rank to collect batches for.
            truncate_long_samples (bool, optional): Whether to truncate long samples
            (True) or drop them (False). Default is True.
            seed (int, optional): Seed for RNG, must be the same on all ranks. Default 0
            max_batches (int, optional): Limit batches exposed per epoch. Used by
                bounded benchmarks to prevent workers from prefetching beyond the
                measured run.
        """
        self.num_replicas = num_replicas
        self.rank = rank
        self.seed = seed
        self.epoch = 0
        self.batch_max_length = batch_max_length
        self.max_batches = max_batches
        self.lengths = np.array(lengths)

        self.valid_indices = np.nonzero(self.lengths <= self.batch_max_length)[0]
        if len(self.valid_indices) < len(self.lengths):
            if truncate_long_samples:
                msg = (
                    f"Found {len(self.lengths) - len(self.valid_indices)}"
                    f"/{len(self.lengths)} samples longer than batch_max_length. "
                    "These samples will be truncated to batch_max_length."
                )
                self.valid_indices = np.arange(len(self.lengths))
                self.lengths = np.clip(self.lengths, 0, self.batch_max_length)
            else:
                msg = (
                    f"Dropping {len(self.lengths) - len(self.valid_indices)}"
                    f"/{len(self.lengths)} samples longer than batch_max_length. Ensure"
                    " that the right max_batch_length is used during data processing."
                )

            if self.rank == 0:
                warnings.warn(msg, stacklevel=1)

        self._cached_generated_batches = (-1, [])

    def __iter__(self):
        batches = self._generate_batches(self.epoch)
        return iter(batches)

    def __len__(self):
        batches = self._generate_batches(self.epoch)
        return len(batches)

    def set_epoch(self, epoch: int):
        self.epoch = epoch

    def _generate_batches(self, epoch: int) -> list[NDArray]:
        """Generate batches for this rank

        Returns:
            list[NDArray]: list of np.arrays containing the indices for each batch on
            this rank
        """
        if self._cached_generated_batches[0] == epoch:
            return self._cached_generated_batches[1]

        rng = np.random.default_rng(seed=self.seed + epoch)
        indices = rng.permutation(self.valid_indices)

        batches = _assign_to_packed_batches(
            self.lengths[indices], self.batch_max_length, self.rank, self.num_replicas
        )

        # The indices in batches are relative to the shuffled self.lengths[indices]
        # Translate them so that they are instead relative to the overall unshuffled
        # self.lengths array.
        batches = [indices[batch] for batch in batches]
        if self.max_batches is not None:
            batches = batches[: self.max_batches]

        # Cache result
        self._cached_generated_batches = (epoch, batches)
        return batches