Distributed loading error with from_pretrained for tp_plan is None
Author: kcz358
Pull request: no
Labels: name: bug

### System Info

- `transformers` version: 4.52.0.dev0
- Platform: Linux-5.15.120.bsk.2-amd64-x86_64-with-glibc2.36
- Python version: 3.10.16
- Huggingface_hub version: 0.30.2
- Safetensors version: 0.5.3
- Accelerate version: 1.6.0
- Accelerate config: 	not found
- DeepSpeed version: not installed
- PyTorch version (GPU?): 2.6.0+cu124 (True)
- 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?: True
- Using GPU in script?: True
- GPU type: NVIDIA H100 80GB HBM3



### Who can help?

@zucchini-nlp @amyeroberts @qubvel 

### Information

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

### Tasks

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

### Reproduction

Minimal Code snippet

```python

import torch.distributed as dist
import torch
from transformers import Qwen2_5OmniForConditionalGeneration, Qwen2_5OmniProcessor

def setup():
    dist.init_process_group(backend="nccl")
    local_rank = dist.get_rank() % dist.get_world_size()
    world_size = dist.get_world_size()
    torch.cuda.set_device(local_rank)
    return local_rank, world_size

def cleanup():
    dist.destroy_process_group()


def main():
    local_rank, world_size = setup()
    model = Qwen2_5OmniForConditionalGeneration.from_pretrained("Qwen/Qwen2.5-Omni-7B", device_map = f"cuda:{local_rank}", torch_dtype="auto")
    cleanup()

if __name__ == "__main__":
    main()
```

Then, running with

```bash
torchrun --nproc_per_node="8" \
    --nnodes="1" \
    --node_rank="0" \
    --master_addr="127.0.0.1" \
    --master_port="12345" \
    reproduce_tp_script.py
```


