🎓 A1Learner
🔍

PyTorch Transformer Gotchas

Transformer gotchas are silent traps where the code runs and the answer is wrong. This lesson reproduces three of them and fixes each with one line.

The silent shape

The default layer expects the sequence first. It accepts a batch-first tensor and returns the same shape, so nothing complains.

import torch
from torch import nn
torch.manual_seed(0)
layer = nn.TransformerEncoderLayer(d_model=8, nhead=2, dim_feedforward=16, dropout=0.0)
layer.eval()
x = torch.randn(2, 4, 8)
with torch.no_grad():
    out = layer(x)
print("Input shape:", tuple(x.shape))
print("Output shape:", tuple(out.shape))

The integer mask

Padding masks built from counts are integer tensors. The layer refuses them with an AssertionError that names the two allowed dtypes.

batch_layer = nn.TransformerEncoderLayer(d_model=8, nhead=2, dim_feedforward=16, dropout=0.0, batch_first=True)
int_mask = torch.tensor([[0, 0, 1, 1], [0, 1, 1, 1]])
try:
    batch_layer(x, src_key_padding_mask=int_mask)
except AssertionError as error:
    print("AssertionError:", error)

Dropout at scoring time

A layer left in train mode drops different units on each pass. Eval mode makes the two passes match.

torch.manual_seed(1)
live = nn.TransformerEncoderLayer(d_model=8, nhead=2, dim_feedforward=16, dropout=0.5, batch_first=True)
live.train()
print("Train mode, two passes match:", torch.equal(live(x), live(x)))
live.eval()
with torch.no_grad():
    print("Eval mode, two passes match:", torch.equal(live(x), live(x)))

Example

import torch
from torch import nn

# Bug 1: a layer built without batch_first reads the batch as the sequence.
torch.manual_seed(0)
wrong_layer = nn.TransformerEncoderLayer(
    d_model=8, nhead=2, dim_feedforward=16, dropout=0.0)
right_layer = nn.TransformerEncoderLayer(
    d_model=8, nhead=2, dim_feedforward=16, dropout=0.0, batch_first=True)
right_layer.load_state_dict(wrong_layer.state_dict())
wrong_layer.eval()
right_layer.eval()
x = torch.randn(2, 4, 8)
with torch.no_grad():
    wrong = wrong_layer(x)
    right = right_layer(x)
    swapped = right_layer(x.transpose(0, 1)).transpose(0, 1)
print("Bug 1: the batch_first mismatch")
print("Runs without an error:", tuple(wrong.shape) == (2, 4, 8))
print("Matches the correct run:", torch.allclose(wrong, right))
print("Matches a run on the transposed input:", torch.allclose(wrong, swapped, atol=1e-6))
print("Fix: build the layer with batch_first=True")
print()

# Bug 2: a padding mask built as integers is refused.
layer = nn.TransformerEncoderLayer(
    d_model=8, nhead=2, dim_feedforward=16, dropout=0.0, batch_first=True)
layer.eval()
int_mask = torch.tensor([[0, 0, 1, 1], [0, 1, 1, 1]])
try:
    layer(x, src_key_padding_mask=int_mask)
    print("Bug 2: no error was raised")
except AssertionError as error:
    print("Bug 2: the mask dtype trap")
    print("AssertionError:", error)
bool_mask = int_mask.bool()
with torch.no_grad():
    fixed = layer(x, src_key_padding_mask=bool_mask)
print("Fix: cast the mask with .bool(), output shape:", tuple(fixed.shape))
print()

# Bug 3: dropout stays live while the model is in train mode.
torch.manual_seed(1)
live = nn.TransformerEncoderLayer(
    d_model=8, nhead=2, dim_feedforward=16, dropout=0.5, batch_first=True)
live.train()
first = live(x)
second = live(x)
print("Bug 3: dropout live at scoring time")
print("Two train-mode passes match:", torch.equal(first, second))
live.eval()
with torch.no_grad():
    first = live(x)
    second = live(x)
print("Fix: call eval(), two passes match:", torch.equal(first, second))

Key points

  • batch_first=False is the default, so a layer reads its input as sequence, batch, width
  • A padding mask must be bool or float, and an integer mask raises AssertionError
  • Dropout stays live in train mode, so two passes over one input can differ
  • Use model.eval() before scoring, so dropout switches off
  • A transformer bug can hide in the argument order, the mask dtype or the mode

In this example, three bugs are reproduced one at a time. The first builds two layers with the same weights, one default and one batch_first. The default layer runs without an error and returns the right shape. Its output still differs from the correct run. It matches a run on the transposed input to within float rounding. The fix builds the layer with batch_first=True. The second bug passes an integer padding mask and catches the AssertionError it raises. The fix casts the mask with .bool(). The third bug leaves a dropout layer in train mode, where two passes over one input differ. The fix calls eval(), and the two passes then match.

Note

A float mask of 0s and 1s runs without error, but it biases the scores instead of hiding positions

Try it Yourself

Edit the code, then press Run.

Running is simulated · Python Playground

Exercise

A transformer layer returns the right shape but wrong values, and its input is (batch, sequence, width). What is the first thing to check?