Pytorch 启动训练过程在阻塞 GPU 0 时卡住[已关闭]

Question

我的服务器中的 GPU0 已被其他进程占用，因此我阻塞了 GPU0 并使用

mp.spawn

来训练我的模型，但未能创建训练进程。

这是可重现的示例：

import torch
import torch.multiprocessing as mp
import torch.nn as nn
import torch.optim as optim
import os
import torch.distributed as dist

def train(rank, gpu_ids):
    device_id = gpu_ids[rank]
    dist.init_process_group(
        backend="nccl", init_method="env://", world_size=len(gpu_ids), rank=device_id)
    torch.cuda.set_device(device_id)
    print(f"Rank {rank} is using device {device_id}")
    model = nn.Linear(10, 1).to(device_id)
    optimizer = optim.SGD(model.parameters(), lr=0.01)
    input_data = torch.randn(5, 10).to(device_id)
    target = torch.randn(5, 1).to(device_id)
    for epoch in range(5):
        optimizer.zero_grad()
        output = model(input_data)
        loss = nn.MSELoss()(output, target)
        loss.backward()
        optimizer.step()
        print(f"Rank {rank}, Epoch {epoch}, Loss: {loss.item()}")

def test_main():
    os.environ['CUDA_VISIBLE_DEVICES'] = '1,2'
    os.environ["MASTER_ADDR"] = '127.0.0.1'
    os.environ["MASTER_PORT"] = '12345'
    gpu_ids = [1]
    world_size = len(gpu_ids)
    mp.spawn(train, args=(gpu_ids,), nprocs=world_size, join=True)

在我的测试中，

train()

方法卡在这里：

dist.init_process_group(backend="nccl", init_method="env://", world_size=len(gpu_ids), rank=device_id)

那么到底发生了什么以及如何解决呢？

Answer 1

单 GPU 设置：如果您不使用多个 GPU 进行分布式训练，请删除分布式设置代码 (dist.init_process_group) 以避免不必要的复杂性。

GPU 冲突：确保 GPU 0 没有被任何其他进程使用。您可以设置 CUDA_VISIBLE_DEVICES='1' 来显式限制 GPU 1 的使用。