如何使用torchrun启动单机多卡DDP并行训练

如何使用torchrun启动单机多卡DDP并行训练

这是一个最近项目中需要使用的方式，新近的数据集大概在40w的规模，而且载入的原始特征都比较大（5～7M），所以准备尝试DistributedDataParallel；

主要目的有两个：

• 1是通过并行训练有效降低训练耗时，之前也使用过单机单卡，以及单机多卡DataParallel的方式；
• 2是通过使用并行训练的方式，倍增有效batch_size的数量，提高单卡对显存的利用率，仍然是为了加速；

类似的实现代码你可以在许多技术博客上检索到，以下是我实际实践的主要代码；

单机单卡到单机多卡的代码修改

# 导入库
import math
# import torch
from torch.utils.data import Dataset, DataLoader
import torch.distributed as dist ## DDP
from torch.utils.data.distributed import DistributedSampler ## DDP
from torch.nn.parallel import DistributedDataParallel as DDP ## DDP
from torch.distributed import init_process_group, destroy_process_group ## DDP

# 设置可见的GPU
# os.environ["CUDA_VISIBLE_DEVICES"] = "0,1,2,3"
# 忽略冗余数据存在的潜在问题
# os.environ["KMP_DUPLICATE_LIB_OK"] = "True"


def train():
    # setup
    local_rank = int(os.environ["LOCAL_RANK"]) ## DDP   
    init_process_group(backend="nccl")
    torch.cuda.set_device(local_rank)

    # ...

    # 
    local_rank = dist.get_rank()
    total_rank = dist.get_world_size()
    device = torch.device("cuda", local_rank)
    
    # random reset
    seed = 1234
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    random.seed(seed)
    np.random.seed(seed)

    # 
    model = Model().to(device)
    if local_rank <= 0:
        print(model)
    # model.load(local_model_path)
    lr *= math.sqrt(total_rank)
    # optimizer = ...
    # scheduler = ...


    #
    # train_datas = Dataset("train")
    sampler = DistributedSampler(train_datas)
    # 或者自定义
    # sampler = MyDistributedSampler(train_data_path)
    train_dataloader = DataLoader(
        train_datas,
        batch_size=batch_size,
        num_workers=num_workers,
        shuffle=(sampler is None),
        sampler=sampler,
        pin_memory=True,
        drop_last=True,
    )
    train_dataloaders = [train_dataloader]

    # warp model
    model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model).to(device)
    model = DDP(model, device_ids=[local_rank], output_device=local_rank, find_unused_parameters=True)


    for epoch in range(max_epoch):
        e_INDEX = train_epoch(epoch, local_rank, total_rank, device, train_dataloaders, model, optimizer, scheduler)

    dist.destroy_process_group()

def train_epoch(epoch, local_rank, total_rank, device, train_dataloaders, model, optimizer, scheduler):
    model.train()
    if local_rank >= 0:
        torch.distributed.barrier()

    # 如果使用了自定义的Sampler 记得实现一个对应set_epoch方法功能的新方法，并由采样器进行调用
    if local_rank >= 0:
        for i, train_dataloader in enumerate(train_dataloaders):
            # train_dataloader.sampler.set_epoch(epoch)
            train_dataloader.sampler.shuffle_data_on_ranks(seed=epoch + i)
    
    e_loss = 0
    e_num = 0
    # ...

    for train_dataloader in train_dataloaders:
        for data in train_dataloader:
            loss = train_step(data, local_rank, total_rank, device, train_dataloaders, model, optimizer, scheduler)

            e_loss += loss[0]
            e_num += loss[1]

    if local_rank >= 0:
        torch.distributed.barrier()
        total = torch.tensor(
            copy.deepcopy([e_loss, e_num]),
            dtype=torch.float32,
            device=device,
        )
        torch.distributed.all_reduce(
            total, torch.distributed.ReduceOp.SUM, async_op=False
        )
        e_loss, e_num = total.tolist()

    e_loss /= e_num

    # 
    if local_rank <= 0:
        if isinstance(model, torch.nn.parallel.DistributedDataParallel):
            model.module.save(opt.notes)   
        else:
            model.save(opt.notes)   

    scheduler.step(e_loss) # if need e_loss
    # scheduler.step()

    # 测试：每个GPU进程输出的模型第一层参数是相同的
    for param in model.parameters():
        print("    GPU{} Model param layer1=>".format(local_rank), param)
        break
    e_INDEX = eval_epcoh(model, val_dataloader, epoch, local_rank)

    return e_INDEX

