File size: 945 Bytes
a544dd6 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 | """GPU preflight diagnostics for Nomen-AI training.
Run before smoke/SFT/DPO to fail early if CUDA or VRAM is unavailable.
"""
import sys
import torch
def main():
print('torch:', torch.__version__)
print('cuda_available:', torch.cuda.is_available())
if not torch.cuda.is_available():
raise SystemExit('ERROR: CUDA is not available. Use a Colab GPU/T4 runtime or Docker with NVIDIA runtime.')
device = torch.cuda.current_device()
name = torch.cuda.get_device_name(device)
props = torch.cuda.get_device_properties(device)
total_gb = props.total_memory / 1e9
print('gpu_name:', name)
print('total_vram_gb:', round(total_gb, 2))
print('compute_capability:', f'{props.major}.{props.minor}')
if total_gb < 14:
raise SystemExit(f'ERROR: VRAM {total_gb:.1f}GB is below expected T4-class 15GB. Use T4/A10G or larger.')
print('GPU_PREFLIGHT_PASS')
if __name__ == '__main__':
main()
|