''' 运行: python gpu_utils/gpu_stress.py --target_utilization 0.8 --duration 4 --num_gpus 8 ''' import torch import torch.nn as nn import torch.nn.functional as F import time import threading import argparse import numpy as np from typing import List, Dict import psutil import os from concurrent.futures import ThreadPoolExecutor import multiprocessing as mp class H100StressTest: def __init__(self, target_utilization: float = 0.5, num_gpus: int = None): self.target_utilization = target_utilization self.running = False self.stats = { 'operations': 0, 'start_time': None, 'peak_memory': 0, 'avg_utilization': 0 } # 检查GPU可用性并确定使用的GPU数量 if torch.cuda.is_available(): available_gpus = torch.cuda.device_count() if num_gpus is None: self.num_gpus = available_gpus else: self.num_gpus = min(num_gpus, available_gpus) self.devices = [torch.device(f"cuda:{i}") for i in range(self.num_gpus)] print(f"检测到 {available_gpus} 个GPU，使用 {self.num_gpus} 个GPU") for i in range(self.num_gpus): print(f"GPU {i}: {torch.cuda.get_device_name(i)}") print(f"GPU {i} 内存: {torch.cuda.get_device_properties(i).total_memory / 1024**3:.1f} GB") print(f"GPU {i} 计算能力: {torch.cuda.get_device_properties(i).major}.{torch.cuda.get_device_properties(i).minor}") else: raise RuntimeError("CUDA不可用，无法进行GPU压力测试") def create_large_transformer_model(self, hidden_size: int = 4096, num_layers: int = 8, device: torch.device = None) -> nn.Module: """创建大型Transformer模型以充分利用H100""" class LargeTransformer(nn.Module): def __init__(self, hidden_size, num_layers): super().__init__() self.hidden_size = hidden_size self.num_layers = num_layers # 多层Transformer编码器 encoder_layer = nn.TransformerEncoderLayer( d_model=hidden_size, nhead=16, # 多头注意力 dim_feedforward=hidden_size * 4, # 前馈网络 dropout=0.1, batch_first=True, norm_first=True ) self.transformer = nn.TransformerEncoder(encoder_layer, num_layers=num_layers) # 额外的线性层 self.input_projection = nn.Linear(hidden_size, hidden_size) self.output_projection = nn.Linear(hidden_size, hidden_size) self.layer_norm = nn.LayerNorm(hidden_size) def forward(self, x): x = self.input_projection(x) x = self.transformer(x) x = self.layer_norm(x) x = self.output_projection(x) return x if device is None: device = self.devices[0] model = LargeTransformer(hidden_size, num_layers).to(device) return model.half() # 使用半精度以节省内存并提高性能 def create_convolutional_model(self, device: torch.device = None) -> nn.Module: """创建卷积神经网络模型""" class ConvNet(nn.Module): def __init__(self): super().__init__() self.conv1 = nn.Conv2d(3, 64, 3, padding=1) self.conv2 = nn.Conv2d(64, 128, 3, padding=1) self.conv3 = nn.Conv2d(128, 256, 3, padding=1) self.conv4 = nn.Conv2d(256, 512, 3, padding=1) self.conv5 = nn.Conv2d(512, 1024, 3, padding=1) self.pool = nn.MaxPool2d(2, 2) self.adaptive_pool = nn.AdaptiveAvgPool2d((1, 1)) self.fc1 = nn.Linear(1024, 2048) self.fc2 = nn.Linear(2048, 4096) self.fc3 = nn.Linear(4096, 2048) self.fc4 = nn.Linear(2048, 1024) self.dropout = nn.Dropout(0.1) def forward(self, x): x = F.relu(self.conv1(x)) x = self.pool(x) x = F.relu(self.conv2(x)) x = self.pool(x) x = F.relu(self.conv3(x)) x = self.pool(x) x = F.relu(self.conv4(x)) x = self.pool(x) x = F.relu(self.conv5(x)) x = self.adaptive_pool(x) x = x.view(x.size(0), -1) x = F.relu(self.fc1(x)) x = self.dropout(x) x = F.relu(self.fc2(x)) x = self.dropout(x) x = F.relu(self.fc3(x)) x = self.dropout(x) x = self.fc4(x) return x if device is None: device = self.devices[0] model = ConvNet().to(device) return model.half() def matrix_operations_intensive(self, size: int = 8192, duration: float = 10.0, device: torch.device = None): """高强度矩阵运算""" if device is None: device = self.devices[0] # 创建大矩阵 a = torch.randn(size, size, device=device, dtype=torch.float16) b = torch.randn(size, size, device=device, dtype=torch.float16) c = torch.randn(size, size, device=device, dtype=torch.float16) start_time = time.time() operations = 0 while time.time() - start_time < duration: # 矩阵乘法链 result = torch.mm(a, b) result = torch.mm(result, c) result = torch.mm(result, a.transpose(0, 1)) # 矩阵求逆（小矩阵） if size <= 2048: try: inv_result = torch.inverse(result[:min(1024, size), :min(1024, size)]) result = torch.mm(result, inv_result) except: pass # 矩阵加法 result = result + b result = result + c # 矩阵转置 result = result.transpose(0, 1) operations += 1 # 每100次操作同步一次 if operations % 100 == 0: torch.cuda.synchronize(device) return operations def transformer_training_loop(self, model, optimizer, criterion, batch_size: int = 16, seq_length: int = 2048, device: torch.device = None): """Transformer训练循环""" if device is None: device = self.devices[0] hidden_size = model.hidden_size # 创建输入数据 x = torch.randn(batch_size, seq_length, hidden_size, device=device, dtype=torch.float16) y = torch.randn(batch_size, seq_length, hidden_size, device=device, dtype=torch.float16) # 前向传播 outputs = model(x) loss = criterion(outputs, y) # 反向传播 optimizer.zero_grad() loss.backward() optimizer.step() return loss.item() def conv_training_loop(self, model, optimizer, criterion, batch_size: int = 32, device: torch.device = None): """卷积网络训练循环""" if device is None: device = self.devices[0] # 创建输入数据 x = torch.randn(batch_size, 3, 224, 224, device=device, dtype=torch.float16) y = torch.randn(batch_size, 1024, device=device, dtype=torch.float16) # 前向传播 outputs = model(x) loss = criterion(outputs, y) # 反向传播 optimizer.zero_grad() loss.backward() optimizer.step() return loss.item() def monitor_gpu_stats(self): """监控所有GPU统计信息""" while self.running: try: total_allocated = 0 total_cached = 0 gpu_stats = [] # 获取所有GPU的内存使用情况 for i in range(self.num_gpus): allocated = torch.cuda.memory_allocated(i) / 1024**3 cached = torch.cuda.memory_reserved(i) / 1024**3 total_allocated += allocated total_cached += cached gpu_stats.append((allocated, cached)) # 更新峰值内存 if total_allocated > self.stats['peak_memory']: self.stats['peak_memory'] = total_allocated # 每10秒打印一次状态 if self.stats['start_time'] and int(time.time() - self.stats['start_time']) % 10 == 0: elapsed = time.time() - self.stats['start_time'] ops_per_sec = self.stats['operations'] / elapsed if elapsed > 0 else 0 print(f"[{elapsed:.0f}s] 操作数: {self.stats['operations']}, " f"操作/秒: {ops_per_sec:.1f}") print(f"总GPU内存: {total_allocated:.2f}GB/{total_cached:.2f}GB, " f"峰值内存: {self.stats['peak_memory']:.2f}GB") # 显示每个GPU的详细状态 for i, (allocated, cached) in enumerate(gpu_stats): print(f" GPU {i}: {allocated:.2f}GB/{cached:.2f}GB") time.sleep(1) except Exception as e: print(f"监控错误: {e}") time.sleep(5) def gpu_worker(self, gpu_id: int, duration_seconds: int, shared_stats: dict, stats_lock: threading.Lock): """单个GPU的工作线程""" device = self.devices[gpu_id] print(f"GPU {gpu_id} 开始工作...") # 为每个GPU创建独立的模型 transformer_model = self.create_large_transformer_model( hidden_size=4096, num_layers=12, device=device ) transformer_optimizer = torch.optim.AdamW( transformer_model.parameters(), lr=1e-4, weight_decay=0.01 ) transformer_criterion = nn.MSELoss() conv_model = self.create_convolutional_model(device=device) conv_optimizer = torch.optim.AdamW( conv_model.parameters(), lr=1e-4, weight_decay=0.01 ) conv_criterion = nn.MSELoss() # 预热GPU for _ in range(5): self.matrix_operations_intensive(size=4096, duration=1.0, device=device) start_time = time.time() local_operations = 0 cycle_count = 0 try: while time.time() - start_time < duration_seconds: cycle_count += 1 current_time = time.time() - start_time # 每30秒轮换不同类型的操作 if int(current_time) % 30 < 10: # 矩阵运算 ops = self.matrix_operations_intensive(size=8192, duration=8, device=device) local_operations += ops elif int(current_time) % 30 < 20: # Transformer训练 for _ in range(5): loss = self.transformer_training_loop( transformer_model, transformer_optimizer, transformer_criterion, batch_size=8, seq_length=2048, device=device ) local_operations += 1 else: # 卷积网络训练 for _ in range(10): loss = self.conv_training_loop( conv_model, conv_optimizer, conv_criterion, batch_size=16, device=device ) local_operations += 1 # 每100个周期清理一次缓存 if cycle_count % 100 == 0: torch.cuda.empty_cache() # 短暂休息以控制利用率 time.sleep(0.1) except Exception as e: print(f"GPU {gpu_id} 发生错误: {e}") finally: # 更新共享统计信息 with stats_lock: shared_stats['operations'] += local_operations print(f"GPU {gpu_id} 完成工作，操作数: {local_operations}") def run_intensive_stress_test(self, duration_hours: float = 2.0): """运行多GPU高强度压力测试""" duration_seconds = int(duration_hours * 3600) print(f"开始多GPU高强度压力测试") print(f"使用GPU数量: {self.num_gpus}") print(f"目标利用率: {self.target_utilization*100:.1f}%") print(f"持续时间: {duration_hours:.1f}小时 ({duration_seconds}秒)") # 启动监控线程 self.running = True self.stats['start_time'] = time.time() monitor_thread = threading.Thread(target=self.monitor_gpu_stats) monitor_thread.daemon = True monitor_thread.start() # 创建共享统计信息和锁 shared_stats = {'operations': 0} stats_lock = threading.Lock() # 使用线程池并行运行所有GPU with ThreadPoolExecutor(max_workers=self.num_gpus) as executor: # 为每个GPU提交任务 futures = [] for gpu_id in range(self.num_gpus): future = executor.submit( self.gpu_worker, gpu_id, duration_seconds, shared_stats, stats_lock ) futures.append(future) try: # 等待所有任务完成 for future in futures: future.result() except KeyboardInterrupt: print("\n用户中断测试") # 取消所有未完成的任务 for future in futures: future.cancel() self.running = False elapsed = time.time() - self.stats['start_time'] print(f"\n多GPU压力测试完成!") print(f"总运行时间: {elapsed:.1f}秒 ({elapsed/3600:.2f}小时)") print(f"总操作数: {shared_stats['operations']}") print(f"平均操作/秒: {shared_stats['operations']/elapsed:.1f}") print(f"峰值GPU内存使用: {self.stats['peak_memory']:.2f}GB") def get_gpu_info(self): """获取所有GPU的详细信息""" if torch.cuda.is_available(): print("=== 多GPU 详细信息 ===") print(f"GPU数量: {torch.cuda.device_count()}") print(f"使用GPU数量: {self.num_gpus}") total_memory = 0 for i in range(self.num_gpus): print(f"\n--- GPU {i} ---") print(f"GPU名称: {torch.cuda.get_device_name(i)}") props = torch.cuda.get_device_properties(i) gpu_memory = props.total_memory / 1024**3 total_memory += gpu_memory print(f"总内存: {gpu_memory:.1f} GB") print(f"多处理器数量: {props.multi_processor_count}") print(f"计算能力: {props.major}.{props.minor}") # 安全获取可能不存在的属性 try: print(f"最大线程数/块: {props.max_threads_per_block}") except AttributeError: print("最大线程数/块: 信息不可用") try: print(f"最大网格尺寸: {props.max_grid_size}") except AttributeError: print("最大网格尺寸: 信息不可用") try: print(f"共享内存/块: {props.shared_memory_per_block / 1024:.0f} KB") except AttributeError: print("共享内存/块: 信息不可用") # 获取当前内存使用情况 allocated = torch.cuda.memory_allocated(i) / 1024**3 cached = torch.cuda.memory_reserved(i) / 1024**3 print(f"已分配内存: {allocated:.2f} GB") print(f"缓存内存: {cached:.2f} GB") print(f"可用内存: {gpu_memory - cached:.2f} GB") print(f"\n总GPU内存: {total_memory:.1f} GB") else: print("CUDA不可用") def main(): parser = argparse.ArgumentParser(description='多GPU高强度压力测试工具') parser.add_argument('--duration', type=float, default=2.0, help='测试持续时间（小时）') parser.add_argument('--target_utilization', type=float, default=0.5, help='目标GPU利用率（0-1）') parser.add_argument('--num_gpus', type=int, default=None, help='使用的GPU数量（默认使用所有可用GPU）') parser.add_argument('--info_only', action='store_true', help='只显示GPU信息，不进行测试') args = parser.parse_args() try: # 创建压力测试实例 stress_test = H100StressTest( target_utilization=args.target_utilization, num_gpus=args.num_gpus ) # 显示GPU信息 stress_test.get_gpu_info() if args.info_only: return # 运行压力测试 stress_test.run_intensive_stress_test(duration_hours=args.duration) except Exception as e: print(f"错误: {e}") return 1 if __name__ == "__main__": exit(main())