pytorch / pytorch/pytorch

`torch.compile` silently succeeds on `TransformerEncoder` with all-masked `src_key_padding_mask` where eager raises RuntimeError

Open
#178,677 2 comments 0 reactions 0 assignees View on GitHub
bot-triaged module: correctness (silent) module: dynamo module: edge cases module: nestedtensor module: nn oncall: pt2 topic: fuzzer triaged
Dominant language
Python
Stars
103k
Forks
29.5k
PR merge metrics
PR metrics pending

Description

### 🐛 Describe the bug

`torch.compile` silently succeeds when `nn.TransformerEncoder` receives a `src_key_padding_mask` with all `True` values (all positions masked), while **eager mode raises `RuntimeError: to_padded_tensor: at least one constituent tensor should have non-zero numel`**.

In PyTorch's convention, `True` in `src_key_padding_mask` means "this position is padding and should be masked." With an all-`True` mask, every position is padding — the eager path attempts to create nested tensors via `torch._nested_tensor_from_mask`, which fails because no constituent tensor has non-zero elements. The compiled path bypasses the nested tensor path entirely and computes a result from all-masked (NaN/zero) data.

### Affected files

| File | Source | Pattern |
|------|--------|---------|
| `randperm_index_pattern-5.py` | E8 (struct+route+thompson), round-4 | `randperm_index_pattern` |

### Root cause

The `TransformerEncoder` in eager mode converts `src_key_padding_mask` to a nested tensor via `torch._nested_tensor_from_mask`. With all positions masked (all `True`), this creates nested tensors with zero elements, and the subsequent `to_padded_tensor` call fails because it requires at least one constituent tensor with non-zero numel.

The compiled path takes a different execution route — Dynamo traces through the transformer layers without going through the nested tensor conversion, sidestepping the check entirely.

This is an **eager vs compiled consistency bug**: both paths should either succeed or fail for the same input.

### Full model-level reproducer

```python
import torch
import torch.nn as nn

class RandomPermutationTransformer(nn.Module):
def __init__(self, vocab_size=10000, d_model=512, nhead=8,
num_layers=6, slice_shape=8):
super().__init__()
self.slice_shape = slice_shape
self.embedding = nn.Embedding(vocab_size, d_model)
self.pos_encoding = nn.Parameter(torch.randn(1, 1000, d_model))
encoder_layer = nn.TransformerEncoderLayer(
d_model=d_model, nhead=nhead,
dim_feedforward=2048, dropout=0.1, batch_first=True
)
self.transformer = nn.TransformerEncoder(
encoder_layer, num_layers=num_layers
)
self.classifier = nn.Sequential(
nn.Linear(d_model, 256), nn.ReLU(),
nn.Dropout(0.2), nn.Linear(256, 2)
)
self.bn = nn.BatchNorm1d(d_model)

def forward(self, x, attention_mask=None):
batch_size, seq_len = x.shape
embedded = self.embedding(x)
embedded = embedded + self.pos_encoding[:, :seq_len, :]
if attention_mask is not None:
transformer_out = self.transformer(
embedded, src_key_padding_mask=attention_mask
)
else:
transformer_out = self.transformer(embedded)
pooled = transformer_out.mean(dim=1)
pooled = self.bn(pooled)
perm = torch.randperm(batch_size, device=pooled.device)
selected_indices = perm[:self.slice_shape]
selected_pooled = pooled[selected_indices]
logits = self.classifier(selected_pooled)
return logits

model = RandomPermutationTransformer(
vocab_size=10000, d_model=512, nhead=8,
num_layers=6, slice_shape=8
).cuda()
model.eval()

input_ids = torch.randint(0, 10000, (16, 128), dtype=torch.long).cuda()
# All positions masked — edge case
attention_mask = torch.ones(16, 128, dtype=torch.bool).cuda()

# Eager: fails
try:
with torch.no_grad():
model(input_ids, attention_mask)
print("eager: OK")
except RuntimeError as e:
print(f"eager: ERROR — {e}")

# Compiled: succeeds (inconsistent)
torch._dynamo.reset()
compiled_model = torch.compile(model)
try:
with torch.no_grad():
out = compiled_model(input_ids, attention_mask)
print(f"compile: OK — shape={out.shape}")
except Exception as e:
print(f"compile: ERROR — {e}")
```

### Behavior summary

| Mode | Result | Output |
|------|--------|--------|
| Eager | RuntimeError | `to_padded_tensor: at least one constituent tensor should have non-zero numel` |
| `torch.compile` | Success (inconsistent) | `torch.Size([8, 2])` — numerical output from all-masked data |

The bug requires `TransformerEncoder` with `src_key_padding_mask=all_True`. It is specific to the edge case where every position is masked. With normal masks (some True, some False), both eager and compiled paths succeed consistently.

### Error logs

**Eager mode** (fails):
```
UserWarning: The PyTorch API of nested tensors is in prototype stage...
RuntimeError: to_padded_tensor: at least one constituent tensor should have non-zero numel
```

**`torch.compile`** (succeeds inconsistently):
```
compile: OK — shape=torch.Size([8, 2])
```

### Versions

```
PyTorch version: 2.12.0.dev20260315+cu126
OS: Ubuntu 22.04.5 LTS (x86_64)
Python version: 3.10.12
GPU: NVIDIA GeForce RTX 3060 Laptop GPU
CUDA: 12.6
```

cc @albanD @mruberry @jbschlosser @walterddr @mikaylagawarecki @cpuhrsch @bhosmer @drisspg @soulitzer @davidberard98 @YuqingJ @chauhang @penguinwu @voznesenskym @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @kadeng @amjames @Lucaskabela @jataylo @ezyang @msaroufim @bdhirsh @anijain2305

topic: fuzzer

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.