Llama 4 failing with mps not supported
Author: SethBurkart123
Pull request: no
Labels: name: bug

### System Info

I'm trying to run the new `meta-llama/Llama-4-Maverick-17B-128E-Instruct` model with hf/transformers version 4.5.1 on my Mac M4 Max. However, I'm running into:
```
Loading checkpoint shards: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████| 55/55 [00:07<00:00,  7.65it/s]
Some parameters are on the meta device because they were offloaded to the disk.
Traceback (most recent call last):
  File "/Users/admin/Documents/AI/llama4-distil/main.py", line 38, in <module>
    outputs = model.generate(
              ^^^^^^^^^^^^^^^
  File "/Users/admin/Documents/AI/llama4-distil/.venv/lib/python3.11/site-packages/torch/utils/_contextlib.py", line 116, in decorate_context
    return func(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/Users/admin/Documents/AI/llama4-distil/.venv/lib/python3.11/site-packages/transformers/generation/utils.py", line 2460, in generate
    result = self._sample(
             ^^^^^^^^^^^^^
  File "/Users/admin/Documents/AI/llama4-distil/.venv/lib/python3.11/site-packages/transformers/generation/utils.py", line 3426, in _sample
    outputs = self(**model_inputs, return_dict=True)
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/admin/Documents/AI/llama4-distil/.venv/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1739, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/admin/Documents/AI/llama4-distil/.venv/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1750, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/admin/Documents/AI/llama4-distil/.venv/lib/python3.11/site-packages/accelerate/hooks.py", line 176, in new_forward
    output = module._old_forward(*args, **kwargs)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/admin/Documents/AI/llama4-distil/.venv/lib/python3.11/site-packages/transformers/models/llama4/modeling_llama4.py", line 1761, in forward
    outputs = self.language_model(
              ^^^^^^^^^^^^^^^^^^^^
  File "/Users/admin/Documents/AI/llama4-distil/.venv/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1739, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/admin/Documents/AI/llama4-distil/.venv/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1750, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/admin/Documents/AI/llama4-distil/.venv/lib/python3.11/site-packages/transformers/models/llama4/modeling_llama4.py", line 1015, in forward
    outputs = self.model(
              ^^^^^^^^^^^
  File "/Users/admin/Documents/AI/llama4-distil/.venv/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1739, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/admin/Documents/AI/llama4-distil/.venv/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1750, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/admin/Documents/AI/llama4-distil/.venv/lib/python3.11/site-packages/transformers/models/llama4/modeling_llama4.py", line 669, in forward
    causal_mask, chunk_causal_mask = self._update_causal_mask(
                                     ^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/admin/Documents/AI/llama4-distil/.venv/lib/python3.11/site-packages/transformers/models/llama4/modeling_llama4.py", line 776, in _update_causal_mask
    chunked_attention_mask = make_flex_block_causal_mask(
                             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/admin/Documents/AI/llama4-distil/.venv/lib/python3.11/site-packages/transformers/integrations/flex_attention.py", line 140, in make_flex_block_causal_mask
    return create_block_causal_mask_flex(
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/admin/Documents/AI/llama4-distil/.venv/lib/python3.11/site-packages/torch/nn/attention/flex_attention.py", line 882, in create_block_mask
    return torch.compile(create_block_mask)(
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/admin/Documents/AI/llama4-distil/.venv/lib/python3.11/site-packages/torch/_dynamo/eval_frame.py", line 574, in _fn
    return fn(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^
  File "/Users/admin/Documents/AI/llama4-distil/.venv/lib/python3.11/site-packages/torch/nn/attention/flex_attention.py", line 823, in create_block_mask
    def create_block_mask(
  File "/Users/admin/Documents/AI/llama4-distil/.venv/lib/python3.11/site-packages/torch/_dynamo/convert_frame.py", line 1380, in __call__
    return self._torchdynamo_orig_callable(
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/admin/Documents/AI/llama4-distil/.venv/lib/python3.11/site-packages/torch/_dynamo/convert_frame.py", line 1164, in __call__
    result = self._inner_convert(
             ^^^^^^^^^^^^^^^^^^^^
  File "/Users/admin/Documents/AI/llama4-distil/.venv/lib/python3.11/site-packages/torch/_dynamo/convert_frame.py", line 547, in __call__
    return _compile(
           ^^^^^^^^^
  File "/Users/admin/Documents/AI/llama4-distil/.venv/lib/python3.11/site-packages/torch/_dynamo/convert_frame.py", line 986, in _compile
    guarded_code = compile_inner(code, one_graph, hooks, transform)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/admin/Documents/AI/llama4-distil/.venv/lib/python3.11/site-packages/torch/_dynamo/convert_frame.py", line 715, in compile_inner
    return _compile_inner(code, one_graph, hooks, transform)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/admin/Documents/AI/llama4-distil/.venv/lib/python3.11/site-packages/torch/_utils_internal.py", line 95, in wrapper_function
    return function(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/admin/Documents/AI/llama4-distil/.venv/lib/python3.11/site-packages/torch/_dynamo/convert_frame.py", line 750, in _compile_inner
    out_code = transform_code_object(code, transform)
               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/admin/Documents/AI/llama4-distil/.venv/lib/python3.11/site-packages/torch/_dynamo/bytecode_transformation.py", line 1361, in transform_code_object
    transformations(instructions, code_options)
  File "/Users/admin/Documents/AI/llama4-distil/.venv/lib/python3.11/site-packages/torch/_dynamo/convert_frame.py", line 231, in _fn
    return fn(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^
  File "/Users/admin/Documents/AI/llama4-distil/.venv/lib/python3.11/site-packages/torch/_dynamo/convert_frame.py", line 662, in transform
    tracer.run()
  File "/Users/admin/Documents/AI/llama4-distil/.venv/lib/python3.11/site-packages/torch/_dynamo/symbolic_convert.py", line 2868, in run
    super().run()
  File "/Users/admin/Documents/AI/llama4-distil/.venv/lib/python3.11/site-packages/torch/_dynamo/symbolic_convert.py", line 1052, in run
    while self.step():
          ^^^^^^^^^^^
  File "/Users/admin/Documents/AI/llama4-distil/.venv/lib/python3.11/site-packages/torch/_dynamo/symbolic_convert.py", line 962, in step
    self.dispatch_table[inst.opcode](self, inst)
  File "/Users/admin/Documents/AI/llama4-distil/.venv/lib/python3.11/site-packages/torch/_dynamo/symbolic_convert.py", line 657, in wrapper
    return handle_graph_break(self, inst, speculation.reason)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/admin/Documents/AI/llama4-distil/.venv/lib/python3.11/site-packages/torch/_dynamo/symbolic_convert.py", line 698, in handle_graph_break
    self.output.compile_subgraph(self, reason=reason)
  File "/Users/admin/Documents/AI/llama4-distil/.venv/lib/python3.11/site-packages/torch/_dynamo/output_graph.py", line 1136, in compile_subgraph
    self.compile_and_call_fx_graph(
  File "/Users/admin/Documents/AI/llama4-distil/.venv/lib/python3.11/site-packages/torch/_dynamo/output_graph.py", line 1382, in compile_and_call_fx_graph
    compiled_fn = self.call_user_compiler(gm)
                  ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/admin/Documents/AI/llama4-distil/.venv/lib/python3.11/site-packages/torch/_dynamo/output_graph.py", line 1432, in call_user_compiler
    return self._call_user_compiler(gm)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/admin/Documents/AI/llama4-distil/.venv/lib/python3.11/site-packages/torch/_dynamo/output_graph.py", line 1483, in _call_user_compiler
    raise BackendCompilerFailed(self.compiler_fn, e).with_traceback(
  File "/Users/admin/Documents/AI/llama4-distil/.venv/lib/python3.11/site-packages/torch/_dynamo/output_graph.py", line 1462, in _call_user_compiler
    compiled_fn = compiler_fn(gm, self.example_inputs())
                  ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/admin/Documents/AI/llama4-distil/.venv/lib/python3.11/site-packages/torch/_dynamo/repro/after_dynamo.py", line 130, in __call__
    compiled_gm = compiler_fn(gm, example_inputs)
                  ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/admin/Documents/AI/llama4-distil/.venv/lib/python3.11/site-packages/torch/__init__.py", line 2340, in __call__
    return compile_fx(model_, inputs_, config_patches=self.config)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/admin/Documents/AI/llama4-distil/.venv/lib/python3.11/site-packages/torch/_inductor/compile_fx.py", line 1863, in compile_fx
    return aot_autograd(
           ^^^^^^^^^^^^^
  File "/Users/admin/Documents/AI/llama4-distil/.venv/lib/python3.11/site-packages/torch/_dynamo/backends/common.py", line 83, in __call__
    cg = aot_module_simplified(gm, example_inputs, **self.kwargs)
         ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/admin/Documents/AI/llama4-distil/.venv/lib/python3.11/site-packages/torch/_functorch/aot_autograd.py", line 1155, in aot_module_simplified
    compiled_fn = dispatch_and_compile()
                  ^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/admin/Documents/AI/llama4-distil/.venv/lib/python3.11/site-packages/torch/_functorch/aot_autograd.py", line 1131, in dispatch_and_compile
    compiled_fn, _ = create_aot_dispatcher_function(
                     ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/admin/Documents/AI/llama4-distil/.venv/lib/python3.11/site-packages/torch/_functorch/aot_autograd.py", line 580, in create_aot_dispatcher_function
    return _create_aot_dispatcher_function(
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/admin/Documents/AI/llama4-distil/.venv/lib/python3.11/site-packages/torch/_functorch/aot_autograd.py", line 830, in _create_aot_dispatcher_function
    compiled_fn, fw_metadata = compiler_fn(
                               ^^^^^^^^^^^^
  File "/Users/admin/Documents/AI/llama4-distil/.venv/lib/python3.11/site-packages/torch/_functorch/_aot_autograd/jit_compile_runtime_wrappers.py", line 203, in aot_dispatch_base
    compiled_fw = compiler(fw_module, updated_flat_args)
                  ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/admin/Documents/AI/llama4-distil/.venv/lib/python3.11/site-packages/torch/_functorch/aot_autograd.py", line 489, in __call__
    return self.compiler_fn(gm, example_inputs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/admin/Documents/AI/llama4-distil/.venv/lib/python3.11/site-packages/torch/_inductor/compile_fx.py", line 1741, in fw_compiler_base
    return inner_compile(
           ^^^^^^^^^^^^^^
  File "/Users/admin/Documents/AI/llama4-distil/.venv/lib/python3.11/site-packages/torch/_inductor/compile_fx.py", line 569, in compile_fx_inner
    return wrap_compiler_debug(_compile_fx_inner, compiler_name="inductor")(
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/admin/Documents/AI/llama4-distil/.venv/lib/python3.11/site-packages/torch/_dynamo/repro/after_aot.py", line 102, in debug_wrapper
    inner_compiled_fn = compiler_fn(gm, example_inputs)
                        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/admin/Documents/AI/llama4-distil/.venv/lib/python3.11/site-packages/torch/_inductor/compile_fx.py", line 685, in _compile_fx_inner
    mb_compiled_graph = fx_codegen_and_compile(
                        ^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/admin/Documents/AI/llama4-distil/.venv/lib/python3.11/site-packages/torch/_inductor/compile_fx.py", line 1129, in fx_codegen_and_compile
    return scheme.codegen_and_compile(gm, example_inputs, inputs_to_check, graph_kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/admin/Documents/AI/llama4-distil/.venv/lib/python3.11/site-packages/torch/_inductor/compile_fx.py", line 1044, in codegen_and_compile
    compiled_fn = graph.compile_to_module().call
                  ^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/admin/Documents/AI/llama4-distil/.venv/lib/python3.11/site-packages/torch/_inductor/graph.py", line 2027, in compile_to_module
    return self._compile_to_module()
           ^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/admin/Documents/AI/llama4-distil/.venv/lib/python3.11/site-packages/torch/_inductor/graph.py", line 2033, in _compile_to_module
    self.codegen_with_cpp_wrapper() if self.cpp_wrapper else self.codegen()
                                                             ^^^^^^^^^^^^^^
  File "/Users/admin/Documents/AI/llama4-distil/.venv/lib/python3.11/site-packages/torch/_inductor/graph.py", line 1962, in codegen
    self.init_wrapper_code()
  File "/Users/admin/Documents/AI/llama4-distil/.venv/lib/python3.11/site-packages/torch/_inductor/graph.py", line 1845, in init_wrapper_code
    wrapper_code_gen_cls is not None
torch._dynamo.exc.BackendCompilerFailed: backend='inductor' raised:
AssertionError: Device mps not supported

Set TORCH_LOGS="+dynamo" and TORCHDYNAMO_VERBOSE=1 for more information


You can suppress this exception and fall back to eager by setting:
    import torch._dynamo
    torch._dynamo.config.suppress_errors = True
```

### Who can help?

_No response_

### 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'm just running it with the sample code here:
```python
from transformers import AutoProcessor, Llama4ForConditionalGeneration
import torch

model_id = "meta-llama/Llama-4-Maverick-17B-128E-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])
```

Here's my installed dependencies with uv:
```
Package            Version
------------------ ---------
accelerate         1.6.0
certifi            2025.1.31
charset-normalizer 3.4.1
filelock           3.18.0
fsspec             2025.3.2
huggingface-hub    0.30.1
idna               3.10
jinja2             3.1.6
markupsafe         3.0.2
mpmath             1.3.0
networkx           3.4.2
numpy              2.2.4
packaging          24.2
pillow             11.1.0
psutil             7.0.0
pyyaml             6.0.2
regex              2024.11.6
requests           2.32.3
safetensors        0.5.3
sympy              1.13.1
tokenizers         0.21.1
torch              2.6.0
torchvision        0.21.0
tqdm               4.67.1
transformers       4.51.0
typing-extensions  4.13.1
urllib3            2.3.0
```

### Expected behavior

The model should be run.

Comments:
You're using a quantization method that is not supported by MPS devices (MacOs).  Use the full weights instead.