def dispatch_model(model: PreTrainedModel, use_v2: bool) -> PreTrainedModel:
    r"""
    Dispatches a pre-trained model to GPUs with balanced memory.
    """
    if torch.cuda.device_count() > 1:
        from accelerate import dispatch_model

        device_map = auto_configure_device_map(torch.cuda.device_count(), use_v2=use_v2)
        model.tie_weights()
        return dispatch_model(model, device_map)
    else:
        return model.cuda()

    Dispatches a pre-trained model to GPUs with balanced memory.
    
src/glmtuner/extras/misc.py