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
  1. 环境变量设置:
    • 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)
  1. 弹性训练:
    • torchrun支持弹性训练,即可以动态地添加或移除训练节点和进程。
    • 要实现弹性训练,需要在训练脚本中处理进程的动态加入和退出,并正确地同步模型参数。

三、注意事项

  1. 网络配置:

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

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

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

欢迎来到由智源人工智能研究院发起的Triton中文社区,这里是一个汇聚了AI开发者、数据科学家、机器学习爱好者以及业界专家的活力平台。我们致力于成为业内领先的Triton技术交流与应用分享的殿堂,为推动人工智能技术的普及与深化应用贡献力量。

更多推荐