gpu_bandwidth.py

# test gpu bandwidth 그냥 멀티 GPU 환경이 이상해서 쓰는 코드
import torch

def create_timing_events():
    """CUDA 타이밍 이벤트 생성"""
    start_time = torch.cuda.Event(enable_timing=True)
    end_time = torch.cuda.Event(enable_timing=True)
    return start_time, end_time

def get_gpu_memory_size(gpu_id):
    """GPU의 가용 메모리 크기 반환 (바이트 단위)"""
    total_memory = torch.cuda.get_device_properties(gpu_id).total_memory
    return int(total_memory * 0.95)  # 95% 사용

def measure_transfer_time(cpu_tensor, gpu_id, test_num=None):
    """특정 GPU에 대한 양방향 전송 시간 및 대역폭 측정"""
    start_time_up, end_time_up = create_timing_events()
    start_time_down, end_time_down = create_timing_events()
    
    torch.cuda.synchronize()  # 시작 전 동기화
    
    # CPU->GPU 전송 시간 측정
    with torch.cuda.device(gpu_id):
        start_time_up.record()
        gpu_tensor = cpu_tensor.cuda(gpu_id, non_blocking=True)  # non_blocking은 유지
        end_time_up.record()
        
        # GPU->CPU 전송 시간 측정
        start_time_down.record()
        cpu_tensor_copy = gpu_tensor.cpu()  # non_blocking 제거
        end_time_down.record()
    
    torch.cuda.synchronize()  # 모든 작업 완료 후 동기화
    
    elapsed_time_up = start_time_up.elapsed_time(end_time_up)
    elapsed_time_down = start_time_down.elapsed_time(end_time_down)
    
    tensor_size = cpu_tensor.nelement() * cpu_tensor.element_size()
    bandwidth_up = tensor_size / (elapsed_time_up / 1000) / (1024**3)
    bandwidth_down = tensor_size / (elapsed_time_down / 1000) / (1024**3)
    
    if test_num is not None:
        print(f"\n테스트 #{test_num}")
        print(f"텐서 크기: {tensor_size / (1024**3):.2f} GB")  # 텐서 크기 출력 추가
    print(f"GPU {gpu_id} - 업로드 시간: {elapsed_time_up:.2f} ms, 대역폭: {bandwidth_up:.2f} GB/s")
    print(f"GPU {gpu_id} - 다운로드 시간: {elapsed_time_down:.2f} ms, 대역폭: {bandwidth_down:.2f} GB/s")
    
    return (elapsed_time_up, bandwidth_up), (elapsed_time_down, bandwidth_down)

def run_bandwidth_test(num_tests=5):
    """모든 GPU에 대한 대역폭 테스트 실행"""
    num_gpus = torch.cuda.device_count()
    print(f"\n=== {num_gpus}개 GPU 대역폭 테스트 시작 ===")
    
    # GPU 정보 출력
    for gpu_id in range(num_gpus):
        memory_size = get_gpu_memory_size(gpu_id)
        print(f"\nGPU {gpu_id}: {torch.cuda.get_device_name(gpu_id)}")
        print(f"총 메모리: {memory_size / (1024**3):.1f} GB (95% 사용)")
    
    results = {gpu_id: {'up': [], 'down': []} for gpu_id in range(num_gpus)}
    
    for gpu_id in range(num_gpus):
        # GPU별로 95% 메모리 크기의 테스트 데이터 생성
        memory_size = get_gpu_memory_size(gpu_id)
        tensor_size = memory_size // 4  # float32는 4바이트
        size_gb = memory_size / (1024**3)
        
        print(f"\nGPU {gpu_id}용 테스트 데이터 생성 중 ({size_gb:.1f} GB)...")
        
        # pinned memory 사용
        cpu_tensor = torch.rand(tensor_size, dtype=torch.float32, pin_memory=True)
        
        # 웜업 단계
        print(f"GPU {gpu_id} 웜업 중...")
        for _ in range(2):
            with torch.no_grad():
                measure_transfer_time(cpu_tensor, gpu_id)
        
        # 실제 테스트 실행
        for test_num in range(num_tests):
            print(f"\n=== GPU {gpu_id} 테스트 라운드 {test_num + 1}/{num_tests} ===")
            (_, bw_up), (_, bw_down) = measure_transfer_time(cpu_tensor, gpu_id, test_num + 1)
            results[gpu_id]['up'].append(bw_up)
            results[gpu_id]['down'].append(bw_down)
        
        # 테스트 완료 후 메모리 해제
        del cpu_tensor
        torch.cuda.empty_cache()
    
    # 결과 통계
    print("\n=== 최종 결과 ===")
    for gpu_id in range(num_gpus):
        print(f"\nGPU {gpu_id} 결과:")
        for direction in ['up', 'down']:
            avg_bw = sum(results[gpu_id][direction]) / len(results[gpu_id][direction])
            max_bw = max(results[gpu_id][direction])
            min_bw = min(results[gpu_id][direction])
            dir_name = "업로드" if direction == 'up' else "다운로드"
            print(f"{dir_name} 평균 대역폭: {avg_bw:.2f} GB/s")
            print(f"{dir_name} 최대 대역폭: {max_bw:.2f} GB/s")
            print(f"{dir_name} 최소 대역폭: {min_bw:.2f} GB/s")

# 대역폭 테스트 실행
if __name__ == "__main__":
    run_bandwidth_test()