[MPS] embedding_bag(include_last_offset=True) ignores the terminal offset
- 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
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