huggingface / huggingface/diffusers

[Bug] Positional Embedding Index Out-of-Bounds in AuraFlow

Đang mở
#12,656 8 bình luận 0 reaction 0 người được giao Xem trên GitHub
bug stale
Ngôn ngữ chính
Python
Star
34.5k
Fork
7.3k
Merge trung bình
3 ngày 3 giờ
Pull request đã merge (30 ngày)
91

Mô tả

### Describe the bug

When training the AuraFlow (Pony V7) lora model (via SimpleTuner), the positional embedding selection function
`pe_selection_index_based_on_dim()` can generate **negative** or **out-of-range** indices for the positional embedding table.

This causes:

```
CUDA error: device-side assert triggered
vectorized_gather_kernel: Assertion `ind >=0 && ind < ind_dim_size` failed.
```

Root cause: the code computes a centered crop of the PE grid, but when the input latent grid
(`h // patch_size`, `w // patch_size`) is **larger** than the pretrained positional embedding grid,
the computed indices become invalid (negative start, or exceeding max grid index).

### Reproduction

```python
import torch

# Simulated model parameters
patch_size = 2
pos_embed_max_size = 96 * 96 # 9216 positions
pos_embed = torch.zeros(1, pos_embed_max_size, 3072)

def pe_selection_index_based_on_dim(h, w):
h_p, w_p = h // patch_size, w // patch_size
h_max = w_max = int(pos_embed_max_size ** 0.5)

starth = h_max // 2 - h_p // 2
startw = w_max // 2 - w_p // 2

rows = torch.arange(starth, starth + h_p)
cols = torch.arange(startw, startw + w_p)
row_idx, col_idx = torch.meshgrid(rows, cols, indexing="ij")

return (row_idx * w_max + col_idx).flatten()

# Example that produces invalid indices (real case)
h, w = 196, 196 # latent spatial dims
idx = pe_selection_index_based_on_dim(h, w)

print(idx.min().item(), idx.max().item()) # → negative & >9215
assert (idx >= 0).all() and (idx < pos_embed_max_size).all()
```

### **Observed Output**

```
-97 9312
AssertionError
```

This reproduces the exact index state that triggers the CUDA `vectorized_gather_kernel` failure during training.

### Logs

```shell

```

### System Info

- 🤗 Diffusers version: 0.35.2
- Platform: Linux-6.6.87.2-microsoft-standard-WSL2-x86_64-with-glibc2.39
- Running on Google Colab?: No
- Python version: 3.12.3
- PyTorch version (GPU?): 2.9.0+cu130 (True)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Huggingface_hub version: 0.36.0
- Transformers version: 4.57.1
- Accelerate version: 1.11.0
- PEFT version: 0.17.1
- Bitsandbytes version: 0.48.2
- Safetensors version: 0.6.2
- xFormers version: not installed
- Accelerator: NVIDIA GeForce RTX 4090, 24564 MiB
- Using GPU in script?: yes
- Using distributed or parallel set-up in script?: no

### Who can help?

_No response_

Hướng dẫn đóng góp

Mở hướng dẫn đóng góp

Hướng nghiên cứu

Tìm pe_selection_index_based_on_dim() trong phần triển khai AuraFlow và chạy reproduction từ issue với h và w được đặt thành 196. Theo dõi cách crop ở giữa được tính toán cho các grid lớn hơn embedding grid được pretrained; hoàn tất khi các chỉ mục được trả về vẫn nằm trong bảng positional embedding và CUDA gather assertion không còn xảy ra.

Do mô hình lập chỉ mục viết ra từ nội dung của issue.

Đánh giá

Công nghệ
python, pytorch
Lĩnh vực
machine-learning
Loại issue
Lỗi
Độ khó
3/5
Thời gian dự kiến
1-2 ngày
Mức độ hoạt động
Đình trệ
Độ rõ ràng
Khá rõ ràng
Mức phù hợp với người mới
42/100

Nhận issue mới trong hộp thư của bạn

Bản tóm tắt ngắn những issue GitHub phù hợp với người mới.