Repository: Lightning-AI/pytorch-lightning
File: src/lightning/pytorch/tuner/batch_size_scaling.py (Python)

Code before:

    Args:
        trainer: A Trainer instance.
        mode: Search strategy to update the batch size:

            - ``'power'``: Keep multiplying the batch size by 2, until we get an OOM error.
            - ``'binsearch'``: Initially keep multiplying by 2 and after encountering an OOM error
                do a binary search between the last successful batch size and the batch size that failed.

        steps_per_trial: number of steps to run with a given batch size.
            Ideally 1 should be enough to test if an OOM error occurs,
            however in practise a few are needed
        init_val: initial batch size to start the search with
        max_trials: max number of increases in batch size done before
           algorithm is terminated
        batch_arg_name: name of the attribute that stores the batch size.
            It is expected that the user has provided a model or datamodule that has a hyperparameter
            with that name. We will look for this attribute name in the following places

            - ``model``
            - ``model.hparams``
            - ``trainer.datamodule`` (the datamodule passed to the tune method)

        margin: Margin to reduce the found batch size by to provide a safety buffer. Only applied when using
            'binsearch' mode. Should be a float between 0 and 1. Defaults to 0.05 (5% reduction).
        max_val: Maximum batch size limit. If provided, the found batch size will not exceed this value.

    """
    if trainer.fast_dev_run:
        rank_zero_warn("Skipping batch size scaler since `fast_dev_run` is enabled.")
        return None

    # Save initial model, that is loaded after batch size is found
    ckpt_path = os.path.join(trainer.default_root_dir, f".scale_batch_size_{uuid.uuid4()}.ckpt")
    trainer.save_checkpoint(ckpt_path)

    # Arguments we adjust during the batch size finder, save for restoring
    params = __scale_batch_dump_params(trainer)

    # Set to values that are required by the algorithm
    __scale_batch_reset_params(trainer, steps_per_trial)

    if trainer.progress_bar_callback:
        trainer.progress_bar_callback.disable()

    try:
        new_size, _ = _adjust_batch_size(trainer, batch_arg_name, value=init_val)

        if mode == "power":
            new_size = _run_power_scaling(trainer, new_size, batch_arg_name, max_trials, params, max_val)
        elif mode == "binsearch":

Review comment (performance):
```suggestion
        max_val: Maximum batch size limit, to prevent overly large or inefficient batch sizes.
```

Code after:

    Args:
        trainer: A Trainer instance.
        mode: Search strategy to update the batch size:

            - ``'power'``: Keep multiplying the batch size by 2, until we get an OOM error.
            - ``'binsearch'``: Initially keep multiplying by 2 and after encountering an OOM error
                do a binary search between the last successful batch size and the batch size that failed.

        steps_per_trial: number of steps to run with a given batch size.
            Ideally 1 should be enough to test if an OOM error occurs,
            however in practise a few are needed
        init_val: initial batch size to start the search with
        max_trials: max number of increases in batch size done before
           algorithm is terminated
        batch_arg_name: name of the attribute that stores the batch size.
            It is expected that the user has provided a model or datamodule that has a hyperparameter
            with that name. We will look for this attribute name in the following places

            - ``model``
            - ``model.hparams``
            - ``trainer.datamodule`` (the datamodule passed to the tune method)

        margin: Margin to reduce the found batch size by to provide a safety buffer. Only applied when using
            'binsearch' mode. Should be a float between 0 and 1. Defaults to 0.05 (5% reduction).
        max_val: Maximum batch size limit, defaults to 8192.
            Helps prevent testing unrealistically large or inefficient batch sizes (e.g., 2**25)
            when running on CPU or when automatic OOM detection is not available.

    """
    assert 0.0 <= margin < 1.0, f"`margin` should be between 0 and 1. Found {margin=}"

    if trainer.fast_dev_run:
        rank_zero_warn("Skipping batch size scaler since `fast_dev_run` is enabled.")
        return None

    # Save initial model, that is loaded after batch size is found
    ckpt_path = os.path.join(trainer.default_root_dir, f".scale_batch_size_{uuid.uuid4()}.ckpt")
    trainer.save_checkpoint(ckpt_path)

    # Arguments we adjust during the batch size finder, save for restoring
    params = __scale_batch_dump_params(trainer)

    # Set to values that are required by the algorithm
    __scale_batch_reset_params(trainer, steps_per_trial)

    if trainer.progress_bar_callback:
        trainer.progress_bar_callback.disable()

    try:
        new_size, _ = _adjust_batch_size(trainer, batch_arg_name, value=init_val)