Python使用cuda报错KeyError: ‘LOCAL_RANK‘

code

rank = int(os.environ['LOCAL_RANK'])
world_size = int(os.environ['WORLD_SIZE'])

报错

Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/conda/Miniconda3/envs/jfs_py38_tts/lib/python3.8/os.py", line 675, in __getitem__
    raise KeyError(key) from None
KeyError: 'LOCAL_RANK'

Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/conda/Miniconda3/envs/jfs_py38_tts/lib/python3.8/os.py", line 675, in __getitem__
    raise KeyError(key) from None
KeyError: 'WORLD_SIZE'

解决方案

在Python脚本调用之前，指定CUDA-visuble

PyTorch 进行分布式训练时常用的命令行指令 from 豆包

这是一个在使用 PyTorch 进行分布式训练时常用的命令行指令。

以下是对其各个部分的解释和整体功能说明：

一、CUDA_VISIBLE_DEVICES=$gpu_ids

1. 解释：

   • CUDA_VISIBLE_DEVICES是一个环境变量，用于控制哪些 GPU 设备对当前进程可见。
   • $gpu_ids是一个变量，通常代表一个或多个 GPU 的编号列表，例如0,1,2表示让当前进程只看到编号为 0、1、2 的 GPU。

2. 功能：

   • 通过设置这个环境变量，可以指定当前任务只使用特定的 GPU 设备，从而有效地管理 GPU 资源分配，避免不必要的 GPU 使用冲突。

二、torchrun --nnodes=1 --nproc_per_node=$num_gpus

1. 解释：

   • torchrun是 PyTorch 提供的用于启动分布式训练的命令行工具。
   • --nnodes参数指定了参与分布式训练的节点数量。在这里设置为1表示只有一个节点参与训练。
   • --nproc_per_node参数指定了每个节点上启动的进程数量。$num_gpus通常是一个变量，表示每个节点上可用的 GPU 数量，这个参数设置每个节点上启动的进程数量等于 GPU 数量，以实现每个 GPU 对应一个进程的分布式训练。

2. 功能：

   • 启动一个分布式训练任务，在单个节点上根据指定的 GPU 数量启动相应数量的进程，每个进程可以独立地执行训练任务，并通过分布式通信机制进行协作和同步。
   • 这种方式可以充分利用多个 GPU 的计算能力，加速深度学习模型的训练过程。

例如， 如果有 4 个 GPU，并且设置gpu_ids=0,1,2,3以及num_gpus=4，那么这个命令会让当前任务只看到编号为 0、1、2、3 的 GPU，并在一个节点上启动 4 个进程，每个进程对应一个 GPU 进行分布式训练。

torchrun基本用法 from 豆包

torchrun是 PyTorch 用于启动分布式训练的命令行工具。以下是它的使用方法：

一、基本用法

1. 启动分布式训练：
   • torchrun --nnodes=<num_nodes> --nproc_per_node=<num_processes_per_node> your_training_script.py [--arg1 value1 --arg2 value2...]
   • <num_nodes>是参与分布式训练的节点数量。
   • <num_processes_per_node>是每个节点上启动的进程数量，通常与节点上的 GPU 数量相等，以实现每个 GPU 对应一个进程。
   • your_training_script.py是要运行的训练脚本文件。
   • [--arg1 value1 --arg2 value2...]是传递给训练脚本的参数。

例如，启动一个包含两个节点，每个节点有 4 个 GPU 的分布式训练：

torchrun --nnodes=2 --nproc_per_node=4 train_script.py --learning_rate=0.001

2. 环境变量设置：
   • CUDA_VISIBLE_DEVICES：可以在命令前设置这个环境变量来控制每个节点上可见的 GPU 设备。例如，CUDA_VISIBLE_DEVICES=0,1,2,3 torchrun --nnodes=1 --nproc_per_node=4 train_script.py将只使用编号为 0、1、2、3 的 GPU。

二、高级用法

1. 分布式数据并行（DDP）：
   • 在训练脚本中，使用torch.distributed模块来设置分布式环境和初始化进程组。
   • 例如：

import torch
import torch.distributed as dist

def setup(rank, world_size):
    os.environ['MASTER_ADDR'] = 'localhost'
    os.environ['MASTER_PORT'] = '12345'
    dist.init_process_group("nccl", rank=rank, world_size=world_size)

def cleanup():
    dist.destroy_process_group()

• 在训练循环中，使用torch.nn.parallel.DistributedDataParallel来包装模型，实现数据并行训练。
• 例如：

from torch.nn.parallel import DistributedDataParallel as DDP

model = YourModel()
model = DDP(model, device_ids=[rank], output_device=rank)

2. 弹性训练：
   • torchrun支持弹性训练，即可以动态地添加或移除训练节点和进程。
   • 要实现弹性训练，需要在训练脚本中处理进程的动态加入和退出，并正确地同步模型参数。

三、注意事项

1. 网络配置：

   • 确保所有参与分布式训练的节点可以相互通信。如果在不同的机器上运行，需要正确设置网络地址和端口号。
   • 可以通过设置环境变量MASTER_ADDR和MASTER_PORT来指定主节点的地址和端口号。

2. 依赖管理：

   • 确保所有节点上的环境和依赖项一致。特别是对于深度学习训练，需要确保相同版本的 PyTorch、CUDA 和其他依赖库在所有节点上安装。

3. 日志和错误处理：

   • 分布式训练中，错误可能在不同的进程中出现，并且可能难以调试。可以使用集中式日志记录或分布式日志工具来帮助诊断问题。
   • 注意观察torchrun的输出和训练脚本的日志，以获取有关训练进度和错误的信息。