Python使用cuda报错KeyError: ‘LOCAL_RANK‘
那么这个命令会让当前任务只看到编号为 0、1、2、3 的 GPU,并在一个节点上启动 4 个进程,每个进程对应一个 GPU 进行分布式训练。是 PyTorch 用于启动分布式训练的命令行工具。这是一个在使用 PyTorch 进行分布式训练时常用的命令行指令。在Python脚本调用之前,指定CUDA-visuble。如果有 4 个 GPU,并且设置。
·
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
-
解释:
CUDA_VISIBLE_DEVICES
是一个环境变量,用于控制哪些 GPU 设备对当前进程可见。$gpu_ids
是一个变量,通常代表一个或多个 GPU 的编号列表,例如0,1,2
表示让当前进程只看到编号为 0、1、2 的 GPU。
-
功能:
- 通过设置这个环境变量,可以指定当前任务只使用特定的 GPU 设备,从而有效地管理 GPU 资源分配,避免不必要的 GPU 使用冲突。
二、torchrun --nnodes=1 --nproc_per_node=$num_gpus
-
解释:
torchrun
是 PyTorch 提供的用于启动分布式训练的命令行工具。--nnodes
参数指定了参与分布式训练的节点数量。在这里设置为1
表示只有一个节点参与训练。--nproc_per_node
参数指定了每个节点上启动的进程数量。$num_gpus
通常是一个变量,表示每个节点上可用的 GPU 数量,这个参数设置每个节点上启动的进程数量等于 GPU 数量,以实现每个 GPU 对应一个进程的分布式训练。
-
功能:
- 启动一个分布式训练任务,在单个节点上根据指定的 GPU 数量启动相应数量的进程,每个进程可以独立地执行训练任务,并通过分布式通信机制进行协作和同步。
- 这种方式可以充分利用多个 GPU 的计算能力,加速深度学习模型的训练过程。
例如, 如果有 4 个 GPU,并且设置gpu_ids=0,1,2,3
以及num_gpus=4
,那么这个命令会让当前任务只看到编号为 0、1、2、3 的 GPU,并在一个节点上启动 4 个进程,每个进程对应一个 GPU 进行分布式训练。
torchrun基本用法 from 豆包
torchrun
是 PyTorch 用于启动分布式训练的命令行工具。以下是它的使用方法:
一、基本用法
- 启动分布式训练:
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
- 环境变量设置:
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。
二、高级用法
- 分布式数据并行(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)
- 弹性训练:
torchrun
支持弹性训练,即可以动态地添加或移除训练节点和进程。- 要实现弹性训练,需要在训练脚本中处理进程的动态加入和退出,并正确地同步模型参数。
三、注意事项
-
网络配置:
- 确保所有参与分布式训练的节点可以相互通信。如果在不同的机器上运行,需要正确设置网络地址和端口号。
- 可以通过设置环境变量
MASTER_ADDR
和MASTER_PORT
来指定主节点的地址和端口号。
-
依赖管理:
- 确保所有节点上的环境和依赖项一致。特别是对于深度学习训练,需要确保相同版本的 PyTorch、CUDA 和其他依赖库在所有节点上安装。
-
日志和错误处理:
- 分布式训练中,错误可能在不同的进程中出现,并且可能难以调试。可以使用集中式日志记录或分布式日志工具来帮助诊断问题。
- 注意观察
torchrun
的输出和训练脚本的日志,以获取有关训练进度和错误的信息。
欢迎来到由智源人工智能研究院发起的Triton中文社区,这里是一个汇聚了AI开发者、数据科学家、机器学习爱好者以及业界专家的活力平台。我们致力于成为业内领先的Triton技术交流与应用分享的殿堂,为推动人工智能技术的普及与深化应用贡献力量。
更多推荐
已为社区贡献1条内容
所有评论(0)