Fix Aria checkpoint key conversion mapping

[sywangyi]

Fixes Aria checkpoint loading by registering Aria with the existing Llava-style checkpoint conversion mapping.

• model loading (from pretrained, etc): @Cyrilvallez

tests/models/aria/test_modeling_aria.py::AriaForConditionalGenerationIntegrationTest::test_batched_generation [FAILED] 14.94s


============================================================================ FAILURES =============================================================================
_______________________________________________ AriaForConditionalGenerationIntegrationTest.test_batched_generation _______________________________________________

self = <tests.models.aria.test_modeling_aria.AriaForConditionalGenerationIntegrationTest testMethod=test_batched_generation>

    @require_torch_large_accelerator
    @require_vision
    @require_bitsandbytes
    def test_batched_generation(self):
        # Skip multihead_attn for 4bit because MHA will read the original weight without dequantize.
        # See https://github.com/huggingface/transformers/pull/37444#discussion_r2045852538.
        model = AriaForConditionalGeneration.from_pretrained(
            "rhymes-ai/Aria",
            quantization_config=BitsAndBytesConfig(load_in_4bit=True, llm_int8_skip_modules=["multihead_attn"]),
        )
        processor = AutoProcessor.from_pretrained("rhymes-ai/Aria")

        prompt1 = "<image>\n<image>\nUSER: What's the difference of two images?\nASSISTANT:"
        prompt2 = "<image>\nUSER: Describe the image.\nASSISTANT:"
        prompt3 = "<image>\nUSER: Describe the image.\nASSISTANT:"
        url1 = "https://images.unsplash.com/photo-1552053831-71594a27632d?q=80&w=3062&auto=format&fit=crop&ixlib=rb-4.0.3&ixid=M3wxMjA3fDB8MHxwaG90by1wYWdlfHx8fGVufDB8fHx8fA%3D%3D"
        url2 = "https://images.unsplash.com/photo-1617258683320-61900b281ced?q=80&w=3087&auto=format&fit=crop&ixlib=rb-4.0.3&ixid=M3wxMjA3fDB8MHxwaG90by1wYWdlfHx8fGVufDB8fHx8fA%3D%3D"
        image1 = Image.open(requests.get(url1, stream=True).raw)
        image2 = Image.open(requests.get(url2, stream=True).raw)

        # Create inputs
        messages = [
            {
                "role": "user",
                "content": [
                    {"type": "image"},
                    {"type": "text", "text": prompt1},
                    {"type": "image"},
                    {"type": "text", "text": prompt2},
                ],
            },
            {
                "role": "user",
                "content": [
                    {"type": "image"},
                    {"type": "text", "text": prompt3},
                ],
            },
        ]

        prompts = [processor.apply_chat_template([message], add_generation_prompt=True) for message in messages]
        images = [[image1, image2], [image2]]
        inputs = processor(text=prompts, images=images, padding=True, return_tensors="pt").to(
            device=model.device, dtype=model.dtype
        )

        EXPECTED_OUTPUTS = Expectations(
            {
                ("cpu", None): [
                    "<|im_start|>user\n<fim_prefix><fim_suffix> <image>\n <image>\n USER: What's the difference of two images?\n ASSISTANT:<fim_prefix><fim_suffix> <image>\n USER: Describe the image.\n ASSISTANT:<|im_end|>\n <|im_start|>assistant\n The first image features a cute, light-colored puppy sitting on a paved surface with",
                    "<|im_start|>user\n<fim_prefix><fim_suffix> <image>\n USER: Describe the image.\n ASSISTANT:<|im_end|>\n <|im_start|>assistant\n The image shows a young alpaca standing on a grassy hill. The alpaca has",
                ],
                ("cuda", None): [
                    "<|im_start|>user\n<fim_prefix><fim_suffix> <image>\n <image>\n USER: What's the difference of two images?\n ASSISTANT:<fim_prefix><fim_suffix> <image>\n USER: Describe the image.\n ASSISTANT:<|im_end|>\n <|im_start|>assistant\n The first image features a cute, light-colored puppy sitting on a paved surface with",
                    "<|im_start|>user\n<fim_prefix><fim_suffix> <image>\n USER: Describe the image.\n ASSISTANT:<|im_end|>\n <|im_start|>assistant\n The image shows a young alpaca standing on a patch of ground with some dry grass. The",
                ],
                ("xpu", 3): [
                    "<|im_start|>user\n<fim_prefix><fim_suffix> <image>\n <image>\n USER: What's the difference of two images?\n ASSISTANT:<fim_prefix><fim_suffix> <image>\n USER: Describe the image.\n ASSISTANT:<|im_end|>\n <|im_start|>assistant\n The first image features a cute, light-colored puppy sitting on a paved surface with",
                    "<|im_start|>user\n<fim_prefix><fim_suffix> <image>\n USER: Describe the image.\n ASSISTANT:<|im_end|>\n <|im_start|>assistant\n The image shows a young alpaca standing on a patch of ground with some dry grass. The",
                ],
                ("rocm", (9, 5)): [
                    "<|im_start|>user\n<fim_prefix><fim_suffix> <image>\n <image>\n USER: What's the difference of two images?\n ASSISTANT:<fim_prefix><fim_suffix> <image>\n USER: Describe the image.\n ASSISTANT:<|im_end|>\n <|im_start|>assistant\n The first image shows a cute golden retriever puppy sitting on a paved surface with a stick",
                    '<|im_start|>user\n<fim_prefix><fim_suffix> <image>\n USER: Describe the image.\n ASSISTANT:<|im_end|>\n <|im_start|>assistant\n The image shows a young llama standing on a patch of ground with some dry grass and dirt. The'
                ],
            }
        )  # fmt: skip
        EXPECTED_OUTPUT = EXPECTED_OUTPUTS.get_expectation()
