scripts/device.py

"""Runtime placement for CUDA, Apple Silicon, and CPU.

8GB-class CUDA boards move one transformer block, or one text-encoder leaf, onto
the accelerator at a time. Apple Silicon loads the same checkpoint onto MPS;
unified memory has to hold it, because CPU and GPU do not have separate pools.
"""
SMALL_CUDA_BYTES = 10 * 1024 ** 3


def choose_runtime(device_type, accelerator_bytes):
    if device_type not in ('cuda', 'mps', 'cpu'):
        raise ValueError(f'Unsupported device type: {device_type}')
    if accelerator_bytes <= 0:
        raise ValueError('Accelerator memory must be positive')
    if device_type == 'cuda' and accelerator_bytes <= SMALL_CUDA_BYTES:
        # One transformer block, or one text-encoder leaf, is on the GPU at a time.
        return {'offload': 'group', 'dtype': 'bfloat16', 'use_stream': False}
    return {'offload': 'model' if device_type == 'cuda' else 'resident',
            'dtype': 'bfloat16', 'use_stream': False}