H100 Mem BW

mem_bw.py

import torch
import pandas as pd
import matplotlib.pyplot as plt
import numpy as np

# Define the range of message sizes in bytes
sizes = [2**i for i in range(0, 35)]  # 1B to 16GB

# Define the number of iterations
iterations = 1000

# Function to measure memory bandwidth
def measure_bandwidth():
    results = []
    for size in sizes:
        num_elements = size // 4 if size >= 4 else 1  # Number of float32 elements (at least 1 element)

        # Allocate tensors
        src = torch.ones(num_elements, dtype=torch.float32, device='cuda')
        dst = torch.zeros(num_elements, dtype=torch.float32, device='cuda')

        # Warm-up
        dst.copy_(src)
        torch.cuda.synchronize()

        elapsed_times = []

        for _ in range(iterations):
            start_event = torch.cuda.Event(enable_timing=True)
            end_event = torch.cuda.Event(enable_timing=True)

            start_event.record()
            dst.copy_(src)
            end_event.record()

            # Wait for the copy to complete
            torch.cuda.synchronize()

            # Calculate elapsed time in milliseconds
            elapsed_time_ms = start_event.elapsed_time(end_event)
            elapsed_times.append(elapsed_time_ms)

        # Average elapsed time
        avg_elapsed_time_ms = sum(elapsed_times) / len(elapsed_times)
        
        # Calculate bandwidth in GB/s
        bandwidth = (size / (avg_elapsed_time_ms * 1e6))  # GB/s
        results.append((size, bandwidth))

        print(f"Size: {size / 1e6} MB, Bandwidth: {bandwidth:.2f} GB/s")

    return results

# Run the measurement
if __name__ == "__main__":
    bandwidth_results = measure_bandwidth()
    
    # Convert results to a DataFrame
    df = pd.DataFrame(bandwidth_results, columns=['Size (bytes)', 'Bandwidth (GB/s)'])
    
    # Export to CSV
    df.to_csv('bandwidth_results.csv', index=False)
    
    # Plot the results
    plt.figure(figsize=(10, 6))
    plt.plot(df['Size (bytes)'], df['Bandwidth (GB/s)'], marker='o')
    plt.xscale('log', base=2)
    plt.xticks(ticks=sizes, labels=[f"{int(size)}B" if size < 1024 else 
                                    f"{int(size / 1024)}KB" if size < 1024**2 else 
                                    f"{int(size / 1024**2)}MB" if size < 1024**3 else 
                                    f"{int(size / 1024**3)}GB" for size in sizes], rotation=45)
    plt.yticks(ticks=np.arange(0, max(df['Bandwidth (GB/s)']) + 100, 100))  # Y-axis ticks every 100 GB/s
    plt.xlabel('Message Size')
    plt.ylabel('Bandwidth (GB/s)')
    plt.title('Memory Bandwidth vs Message Size')
    plt.grid(True)
    plt.tight_layout()
    plt.savefig('bandwidth_plot.png')
    plt.show()

results.csv