gpu_zombie.py

gpu_zombie.py

import torch

ngpus = torch.cuda.device_count()
N, M, K = 1280, 2560, 5120

def matmul():
    A = []
    B = torch.randn(M, K, device='cuda:0')

    for i in range(ngpus):
        A.append(torch.randn(N // ngpus, M, device='cuda:'+str(i)))

    B_ = [B]
    for i in range(ngpus):
        if i != 0:
            B_.append(B.to('cuda:'+str(i)))

    C_ = []
    for i in range(ngpus):
        C_.append(torch.matmul(A[i], B_[i]))

    C = torch.empty(N, K)
    for i in range(ngpus):
        start_index = i * (N//ngpus)
        C[start_index:start_index + N//ngpus, :].copy_(C_[i])

while True:
    matmul()