File size: 959 Bytes
9116984 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 | """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}
|