huggingface / huggingface/pytorch-image-models

[BUG] `PatchEmbed` returns a view tensor, triggering an FSDP2 warning

Closed
#2,766 6 comments 0 reactions 1 assignee Claimed by @rwightman View on GitHub
bug
Dominant language
Python
Stars
37.1k
Forks
5.2k
Avg merge
1d 11h
Merged PRs (30d)
37

Description

**Describe the bug**

```shell
/usr/local/lib/python3.12/dist-packages/torch/distributed/fsdp/_fully_shard/_fsdp_state.py:336: UserWarning: FSDP2-wrapped module (FSDPPatchEmbed) returned a view tensor. An in-place op on this view (e.g., `x += y`) will silently drop the pre-backward hook and skip the all-gather, which can cause backward to fail or produce wrong gradients. Use out-of-place ops (`out = out + y`, not `out += y`) or `.clone()` the output before any in-place op.
output = self._register_pre_backward_hook(output)
```

This warning was introduced to pytorch in https://github.com/pytorch/pytorch/pull/181850

**To Reproduce**
Steps to reproduce the behavior:
1. install torch 2.14.0+cu126
2. save the following to `timmpe_fsdp2warn.py`

```python
import torch
import torch.distributed as dist
from torch.distributed.fsdp import fully_shard
from timm.layers import PatchEmbed

deviceid = torch.cuda.current_device()
device = torch.device(f"cuda:{deviceid}")
backend = torch.distributed.get_default_backend_for_device(device)
dist.init_process_group(
backend=backend,
init_method="env://",
device_id=device,
)

pe = PatchEmbed()
pe.cuda()
fully_shard(pe)

x = torch.randn(2, 3, 224, 224, device="cuda")

out = pe(x)
```
3. run

```shell
$ torchrun --standalone --nproc-per-node gpu timmpe_fsdp2warn.py
/usr/local/lib/python3.12/dist-packages/torch/distributed/fsdp/_fully_shard/_fsdp_state.py:336: UserWarning: FSDP2-wrapped module (FSDPPatchEmbed) returned a view tensor. An in-place op on this view (e.g., `x += y`) will silently drop the pre-backward hook and skip the all-gather, which can cause backward to fail or produce wrong gradients. Use out-of-place ops (`out = out + y`, not `out += y`) or `.clone()` the output before any in-place op.
output = self._register_pre_backward_hook(output)
```

**Expected behavior**
Avoid the warning.

Should not return a view tensor

https://github.com/huggingface/pytorch-image-models/blob/a6940450df36e0ef7a0a52b2f9b08690d0c89a23/timm/layers/patch_embed.py#L137-L138

**Desktop (please complete the following information):**
- OS: Ubuntu 22.04
- timm v1.0.29
- torch==2.14.0+cu126

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.