[llama4] Inference with multiple GPU torch.distributed.DistStoreError
Author: GilgameshD
Pull request: no
Labels: name: bug

### System Info

- `transformers` version: 4.52.0.dev0
- Platform: Linux-5.15.0-1032-oracle-x86_64-with-glibc2.31
- Python version: 3.10.16
- Huggingface_hub version: 0.30.1
- Safetensors version: 0.5.2
- Accelerate version: 1.3.0
- Accelerate config: 	not found
- DeepSpeed version: not installed
- PyTorch version (GPU?): 2.6.0+cu124 (False)
- 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?: Yes
- Using GPU in script?: Yes
- GPU type: NVIDIA A100-SXM4-80GB x 8

### Who can help?

@ArthurZucker 

### Information

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

### Tasks

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

### Reproduction

I am using the official example of llama4 with the latest transformer code:

```
from transformers import AutoProcessor, Llama4ForConditionalGeneration
import torch


model_id = "meta-llama/Llama-4-Scout-17B-16E-Instruct"

processor = AutoProcessor.from_pretrained(model_id)
model = Llama4ForConditionalGeneration.from_pretrained(
    model_id,
    attn_implementation="flex_attention",
    device_map="auto",
    torch_dtype=torch.bfloat16,
)

url1 = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/0052a70beed5bf71b92610a43a52df6d286cd5f3/diffusers/rabbit.jpg"
url2 = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/datasets/cat_style_layout.png"
messages = [
    {
        "role": "user",
        "content": [
            {"type": "image", "url": url1},
            {"type": "image", "url": url2},
            {
                "type": "text",
                "text": "Can you describe how these two images are similar, and how they differ?",
            },
        ],
    },
]

inputs = processor.apply_chat_template(
    messages,
    add_generation_prompt=True,
    tokenize=True,
    return_dict=True,
    return_tensors="pt",
).to(model.device)

outputs = model.generate(
    **inputs,
    max_new_tokens=256,
)

response = processor.batch_decode(outputs[:, inputs["input_ids"].shape[-1] :])[0]
print(response)
print(outputs[0])
```

I got the following error:
```
torch.distributed.DistStoreError: Timed out after 601 seconds waiting for clients. 1/8 clients joined. 
OSError: We tried to initialize torch.distributed for you, but it failed, makesure you init torch distributed in your script to use `tp_plan='auto'`
```

### Expected behavior

Text output from the model

Comments:
Can you show how you ran the script? 

I got the above error when I used `python script.py`.
Then I use `accelerate launch script.py` to run the script and it seems the distributed part works. But got a new error:
```
[rank0]: ValueError: block_mask was created for block_mask.shape=(1, 1, 4959, tensor(8192, device='cuda:0')) but got q_len=4959 and kv_len=4959. As the block mask was created for a larger length than you're using it for, you can either 1. create a new block mask with the correct length, or 2. 'adjust' the existing block mask to the correct length by calling block_mask._adjust(q_len, kv_len). This essentially 'crops' the block mask to the upper left corner, which does not work for all mask_mods! 
```

`torch.distributed.DistStoreError` should only happen with torchrun 


The new error is probably related to #37351 - it's WIP :)

```diff
outputs = model.generate(
    **inputs,
    max_new_tokens=256,
+    cache_implementation="hybrid"
)
```


Also `model.eval()`

Fixed by #37353

Thank you both!

Adding `cache_implementation="hybrid"` and running with `accelerate launch script.py` works. But I still don't know why using `python script.py` leads to `torch.distributed.DistStoreError`. Maybe due to some system-level settings 

Hello @ArthurZucker : Just a quick question, for inference with llama-4, do we need this line or is the default cache option sufficient? Thanks.
cache_implementation="hybrid"

The default is now `"hybrid"` 😉