pytorch / pytorch/pytorch

[MPS] embedding_bag(include_last_offset=True) ignores the terminal offset

Open
#190,060 0 comments 0 reactions 0 assignees View on GitHub
module: correctness (silent) module: mps triaged
Dominant language
Python
Stars
103k
Forks
29.6k
PR merge metrics
PR metrics pending

Description

### 🐛 Describe the bug

With `include_last_offset=True`, the final entry of `offsets` must end the
last bag. The MPS host decrements `num_bags` accordingly, but the kernel
always ends the final bag at `num_indices`, so the supplied terminal offset is
never read and trailing indices are wrongly summed into the last bag.

```python
import torch

W = torch.arange(16.0).reshape(4, 4)
idx = torch.tensor([0, 1, 2])
offs = torch.tensor([0, 2]) # one bag: indices [0, 2)
ref = torch.nn.functional.embedding_bag(idx, W, offs, mode="sum", include_last_offset=True)
out = torch.nn.functional.embedding_bag(idx.to("mps"), W.to("mps"), offs.to("mps"),
mode="sum", include_last_offset=True)
print("cpu:", ref.tolist())
print("mps:", out.cpu().tolist())
```

Output on a source build of current main (2.14.0a0+git2662808); also
reproduces on 2.13.0 stable:

```text
cpu: [[4.0, 6.0, 8.0, 10.0]]
mps: [[12.0, 15.0, 18.0, 21.0]] # row 2 wrongly included
```

The host decrement is at
https://github.com/pytorch/pytorch/blob/26628087023e5102381849e9a27c3f09b6d485a4/aten/src/ATen/native/mps/operations/EmbeddingBag.mm#L70-L73
and the kernel's last-bag end at
https://github.com/pytorch/pytorch/blob/26628087023e5102381849e9a27c3f09b6d485a4/aten/src/ATen/native/mps/kernels/EmbeddingBag.metal#L172-L176
uses `num_indices` unconditionally. Backward gradients for the trailing
indices are correspondingly wrong. #52851 documents the expected
`include_last_offset` contract generally (it predates the MPS backend).

### Versions

PyTorch 2.14.0a0+git2662808 (source build of main) and 2.13.0; macOS 26.5.2 (arm64), Apple M5 Max, Python 3.12.13.

cc @ezyang @gchanan @kadeng @msaroufim @kulinseth @malfet @DenisVieriu97 @jhavukainen @aditvenk @Isalia20

Contributor guide

Open the contributing guide

Research direction

Start with the host logic in aten/src/ATen/native/mps/operations/EmbeddingBag.mm at lines 70–73 and the last-bag handling in aten/src/ATen/native/mps/kernels/EmbeddingBag.metal at lines 172–176. Run the supplied include_last_offset reproduction and compare CPU and MPS results. Done means the terminal offset bounds the final bag and trailing-index forward and backward results match the expected contract.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
backend, machine-learning
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Quiet
Clarity
Mostly clear
Newbie friendliness
55/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.