ViTForMaskedImageModeling (vit-base-patch16-224-in21k) distributed training warning
Author: willxxy
Pull request: no
Labels: name: bug

name: Vision

### System Info

I am having some strange errors/warnings when running the vit model. I am on version `4.56.0`. 2 A6000 NVIDIA GPUS. Ubuntu 20.04.6 LTS. Python 3.10.18. torch version 2.6.0+cu124.

### Who can help?

@molbap @yonigozlan 
1. When I use the `AutoModel.from_pretrained` class for the `vit-base-patch16-224-in21k` checkpoint and wrap the model with `DDP` and run with `torchrun`, I get this error 

```
[rank1]:   File "/home/william/miniconda3/envs/ecgv2/lib/python3.10/site-packages/transformers/utils/generic.py", line 1064, in wrapper
[rank1]:     outputs = func(self, *args, **kwargs)
[rank1]:   File "/home/william/miniconda3/envs/ecgv2/lib/python3.10/site-packages/transformers/models/vit/modeling_vit.py", line 482, in forward
[rank1]:     embedding_output = self.embeddings(
[rank1]:   File "/home/william/miniconda3/envs/ecgv2/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1739, in _wrapped_call_impl
[rank1]:     return self._call_impl(*args, **kwargs)
[rank1]:   File "/home/william/miniconda3/envs/ecgv2/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1750, in _call_impl
[rank1]:     return forward_call(*args, **kwargs)
[rank1]:   File "/home/william/miniconda3/envs/ecgv2/lib/python3.10/site-packages/transformers/models/vit/modeling_vit.py", line 112, in forward
[rank1]:     mask_tokens = self.mask_token.expand(batch_size, seq_length, -1)
[rank1]: AttributeError: 'NoneType' object has no attribute 'expand'
```

However, when I use the `ViTForMaskedImageModeling.from_pretrained` class I get this warning

```
/home/william/miniconda3/envs/ecgv2/lib/python3.10/site-packages/torch/autograd/graph.py:823: UserWarning: Grad strides do not match bucket view strides. This may indicate grad was not created according to the gradient layout contract, or that the param's strides changed since DDP was constructed.  This is not an error, but may impair performance.
grad.sizes() = [768, 768, 1, 1], strides() = [768, 1, 768, 768]
bucket_view.sizes() = [768, 768, 1, 1], strides() = [768, 1, 1, 1] (Triggered internally at /pytorch/torch/csrc/distributed/c10d/reducer.cpp:327.)
  return Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass
/home/william/miniconda3/envs/ecgv2/lib/python3.10/site-packages/torch/autograd/graph.py:823: UserWarning: Grad strides do not match bucket view strides. This may indicate grad was not created according to the gradient layout contract, or that the param's strides changed since DDP was constructed.  This is not an error, but may impair performance.
grad.sizes() = [768, 768, 1, 1], strides() = [768, 1, 768, 768]
bucket_view.sizes() = [768, 768, 1, 1], strides() = [768, 1, 1, 1] (Triggered internally at /pytorch/torch/csrc/distributed/c10d/reducer.cpp:327.)
  return Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass
```

