/ludwig/utils/batch_size_tuner.py

import gc
import logging
import statistics
import time
from abc import ABC

import torch

from ludwig.api_annotations import DeveloperAPI
from ludwig.constants import MAX_BATCH_SIZE_DATASET_FRACTION, MIN_POSSIBLE_BATCH_SIZE

logger = logging.getLogger(__name__)

TOTAL_STEPS = 5


@DeveloperAPI
class BatchSizeEvaluator(ABC):
    def select_best_batch_size(
        self,
        dataset_len: int,
        max_batch_size: int | None = None,
        max_trials: int = 20,
        is_coordinator: bool | None = True,
        global_max_sequence_length: int | None = None,
    ) -> int:
        """Returns optimal batch size as measured by throughput (samples / sec)."""
        logger.info("Tuning batch size...")

        max_batch_size = max_batch_size or dataset_len

        def _is_valid_batch_size(batch_size):
            # make sure that batch size is valid (e.g. less than 20% of ds size and max_batch_size)
            is_smaller_than_training_set = batch_size <= MAX_BATCH_SIZE_DATASET_FRACTION * dataset_len
            is_under_max_batch_size = batch_size <= max_batch_size
            is_valid = is_smaller_than_training_set and is_under_max_batch_size
            if not is_valid and is_coordinator:
                logger.info(
                    f"Batch size {batch_size} is invalid, must be less than or equal to "
                    f"{MAX_BATCH_SIZE_DATASET_FRACTION * 100}% dataset size "
                    f"({int(MAX_BATCH_SIZE_DATASET_FRACTION * dataset_len)} samples "
                    f"of {dataset_len}) and less than or equal to max batch size {max_batch_size}"
                )
            return is_valid

        batch_size = MIN_POSSIBLE_BATCH_SIZE
        best_samples_per_sec = 0
        best_batch_size = None
        count = 0
        while count < max_trials and _is_valid_batch_size(batch_size):
            if is_coordinator:
                logger.info(f"Exploring batch_size={batch_size}")
            gc.collect()

            try:
                samples_per_sec = self.evaluate(
                    batch_size, total_steps=TOTAL_STEPS, global_max_sequence_length=global_max_sequence_length
                )
                if is_coordinator:
                    logger.info(f"Throughput at batch_size={batch_size}: {samples_per_sec:.5f} samples/s")
                if samples_per_sec < best_samples_per_sec:
                    # We assume that once the throughput starts degrading, it won't go up again
                    if is_coordinator:
                        logger.info(f"Throughput decrease at batch_size={batch_size}")
                    break

                best_samples_per_sec = samples_per_sec
                best_batch_size = batch_size
                count += 1

                # double batch size
                batch_size *= 2
            except RuntimeError as e:
                # PyTorch only generates Runtime errors for CUDA OOM.
                gc.collect()
                if "CUDA out of memory" in str(e) or isinstance(e, torch.cuda.OutOfMemoryError):
                    if is_coordinator:
                        logger.info(f"OOM at batch_size={batch_size}")
                else:
                    # Not a CUDA error
                    raise
                break

        # Ensure that some batch size is found.
        # `best_batch_size` can be None if the first batch size is invalid.
        if best_batch_size is None:
            if is_coordinator:
                logger.info(f"Could not tune batch size, using minimum batch size of {MIN_POSSIBLE_BATCH_SIZE}")
            best_batch_size = MIN_POSSIBLE_BATCH_SIZE

        if is_coordinator:
            logger.info(f"Selected batch_size={best_batch_size}")
        return best_batch_size

    def evaluate(self, batch_size: int, total_steps: int = 5, global_max_sequence_length: int | None = None) -> float:
        """Evaluates throughput of the given batch size.

        Return:
            Median throughput in samples / sec.
        """
        durations = []
        for _ in range(total_steps):
            self.reset()
            start_ts = time.time()
            self.step(batch_size, global_max_sequence_length=global_max_sequence_length)
            durations.append(time.time() - start_ts)

        med_duration_s = statistics.median(durations)
        if med_duration_s == 0.0:
            return float("inf")

        return batch_size / med_duration_s

    def reset(self):
        """Called at the beginning of each evaluation step."""

    def step(self, batch_size: int, global_max_sequence_length: int | None = None):
        """Called each step to evaluate the given batch size."""
        raise NotImplementedError("`step` must be implemented by concrete evaluator.")


class BaseLLMBatchSizeEvaluator(BatchSizeEvaluator):
    """Base class for batch size evaluators for LLM models."""

    def __init__(self, trainer):
        self.trainer = trainer
        self.input_feature_name, self.input_feature = list(trainer.model.input_features.items())[0]
        self.output_feature_name, self.output_feature = list(trainer.model.output_features.items())[0]

        # Get the length of the longest input sequence from the training data
        self.input_msl = self.input_feature.input_shape[0]
        if trainer.model.config_obj.input_features[0].preprocessing.max_sequence_length:
            self.input_msl = trainer.model.config_obj.input_features[0].preprocessing.max_sequence_length

        # Get the length of the longest output sequence from the training data
        self.output_msl = self.output_feature.output_shape[0]
        if trainer.model.config_obj.output_features[0].preprocessing.max_sequence_length:
            self.output_msl = trainer.model.config_obj.output_features[0].preprocessing.max_sequence_length

        # This is useful to create the synthetic input and target data which will be a
        # random sequence of integers between 0 and vocab_size
        self.vocab_size = len(trainer.model.config_obj.input_features[0].encoder.vocab)

    def reset(self):
        self.trainer.model.reset_metrics()
        self.trainer.optimizer.zero_grad()

    def step(self, batch_size: int, global_max_sequence_length: int | None = None):
        if global_max_sequence_length and self.input_msl + self.output_msl > global_max_sequence_length:
            # In this case, we just need to make sure that the length of the synthetic data exceeds
            # max_sequence_length by at most a small amount
            self.input_msl = global_max_sequence_length // 2 + 1
            self.output_msl = global_max_sequence_length // 2 + 1

        inputs = {
            self.input_feature_name: torch.randint(0, self.vocab_size, size=(batch_size, self.input_msl))
            .to(self.input_feature.input_dtype)
            .to(self.trainer.device)
        }
        targets = {
            self.output_feature_name: torch.randint(0, self.vocab_size, size=(batch_size, self.output_msl))
            .to(self.output_feature.get_output_dtype())
            .to(self.trainer.device)
        }

        self.perform_step(inputs, targets)

    def perform_step(self, inputs, targets):
        raise NotImplementedError("perform_step method must be implemented in subclasses")


class LLMFinetuneTrainerBatchSizeEvaluator(BaseLLMBatchSizeEvaluator):
    """Batch size evaluator for training batch size for LLM finetuning."""

    def perform_step(self, inputs, targets):
        self.trainer.train_step(inputs, targets)


class LLMFinetunePredictBatchSizeEvaluator(BaseLLMBatchSizeEvaluator):
    """Batch size evaluator for prediction/evaluation batch size for LLM finetuning."""

    def perform_step(self, inputs, targets):
        with torch.no_grad():
            self.trainer.dist_model((inputs, targets))