[muse-glimmer] interpolate_pos_emb returns corrupted bf16 values on CPU and can segfault
Mantenedores costumam responder em até 1 dia
Avaliação
Esta issue ainda não foi avaliada.
Descrição
🐛 Describe the bug
interpolate_pos_emb in examples/models/muse-glimmer/vision/precompute.py returns wrong values for a bf16 positional table on CPU, and it is the likely source of the intermittent segfault in the pull / unittest / macos / macos-job job.
In CI. The job crashes with Fatal Python error: Segmentation fault in torch.nn.functional.grid_sample, called from interpolate_pos_emb, while running examples/models/muse-glimmer/tests/test_vision_precompute.py. I found two occurrences, both with torch 2.15.0a0+gite91899c on osx-arm64: on main at 69eeed2 (run 38405593574, in PrecomputeBundleTest::test_all_inputs_consistent) and on #22373 (run 38325606287, in PosEmbInterpTest::test_downscaled_shape). The crashing thread in the second one:
Fatal Python error: Segmentation fault
Current thread 0x00000001f0e7a140 (most recent call first):
File ".../site-packages/torch/nn/functional.py", line 5701 in grid_sample
File ".../site-packages/executorch/examples/models/muse_glimmer/vision/precompute.py", line 114 in interpolate_pos_emb
File ".../examples/models/muse-glimmer/tests/test_vision_precompute.py", line 72 in test_downscaled_shape
Values. The output is wrong on every run, not only when the process crashes. Interpolating to the native 32x32 grid puts every sample on a texel center, so it should return the table exactly:
import torch
from executorch.examples.models.muse_glimmer.vision.precompute import (
MuseGlimmerVisionConfig,
interpolate_pos_emb,
)
torch.manual_seed(0)
cfg = MuseGlimmerVisionConfig()
table = torch.randn(cfg.pos_emb_grid_h * cfg.pos_emb_grid_w, cfg.latent_dim, dtype=torch.bfloat16)
pos = interpolate_pos_emb(table, 32, 32, cfg) # native 32x32 grid: should return the table
print("mismatched:", (pos[0] != table).float().mean().item(), "NaN:", int(pos.isnan().sum()))
mismatched: 0.9780623316764832 NaN: 826
Cause. interpolate_pos_emb passes the permuted view pos_embed_table.view(32, 32, 1536).permute(2, 0, 1).unsqueeze(0), whose strides are (1536, 1, 49152, 1536), to F.grid_sample. For Half and BFloat16, the vectorized CPU kernel computes its gather offsets in int_same_size_t<scalar_t>, that is int16_t (GridSamplerKernel.cpp, lines 468 to 469, offsets at lines 568 to 571), and grid_sampler_2d_cpu guards against offset overflow only for float32 (GridSampler.cpp, line 1101). For this view the offsets go up to 31 * 49152 + 31 * 1536 = 1571328, far above the int16 maximum of 32767, so they wrap. An emulation of that int16 arithmetic puts about half of the gathered reads before the start of the storage: they return garbage or NaN, and can fault when that memory is unmapped, which is consistent with the CI crash. On the contiguous tensor, the same 32x32 call returns the table exactly, and an 8x8 downscale matches a float32 reference within bf16 rounding. The kernel code is the same on PyTorch main (afa81a4) and at the CI pin (e91899c); the 16-bit dtypes reached this kernel with pytorch/pytorch#134812, first released in 2.6.0.
Impact. precompute_vision_inputs is called by inference.py (generate_with_image, whose vision encoder adds pos_emb to the patch embeddings) and by export/common.py (sample_vision_inputs, example inputs for export). The C++ runtime does not go through this code: vision/preprocess.h interpolates from the float32 table. The existing tests check only shapes and dtypes, which is why the wrong values went unnoticed.
Proposed fix. Make the view contiguous before grid_sample, which keeps the computation in the table's dtype and brings the largest gather offset to 1023, and add value checks to the tests. I have a PR ready with this change.
A separate observation, not addressed by that fix: the Python code computes the sampling coordinates in bf16 (it builds the grid in bf16, and grid_sample, which requires the grid in the input dtype, also unnormalizes it in bf16), while vision/preprocess.h computes them in double. On grids that are not powers of two the results then diverge: on 74x74, a Python transcription of detail::interpolate_pos_emb (I did not build the C++) differs from the corrected Python output on about 97% of the values, by up to about 0.75. The docstring says the interpolation stays in the table's dtype to preserve reference numerics, so this may be intended. Which of the two paths is meant to match the reference?
Versions
Local reproduction: Windows 11, x86_64 (AVX2), Python 3.12.11, torch 2.14.0+cpu, executorch 1.5.1 (its precompute.py is identical to the one on main at 060075b). A pure PyTorch transcription of the reproduction also fails with torch 2.11.0+cu128 on Python 3.14.2 (97.5% of the values mismatched, 826 to 842 NaN depending on the run).
CI: osx-arm64, torch 2.15.0a0+gite91899c.
cc @mergennachin @iseeyuan @lucylq @helunwencser @tarun292 @kimishpatel @jackzhxng
- Linguagem predominante
- Python
- Estrelas
- 5.1k
- Forks
- 1.2k
- Merge médio
- 1d 23h
- PRs com merge (30d)
- 713
Preparar o ambiente
- Sem Dockerfile nem arquivo Docker Compose
- Tem um modelo de pull request
- Ler o guia de contribuição
Primeiros passos
- Leia a issue inteira e depois o guia de contribuição do projeto.
- Comente na issue dizendo que vai assumir — evita que duas pessoas façam o mesmo trabalho.
- Faça um fork do repositório e trabalhe em uma branch.
- Abra um pull request que referencie o número da issue.
Mais de pytorch/executorch
-
Portable remainder kernel: wrong sign for integers and float64 rounded to floatTalvez já em andamento @JakeStevens assumiu há 4 dias. Abertabug module: kernels triaged
Dificuldade 2/5 1-3 horas Facilidade para iniciantes 82/100
pytorch/executorch#23682 · 2 responsáveis ·
Mantenedores costumam responder em até 1 dia
-
enhancement triaged
Dificuldade 2/5 Meio dia Facilidade para iniciantes 68/100
pytorch/executorch#21640 ·
Mantenedores costumam responder em até 1 dia
-
bug module: microcontrollers partner: arm
Dificuldade 4/5 3-5 dias Facilidade para iniciantes 38/100
pytorch/executorch#23812 · 2 comentários · 1 reação ·
Mantenedores costumam responder em até 1 dia
-
macOS Python wheel _portable_lib has non-relocatable libtorch rpathTalvez já em andamento @JacobSzwejbka assumiu há 1 dia. Aberta
pytorch/executorch#23742 · 1 comentário · 2 responsáveis ·
Mantenedores costumam responder em até 1 dia
-
Feedback about ToolsAbertainvalid module: doc
Dificuldade 3/5 1-2 dias Facilidade para iniciantes 35/100
pytorch/executorch#23706 · 1 comentário ·
Mantenedores costumam responder em até 1 dia
Todas as issues de pytorch/executorch
Issues semelhantes
-
Dificuldade 1/5 Menos de uma hora Facilidade para iniciantes 88/100
-
Dificuldade 2/5 1-3 horas Facilidade para iniciantes 72/100
aws-samples/sample-per-model-bedrock#261 ·
Mantenedores costumam responder em até 1 dia
-
[UI/Accessibility] Low contrast text visibility in country selection dropdown on settings pageAberta
Dificuldade 2/5 1-3 horas Facilidade para iniciantes 65/100
pwncollege/dojo#1188 ·
Mantenedores costumam responder em até 1 dia
-
bug launcher packaging
Dificuldade 1/5 Menos de uma hora Facilidade para iniciantes 88/100
tenstorrent/tt-studio#1426 ·
Mantenedores costumam responder em até 1 dia
-
Dificuldade 2/5 1-3 horas Facilidade para iniciantes 72/100
python-attrs/cattrs#796 ·
Mantenedores costumam responder em até 1 dia