烈火灼冰 - gradient_accumulation
type: channel
link: https://blog.4c43.work/index.php/tag/gradient_accumulation/
items: 【学习/记录】torchrun/DDP使用方法与注意事项

【学习/记录】torchrun/DDP使用方法与注意事项
type: item
date: 2024-08-19T12:20:00+00:00
description: 概述 基于pytorch自身的分布式训练，通过调用nccl等backend实现多机多卡并行训练。 面向 DDP手动启动 的实现 启动方式 python main.py 面向 torchrun 的实现 启动方式 torchrun main.py 重要环境变量参数说明 变量名 定义及说明 LOCAL_RANK 本地排名 RANK 全局排名 GROUP_RANK worker 组的排名。一个介于 0 和 max_nnodes 之间的数字。当每个节点运行一个 worker 组时，这是节点的排名 ROLE_RANK 所有具有相同角色的 worker 中的 worker 的排名。worker 的角色在 WorkerSpec 中指定 LOCAL_WORLD_SIZE 本地世界大小（例如本地运行的 worker 数）；等于在 torchrun 上指定的 --nproc-per-node WORLD_SIZE 世界大小（作业中 worker 的总数） ROLE_WORLD_SIZE 使用在 WorkerSpec 中指定的相同角色启动的 worker 的总数 MASTER_ADDR 运行排名为 0 的 worker 的主机的 FQDN；用于初始化 Torch 分布式后端 MASTER_PORT MASTER_ADDR 上的端口，可用于托管 C10d TCP 存储 TORCHELASTIC_RESTART_COUNT 到目前为止的 worker 组重启次数 TORCHELASTIC_MAX_RESTARTS 配置的最大重启次数 TORCHELASTIC_RUN_ID 等于 rendezvous run_id（例如唯一的作业 ID） PYTHON_EXEC 系统可执行文件覆盖。如果提供，python 用户脚本将使用 PYTHON_EXEC 的值作为可执行文件。默认情况下使用 sys.executable rank问题 端口问题 个人模版 参考： Pytorch - DDP教程 Pytorch - torchrun 弹性启动 import argparse import logging import os import sys from importlib import reload import datasets import torch import torch.distributed as dist import torch.nn as nn from einops import rearrange from torch.distributed.elastic.multiprocessing.errors import record from torch.nn.parallel import DistributedDataParallel as DDP from torch.optim.adamw import AdamW from torch.utils.data import DataLoader, Dataset from torch.utils.data.distributed import DistributedSampler from torchvision import transforms class CustomModel(nn.Module): def __init__(self, dim, heads=4, dim_head=32): super().__init__() self.heads = heads hidden_dim = dim_head * heads self.to_qkv = nn.Conv2d(dim, hidden_dim * 3, 1, bias=False) self.to_out = nn.Conv2d(hidden_dim, dim, 1) def forward(self, x): b, c, h, w = x.shape qkv = self.to_qkv(x) q, k, v = rearrange( qkv, "b (qkv heads c) h w -> qkv b heads c (h w)", heads=self.heads, qkv=3 ) k = k.softmax(dim=-1) context = torch.einsum("bhdn,bhen->bhde", k, v) out = torch.einsum("bhde,bhdn->bhen", context, q) out = rearrange( out, "b heads c (h w) -> b (heads c) h w", heads=self.heads, h=h, w=w ) return self.…