ddp多机多卡训练主节点自动启动其余节点

### 基于DDP的多机多卡训练中主节点自动启动其余工作节点的配置方法 在分布式深度学习训练中，使用 `DistributedDataParallel`（简称 DDP）可以有效提升模型训练的速度和效率。为了实现多机多卡训练，通常需要一个主节点负责协调整个集群的工作流程，并自动启动其他工作节点。以下是关于如何配置主节点以自动启动其余工作节点的方法。 #### 1. 初始化环境变量 在多机多卡训练中，每台机器都需要设置一些必要的环境变量以便进行通信。这些变量可以通过命令行参数传递给脚本或者直接写入代码中。常见的环境变量包括： - **MASTER_ADDR**: 指定主节点的 IP 地址。 - **MASTER_PORT**: 指定主节点监听的端口号。 - **WORLD_SIZE**: 表示参与训练的所有 GPU 总数。 - **RANK**: 当前进程在全球范围内的排名编号，从 0 开始计数，其中 Rank=0 的设备为主节点。 ```bash export MASTER_ADDR="192.168.1.1" export MASTER_PORT="12345" export WORLD_SIZE=8 ``` 上述命令需根据实际部署情况调整[^1]。 #### 2. 使用 `torch.distributed.launch` 或 `torchrun` PyTorch 提供了两种工具可以帮助简化多机多卡训练过程中的节点管理问题——分别是 `torch.distributed.launch` 脚本以及更现代化替代品 `torchrun` 工具。它们能够帮助开发者轻松指定运行时所需的资源数量及其分布位置。 ##### 方法一：通过 `torch.distributed.launch` 如果采用这种方式，则可以在 Python 文件开头部分加入如下内容来定义入口函数: ```python import os import torch from torch.nn.parallel import DistributedDataParallel as DDP def main(rank, world_size): setup(rank, world_size) if __name__ == "__main__": free_gpu_ids = ... # 获取可用GPU列表逻辑省略 n_gpus_per_node = len(free_gpu_ids) total_world_size = int(os.environ['WORLD_SIZE']) node_rank = int(os.environ.get('NODE_RANK', '0')) num_nodes = total_world_size // n_gpus_per_node mp.spawn(main, args=(n_gpus_per_node,), nprocs=n_gpus_per_node, join=True) ``` 然后利用 shell script 启动不同服务器上的 python 进程即可完成全部操作[^2]. ##### 方法二：借助 `torchrun` 相比前者而言更为简洁明了一些，在终端执行下面这样的指令就可以让程序自行处理好一切细节方面的事情啦！ ```bash torchrun --nnodes=$NUM_NODES \ --node-rank=$CURRENT_NODE_RANK \ --master_addr $MAIN_SERVER_IP \ --master_port $PORT_NUMBER \ --nproc_per_node=$GPUS_PER_MACHINE train.py ... ``` 这里需要注意的是 `$VAR_NAME` 形式的占位符应该替换成为具体的数值或者是路径字符串等等[^3]。 #### 3. 编写自定义启动脚本 对于更加复杂的场景来说可能还需要额外编写专门用来控制各个子任务之间相互关系的一个总控类别的 Shell/Bash/Python 等形式的小型应用程序才行哦~ 它的主要职责就是按照预定计划依次唤醒目标主机上面对应的 worker processes 并传送给他们各自专属的任务描述信息包咯~ --- ### 示例代码片段展示 以下是一个简单的例子展示了如何在一个项目里集成以上提到的功能特性之一的部分源码结构布局样式仅供参考借鉴之用而已哈😊 ```python # utils/distribute_utils.py import os import torch def init_process_group(local_rank=None): """ Initialize the process group. Args: local_rank (int): Local rank of current device. Returns: None """ backend = "nccl" if torch.cuda.is_available() else "gloo" master_ip = os.getenv("MASTER_ADDR", default="localhost") port = os.getenv("MASTER_PORT", default="12345") dist_url = f"tcp://{master_ip}:{port}" rank = int(os.getenv("RANK", "0")) world_size = int(os.getenv("WORLD_SIZE", "1")) torch.distributed.init_process_group( backend=backend, init_method=dist_url, world_size=world_size, rank=rank ) if local_rank is not None: torch.cuda.set_device(local_rank) class DDPLauncher(object): @staticmethod def launch(fn, args_tuple, gpus_list): ... ``` ---