check_dtype.py

"""检查训练时实际使用的数据类型"""
import torch

print("=" * 60)
print("PyTorch & CUDA 环境")
print("=" * 60)
print(f"PyTorch 版本: {torch.__version__}")
print(f"CUDA 可用: {torch.cuda.is_available()}")
if torch.cuda.is_available():
    print(f"CUDA 版本: {torch.version.cuda}")
    print(f"显卡名称: {torch.cuda.get_device_name(0)}")
    print(f"显卡算力: {torch.cuda.get_device_capability(0)}")
    print(f"显存总量: {torch.cuda.get_device_properties(0).total_memory / 1024**3:.2f} GB")

print("\n" + "=" * 60)
print("BF16 支持检查")
print("=" * 60)
print(f"BF16 支持 (CPU): {torch.cpu.amp.autocast.is_enabled()}")
if torch.cuda.is_available():
    # 检查 BF16 是否可用
    try:
        device = torch.device("cuda")
        x = torch.randn(10, 10, device=device, dtype=torch.bfloat16)
        y = torch.randn(10, 10, device=device, dtype=torch.bfloat16)
        z = torch.matmul(x, y)
        print(f"BF16 支持 (GPU): ✅ 正常 (结果 dtype={z.dtype})")
    except Exception as e:
        print(f"BF16 支持 (GPU): ❌ 失败 - {e}")
    
    # 检查 TF32
    print(f"TF32 for matmul: {torch.backends.cuda.matmul.allow_tf32}")
    print(f"TF32 for cuDNN: {torch.backends.cudnn.allow_tf32}")

print("\n" + "=" * 60)
print("Tensor Core 测试 (BF16 vs FP32)")
print("=" * 60)
if torch.cuda.is_available():
    import time
    device = torch.device("cuda")
    size = 4096
    iterations = 100
    
    # BF16 测试
    torch.cuda.synchronize()
    x_bf16 = torch.randn(size, size, device=device, dtype=torch.bfloat16)
    y_bf16 = torch.randn(size, size, device=device, dtype=torch.bfloat16)
    torch.cuda.synchronize()
    
    start = time.time()
    for _ in range(iterations):
        z_bf16 = torch.matmul(x_bf16, y_bf16)
    torch.cuda.synchronize()
    bf16_time = time.time() - start
    
    # FP32 测试
    torch.cuda.synchronize()
    x_fp32 = torch.randn(size, size, device=device, dtype=torch.float32)
    y_fp32 = torch.randn(size, size, device=device, dtype=torch.float32)
    torch.cuda.synchronize()
    
    start = time.time()
    for _ in range(iterations):
        z_fp32 = torch.matmul(x_fp32, y_fp32)
    torch.cuda.synchronize()
    fp32_time = time.time() - start
    
    print(f"BF16 矩阵乘法 ({iterations}次, {size}x{size}): {bf16_time:.4f}s")
    print(f"FP32 矩阵乘法 ({iterations}次, {size}x{size}): {fp32_time:.4f}s")
    print(f"加速比: {fp32_time/bf16_time:.2f}x")
    
    if bf16_time > fp32_time * 0.9:
        print("⚠️  警告: BF16 没有明显加速,可能未使用 Tensor Core!")
    else:
        print("✅ BF16 加速正常,Tensor Core 工作正常")

print("\n" + "=" * 60)
print("建议")
print("=" * 60)
print("1. 在两台机器上分别运行此脚本")
print("2. 对比 '加速比' 和 '显卡算力'")
print("3. 如果 BF16 加速比 < 1.5x,说明可能有回退")
print("4. RTX 4060 Ti 的算力应该是 (8, 9),RTX 3060 是 (8, 6)")