File size: 5,137 Bytes
abd39d3 | 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 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 | """Tự động kiểm tra, tải và cài đặt PyTorch CUDA + thư viện cuBLAS 12 cho GPU (Colab T4 / Linux)."""
import os
import sys
import shutil
import subprocess
from pathlib import Path
ROOT = Path(__file__).resolve().parent.parent
sys.path.insert(0, str(ROOT))
if hasattr(sys.stdout, 'reconfigure'):
sys.stdout.reconfigure(encoding='utf-8')
from translation.core.gpu_setup import (
choose_cuda_channel,
torch_install_command,
verify_torch_cuda,
verify_ct2_cuda,
install_torch_cuda,
_TORCH_CUDA_CHANNELS,
)
from translation.core.ct2_safe_import import preload_cuda_libraries
def get_driver_cuda() -> tuple[bool, str | None, str | None, float | None]:
"""Lấy thông tin GPU NVIDIA và CUDA version từ nvidia-smi."""
if shutil.which("nvidia-smi") is None:
return False, None, None, None
try:
header = subprocess.run(["nvidia-smi"], capture_output=True, text=True, timeout=10)
gpu_info = subprocess.run(
["nvidia-smi", "--query-gpu=name,memory.total", "--format=csv,noheader,nounits"],
capture_output=True,
text=True,
timeout=10,
)
except Exception:
return False, None, None, None
gpu_name = None
vram_gb = None
cuda_ver = None
if gpu_info.returncode == 0 and gpu_info.stdout.strip():
first_line = gpu_info.stdout.strip().splitlines()[0]
parts = [p.strip() for p in first_line.split(",", 1)]
gpu_name = parts[0]
if len(parts) > 1:
try:
vram_gb = float(parts[1]) / 1024.0
except ValueError:
pass
if header.returncode == 0:
import re
m = re.search(r"CUDA Version:\s*([0-9]+\.[0-9]+)", header.stdout)
if m:
cuda_ver = m.group(1)
return True, gpu_name, cuda_ver, vram_gb
def install_nvidia_libs_if_needed():
"""Đảm bảo các thư viện cuBLAS 12 và cuDNN đã được cài đặt trên Linux."""
if not sys.platform.startswith("linux"):
return
try:
import nvidia.cublas # noqa: F401
except ImportError:
print("📦 Đang tải và cài đặt thư viện runtime nvidia-cublas-cu12 & nvidia-cudnn-cu12...")
subprocess.run(
[sys.executable, "-m", "pip", "install", "-q", "nvidia-cublas-cu12", "nvidia-cudnn-cu12"],
check=False
)
def main():
print("=" * 65)
print("🚀 NOVELREADER — TRÌNH KIỂM TRA & KÍCH HOẠT GPU T4 / CUDA")
print("=" * 65)
has_gpu, gpu_name, driver_cuda, vram_gb = get_driver_cuda()
if not has_gpu:
print("❌ Không tìm thấy GPU NVIDIA vật lý qua nvidia-smi.")
print(" Nếu đang chạy trên Google Colab:")
print(" Vào menu: Thời gian chạy -> Thay đổi loại thời gian chạy -> T4 GPU -> Lưu.")
sys.exit(1)
vram_str = f"{vram_gb:.1f} GB" if vram_gb else "không xác định"
print(f"✅ Phát hiện GPU: {gpu_name} (VRAM: {vram_str})")
print(f"📌 CUDA Driver Version: {driver_cuda or 'N/A'}")
# 1. Cài đặt các gói thư viện cuBLAS 12 và cuDNN nếu là Linux (chứa libcublas.so.12)
install_nvidia_libs_if_needed()
# 2. Kiểm tra PyTorch CUDA hiện tại
torch_ok, msg = verify_torch_cuda()
channel = choose_cuda_channel(driver_cuda) or "cu124"
if torch_ok:
print(f"✅ PyTorch đã nhận GPU thành công: {msg}")
else:
print(f"\n⚠️ PyTorch hiện tại chưa hỗ trợ GPU: {msg}")
print(f"📦 Đang tự động tải và cài đặt PyTorch CUDA tương thích (kênh {channel})...")
success, res_msg = install_torch_cuda(driver_cuda, on_log=print)
if not success:
print(f"❌ Cài đặt PyTorch CUDA thất bại: {res_msg}")
sys.exit(2)
print(f"✅ {res_msg}")
# 3. Preload các thư viện CUDA (libcublas.so.12, libcublasLt.so.12) vào bộ nhớ tiến trình
preload_cuda_libraries()
# 4. Kiểm tra CTranslate2 CUDA
ct2_ok, ct2_msg = verify_ct2_cuda()
if ct2_ok:
print(f"✅ CTranslate2 CUDA: {ct2_msg}")
else:
print(f"⚠️ CTranslate2 chưa nhận CUDA: {ct2_msg}")
print(" Đang thiết lập cờ HACHIMIMT_FORCE_CT2_CUDA=1...")
os.environ["HACHIMIMT_FORCE_CT2_CUDA"] = "1"
# 5. Đánh thức GPU và kiểm tra tensor
try:
import torch
if torch.cuda.is_available():
x = torch.zeros((2048, 2048), device="cuda") @ torch.zeros((2048, 2048), device="cuda")
allocated = torch.cuda.memory_allocated(0) / (1024**2)
reserved = torch.cuda.memory_reserved(0) / (1024**2)
del x
torch.cuda.empty_cache()
print(f"🔥 Kích hoạt GPU Tensor thành công! (Allocated: {allocated:.1f}MB, Reserved: {reserved:.1f}MB)")
print("\n🎉 GPU T4 / CUDA ĐÃ SẴN SÀNG VỚI ĐẦY ĐỦ TÀI NGUYÊN!")
except Exception as exc:
print(f"⚠️ Lỗi khởi động tensor GPU: {exc}")
if __name__ == "__main__":
main()
|