dbvb2k/GRPOTraining

1import os2import sys3from pathlib import Path4import ctypes5import torch6import glob7 8def find_cuda_windows():9    """Find CUDA installation on Windows."""10    # Common CUDA paths11    possible_cuda_paths = [12        "D:/Program Files/NVIDIA GPU Computing Toolkit/CUDA/V12.8",  # Your actual CUDA path13        "C:/Program Files/NVIDIA GPU Computing Toolkit/CUDA/v12.8",14        "D:/Program Files/NVIDIA GPU Computing Toolkit/CUDA/v12.8"15    ]16    17    # Check which path exists18    for cuda_path in possible_cuda_paths:19        if Path(cuda_path).exists():20            return cuda_path21    return None22 23def find_cudart_dll(cuda_path):24    """Find the CUDA runtime DLL."""25    # Check bin directory first26    bin_pattern = os.path.join(cuda_path, "bin", "cudart64_*.dll")27    bin_matches = glob.glob(bin_pattern)28    if bin_matches:29        return bin_matches[0]30    31    # Check lib/x64 directory as fallback32    lib_pattern = os.path.join(cuda_path, "lib", "x64", "cudart64_*.dll")33    lib_matches = glob.glob(lib_pattern)34    return lib_matches[0] if lib_matches else None35 36def set_cuda_paths():37    """Set up CUDA paths for Windows."""38    if not sys.platform.startswith('win'):39        return True40 41    # Get CUDA path42    cuda_path = find_cuda_windows()43    if not cuda_path:44        print("Could not find CUDA installation. Please install CUDA Toolkit.")45        return False46 47    # Add CUDA paths to environment variables48    cuda_bin = str(Path(cuda_path) / "bin")49    cuda_lib = str(Path(cuda_path) / "lib/x64")50 51    print(f"CUDA path: {cuda_path}")52    print(f"CUDA bin: {cuda_bin}")53    print(f"CUDA lib: {cuda_lib}")54    55    # Update PATH56    if cuda_bin not in os.environ['PATH']:57        os.environ['PATH'] = f"{cuda_bin};{os.environ['PATH']}"58    if cuda_lib not in os.environ['PATH']:59        os.environ['PATH'] = f"{cuda_lib};{os.environ['PATH']}"60    61    # Set CUDA_HOME62    os.environ['CUDA_HOME'] = cuda_path63    64    # Try to load CUDA DLLs explicitly65    try:66        cudart_path = find_cudart_dll(cuda_path)67        if cudart_path:68            print(f"Loading CUDA runtime from: {cudart_path}")69            ctypes.CDLL(cudart_path)70            return True71        else:72            print("Could not find CUDA runtime DLL")73            return False74    except Exception as e:75        print(f"Warning: Could not load CUDA runtime DLL: {e}")76        return False77 78def verify_cuda():79    """Verify CUDA is working."""80    if not torch.cuda.is_available():81        print("CUDA is not available in PyTorch!")82        return False83    84    try:85        # Try a simple CUDA operation86        x = torch.tensor([1.0, 2.0, 3.0], device="cuda")87        y = x * 288        print("Successfully performed CUDA operation")89        return True90    except Exception as e:91        print(f"CUDA test failed: {e}")92        return False93 94if __name__ == "__main__":95    if set_cuda_paths():96        print("CUDA paths set successfully")97        if verify_cuda():98            print("CUDA is working properly")99        else:100            print("CUDA verification failed")101    else:102        print("Failed to set CUDA paths")