def train_step(data, local_rank, total_rank, device, train_dataloaders, model, optimizer, scheduler):

    data_a,label_a = data

    data_a = data_a.to(device)
    label_a = label_a.to(device)

    # 这里要注意，如果使用的是自定义api 如这里的compute_loss该方法中 包含了模型的推理和损失计算过程，那么就需要使用model.module进行调用
    # 如果直接使用模型推理 那么直接调用 forward函数即可，如model(data_a)
    if isinstance(model, torch.nn.parallel.DistributedDataParallel):
        loss, _ = model.module.compute_loss(data_a, label_a)
    else:
        loss, _loss_dict_ = model.compute_loss(data_a, label_a)   

    optimizer.zero_grad() 
    loss.backward() 
    optimizer.step() 

    return loss.item(), label_a.shape[0]

对于测试集的验证，可以不使用Sampler（验证集/测试集在每个GPU进程中都相同），仅在某一个GPU中，完整测试即可；

启动方式：

• torchrun -m --nnodes=1 --nproc_per_node=4 main
• nohup torchrun -m --nnodes=1 --nproc_per_node=4 main 127.0.0.1 5301 > /dev/null 2>&1 & disown

测试模型的第一层输出（part），可以看到它们都是相同的：

    GPU0 Model param layer1=> Parameter containing:
