兄弟，我看了你的代码和问题描述，一眼就发现了核心问题——你同时混用了两种分布式训练的启动方式，直接导致进程和设备分配彻底混乱，才会报出张量设备不统一的错误！

问题根源 #

你的代码是基于torch.multiprocessing.spawn来手动启动多进程的，但你却用torchrun命令来运行脚本，这两种方式完全不能同时使用：

• mp.spawn是PyTorch提供的手动多进程启动工具，直接执行python script.py就能让它自动创建指定数量的训练进程，rank参数会自动传递给你的train函数。
• torchrun（包括旧版的torch.distributed.launch）是官方推荐的分布式启动器，它会自动帮你管理进程、设置RANK和WORLD_SIZE环境变量，不需要你再用mp.spawn二次启动进程。

你现在的操作相当于：用torchrun启动了4个进程，每个进程又通过mp.spawn再启动4个进程，总共16个进程！这直接导致进程组的rank完全错乱，模型和数据的设备分配逻辑崩溃，自然会出现“张量在cuda:0和CPU”的错误。

两种修复方案（二选一即可） #

方案1：保留mp.spawn，直接用Python运行脚本

不需要修改代码，把启动命令改成：

python script.py

mp.spawn会自动创建4个训练进程，每个进程的rank参数正确传递，模型和数据的设备分配就能正常工作，那个RuntimeError应该会立刻消失。

方案2：改用torchrun启动（推荐正式训练使用）

这种方式更简洁，也支持弹性训练等高级功能，需要修改代码移除mp.spawn的逻辑，改成从环境变量获取rank和world_size：

import os
# 其他导入保持不变

class SimpleModel(nn.Module):
    def __init__(self):
        super(SimpleModel, self).__init__()
        self.fc = nn.Linear(10, 10)
    
    def forward(self, x):
        return self.fc(x)

def setup():
    # 从环境变量获取rank和world_size
    rank = int(os.environ["RANK"])
    world_size = int(os.environ["WORLD_SIZE"])
    dist.init_process_group("nccl", rank=rank, world_size=world_size)
    torch.cuda.set_device(rank)

def train():
    setup()
    rank = int(os.environ["RANK"])
    world_size = int(os.environ["WORLD_SIZE"])
    
    model = SimpleModel().to(rank)
    model = DDP(model, device_ids=[rank])
    
    optimizer = torch.optim.Adam(model.parameters())
    
    # Dummy data loader（实际训练建议用DistributedSampler）
    inputs = torch.randn(4, 10).to(rank)
    labels = torch.randint(0, 10, (4,)).to(rank)
    data_loader = [(inputs, labels)]
    
    for epoch in range(1):
        for batch in data_loader:
            inputs, labels = batch
            optimizer.zero_grad()
            outputs = model(inputs)
            loss = nn.CrossEntropyLoss()(outputs, labels)
            loss.backward()
            optimizer.step()
    
    dist.destroy_process_group()

if __name__ == "__main__":
    train()

启动命令还是你原来的：

torchrun --nproc_per_node=4 script.py

额外训练优化建议 #

如果之后要做正式训练，记得用DistributedSampler来拆分数据集，让每个进程只处理自己的专属数据子集，避免重复计算：

from torch.utils.data import Dataset, DataLoader, DistributedSampler

# 定义自定义数据集
class DummyDataset(Dataset):
    def __len__(self):
        return 16  # 总样本数
    def __getitem__(self, idx):
        return torch.randn(10), torch.randint(0, 10, ())

# 在train函数里替换数据加载部分
dataset = DummyDataset()
sampler = DistributedSampler(dataset)  # 自动拆分数据
data_loader = DataLoader(dataset, batch_size=4, sampler=sampler)

# 训练循环里需要设置sampler的epoch
for epoch in range(epochs):
    sampler.set_epoch(epoch)  # 保证每个epoch数据打乱的一致性
    for batch in data_loader:
        # 训练逻辑...

先试试方案1，应该能快速解决你的设备不统一问题～如果之后要做大规模训练，再切换到方案2配合DistributedSampler，训练效率会更高！