如何在 8 个 GPU 上并行化 Transformer 模型进行机器翻译？

我正在尝试使用变压器模型以几乎与原始文章相同的方式执行机器翻译。虽然该模型运行得相当好，但它需要更多的计算资源。为了解决这个问题，我在一台具有 8 个 GPU 处理器的计算机上运行了该模型，但我缺乏这方面的经验。 我尝试对并行化进行必要的调整：

transformer = nn.DataParallel(transformer)
transformer = transformer.to(DEVICE)

但是，由于我缺乏经验，事情运行得并不顺利。具体来说，我已经被以下错误消息困扰了很长时间：

文件 “C:\Projects\MT005.venv\Lib\site-packages\torch\nn\function.py”， 第 5382 行，在 multi_head_attention_forward 中引发 RuntimeError(f"The 2D attn_mask 的形状是 {attn_mask.shape}，但应该是 { Correct_2d_size}。") RuntimeError: 2D attn_mask 的形状是 torch.Size([8, 64])，但应该是 (4, 4)。

有人可以帮我解决这个问题并让模型在所有 8 个 GPU 上运行吗？

遇到的错误消息“RuntimeError: 2D attn_mask 的形状是 torch.Size([8, 64])，但应该是 (4, 4)”表明注意力掩码的尺寸不正确。这很可能是因为 nn.DataParallel 在批处理维度（维度 0）上拆分数据，并且注意力掩码需要针对每个设备上的拆分数据进行调整。

以下是如何在 8 个 GPU 上并行化 Transformer 模型以进行机器翻译并解决错误的方法：

1. 调整注意力掩码： 需要修改代码以正确处理注意力掩码。由于 nn.DataParallel 沿批处理维度拆分数据，因此需要确保注意力掩码也已拆分并放置在正确的设备上。

```python import torch

class Transformer(nn.Module): # ... 的 Transformer 模型定义 ...

# ... 在的训练循环中 ... for batch in train_dataloader: # 将数据移动到设备 src = batch.src.to(DEVICE) trg = batch.trg.to(DEVICE) src_mask = batch.src_mask.to(DEVICE) trg_mask = batch.trg_mask.to(DEVICE)

   # 创建注意力掩码
   src_attn_mask = (src_mask.unsqueeze(1) @ src_mask.unsqueeze(2)).bool()
   trg_attn_mask = (trg_mask.unsqueeze(1) @ trg_mask.unsqueeze(2)).bool()

   # ... 的其余训练循环 ...

   # 将注意力掩码传递给模型
   output = transformer(src, trg, src_mask=src_attn_mask, trg_mask=trg_attn_mask)

```

1. 优化并行化： 虽然 nn.DataParallel 是一个很好的起点，但可以考虑使用更高级的技术来进一步优化并行化，尤其是在大型模型和数据集的情况下。以下是一些选项：

2. 使用更大的批大小： 增加批大小可以提高 GPU 利用率。尝试不同的批大小以找到最佳平衡点。

3. 梯度累积： 如果无法在所有 8 个 GPU 上安装更大的批大小，请模拟更大的批大小，方法是在多个步骤中累积梯度，然后更新模型参数。

4. 模型并行化： 对于非常大的模型，可以探索模型并行化，即将模型的不同部分分布在不同的 GPU 上。这需要更复杂的实现。

5. 性能监控： 使用诸如 nvidia-smi 之类的工具监控 GPU 利用率，以确保所有 8 个 GPU 都得到充分利用。

通过实施这些更改，应该能够解决错误并在所有 8 个 GPU 上有效地并行化 Transformer 模型以进行机器翻译。请记住，根据的特定硬件和模型大小微调并行化策略以获得最佳性能至关重要。