Error is like this:
```bash
[W422 02:54:33.854949743 Utils.hpp:165] Warning: Environment variable NCCL_BLOCKING_WAIT is deprecated; use TORCH_NCCL_BLOCKING_WAIT instead (function operator())
[W422 02:54:33.872220475 Utils.hpp:165] Warning: Environment variable NCCL_BLOCKING_WAIT is deprecated; use TORCH_NCCL_BLOCKING_WAIT instead (function operator())
[W422 02:54:33.046816487 Utils.hpp:165] Warning: Environment variable NCCL_BLOCKING_WAIT is deprecated; use TORCH_NCCL_BLOCKING_WAIT instead (function operator())
[W422 02:54:33.118141019 Utils.hpp:165] Warning: Environment variable NCCL_BLOCKING_WAIT is deprecated; use TORCH_NCCL_BLOCKING_WAIT instead (function operator())
[W422 02:54:33.245757189 Utils.hpp:165] Warning: Environment variable NCCL_BLOCKING_WAIT is deprecated; use TORCH_NCCL_BLOCKING_WAIT instead (function operator())
[W422 02:54:33.367135370 Utils.hpp:165] Warning: Environment variable NCCL_BLOCKING_WAIT is deprecated; use TORCH_NCCL_BLOCKING_WAIT instead (function operator())
[W422 02:54:33.496151955 Utils.hpp:165] Warning: Environment variable NCCL_BLOCKING_WAIT is deprecated; use TORCH_NCCL_BLOCKING_WAIT instead (function operator())
[W422 02:54:33.503352722 Utils.hpp:165] Warning: Environment variable NCCL_BLOCKING_WAIT is deprecated; use TORCH_NCCL_BLOCKING_WAIT instead (function operator())
Unrecognized keys in `rope_scaling` for 'rope_type'='default': {'mrope_section'}
Qwen2_5OmniToken2WavModel does not support eager attention implementation, fall back to sdpa
Unrecognized keys in `rope_scaling` for 'rope_type'='default': {'mrope_section'}

Loading checkpoint shards:   0%|          | 0/5 [00:00<?, ?it/s][rank6]: Traceback (most recent call last):
[rank6]:   File "/opt/tiger/dev/lmms-eval/scripts/reproduce_tp_script.py", line 23, in <module>
[rank6]:     main()
[rank6]:   File "/opt/tiger/dev/lmms-eval/scripts/reproduce_tp_script.py", line 19, in main
[rank6]:     model = Qwen2_5OmniForConditionalGeneration.from_pretrained("Qwen/Qwen2.5-Omni-7B", device_map = f"cuda:{local_rank}", torch_dtype="auto")
[rank6]:   File "/home/tiger/miniconda3/envs/reproduce/lib/python3.10/site-packages/transformers/models/qwen2_5_omni/modeling_qwen2_5_omni.py", line 4421, in from_pretrained
[rank6]:     model = super().from_pretrained(
[rank6]:   File "/home/tiger/miniconda3/envs/reproduce/lib/python3.10/site-packages/transformers/modeling_utils.py", line 282, in _wrapper
[rank6]:     return func(*args, **kwargs)
[rank6]:   File "/home/tiger/miniconda3/envs/reproduce/lib/python3.10/site-packages/transformers/modeling_utils.py", line 4470, in from_pretrained
[rank6]:     ) = cls._load_pretrained_model(
[rank6]:   File "/home/tiger/miniconda3/envs/reproduce/lib/python3.10/site-packages/transformers/modeling_utils.py", line 4869, in _load_pretrained_model
[rank6]:     caching_allocator_warmup(model_to_load, expanded_device_map, factor=2 if hf_quantizer is None else 4)
[rank6]:   File "/home/tiger/miniconda3/envs/reproduce/lib/python3.10/site-packages/transformers/modeling_utils.py", line 5901, in caching_allocator_warmup
[rank6]:     re.compile("|".join([re.escape(plan) for plan in model._tp_plan]))
[rank6]: TypeError: 'NoneType' object is not iterable
Qwen2_5OmniToken2WavModel does not support eager attention implementation, fall back to sdpa

Loading checkpoint shards:   0%|          | 0/5 [00:00<?, ?it/s]
Unrecognized keys in `rope_scaling` for 'rope_type'='default': {'mrope_section'}
Unrecognized keys in `rope_scaling` for 'rope_type'='default': {'mrope_section'}
Qwen2_5OmniToken2WavModel does not support eager attention implementation, fall back to sdpa
Qwen2_5OmniToken2WavModel does not support eager attention implementation, fall back to sdpa

Loading checkpoint shards:   0%|          | 0/5 [00:00<?, ?it/s]Unrecognized keys in `rope_scaling` for 'rope_type'='default': {'mrope_section'}
[rank5]: Traceback (most recent call last):
[rank5]:   File "/opt/tiger/dev/lmms-eval/scripts/reproduce_tp_script.py", line 23, in <module>
[rank5]:     main()
[rank5]:   File "/opt/tiger/dev/lmms-eval/scripts/reproduce_tp_script.py", line 19, in main
[rank5]:     model = Qwen2_5OmniForConditionalGeneration.from_pretrained("Qwen/Qwen2.5-Omni-7B", device_map = f"cuda:{local_rank}", torch_dtype="auto")
[rank5]:   File "/home/tiger/miniconda3/envs/reproduce/lib/python3.10/site-packages/transformers/models/qwen2_5_omni/modeling_qwen2_5_omni.py", line 4421, in from_pretrained
[rank5]:     model = super().from_pretrained(
[rank5]:   File "/home/tiger/miniconda3/envs/reproduce/lib/python3.10/site-packages/transformers/modeling_utils.py", line 282, in _wrapper
[rank5]:     return func(*args, **kwargs)
[rank5]:   File "/home/tiger/miniconda3/envs/reproduce/lib/python3.10/site-packages/transformers/modeling_utils.py", line 4470, in from_pretrained
[rank5]:     ) = cls._load_pretrained_model(
[rank5]:   File "/home/tiger/miniconda3/envs/reproduce/lib/python3.10/site-packages/transformers/modeling_utils.py", line 4869, in _load_pretrained_model
[rank5]:     caching_allocator_warmup(model_to_load, expanded_device_map, factor=2 if hf_quantizer is None else 4)
[rank5]:   File "/home/tiger/miniconda3/envs/reproduce/lib/python3.10/site-packages/transformers/modeling_utils.py", line 5901, in caching_allocator_warmup
[rank5]:     re.compile("|".join([re.escape(plan) for plan in model._tp_plan]))
[rank5]: TypeError: 'NoneType' object is not iterable
Unrecognized keys in `rope_scaling` for 'rope_type'='default': {'mrope_section'}
Qwen2_5OmniToken2WavModel does not support eager attention implementation, fall back to sdpa
Qwen2_5OmniToken2WavModel does not support eager attention implementation, fall back to sdpa

Loading checkpoint shards:   0%|          | 0/5 [00:00<?, ?it/s]

Loading checkpoint shards:   0%|          | 0/5 [00:00<?, ?it/s][rank7]: Traceback (most recent call last):
[rank7]:   File "/opt/tiger/dev/lmms-eval/scripts/reproduce_tp_script.py", line 23, in <module>
[rank7]:     main()
[rank7]:   File "/opt/tiger/dev/lmms-eval/scripts/reproduce_tp_script.py", line 19, in main
[rank7]:     model = Qwen2_5OmniForConditionalGeneration.from_pretrained("Qwen/Qwen2.5-Omni-7B", device_map = f"cuda:{local_rank}", torch_dtype="auto")
[rank7]:   File "/home/tiger/miniconda3/envs/reproduce/lib/python3.10/site-packages/transformers/models/qwen2_5_omni/modeling_qwen2_5_omni.py", line 4421, in from_pretrained
[rank7]:     model = super().from_pretrained(
[rank7]:   File "/home/tiger/miniconda3/envs/reproduce/lib/python3.10/site-packages/transformers/modeling_utils.py", line 282, in _wrapper
[rank7]:     return func(*args, **kwargs)
[rank7]:   File "/home/tiger/miniconda3/envs/reproduce/lib/python3.10/site-packages/transformers/modeling_utils.py", line 4470, in from_pretrained
[rank7]:     ) = cls._load_pretrained_model(
[rank7]:   File "/home/tiger/miniconda3/envs/reproduce/lib/python3.10/site-packages/transformers/modeling_utils.py", line 4869, in _load_pretrained_model
[rank7]:     caching_allocator_warmup(model_to_load, expanded_device_map, factor=2 if hf_quantizer is None else 4)
[rank7]:   File "/home/tiger/miniconda3/envs/reproduce/lib/python3.10/site-packages/transformers/modeling_utils.py", line 5901, in caching_allocator_warmup
[rank7]:     re.compile("|".join([re.escape(plan) for plan in model._tp_plan]))
[rank7]: TypeError: 'NoneType' object is not iterable

Loading checkpoint shards:   0%|          | 0/5 [00:00<?, ?it/s][rank1]: Traceback (most recent call last):
[rank1]:   File "/opt/tiger/dev/lmms-eval/scripts/reproduce_tp_script.py", line 23, in <module>
[rank1]:     main()
[rank1]:   File "/opt/tiger/dev/lmms-eval/scripts/reproduce_tp_script.py", line 19, in main
[rank1]:     model = Qwen2_5OmniForConditionalGeneration.from_pretrained("Qwen/Qwen2.5-Omni-7B", device_map = f"cuda:{local_rank}", torch_dtype="auto")
[rank1]:   File "/home/tiger/miniconda3/envs/reproduce/lib/python3.10/site-packages/transformers/models/qwen2_5_omni/modeling_qwen2_5_omni.py", line 4421, in from_pretrained
[rank1]:     model = super().from_pretrained(
[rank1]:   File "/home/tiger/miniconda3/envs/reproduce/lib/python3.10/site-packages/transformers/modeling_utils.py", line 282, in _wrapper
[rank1]:     return func(*args, **kwargs)
[rank1]:   File "/home/tiger/miniconda3/envs/reproduce/lib/python3.10/site-packages/transformers/modeling_utils.py", line 4470, in from_pretrained
[rank1]:     ) = cls._load_pretrained_model(
[rank1]:   File "/home/tiger/miniconda3/envs/reproduce/lib/python3.10/site-packages/transformers/modeling_utils.py", line 4869, in _load_pretrained_model
[rank1]:     caching_allocator_warmup(model_to_load, expanded_device_map, factor=2 if hf_quantizer is None else 4)
[rank1]:   File "/home/tiger/miniconda3/envs/reproduce/lib/python3.10/site-packages/transformers/modeling_utils.py", line 5901, in caching_allocator_warmup
[rank1]:     re.compile("|".join([re.escape(plan) for plan in model._tp_plan]))
[rank1]: TypeError: 'NoneType' object is not iterable

Loading checkpoint shards:   0%|          | 0/5 [00:00<?, ?it/s]

Loading checkpoint shards:   0%|          | 0/5 [00:00<?, ?it/s][rank2]: Traceback (most recent call last):
[rank2]:   File "/opt/tiger/dev/lmms-eval/scripts/reproduce_tp_script.py", line 23, in <module>
[rank2]:     main()
[rank2]:   File "/opt/tiger/dev/lmms-eval/scripts/reproduce_tp_script.py", line 19, in main
[rank2]:     model = Qwen2_5OmniForConditionalGeneration.from_pretrained("Qwen/Qwen2.5-Omni-7B", device_map = f"cuda:{local_rank}", torch_dtype="auto")
[rank2]:   File "/home/tiger/miniconda3/envs/reproduce/lib/python3.10/site-packages/transformers/models/qwen2_5_omni/modeling_qwen2_5_omni.py", line 4421, in from_pretrained
[rank2]:     model = super().from_pretrained(
[rank2]:   File "/home/tiger/miniconda3/envs/reproduce/lib/python3.10/site-packages/transformers/modeling_utils.py", line 282, in _wrapper
[rank2]:     return func(*args, **kwargs)
[rank2]:   File "/home/tiger/miniconda3/envs/reproduce/lib/python3.10/site-packages/transformers/modeling_utils.py", line 4470, in from_pretrained
[rank2]:     ) = cls._load_pretrained_model(
[rank2]:   File "/home/tiger/miniconda3/envs/reproduce/lib/python3.10/site-packages/transformers/modeling_utils.py", line 4869, in _load_pretrained_model
[rank2]:     caching_allocator_warmup(model_to_load, expanded_device_map, factor=2 if hf_quantizer is None else 4)
[rank2]:   File "/home/tiger/miniconda3/envs/reproduce/lib/python3.10/site-packages/transformers/modeling_utils.py", line 5901, in caching_allocator_warmup
[rank2]:     re.compile("|".join([re.escape(plan) for plan in model._tp_plan]))
[rank2]: TypeError: 'NoneType' object is not iterable

Loading checkpoint shards:   0%|          | 0/5 [00:00<?, ?it/s][rank3]: Traceback (most recent call last):
[rank3]:   File "/opt/tiger/dev/lmms-eval/scripts/reproduce_tp_script.py", line 23, in <module>
[rank3]:     main()
[rank3]:   File "/opt/tiger/dev/lmms-eval/scripts/reproduce_tp_script.py", line 19, in main
[rank3]:     model = Qwen2_5OmniForConditionalGeneration.from_pretrained("Qwen/Qwen2.5-Omni-7B", device_map = f"cuda:{local_rank}", torch_dtype="auto")
[rank3]:   File "/home/tiger/miniconda3/envs/reproduce/lib/python3.10/site-packages/transformers/models/qwen2_5_omni/modeling_qwen2_5_omni.py", line 4421, in from_pretrained
[rank3]:     model = super().from_pretrained(
[rank3]:   File "/home/tiger/miniconda3/envs/reproduce/lib/python3.10/site-packages/transformers/modeling_utils.py", line 282, in _wrapper
[rank3]:     return func(*args, **kwargs)
[rank3]:   File "/home/tiger/miniconda3/envs/reproduce/lib/python3.10/site-packages/transformers/modeling_utils.py", line 4470, in from_pretrained
[rank3]:     ) = cls._load_pretrained_model(
[rank3]:   File "/home/tiger/miniconda3/envs/reproduce/lib/python3.10/site-packages/transformers/modeling_utils.py", line 4869, in _load_pretrained_model
[rank3]:     caching_allocator_warmup(model_to_load, expanded_device_map, factor=2 if hf_quantizer is None else 4)
[rank3]:   File "/home/tiger/miniconda3/envs/reproduce/lib/python3.10/site-packages/transformers/modeling_utils.py", line 5901, in caching_allocator_warmup
[rank3]:     re.compile("|".join([re.escape(plan) for plan in model._tp_plan]))
[rank3]: TypeError: 'NoneType' object is not iterable

Loading checkpoint shards:   0%|          | 0/5 [00:00<?, ?it/s]
W0422 02:54:36.514000 100787 site-packages/torch/distributed/elastic/multiprocessing/api.py:897] Sending process 100799 closing signal SIGTERM
W0422 02:54:36.515000 100787 site-packages/torch/distributed/elastic/multiprocessing/api.py:897] Sending process 100800 closing signal SIGTERM
W0422 02:54:36.516000 100787 site-packages/torch/distributed/elastic/multiprocessing/api.py:897] Sending process 100801 closing signal SIGTERM
W0422 02:54:36.516000 100787 site-packages/torch/distributed/elastic/multiprocessing/api.py:897] Sending process 100802 closing signal SIGTERM
W0422 02:54:36.517000 100787 site-packages/torch/distributed/elastic/multiprocessing/api.py:897] Sending process 100803 closing signal SIGTERM
W0422 02:54:36.517000 100787 site-packages/torch/distributed/elastic/multiprocessing/api.py:897] Sending process 100804 closing signal SIGTERM
W0422 02:54:36.518000 100787 site-packages/torch/distributed/elastic/multiprocessing/api.py:897] Sending process 100806 closing signal SIGTERM
E0422 02:54:37.578000 100787 site-packages/torch/distributed/elastic/multiprocessing/api.py:869] failed (exitcode: 1) local_rank: 6 (pid: 100805) of binary: /home/tiger/miniconda3/envs/reproduce/bin/python3
Traceback (most recent call last):
  File "/home/tiger/miniconda3/envs/reproduce/bin/torchrun", line 8, in <module>
    sys.exit(main())
  File "/home/tiger/miniconda3/envs/reproduce/lib/python3.10/site-packages/torch/distributed/elastic/multiprocessing/errors/__init__.py", line 355, in wrapper
    return f(*args, **kwargs)
  File "/home/tiger/miniconda3/envs/reproduce/lib/python3.10/site-packages/torch/distributed/run.py", line 918, in main
    run(args)
  File "/home/tiger/miniconda3/envs/reproduce/lib/python3.10/site-packages/torch/distributed/run.py", line 909, in run
    elastic_launch(
  File "/home/tiger/miniconda3/envs/reproduce/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/tiger/miniconda3/envs/reproduce/lib/python3.10/site-packages/torch/distributed/launcher/api.py", line 269, in launch_agent
    raise ChildFailedError(
torch.distributed.elastic.multiprocessing.errors.ChildFailedError: 
============================================================
scripts/reproduce_tp_script.py FAILED
------------------------------------------------------------
Failures:
  <NO_OTHER_FAILURES>
------------------------------------------------------------
Root Cause (first observed failure):
[0]:
  time      : 2025-04-22_02:54:36
  host      : n124-174-015.byted.org
  rank      : 6 (local_rank: 6)
  exitcode  : 1 (pid: 100805)
  error_file: <N/A>
  traceback : To enable traceback see: https://pytorch.org/docs/stable/elastic/errors.html
============================================================
```