tensor([[[[ 0.1158,  0.0723, -0.0191],
          [ 0.0715, -0.1204, -0.0295],
          [-0.0790,  0.0723,  0.0963],

    GPU3 Model param layer1=> Parameter containing:
tensor([[[[ 0.1158,  0.0723, -0.0191],
          [ 0.0715, -0.1204, -0.0295],
          [-0.0790,  0.0723,  0.0963],

    GPU1 Model param layer1=> Parameter containing:
tensor([[[[ 0.1158,  0.0723, -0.0191],
          [ 0.0715, -0.1204, -0.0295],
          [-0.0790,  0.0723,  0.0963],

    GPU2 Model param layer1=> Parameter containing:
tensor([[[[ 0.1158,  0.0723, -0.0191],
          [ 0.0715, -0.1204, -0.0295],
          [-0.0790,  0.0723,  0.0963],                    

一点经验

梯度同步后的模型参数不一致

其实主要的代码并不复杂，但是这里我要提到的一点经验是：如果你的模型参数在每个epoch结束之后（梯度同步更新完成）的打印结果，并不一致，并且你的DDP主要代码也与上述代码大同小异，那么很有可能是你的模型有问题（并不是不正确，只不过无法得到DDP的支持）；

这也是我的项目所面临的问题：我使用了一个相对简单的模型，可以很好的实践DDP，但另外一个复杂许多的模型就无法通过测试（每个epoch后模型参数在各个GPU进程间会变的不同），具体的后续补充；

整理附上GPT的回答，仅供参考：

我正在使用一个以Conformer为主要网络结构的模型，训练一个音频分类模型，但是我发现这个模型并不能在DP和DDP中得到很好的支持，主要表现在每个epoch训练梯度同步之后，各个GPU进程中的模型参数不一致，而如果我换成一个简单的卷积网实现分类模型，就可以被DP和DDP很好的支持，它们会在每个epoch训练梯度同步之后，使得各个GPU进程中的模型参数完全一致；请根据我所描述的情况，分析为什么这个基于Conformer的模型结构无法被DP和DDP支持？

R:
在使用 DataParallel（DP）和 DistributedDataParallel（DDP）时，确保模型能够正确同步梯度并保持参数一致是关键。对于复杂的模型结构，如基于 Conformer 的音频分类模型，可能会出现一些挑战，导致同步问题。以下是可能的原因：

1. **自定义操作或模块：** 如果你的 Conformer 模型包含自定义的操作或模块，这可能导致 DP 或 DDP 的同步问题。确保所有操作都能够正确地在 GPU 上执行，并且它们的梯度能够被正确地计算和同步。

2. **非标准的层：** Conformer 模型可能包含一些非标准的层，如深度卷积块、多头自注意力等。确保这些层在 DP 和 DDP 中能够正确地进行梯度同步。你可能需要在这些层的实现中查看是否有专门为 DP 和 DDP 提供支持的选项。

3. **权重初始化：** DP 和 DDP 要求每个 GPU 上的模型初始权重是一致的。确保你的模型在每个 GPU 上的权重初始化是相同的，可以考虑在初始化之后显式地同步一次模型参数。

4. **梯度累积：** 在某些情况下，可能会发生梯度累积的情况，导致模型参数不一致。确保在 DP 和 DDP 中使用梯度累积时，模型的参数同步仍然能够正确进行。

5. **DDP 启动方式：** 使用 DDP 时，确保在启动时每个 GPU 上的进程都能够正确初始化。这可能需要查看 DDP 初始化的方式，确保每个进程都按预期地启动。

6. **PyTorch 版本：** 有时，特定版本的 PyTorch 可能对某些模型结构的 DP 和 DDP 支持更好。确保你使用的是最新版本的 PyTorch。

在调试时，可以尝试通过在每个 GPU 上输出和比较梯度、权重等信息，以找出在同步中出现问题的具体位置。此外，查看 PyTorch 的官方文档和社区讨论，以获取与你的模型结构相关的最新支持和最佳实践。

NCCL timeout

Watchdog caught collective operation timeout: WorkNCCL(SeqNum=1363168, OpType=BROADCAST, Timeout(ms)=1800000) ran for 1805862 milliseconds before timing out.

我在测试的时候，每个GPU进程都做了验证集的val，这可能并不必要，有效验证应该仅在Rank0进程进行即可；

当然还有可能是其他原因造成的这个异常，可=>参考链接；

LR的设置

可以是简单的线性增大，我查了一些资料因此有上述代码的实现，但要注意的是，具体增大多少要结合具体的任务，也就是要调参，LR保持不变或更小可能会使得收敛十分缓慢，如果太大又可能使训练崩溃；

一些参考链接：

• 参考链接1
• 参考链接2
• 参考链接3
• 参考链接4
• 参考链接5
• 参考链接6
• 参考链接7
• 参考链接8

关于MyDistributedSampler

from torch.utils.data.sampler import Sampler


class MyDistributedSampler(Sampler):
    def __init__(self, data_path, m, batch_size, local_rank, total_rank, device, distribute=False):
        # 这里的读取 与 Dataset中 读取原始数据是一致的（我们使用的文件是excel）
        data_lines = pd.read_excel(data_path)
        self.distribute = distribute

        self.total_rank = -1
        self.local_rank = -1
        self.device = device

        if distribute:
            if not dist.is_available():
                raise RuntimeError("Requires distributed package to be available")
            self.total_rank = total_rank
            self.local_rank = local_rank

        self._m_per_class = m
        self._batch_size = batch_size

        # 拿到与单卡一致的所有数据
        self._labels_to_indices = self._get_labels_to_indices(data_lines)
        # 拿到与单卡一致的索引数据
        self._global_labels = (
            self._labels_to_indices.reset_index()["a_cls"]
            .drop_duplicates()
            .values.tolist()
        )

        if not self.distribute:
            # 单卡情况 待操作数据 就是 所有的索引数据
            self.labels = self._global_labels
        else:
            # 多卡时，则要按照 当前的local_rank 进行划分
            self.labels = self._split_label_randoms(seed=0)[self.local_rank]
        assert (
            self._batch_size % self._m_per_class
        ) == 0, "m_per_class must divide batch_size without any remainder"
        # 采样器的样本数
        self._sample_length = self._get_sample_length()

        print(
            "Init Sampler with Mper with {} items, and m = {}, batch_num = {}"
            "\n".format(self._sample_length, m, self.num_iters())
        )

        return

    def __iter__(self):
        # 创建一个迭代器：迭代长度 与样本数量一致 并初始化为0
        idx_list = [0] * self._sample_length
        i = 0
        num_iters = self.num_iters()
        for _ in range(num_iters):
            random.shuffle(self.labels)
            # 为使用HardTripletLoss计算损失，这里随机 固定个数的分类 以batch_size=32 m=4为例，则这里去的个数为8
            # 对于我们的数据集：实际count的最小关联数是6 因此将m设为4是合理的
            curr_label_set = self.labels[: self._batch_size // self._m_per_class]
            for label in curr_label_set:
                if len(self._labels_to_indices.loc[label]) >= self._m_per_class:
                    items = self._labels_to_indices.loc[label].sample(
                        self._m_per_class
                    ) 
                else:
                    items = self._labels_to_indices.loc[label].sample(
                        self._m_per_class,
                        replace=True
                    )  
                t = []
                for item in items.index.values:
                    cur_id = np.random.choice(
                        self._labels_to_indices.loc[label, item]["ids"],
                        1,
                        replace=False,
                    )[0]
                    t.append(cur_id)
                # 确定采样规范 并返回，可以看到这里已经是一个符合要求的完整序列
                # 注意：对应的 Dataloader 不可以使用shuffle
                idx_list[i : i + self._m_per_class] = t[: self._m_per_class]
                i += self._m_per_class
        return iter(idx_list)

    # set_epoch的实现: 仅多卡的时候才需要调用
    def shuffle_data_on_ranks(self, seed, display_details=False):
        # 每个epoch都要重新划分各个GPU进程所分的的数据（它们彼此不重复）
        self.labels = self._split_label_randoms(seed)[self.local_rank]
        self._sample_length = self._get_sample_length()
        if display_details:
            display = sorted(self.labels)[:5]
            print(
                "Shuffle sampler with {} items, rank = {}, total rank = {}, batch_num = {}, label of head{} is {}".format(
                    self._sample_length,
                    self.local_rank,
                    self.total_rank,
                    self.num_iters(),
                    len(display),
                    display,
                )
            )

    def num_iters(self):
        # 样本数 整除 batch_size的 值，即整个训练集被迭代一次 所需要的次数
        return self._sample_length // self._batch_size

    def _get_sample_length(self):
        # 这就是实际要加载的样本数，由于这里的self.labels 已经根据local_rank完成了划分，因此这个sample_len就是每个GPU进程所分得得样本数
        sample_length = sum([self._labels_to_indices.loc[k]["count"].sum() for k in self.labels])

        # 由于梯度同步时，需要每个GPU进程都有数据，因此如下这段同步只要是为了 保证所有GPU进程中的采样数量是一致的（与最小值相同）（如果你的数据很规整 这段同步实际用不到）
        if self.total_rank >= 0:
            torch.distributed.barrier()
            total = torch.tensor(
                [sample_length],
                dtype=torch.float32,
                device=self.device,
            )
            torch.distributed.all_reduce(
                total, torch.distributed.ReduceOp.MIN, async_op=True
            )
            sample_length = int(total.tolist()[0])

        sample_length -= sample_length % self._batch_size
        return sample_length
    
    def _split_label_randoms(self, seed):
        split_label = []
        global_label = self._global_labels.copy()
        random.Random(seed).shuffle(global_label)
        for i in range(self.total_rank):
            # 可以看到 索引数据会根据 total_rank 间隔取值
            split_label.append(global_label[i :: self.total_rank])
        return split_label

    # @staticmethod
    def _get_labels_to_indices(self, data_lines):
        df = data_lines.groupby(["a_cls", "a_item"]).agg(
            count=("a_id", "count"), ids=("a_id", "unique")
        )
        return df

    def __len__(self):
        return self._sample_length