>       generate_ids = model.generate(**inputs, max_new_tokens=20)
                       ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^

tests/models/aria/test_modeling_aria.py:484:
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _
/opt/venv/lib/python3.12/site-packages/torch/utils/_contextlib.py:124: in decorate_context
    return func(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^
src/transformers/generation/utils.py:2645: in generate
    result = decoding_method(
src/transformers/generation/utils.py:2845: in _sample
    outputs = self._prefill(
src/transformers/generation/utils.py:3896: in _prefill
    return self(**model_inputs, return_dict=True)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
/opt/venv/lib/python3.12/site-packages/torch/nn/modules/module.py:1778: in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
/opt/venv/lib/python3.12/site-packages/torch/nn/modules/module.py:1789: in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
src/transformers/utils/generic.py:911: in wrapper
    output = func(self, *args, **kwargs)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^
src/transformers/models/aria/modeling_aria.py:1128: in forward
    outputs = self.model(
/opt/venv/lib/python3.12/site-packages/torch/nn/modules/module.py:1778: in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
/opt/venv/lib/python3.12/site-packages/torch/nn/modules/module.py:1789: in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
src/transformers/utils/generic.py:911: in wrapper
    output = func(self, *args, **kwargs)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^
src/transformers/models/aria/modeling_aria.py:982: in forward
    outputs = self.language_model(
/opt/venv/lib/python3.12/site-packages/torch/nn/modules/module.py:1778: in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
/opt/venv/lib/python3.12/site-packages/torch/nn/modules/module.py:1789: in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
src/transformers/utils/generic.py:1040: in wrapper
    output = func(self, *args, **kwargs)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^
src/transformers/utils/output_capturing.py:252: in wrapper
    outputs = func(self, *args, **kwargs)
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^
src/transformers/models/aria/modeling_aria.py:742: in forward
    hidden_states = decoder_layer(
src/transformers/modeling_layers.py:93: in __call__
    return super().__call__(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
/opt/venv/lib/python3.12/site-packages/torch/nn/modules/module.py:1778: in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
/opt/venv/lib/python3.12/site-packages/torch/nn/modules/module.py:1789: in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
src/transformers/models/aria/modeling_aria.py:553: in forward
    hidden_states, _ = self.self_attn(
/opt/venv/lib/python3.12/site-packages/torch/nn/modules/module.py:1778: in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
/opt/venv/lib/python3.12/site-packages/torch/nn/modules/module.py:1789: in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
src/transformers/models/aria/modeling_aria.py:488: in forward
    query_states = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^
/opt/venv/lib/python3.12/site-packages/torch/nn/modules/module.py:1778: in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
/opt/venv/lib/python3.12/site-packages/torch/nn/modules/module.py:1789: in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
/opt/venv/lib/python3.12/site-packages/bitsandbytes/nn/modules.py:529: in forward
    fix_4bit_weight_quant_state_from_module(self)
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _

module = Linear4bit(in_features=2560, out_features=2560, bias=False)

    def fix_4bit_weight_quant_state_from_module(module: Union["Embedding4bit", "Linear4bit"]):
        if getattr(module.weight, "quant_state", None) is not None:
            return

        if getattr(module, "quant_state", None) is None:
            warnings.warn(
                "FP4 quantization state not initialized. Please call .cuda() or .to(device) on the LinearFP4 layer first.",
            )

        # the quant state got lost when the parameter got converted. This happens for example for fsdp
        # since we registered the module, we can recover the state here
>       assert module.weight.shape[1] == 1
               ^^^^^^^^^^^^^^^^^^^^^^^^^^^
E       AssertionError

/opt/venv/lib/python3.12/site-packages/bitsandbytes/nn/modules.py:415: AssertionError
------------------------------------------------------------------------ Captured log call ------------------------------------------------------------------------
WARNING  transformers.modeling_utils:loading_report.py:269 AriaForConditionalGeneration LOAD REPORT from: rhymes-ai/Aria
Key                                                                      | Status     |
-------------------------------------------------------------------------+------------+-
language_model.model.layers.{0...27}.mlp.experts.fc2.weight              | UNEXPECTED |
language_model.model.layers.{0...27}.mlp.shared_experts.up_proj.weight   | UNEXPECTED |
language_model.model.layers.{0...27}.post_attention_layernorm.weight     | UNEXPECTED |
language_model.model.layers.{0...27}.mlp.shared_experts.down_proj.weight | UNEXPECTED |
language_model.model.layers.{0...27}.mlp.experts.fc1.weight              | UNEXPECTED |
language_model.model.layers.{0...27}.self_attn.v_proj.weight             | UNEXPECTED |
language_model.model.layers.{0...27}.mlp.router.weight                   | UNEXPECTED |
language_model.model.layers.{0...27}.self_attn.k_proj.weight             | UNEXPECTED |
language_model.model.layers.{0...27}.mlp.shared_experts.gate_proj.weight | UNEXPECTED |
language_model.model.layers.{0...27}.input_layernorm.weight              | UNEXPECTED |
language_model.model.layers.{0...27}.self_attn.q_proj.weight             | UNEXPECTED |
language_model.model.layers.{0...27}.self_attn.o_proj.weight             | UNEXPECTED |
language_model.lm_head.weight                                            | UNEXPECTED |
language_model.model.norm.weight                                         | UNEXPECTED |
language_model.model.embed_tokens.weight                                 | UNEXPECTED |
model.language_model.layers.{0...27}.self_attn.q_proj.weight             | MISSING    |
model.language_model.layers.{0...27}.self_attn.k_proj.weight             | MISSING    |
model.language_model.layers.{0...27}.mlp.shared_experts.up_proj.weight   | MISSING    |
model.language_model.layers.{0...27}.post_attention_layernorm.weight     | MISSING    |
model.language_model.layers.{0...27}.mlp.experts.fc1.weight              | MISSING    |
model.language_model.layers.{0...27}.mlp.shared_experts.down_proj.weight | MISSING    |
model.language_model.layers.{0...27}.mlp.experts.fc2.weight              | MISSING    |
model.language_model.layers.{0...27}.mlp.router.weight                   | MISSING    |
model.language_model.layers.{0...27}.self_attn.o_proj.weight             | MISSING    |
model.language_model.layers.{0...27}.mlp.shared_experts.gate_proj.weight | MISSING    |
model.language_model.layers.{0...27}.self_attn.v_proj.weight             | MISSING    |
model.language_model.layers.{0...27}.input_layernorm.weight              | MISSING    |
model.language_model.norm.weight                                         | MISSING    |
model.language_model.embed_tokens.weight                                 | MISSING    |
lm_head.weight                                                           | MISSING    |

[sywangyi]

@ydshieh @IlyasMoutawwakil please help review as well

[github-actions]

[For maintainers] Suggested jobs to run (before merge)

run-slow: aria

[github-actions]

---

CI recap

Dashboard: View test results in Grafana
Latest run: 28919855953:2
Result: success | Jobs: 15 | Tests: 172,005 | Failures: 0 | Duration: 18h 14m

[zucchini-nlp]

Wow, how did we even miss this. Thanks @sywangyi for fixing, aria is not used much so we didn't see any GH issues

I suppose Aria doesn't run on nightly since it is too big @ydshieh? Do you think we need a minified version with less layers, still deterministic logits but output is garbage. That is what we did with diffusion gemma

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[sywangyi]

could you help merge it @zucchini-nlp ?

[zucchini-nlp]

oh yeah, let's merge, got lost in my GH