在使用多 GPU 的 Transformer 中进行训练时，掩码的形状会除以 GPU 的数量。为什么？

Question

那里 我正在用多 GPU 训练 Transformer，但我遇到了问题。 我正在使用 Pytorch 并使用

model = Transformer(
src_tokens=src_tokens, tgt_tokens=tgt_tokens, dim_model=dim_model, num_heads=num_heads,
num_encoder_layers=num_encoder_layers, num_decoder_layers=num_decoder_layers, dropout_p=0.1)
model = nn.DataParallel(model, device_ids=device_ids)
model.to(device)

训练过程是这样的：

def train_loop(model, opt, loss_fn, dataloader):
    model.train()
    total_loss = 0

    for X, y in dataloader:
        X, y = X.t().to(device), y.t().to(device)

        y_input = y[:, :-1]
        y_expected = y[:, 1:]

        sequence_length = y_input.size(1)
        src_pad_mask = create_pad_mask(X, 1)
        tgt_pad_mask = create_pad_mask(y_input, 1)
        tgt_mask = get_tgt_mask(sequence_length)
        pred = model(X, y_input, tgt_mask=tgt_mask, src_pad_mask=src_pad_mask, tgt_pad_mask=tgt_pad_mask)
        # Permute pred to have batch size first again

        pred = pred.permute(1, 2, 0)
        loss = loss_fn(pred, y_expected)
        opt.zero_grad()
        loss.backward()
        opt.step()
        total_loss += loss.detach().item()
    return total_loss / len(dataloader)

我的model.py是这样的：

class Transformer(nn.Module):
    """
    Model from "A detailed guide to Pytorch's nn.Transformer() module.", by
    Daniel Melchor: https://medium.com/@danielmelchor/a-detailed-guide-to-pytorchs-nn-transformer-module-c80afbc9ffb1
    """

    # Constructor
    def __init__(
            self,
            src_tokens,
            tgt_tokens,
            dim_model,
            num_heads,
            num_encoder_layers,
            num_decoder_layers,
            dropout_p,
    ):
        super().__init__()

        # INFO
        self.model_type = "Transformer"
        self.dim_model = dim_model

        # LAYERS
        self.positional_encoder = PositionalEncoding(
            dim_model=dim_model, dropout_p=dropout_p, max_len=5000
        )
        self.src_embedding = nn.Embedding(src_tokens, dim_model)
        self.tgt_embedding = nn.Embedding(tgt_tokens, dim_model)
        self.transformer = nn.Transformer(
            d_model=dim_model,
            nhead=num_heads,
            num_encoder_layers=num_encoder_layers,
            num_decoder_layers=num_decoder_layers,
            dropout=dropout_p,
        )
        self.out = nn.Linear(dim_model, tgt_tokens)

    def forward(self, src, tgt, tgt_mask=None, src_pad_mask=None, tgt_pad_mask=None):
        # Src size must be (batch_size, src sequence length)
        # Tgt size must be (batch_size, tgt sequence length)

        # Embedding + positional encoding - Out size = (batch_size, sequence length, dim_model)

        src = self.src_embedding(src) * math.sqrt(self.dim_model)
        tgt = self.tgt_embedding(tgt) * math.sqrt(self.dim_model)
        src = self.positional_encoder(src)
        tgt = self.positional_encoder(tgt)

        # We could use the parameter batch_first=True, but our KDL version doesn't support it yet, so we permute
        # to obtain size (sequence length, batch_size, dim_model),
        src = src.permute(1, 0, 2)
        tgt = tgt.permute(1, 0, 2)
        print('src_pad_mask: '+str(src_pad_mask.shape)+'  tgt_pad_mask: '+str(tgt_pad_mask.shape)+'  tgt_mask: '+str(tgt_mask.shape))
        print('++++++++++++++++++++++++++++++++++++++++++++++++++++++++')

        # Transformer blocks - Out size = (sequence length, batch_size, num_tokens)
        transformer_out = self.transformer(src, tgt, tgt_mask=tgt_mask, src_key_padding_mask=src_pad_mask,
                                           tgt_key_padding_mask=tgt_pad_mask)
        out = self.out(transformer_out)

        return out

我收到此错误：

