src/vocalrender/training/dynamic_batch.py

"""
Dynamic batch sampler for VoxCPM training.

Provides :class:`DynamicBatchSampler`, a :class:`torch.utils.data.Sampler`
that groups samples so the total token count per batch stays within a
configurable budget.  For distributed training it plans *global* steps
first (assigning samples to ranks greedily to balance cost) and then
partitions per rank.

Extracted from ``svs_data.py`` to keep batch-planning infrastructure
separate from dataset / collation logic.
"""

import random
from typing import List

import torch
import torch.distributed as dist


class DynamicBatchSampler(torch.utils.data.Sampler):
    """
    Dynamic batch sampler that groups samples based on their total_length.

    Instead of using a fixed batch_size, this sampler creates batches where
    the total number of tokens (sum of all sample lengths) doesn't exceed
    max_batch_tokens. This helps:
    1. Avoid OOM errors on long sequences
    2. Better GPU utilization by packing similar-length samples together

    For distributed training, batches are partitioned across GPUs using
    rank and world_size parameters.

    Implementation note:
        In distributed training this sampler acts as a global-step planner.
        It first constructs a *distributed step* as ``world_size`` local
        micro-batches together, then assigns samples greedily so the costs of
        those local micro-batches stay close. Only after the global steps are
        planned are they shuffled as whole units. This avoids the common
        failure mode where each rank independently gets a valid local batch
        but the same training step is badly imbalanced across ranks.

    Args:
        lengths: Pre-computed list of sequence lengths for each sample
        max_batch_tokens: Maximum total tokens per batch
        max_batch_size: Maximum number of samples per batch (optional cap)
        shuffle: Whether to shuffle samples each epoch
        drop_last: Whether to drop the last incomplete batch
        rank: Process rank for distributed training (0 for single GPU)
        world_size: Total number of processes (1 for single GPU)
        seed: Random seed for reproducible shuffling across processes
    """

    def __init__(
        self,
        lengths: List[int],
        max_batch_tokens: int,
        max_batch_size: int = 64,
        shuffle: bool = True,
        drop_last: bool = False,
        rank: int = 0,
        world_size: int = 1,
        seed: int = 42,
    ):
        self.lengths = lengths
        self.max_batch_tokens = max_batch_tokens
        self.max_batch_size = max_batch_size
        self.shuffle = shuffle
        self.drop_last = drop_last
        self.rank = rank
        self.world_size = world_size
        self.seed = seed
        self.epoch = 0
        self.batches: List[List[int]] | None = None
        self._start_batch = 0

        self._create_batches()

    def set_start_batch(self, n: int) -> None:
        """Skip the first ``n`` batches on the next ``__iter__`` (consumed once).

        Used by the resume path to fast-forward the dataloader to a saved
        position without materializing (decoding) the skipped batches. The
        offset is reset to 0 once an iterator is created, so subsequent epochs
        start from the beginning.
        """
        if self.batches is not None and n > len(self.batches):
            raise RuntimeError(
                f"start_batch={n} exceeds available batches ({len(self.batches)}) "
                f"for epoch={self.epoch}."
            )
        self._start_batch = max(0, int(n))

    def set_epoch(self, epoch: int):
        """Set the epoch for reproducible shuffling.

        Skips replanning when the requested epoch matches the currently
        planned one — resume paths call ``set_epoch(data_epoch)`` right
        after construction, which would otherwise redo an O(N·world_size²)
        pure-Python plan for nothing.
        """
        if epoch == self.epoch and self.batches is not None:
            return
        self.epoch = epoch
        self._create_batches()

    def _batch_max_len(self, batch: List[int]) -> int:
        if not batch:
            return 0
        return max(self.lengths[idx] for idx in batch)

    def _batch_cost(self, batch: List[int]) -> int:
        if not batch:
            return 0
        # Attention/activation cost tracks sequence length worse than a simple
        # token budget. Use B * L^2 as a better proxy so that the same
        # distributed step sees similarly expensive batches on every rank.
        max_len = self._batch_max_len(batch)
        return len(batch) * max_len * max_len

    @staticmethod
    def _batch_tokens_from_stats(count: int, max_len: int) -> int:
        if count <= 0 or max_len <= 0:
            return 0
        return count * max_len

    @staticmethod
    def _batch_cost_from_stats(count: int, max_len: int) -> int:
        if count <= 0 or max_len <= 0:
            return 0
        return count * max_len * max_len

    def _order_indices(self, rng: random.Random) -> List[int]:
        indices = list(range(len(self.lengths)))
        if not self.shuffle:
            return indices

        # Sort by length first, then shuffle within relatively small
        # sortish buckets. Smaller buckets keep lengths tighter than the
        # previous coarse bucketing, which reduces bad local orderings that
        # later force imbalanced distributed steps.
        sorted_indices = sorted(indices, key=lambda i: self.lengths[i])
        chunk_size = max(64, self.world_size * max(1, self.max_batch_size) * 8)
        chunk_size = min(chunk_size, 512)
        chunks = [
            sorted_indices[i:i + chunk_size]
            for i in range(0, len(sorted_indices), chunk_size)
        ]
        for chunk in chunks:
            rng.shuffle(chunk)
        rng.shuffle(chunks)
        return [idx for chunk in chunks for idx in chunk]

    def _candidate_rank_key(
        self,
        *,
        rank: int,
        sample_len: int,
        batch_sizes: List[int],
        batch_max_lens: List[int],
        step_idx: int,
    ):
        count = batch_sizes[rank]
        max_len = batch_max_lens[rank]
        new_count = count + 1
        new_max_len = max(max_len, sample_len)
        projected_tokens = new_count * new_max_len

        # Keep singleton overflow behaviour compatible with the previous
        # implementation: if one sample alone exceeds the budget, still allow
        # it as a batch of size 1 so the sample is not dropped.
        if count > 0 and (
            new_count > self.max_batch_size or projected_tokens > self.max_batch_tokens
        ):
            return None

        projected_cost = self._batch_cost_from_stats(new_count, new_max_len)

        simulated_costs = [
            self._batch_cost_from_stats(existing_count, existing_max_len)
            for existing_count, existing_max_len in zip(batch_sizes, batch_max_lens)
        ]
        simulated_costs[rank] = projected_cost
        cost_spread = max(simulated_costs) - min(simulated_costs)
        max_cost = max(simulated_costs)
        total_cost = sum(simulated_costs)

        simulated_tokens = [
            self._batch_tokens_from_stats(existing_count, existing_max_len)
            for existing_count, existing_max_len in zip(batch_sizes, batch_max_lens)
        ]
        simulated_tokens[rank] = projected_tokens
        token_spread = max(simulated_tokens) - min(simulated_tokens)

        # Rotate tie-breaking so the same rank is not always preferred when
        # costs are equal.
        rotated_rank = (rank - step_idx) % max(1, self.world_size)
        return (
            cost_spread,
            max_cost,
            total_cost,
            token_spread,
            projected_cost,
            projected_tokens,
            new_count,
            rotated_rank,
            rank,
        )

    def _plan_global_steps(
        self,
        ordered_indices: List[int],
        rng: random.Random,
    ) -> List[List[List[int]]]:
        if not ordered_indices:
            return []

        global_steps: List[List[List[int]]] = []
        next_index = 0
        step_idx = 0

        while next_index < len(ordered_indices):
            step_batches = [[] for _ in range(self.world_size)]
            batch_sizes = [0] * self.world_size
            batch_max_lens = [0] * self.world_size
            assigned_any = False

            while next_index < len(ordered_indices):
                sample_idx = ordered_indices[next_index]
                sample_len = self.lengths[sample_idx]

                candidates = []
                for rank in range(self.world_size):
                    candidate_key = self._candidate_rank_key(
                        rank=rank,
                        sample_len=sample_len,
                        batch_sizes=batch_sizes,
                        batch_max_lens=batch_max_lens,
                        step_idx=step_idx,
                    )
                    if candidate_key is not None:
                        candidates.append(candidate_key)

                if not candidates:
                    break

                chosen_rank = min(candidates)[-1]
                step_batches[chosen_rank].append(sample_idx)
                batch_sizes[chosen_rank] += 1
                batch_max_lens[chosen_rank] = max(batch_max_lens[chosen_rank], sample_len)
                assigned_any = True
                next_index += 1

            if not assigned_any:
                # Defensive fallback: this should only happen if the dataset is
                # empty, but avoid an infinite loop if constraints are invalid.
                break

            non_empty_ranks = sum(1 for batch in step_batches if batch)
            if self.drop_last and non_empty_ranks < self.world_size:
                break

            global_steps.append(step_batches)
            step_idx += 1

        if self.shuffle:
            rng.shuffle(global_steps)

        return global_steps

    def _create_batches(self):
        """Create per-rank batches via global distributed-step planning.

        The planner is fully deterministic given ``(seed, epoch)`` and its
        output is identical on every rank — so in distributed runs we let
        global rank 0 do the O(N·world_size²) Python plan once and
        broadcast the result. Every rank then extracts its own slice via
        ``self.rank`` (which is the data-parallel rank, not the global
        one), so correctness under FULL_SHARD and HYBRID_SHARD is
        preserved: replicas that share a dp_rank simply pick the same
        slice.
        """
        use_broadcast = dist.is_available() and dist.is_initialized() and dist.get_world_size() > 1

        if use_broadcast:
            if dist.get_rank() == 0:
                rng = random.Random(self.seed + self.epoch)
                ordered_indices = self._order_indices(rng)
                global_steps = self._plan_global_steps(ordered_indices, rng)
                payload = [global_steps]
            else:
                payload = [None]
            dist.broadcast_object_list(payload, src=0)
            global_steps = payload[0]
        else:
            rng = random.Random(self.seed + self.epoch)
            ordered_indices = self._order_indices(rng)
            global_steps = self._plan_global_steps(ordered_indices, rng)

        self.batches = [
            step[self.rank]
            for step in global_steps
            if self.rank < len(step) and step[self.rank]
        ]

    def __iter__(self):
        # Note: batches are already shuffled in _create_batches() with deterministic seed.
        # _start_batch (set by set_start_batch for resume) is consumed once, then reset.
        start = self._start_batch
        self._start_batch = 0
        for batch in self.batches[start:]:
            yield batch

    def __len__(self):
        return len(self.batches)