More modes for model opt quantization than halving the batch size

[michaelfeil]

It would be exciting to have an additional option besides:

currently mode (e.g) batch_size=32

batch_size=32 -> if OOM -> batch_size=16 -> if OOM raise ValueError.

Proposed: descending in steps of two.

batch_size=32 -> if OOM -> batch_size=16 -> if OOM -> batch_size=8 -> if OOM -> batch_size=4 -> if OOM -> batch_size=4 -> if OOM -> batch_size=2 -> if OOM -> batch_size=1   if OOM -> raise ValueError

Advantages:

• hard to find out a suitable batch size

Disadvantage:

• for extremely long sequences -> padding would be added -> compute overhead if e.g. batch_size of 4 consist of [2048, 130, 140, 150] and only option for not OOM is batch_size=1. I expect this to be negligible in practice. Also a user error is printed, so a more performant batch_size can be set.

Last observed in 0.21 - just found the option: https://github.com/NVIDIA/TensorRT-Model-Optimizer/blob/5c9390c215abcf996594397c76f281a6e4109bd2/modelopt/torch/utils/dataset_utils.py#L191

[michaelfeil]

Code from v0.21.0

def calibrate_loop():
        if calib_dataloader is None:
            return
        with torch.no_grad():
            low_mem_mode = False
            for _, data in enumerate(calib_dataloader):
                batch_size = data[list(data.keys())[0]].shape[0]
                if batch_size == 1:
                    model(**data)
                elif not low_mem_mode:
                    # Try running the forward once.
                    # If output memory, we try running inference with split input tensors
                    try:
                        model(**data)
                    except torch.OutOfMemoryError:
                        print(
                            "Warning: torch.OutOfMemoryError detected, try reducing the batch size..."
                        )
                        low_mem_mode = True

                if low_mem_mode:
                    split_data_1 = {
                        key: data[key][:batch_size // 2, ...]
                        for key in data
                    }
                    model(**split_data_1)

                    split_data_2 = {
                        key: data[key][batch_size // 2:, ...]
                        for key in data
                    }
                    model(**split_data_2)

[kevalmorabia97]

Closing as #137 is merged