Why does this happen? I found maybe one related issue [here](https://github.com/huggingface/transformers/issues/32614). I found [this too](https://github.com/huggingface/transformers/issues/23832) but I don't think this one is related.

### Information

- [ ] The official example scripts
- [ ] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [ ] My own task or dataset (give details below)

### Reproduction

run `CUDA_VISIBLE_DEVICES=1,2 torchrun --nproc_per_node=2 --master_port=12355 ecg_bench/test.py`

```
import os
import torch
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
from transformers import ViTForMaskedImageModeling, AutoModel, AutoImageProcessor
from ecg_bench.configs.constants import HF_CACHE_DIR

def setup_ddp():
    if not dist.is_initialized():
        dist.init_process_group(backend="nccl")
    local_rank = int(os.environ.get("LOCAL_RANK", 0))
    torch.cuda.set_device(local_rank)
    return local_rank

def test_model(model_class):
    model = model_class.from_pretrained("google/vit-base-patch16-224-in21k", cache_dir=HF_CACHE_DIR)
    processor = AutoImageProcessor.from_pretrained("google/vit-base-patch16-224-in21k", cache_dir=HF_CACHE_DIR)
    
    local_rank = setup_ddp()
    model = DDP(model.to(f"cuda:{local_rank}"), device_ids=[local_rank])
    
    images = torch.rand(2, 3, 224, 224)
    inputs = processor(images=images, return_tensors="pt")
    pixel_values = inputs["pixel_values"].to(f"cuda:{local_rank}")
    
    bool_masked_pos = torch.zeros(2, 196, dtype=torch.bool, device=f"cuda:{local_rank}")
    for i in range(2):
        bool_masked_pos[i, torch.randperm(196)[:29]] = True
    outputs = model(pixel_values=pixel_values, bool_masked_pos=bool_masked_pos)
    loss = outputs.loss
    
    loss.backward()
    print(f"Loss: {loss.item()}")

if __name__ == "__main__":
    test_model(ViTForMaskedImageModeling)
    test_model(AutoModel)
```



### Expected behavior

```
(ecgv2) william@gpu2:~/ecg_dev/ECG-Bench-v2$ bash scripts/test_ddp_reproduction.sh 
Running minimal reproduction test with DDP...
W1007 03:11:37.695000 3870938 site-packages/torch/distributed/run.py:792] 
W1007 03:11:37.695000 3870938 site-packages/torch/distributed/run.py:792] *****************************************
W1007 03:11:37.695000 3870938 site-packages/torch/distributed/run.py:792] Setting OMP_NUM_THREADS environment variable for each process to be 1 in default, to avoid your system being overloaded, please further tune the variable for optimal performance in your application as needed. 
W1007 03:11:37.695000 3870938 site-packages/torch/distributed/run.py:792] *****************************************
Some weights of ViTForMaskedImageModeling were not initialized from the model checkpoint at google/vit-base-patch16-224-in21k and are newly initialized: ['decoder.0.bias', 'decoder.0.weight', 'vit.embeddings.mask_token']
You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.
Some weights of ViTForMaskedImageModeling were not initialized from the model checkpoint at google/vit-base-patch16-224-in21k and are newly initialized: ['decoder.0.bias', 'decoder.0.weight', 'vit.embeddings.mask_token']
You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.
Fetching 1 files: 100%|██████████████████████████████████████████| 1/1 [00:00<00:00, 16070.13it/s]
Fetching 1 files: 100%|██████████████████████████████████████████| 1/1 [00:00<00:00, 19239.93it/s]
Fast image processor class <class 'transformers.models.vit.image_processing_vit_fast.ViTImageProcessorFast'> is available for this model. Using slow image processor class. To use the fast image processor class set `use_fast=True`.
Fast image processor class <class 'transformers.models.vit.image_processing_vit_fast.ViTImageProcessorFast'> is available for this model. Using slow image processor class. To use the fast image processor class set `use_fast=True`.
Fetching 1 files: 100%|██████████████████████████████████████████| 1/1 [00:00<00:00, 13443.28it/s]
Fetching 1 files: 100%|███████████████████████████████████████████| 1/1 [00:00<00:00, 8473.34it/s]
It looks like you are trying to rescale already rescaled images. If the input images have pixel values between 0 and 1, set `do_rescale=False` to avoid rescaling them again.
It looks like you are trying to rescale already rescaled images. If the input images have pixel values between 0 and 1, set `do_rescale=False` to avoid rescaling them again.
/home/william/miniconda3/envs/ecgv2/lib/python3.10/site-packages/torch/autograd/graph.py:823: UserWarning: Grad strides do not match bucket view strides. This may indicate grad was not created according to the gradient layout contract, or that the param's strides changed since DDP was constructed.  This is not an error, but may impair performance.
grad.sizes() = [768, 768, 1, 1], strides() = [768, 1, 768, 768]
bucket_view.sizes() = [768, 768, 1, 1], strides() = [768, 1, 1, 1] (Triggered internally at /pytorch/torch/csrc/distributed/c10d/reducer.cpp:327.)
  return Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass
/home/william/miniconda3/envs/ecgv2/lib/python3.10/site-packages/torch/autograd/graph.py:823: UserWarning: Grad strides do not match bucket view strides. This may indicate grad was not created according to the gradient layout contract, or that the param's strides changed since DDP was constructed.  This is not an error, but may impair performance.
grad.sizes() = [768, 768, 1, 1], strides() = [768, 1, 768, 768]
bucket_view.sizes() = [768, 768, 1, 1], strides() = [768, 1, 1, 1] (Triggered internally at /pytorch/torch/csrc/distributed/c10d/reducer.cpp:327.)
  return Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass
Loss: 1.00250244140625Loss: 1.002082109451294

Fetching 1 files: 100%|███████████████████████████████████████████| 1/1 [00:00<00:00, 9664.29it/s]
Fetching 1 files: 100%|███████████████████████████████████████████| 1/1 [00:00<00:00, 6364.65it/s]
Fast image processor class <class 'transformers.models.vit.image_processing_vit_fast.ViTImageProcessorFast'> is available for this model. Using slow image processor class. To use the fast image processor class set `use_fast=True`.
Fast image processor class <class 'transformers.models.vit.image_processing_vit_fast.ViTImageProcessorFast'> is available for this model. Using slow image processor class. To use the fast image processor class set `use_fast=True`.
Fetching 1 files: 100%|██████████████████████████████████████████| 1/1 [00:00<00:00, 18001.30it/s]
Fetching 1 files: 100%|██████████████████████████████████████████| 1/1 [00:00<00:00, 22192.08it/s]
[rank1]: Traceback (most recent call last):
[rank1]:   File "/home/william/ecg_dev/ECG-Bench-v2/ecg_bench/test.py", line 37, in <module>
[rank1]:     test_model(AutoModel)
[rank1]:   File "/home/william/ecg_dev/ECG-Bench-v2/ecg_bench/test.py", line 29, in test_model
[rank1]:     outputs = model(pixel_values=pixel_values, bool_masked_pos=bool_masked_pos)
[rank1]:   File "/home/william/miniconda3/envs/ecgv2/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1739, in _wrapped_call_impl
[rank1]:     return self._call_impl(*args, **kwargs)
[rank1]:   File "/home/william/miniconda3/envs/ecgv2/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1750, in _call_impl
[rank1]:     return forward_call(*args, **kwargs)
[rank1]:   File "/home/william/miniconda3/envs/ecgv2/lib/python3.10/site-packages/torch/nn/parallel/distributed.py", line 1643, in forward
[rank1]:     else self._run_ddp_forward(*inputs, **kwargs)
[rank1]:   File "/home/william/miniconda3/envs/ecgv2/lib/python3.10/site-packages/torch/nn/parallel/distributed.py", line 1459, in _run_ddp_forward
[rank1]:     return self.module(*inputs, **kwargs)  # type: ignore[index]
[rank1]:   File "/home/william/miniconda3/envs/ecgv2/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1739, in _wrapped_call_impl
[rank1]:     return self._call_impl(*args, **kwargs)
[rank1]:   File "/home/william/miniconda3/envs/ecgv2/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1750, in _call_impl
[rank1]:     return forward_call(*args, **kwargs)
[rank1]:   File "/home/william/miniconda3/envs/ecgv2/lib/python3.10/site-packages/transformers/utils/generic.py", line 1064, in wrapper
[rank1]:     outputs = func(self, *args, **kwargs)
[rank1]:   File "/home/william/miniconda3/envs/ecgv2/lib/python3.10/site-packages/transformers/models/vit/modeling_vit.py", line 482, in forward
[rank1]:     embedding_output = self.embeddings(
[rank1]:   File "/home/william/miniconda3/envs/ecgv2/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1739, in _wrapped_call_impl
[rank1]:     return self._call_impl(*args, **kwargs)
[rank1]:   File "/home/william/miniconda3/envs/ecgv2/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1750, in _call_impl
[rank1]:     return forward_call(*args, **kwargs)
[rank1]:   File "/home/william/miniconda3/envs/ecgv2/lib/python3.10/site-packages/transformers/models/vit/modeling_vit.py", line 112, in forward
[rank1]:     mask_tokens = self.mask_token.expand(batch_size, seq_length, -1)
[rank1]: AttributeError: 'NoneType' object has no attribute 'expand'
[rank0]: Traceback (most recent call last):
[rank0]:   File "/home/william/ecg_dev/ECG-Bench-v2/ecg_bench/test.py", line 37, in <module>
[rank0]:     test_model(AutoModel)
[rank0]:   File "/home/william/ecg_dev/ECG-Bench-v2/ecg_bench/test.py", line 29, in test_model
[rank0]:     outputs = model(pixel_values=pixel_values, bool_masked_pos=bool_masked_pos)
[rank0]:   File "/home/william/miniconda3/envs/ecgv2/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1739, in _wrapped_call_impl
[rank0]:     return self._call_impl(*args, **kwargs)
[rank0]:   File "/home/william/miniconda3/envs/ecgv2/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1750, in _call_impl
[rank0]:     return forward_call(*args, **kwargs)
[rank0]:   File "/home/william/miniconda3/envs/ecgv2/lib/python3.10/site-packages/torch/nn/parallel/distributed.py", line 1643, in forward
[rank0]:     else self._run_ddp_forward(*inputs, **kwargs)
[rank0]:   File "/home/william/miniconda3/envs/ecgv2/lib/python3.10/site-packages/torch/nn/parallel/distributed.py", line 1459, in _run_ddp_forward
[rank0]:     return self.module(*inputs, **kwargs)  # type: ignore[index]
[rank0]:   File "/home/william/miniconda3/envs/ecgv2/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1739, in _wrapped_call_impl
[rank0]:     return self._call_impl(*args, **kwargs)
[rank0]:   File "/home/william/miniconda3/envs/ecgv2/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1750, in _call_impl
[rank0]:     return forward_call(*args, **kwargs)
[rank0]:   File "/home/william/miniconda3/envs/ecgv2/lib/python3.10/site-packages/transformers/utils/generic.py", line 1064, in wrapper
[rank0]:     outputs = func(self, *args, **kwargs)
[rank0]:   File "/home/william/miniconda3/envs/ecgv2/lib/python3.10/site-packages/transformers/models/vit/modeling_vit.py", line 482, in forward
[rank0]:     embedding_output = self.embeddings(
[rank0]:   File "/home/william/miniconda3/envs/ecgv2/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1739, in _wrapped_call_impl
[rank0]:     return self._call_impl(*args, **kwargs)
[rank0]:   File "/home/william/miniconda3/envs/ecgv2/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1750, in _call_impl
[rank0]:     return forward_call(*args, **kwargs)
[rank0]:   File "/home/william/miniconda3/envs/ecgv2/lib/python3.10/site-packages/transformers/models/vit/modeling_vit.py", line 112, in forward
[rank0]:     mask_tokens = self.mask_token.expand(batch_size, seq_length, -1)
[rank0]: AttributeError: 'NoneType' object has no attribute 'expand'
[rank0]:[W1007 03:11:47.275383972 ProcessGroupNCCL.cpp:1496] Warning: WARNING: destroy_process_group() was not called before program exit, which can leak resources. For more info, please see https://pytorch.org/docs/stable/distributed.html#shutdown (function operator())
W1007 03:11:48.425000 3870938 site-packages/torch/distributed/elastic/multiprocessing/api.py:897] Sending process 3871025 closing signal SIGTERM
E1007 03:11:48.589000 3870938 site-packages/torch/distributed/elastic/multiprocessing/api.py:869] failed (exitcode: 1) local_rank: 1 (pid: 3871026) of binary: /home/william/miniconda3/envs/ecgv2/bin/python3
Traceback (most recent call last):
  File "/home/william/miniconda3/envs/ecgv2/bin/torchrun", line 10, in <module>
    sys.exit(main())
  File "/home/william/miniconda3/envs/ecgv2/lib/python3.10/site-packages/torch/distributed/elastic/multiprocessing/errors/__init__.py", line 355, in wrapper
    return f(*args, **kwargs)
  File "/home/william/miniconda3/envs/ecgv2/lib/python3.10/site-packages/torch/distributed/run.py", line 918, in main
    run(args)
  File "/home/william/miniconda3/envs/ecgv2/lib/python3.10/site-packages/torch/distributed/run.py", line 909, in run
    elastic_launch(
  File "/home/william/miniconda3/envs/ecgv2/lib/python3.10/site-packages/torch/distributed/launcher/api.py", line 138, in __call__
    return launch_agent(self._config, self._entrypoint, list(args))
  File "/home/william/miniconda3/envs/ecgv2/lib/python3.10/site-packages/torch/distributed/launcher/api.py", line 269, in launch_agent
    raise ChildFailedError(
torch.distributed.elastic.multiprocessing.errors.ChildFailedError: 
============================================================
ecg_bench/test.py FAILED
------------------------------------------------------------
Failures:
  <NO_OTHER_FAILURES>
------------------------------------------------------------
Root Cause (first observed failure):
[0]:
  time      : 2025-10-07_03:11:48
  host      : gpu2
  rank      : 1 (local_rank: 1)
  exitcode  : 1 (pid: 3871026)
  error_file: <N/A>
  traceback : To enable traceback see: https://pytorch.org/docs/stable/elastic/errors.html
============================================================
Test completed!
```

Comments:
The model id `"google/vit-base-patch16-224-in21k"` when used with `AutoModel` will initialize the [`ViTModel`](https://huggingface.co/google/vit-base-patch16-224-in21k/blob/main/config.json#L4). This class when used with [`use_mask_token=True`](https://github.com/huggingface/transformers/blob/caa14e7dabb086f167c14b7eecadc2ba9db25eb6/src/transformers/models/vit/modeling_vit.py#L422) will let you execute the forward pass as you wanted.

Other than that using the `ViTForMaskedImageModeling` would make it specific to your usecase and hence one works and the other complains.

I could not reproduce the same warning or issue with the `ViTForMaskedImageModeling` class. But I will leave this discussion open to the community.


@ariG23498 Thanks for the input. I followed your advice and passed `use_mask_token=True` to the `AutoModel` class but now there is no loss. I tried setting `return_loss=True` as well but this does not seem to do anything. I think the `ViTForMaskedImageModeling` class warning is due to some unsafe things happening in the `ViTForMaskedImageModeling ` source code under distributed setting. I will take a look more when I have time but wanted to see if anyone has experienced it before. 

As I understand this, `ViTModel` is not meant to provide `loss`. The output of the `ViTModel`'s forward is [`BaseModelOutputWithPooling`](https://github.com/huggingface/transformers/blob/3553f0bc23011a5aa459c7b0cd282f570ecd7ac5/src/transformers/modeling_outputs.py#L71) which does not have the `loss` attribute.

On the other hand using the `ViTForMaskedImageModeling` is built to have a head on top of the `ViTModel` which is used to train the model for masked image modeling. I for one did not have any issues with the DDP setup that you have shared when using the `ViTForMaskedImageModeling`

@ariG23498 Thank you for the clarifications. Would you please let me know the torch + nvcc and transformers versions you are using? Maybe I can try updating the versions.

Sure!

Here is the output for `transformers evn`
```shell
- `transformers` version: 4.56.2
- Platform: Linux-5.4.0-216-generic-x86_64-with-glibc2.31
- Python version: 3.12.7
- Huggingface_hub version: 0.35.1
- Safetensors version: 0.6.2
- Accelerate version: 1.10.1
- Accelerate config: 	not found
- DeepSpeed version: not installed
- PyTorch version (accelerator?): 2.8.0+cu128 (CUDA)
- Tensorflow version (GPU?): not installed (NA)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using distributed or parallel set-up in script?: <fill in>
- Using GPU in script?: <fill in>
- GPU type: NVIDIA A100-SXM4-80GB
```

Let me know if updating solves the issue.

it seems at least the transformers and pytorch versions is not the cause of the issue. 

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.