File size: 3,786 Bytes
08bf07d |
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 |
import torch
import time
import argparse
from threading import Thread
def gpu_worker(gpu_id, duration, tensor_size):
"""单个GPU的工作线程，负责持续进行张量运算"""
try:
# 设置当前线程使用的GPU
device = torch.device(f"cuda:{gpu_id}")
torch.cuda.set_device(device)
# 打印GPU信息
gpu_name = torch.cuda.get_device_name(device)
print(f"GPU {gpu_id} 启动: {gpu_name}")
# 创建大随机张量
tensor_a = torch.randn(tensor_size, tensor_size, device=device)
tensor_b = torch.randn(tensor_size, tensor_size, device=device)
# 预热GPU
for _ in range(10):
result = torch.matmul(tensor_a, tensor_b)
torch.cuda.synchronize(device)
# 开始持续运算
start_time = time.time()
iterations = 0
while time.time() - start_time < duration:
# 矩阵乘法运算
result = torch.matmul(tensor_a, tensor_b)
# 定期更新张量避免优化
if iterations % 100 == 0:
tensor_a = 0.999 * tensor_a + 0.001 * torch.randn_like(tensor_a)
tensor_b = 0.999 * tensor_b + 0.001 * torch.randn_like(tensor_b)
iterations += 1
# 每10秒打印一次状态
if iterations % 1000 == 0:
elapsed = time.time() - start_time
print(f"GPU {gpu_id}: 已运行 {elapsed:.1f} 秒, 完成 {iterations} 次迭代")
# 短暂同步确保计算完成
if iterations % 100 == 0:
torch.cuda.synchronize(device)
# 计算结束统计
elapsed = time.time() - start_time
print(f"GPU {gpu_id} 完成: 总时间 {elapsed:.1f} 秒, 总迭代 {iterations} 次, "
f"平均每秒 {iterations/elapsed:.2f} 次")
except Exception as e:
print(f"GPU {gpu_id} 出错: {str(e)}")
finally:
# 清理内存
torch.cuda.empty_cache()
def multi_gpu_stress_test(duration, tensor_size, use_gpus=None):
"""多GPU压力测试主函数"""
# 检查可用GPU数量
available_gpus = torch.cuda.device_count()
if available_gpus == 0:
print("错误: 未检测到可用GPU")
return
# 确定要使用的GPU
if use_gpus is None:
use_gpus = list(range(available_gpus))
else:
# 验证GPU ID有效性
use_gpus = [g for g in use_gpus if 0 <= g < available_gpus]
if not use_gpus:
print("错误: 没有有效的GPU ID")
return
print(f"检测到 {available_gpus} 张GPU，将使用 {len(use_gpus)} 张: {use_gpus}")
# 为每张GPU创建并启动线程
threads = []
for gpu_id in use_gpus:
thread = Thread(target=gpu_worker, args=(gpu_id, duration, tensor_size))
threads.append(thread)
thread.start()
# 等待所有线程完成
for thread in threads:
thread.join()
print("所有GPU测试完成")
if __name__ == "__main__":
parser = argparse.ArgumentParser(description='多GPU压力测试程序')
parser.add_argument('--duration', type=int, default=6000000,
help='测试持续时间(秒)，默认60秒')
parser.add_argument('--size', type=int, default=4096,
help='每张GPU上的张量大小，默认4096x4096')
parser.add_argument('--gpus', type=int, nargs='+',
help=f'指定要使用的GPU ID，如 --gpus 0 1 2 3 4 5 6 7')
args = parser.parse_args()
# 运行多GPU测试
multi_gpu_stress_test(args.duration, args.size, args.gpus)
|