File size: 2,309 Bytes
31f7037 | 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 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 | """CUDA configuration and explicit environment diagnostics for the L40S path."""
import importlib.metadata
import json
import os
import platform
import torch
def configure(device_name='cuda', precision='bf16', seed=42):
device = torch.device(device_name)
if device.type == 'cuda' and device.index is None:
device = torch.device('cuda', 0)
if device.type == 'cuda' and not torch.cuda.is_available():
raise RuntimeError('CUDA requested but unavailable. Run nvidia-smi and install the CUDA PyTorch wheel; CPU fallback is disabled.')
if precision == 'bf16' and device.type == 'cuda' and not torch.cuda.is_bf16_supported():
raise RuntimeError('This CUDA device does not support BF16; select --precision fp32.')
torch.set_num_threads(int(os.environ.get('JULIA_CPU_THREADS', '4')))
torch.manual_seed(seed)
if device.type == 'cuda':
torch.cuda.set_device(device)
torch.cuda.manual_seed_all(seed)
torch.backends.cuda.matmul.allow_tf32 = True
torch.backends.cudnn.allow_tf32 = True
torch.cuda.reset_peak_memory_stats(device)
torch.set_float32_matmul_precision('high')
return device
def environment(device):
result = dict(python=platform.python_version(), torch=str(torch.__version__), cuda=torch.version.cuda,
transformers=importlib.metadata.version('transformers'), device=str(device))
if device.type == 'cuda':
prop = torch.cuda.get_device_properties(device)
result.update(gpu=prop.name, vram_bytes=prop.total_memory,
compute_capability=list(torch.cuda.get_device_capability(device)))
return result
def move(batch, device):
return {k: v.to(device, non_blocking=device.type == 'cuda') for k, v in batch.items()}
def main():
device = configure()
x = torch.randn(512, 512, device=device, dtype=torch.bfloat16, requires_grad=True)
(x @ x.T).float().square().mean().backward()
torch.cuda.synchronize()
report = environment(device)
report['bf16_forward_backward_finite'] = bool(torch.isfinite(x.grad).all())
print(json.dumps(report, indent=2))
if not report['bf16_forward_backward_finite']:
raise RuntimeError('CUDA BF16 diagnostic produced nonfinite gradients')
if __name__ == '__main__':
main()
|