### Expected behavior

Under the example scripts in [here](Qwen/Qwen2.5-Omni-7B), the model is loaded without torch distributed so everything is great. However, under the distributed cases, with a model with `tp_plan` is None, the `caching_allocator_warmup` here
https://github.com/huggingface/transformers/blob/fee1190601b5d04ec6d3f7f58fd22788d7f3236d/src/transformers/modeling_utils.py#L5874-L5904

will try to iterate over a NoneType _tp_plan and cause this error.

This bugs is potentially exists in all distributed environment with no tp_plan and every model such as Qwen2Audio with a tp_plan does not effect from this bugs. To fix this bug I believe is simple, just add a check before iteration and assign _tp_plan with an empty list would solve the bug.

My temporary solution is to manually hack this line before the `from_pretrained`

```python
Qwen2_5OmniForConditionalGeneration._tp_plan = []
```

Comments:
Hey @kcz358, that's right, thanks for reporting ! Can you open a small pr to fix the issue ? You can just add the `_tp_plan = {}` in the `Qwen2_5OmniForConditionalGeneration` class for example

Sure, no problem. Do you think it is better if I add a small check in the `caching_allocator_warmup` function? Since this can solve the problem for all model with `_tp_plan is None`

Say a small check like this
```python
model._tp_plan = {} if model._tp_plan is None else model._tp_plan
```
in `caching_allocator_warmup`

Yes makes sense 👍 

Hi @MekkCyber , I have opened a PR here https://github.com/huggingface/transformers/pull/37708