torch.distributed.get_world_size与cuda.device_count不一致致设备序号错误

在Slurm管理的HPC环境中，使用PyTorch 2.1.0+cu121尝试将模型张量并行拆分到2张NVIDIA A100 GPU上，运行测试代码后触发RuntimeError: CUDA error: invalid device ordinal错误。

测试代码 #

from torch.distributed.tensor.parallel import (
    PairwiseParallel,
    parallelize_module,
)

import torch
from torch import nn
import torch.distributed as dist
from torch.distributed._tensor import DeviceMesh

from torch.testing._internal.common_distributed import (
    spawn_threads_and_init_comms,
)

class ToyModel(nn.Module):
    def __init__(self, in_channels, hidden_channels):
        super(ToyModel, self).__init__()
        self.dummy_param = nn.Parameter(torch.empty(0))
        self.net1 = nn.Linear(in_channels, hidden_channels)
        self.relu = nn.ReLU()
        self.net2 = nn.Linear(hidden_channels, in_channels)

    def forward(self, x):
        return self.net2(self.relu(self.net1(x)))
        
@spawn_threads_and_init_comms
def demo_world_rank_mismatch(world_size):
    rank = dist.get_rank()
    print(f'Rank: {rank}')

    print(f"World size: {dist.get_world_size()}")

    print("Create a sharding plan based on the given world_size", world_size)

    mesh = torch.arange(world_size)
    device_mesh = DeviceMesh(
        "cuda",
        mesh,
    )
    in_dim = 1024
    hidden_dim = 4 * in_dim

    model = ToyModel(in_dim, hidden_dim).to(rank)

    LR = 0.25
    optimizer = torch.optim.SGD(model.parameters(), lr=LR)

    lparallel=True
    if lparallel:
        print("Parallelize the module based on the given Parallel Style", rank)
        model = parallelize_module(model, device_mesh, PairwiseParallel())
    print(model)

    print(f"model of rank {rank} on {model.dummy_param.device}")

    # 原代码未定义此变量
    for i in range(ITER_TIME):
        inp = torch.rand(10000, in_dim).to(rank)
        output = model(inp)
        output.sum().backward()
        optimizer.step()
    
    print(f'{rank}, max memory alloc: {torch.cuda.max_memory_allocated(device=rank)}')


print(f"Device count: {torch.cuda.device_count()}")
demo_world_rank_mismatch(torch.cuda.device_count())

运行输出与错误 #

Device count: 2
Rank: 1
World size: 4
Create a sharding plan based on the given world_size 2
Rank: 2
World size: 4
Create a sharding plan based on the given world_size 2
Rank: 3
World size: 4
Create a sharding plan based on the given world_size 2
Rank: 0
World size: 4
Create a sharding plan based on the given world_size 2
Parallelize the module based on the given Parallel Style 0
Parallelize the module based on the given Parallel Style 1

随后触发错误：

RuntimeError: CUDA error: invalid device ordinal
CUDA kernel errors might be asynchronously reported at some other API call, so the stacktrace below might be incorrect.
For debugging consider passing CUDA_LAUNCH_BLOCKING=1.
Compile with `TORCH_USE_CUDA_DSA` to enable device-side assertions.

错误原因分析 #

1. 进程数与GPU数不匹配：spawn_threads_and_init_comms是PyTorch内部测试工具，默认创建4个进程（world size=4），但实际只有2个GPU。rank=2、3的进程尝试访问序号为2、3的CUDA设备，超出系统仅有的0、1设备范围，触发设备序号无效错误。
2. DeviceMesh配置错误：直接用torch.arange(world_size)创建DeviceMesh，当world_size与实际GPU数不一致时，会导致设备映射逻辑混乱。
3. 未定义变量：代码中ITER_TIME未声明，会导致后续训练循环无法执行。

解决方案 #

1. 替换测试用初始化工具，使用标准分布式流程 #

放弃spawn_threads_and_init_comms，改用PyTorch标准分布式初始化方法，适配Slurm环境的进程管理逻辑。

2. 正确映射Rank与CUDA设备 #

通过Slurm提供的LOCAL_RANK环境变量，将每个进程绑定到节点内对应的GPU设备上。

3. 修复未定义变量 #

补充ITER_TIME的定义，确保训练循环正常执行。

修改后的完整代码 #

from torch.distributed.tensor.parallel import (
    PairwiseParallel,
    parallelize_module,
)

import torch
from torch import nn
import torch.distributed as dist
from torch.distributed._tensor import DeviceMesh
import os

class ToyModel(nn.Module):
    def __init__(self, in_channels, hidden_channels):
        super(ToyModel, self).__init__()
        self.dummy_param = nn.Parameter(torch.empty(0))
        self.net1 = nn.Linear(in_channels, hidden_channels)
        self.relu = nn.ReLU()
        self.net2 = nn.Linear(hidden_channels, in_channels)

    def forward(self, x):
        return self.net2(self.relu(self.net1(x)))

def main():
    # 初始化分布式环境（适配Slurm）
    dist.init_process_group(backend="nccl")
    rank = dist.get_rank()
    local_rank = int(os.environ.get("LOCAL_RANK", 0))
    world_size = dist.get_world_size()

    # 绑定当前进程到对应CUDA设备
    torch.cuda.set_device(local_rank)
    device = torch.device(f"cuda:{local_rank}")

    print(f'Rank: {rank}, Local Rank: {local_rank}, World size: {world_size}')

    # 创建匹配实际GPU数的DeviceMesh
    device_mesh = DeviceMesh("cuda", list(range(world_size)))
    in_dim = 1024
    hidden_dim = 4 * in_dim

    model = ToyModel(in_dim, hidden_dim).to(device)

    LR = 0.25
    optimizer = torch.optim.SGD(model.parameters(), lr=LR)

    lparallel = True
    if lparallel:
        print(f"Parallelize the module on rank {rank}")
        model = parallelize_module(model, device_mesh, PairwiseParallel())
    print(model)

    print(f"model of rank {rank} on {model.dummy_param.device}")

    ITER_TIME = 3
    for i in range(ITER_TIME):
        inp = torch.rand(10000, in_dim).to(device)
        output = model(inp)
        output.sum().backward()
        optimizer.step()
    
    print(f'{rank}, max memory alloc: {torch.cuda.max_memory_allocated(device=device)}')

    # 清理分布式环境
    dist.destroy_process_group()

if __name__ == "__main__":
    print(f"Device count: {torch.cuda.device_count()}")
    main()

4. Slurm环境启动命令 #

使用srun启动2个进程（对应2张GPU）：

srun --gres=gpu:2 --cpus-per-task=8 --nodes=1 --ntasks=2 python your_script.py

内容的提问来源于stack exchange，提问作者DLS