root@3b:/koi/transformer-multi# python train.py
src_pad_mask: torch.Size([1, 4784])  tgt_pad_mask: torch.Size([1, 3225])  tgt_mask: torch.Size([538, 3225])
++++++++++++++++++++++++++++++++++++++++++++++++++++++++
src_pad_mask: torch.Size([1, 4784])  tgt_pad_mask: torch.Size([1, 3225])  tgt_mask: torch.Size([538, 3225])
++++++++++++++++++++++++++++++++++++++++++++++++++++++++
src_pad_mask: torch.Size([1, 4784])  tgt_pad_mask: torch.Size([1, 3225])  tgt_mask: torch.Size([538, 3225])
++++++++++++++++++++++++++++++++++++++++++++++++++++++++
src_pad_mask: torch.Size([1, 4784])  tgt_pad_mask: torch.Size([1, 3225])  tgt_mask: torch.Size([538, 3225])
++++++++++++++++++++++++++++++++++++++++++++++++++++++++
src_pad_mask: torch.Size([1, 4784])  tgt_pad_mask: torch.Size([1, 3225])  tgt_mask: torch.Size([538, 3225])
++++++++++++++++++++++++++++++++++++++++++++++++++++++++
src_pad_mask: torch.Size([1, 4784])  tgt_pad_mask: torch.Size([1, 3225])  tgt_mask: torch.Size([535, 3225])
++++++++++++++++++++++++++++++++++++++++++++++++++++++++
Traceback (most recent call last):
  File "/koi/transformer-multi/train.py", line 160, in <module>
    train_loss = train_loop(model, opt, loss_fn, trn_loader)
  File "/koi/transformer-multi/train.py", line 121, in train_loop
    pred = model(X, y_input, tgt_mask=tgt_mask, src_pad_mask=src_pad_mask, tgt_pad_mask=tgt_pad_mask)
  File "/root/anaconda3/envs/koi/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1102, in _call_impl
    return forward_call(*input, **kwargs)
  File "/root/anaconda3/envs/koi/lib/python3.9/site-packages/torch/nn/parallel/data_parallel.py", line 168, in forward
    outputs = self.parallel_apply(replicas, inputs, kwargs)
  File "/root/anaconda3/envs/koi/lib/python3.9/site-packages/torch/nn/parallel/data_parallel.py", line 178, in parallel_apply
    return parallel_apply(replicas, inputs, kwargs, self.device_ids[:len(replicas)])
  File "/root/anaconda3/envs/koi/lib/python3.9/site-packages/torch/nn/parallel/parallel_apply.py", line 86, in parallel_apply
    output.reraise()
  File "/root/anaconda3/envs/koi/lib/python3.9/site-packages/torch/_utils.py", line 434, in reraise
    raise exception
RuntimeError: Caught RuntimeError in replica 0 on device 0.
Original Traceback (most recent call last):
  File "/root/anaconda3/envs/koi/lib/python3.9/site-packages/torch/nn/parallel/parallel_apply.py", line 61, in _worker
    output = module(*input, **kwargs)
  File "/root/anaconda3/envs/koi/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1102, in _call_impl
    return forward_call(*input, **kwargs)
  File "/koi/transformer-multi/model.py", line 93, in forward
    transformer_out = self.transformer(src, tgt, tgt_mask=tgt_mask, src_key_padding_mask=src_pad_mask,
  File "/root/anaconda3/envs/koi/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1102, in _call_impl
    return forward_call(*input, **kwargs)
  File "/root/anaconda3/envs/koi/lib/python3.9/site-packages/torch/nn/modules/transformer.py", line 142, in forward
    output = self.decoder(tgt, memory, tgt_mask=tgt_mask, memory_mask=memory_mask,
  File "/root/anaconda3/envs/koi/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1102, in _call_impl
    return forward_call(*input, **kwargs)
  File "/root/anaconda3/envs/koi/lib/python3.9/site-packages/torch/nn/modules/transformer.py", line 248, in forward
    output = mod(output, memory, tgt_mask=tgt_mask,
  File "/root/anaconda3/envs/koi/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1102, in _call_impl
    return forward_call(*input, **kwargs)
  File "/root/anaconda3/envs/koi/lib/python3.9/site-packages/torch/nn/modules/transformer.py", line 451, in forward
    x = self.norm1(x + self._sa_block(x, tgt_mask, tgt_key_padding_mask))
  File "/root/anaconda3/envs/koi/lib/python3.9/site-packages/torch/nn/modules/transformer.py", line 460, in _sa_block
    x = self.self_attn(x, x, x,
  File "/root/anaconda3/envs/koi/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1102, in _call_impl
    return forward_call(*input, **kwargs)
  File "/root/anaconda3/envs/koi/lib/python3.9/site-packages/torch/nn/modules/activation.py", line 1003, in forward
    attn_output, attn_output_weights = F.multi_head_attention_forward(
  File "/root/anaconda3/envs/koi/lib/python3.9/site-packages/torch/nn/functional.py", line 5011, in multi_head_attention_forward
    raise RuntimeError(f"The shape of the 2D attn_mask is {attn_mask.shape}, but should be {correct_2d_size}.")
RuntimeError: The shape of the 2D attn_mask is torch.Size([538, 3225]), but should be (3225, 3225).

我测试了好几次，每次都改变 GPU 的数量。掩码的形状将除以 GPU 的数量。 我不知道如何解决这个问题。

Answer 1

Dataparallel 通过假设它大约是批量大小来拆分第一个维度，因此它将您的 tgt_mask 拆分为 6 个张量。围绕这个问题进行了讨论，但我确定现在是否有解决方案。 https://discuss.pytorch.org/t/avoid-tensor-split-with-nn-dataparallel/18293

您可以重复 tgt_mask [3225, 3225] 到 [N, 3225, 3225] 并将其从 [1, 3225, 3225] 重新调整为 [3225, 3225]，然后再将其传递给变压器模型。