| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import pytest |
| 5 | import torch |
| 6 | from transformers import AriaTextConfig |
| 7 | from transformers.models.aria.modeling_aria import AriaTextMoELayer as HFMoE |
| 8 | |
| 9 | from vllm.config import VllmConfig |
| 10 | from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors import ( # noqa: E501 |
| 11 | CompressedTensorsConfig, |
| 12 | ) |
| 13 | from vllm.model_executor.models.aria import AriaTextModel, AriaTextMoELayer |
| 14 | from vllm.model_executor.models.utils import AutoWeightsLoader |
| 15 | |
| 16 | |
| 17 | @pytest.mark.cpu_test |
| 18 | @pytest.mark.usefixtures("dist_init") |
| 19 | @pytest.mark.parametrize( |
| 20 | "checkpoint_name,loaded_name,param_name", |
| 21 | [ |
| 22 | ( |
| 23 | "experts.fc1.weight", |
| 24 | "experts.w13_weight", |
| 25 | "experts.routed_experts.w13_weight", |
| 26 | ), |
| 27 | ( |
| 28 | "experts.fc2.weight", |
| 29 | "experts.w2_weight", |
| 30 | "experts.routed_experts.w2_weight", |
| 31 | ), |
| 32 | ( |
| 33 | "shared_experts.down_proj.weight", |
| 34 | "shared_experts.down_proj.weight", |
| 35 | "shared_experts.down_proj.weight", |
| 36 | ), |
| 37 | ], |
| 38 | ) |
| 39 | def test_aria_expert_weights_load_with_checkpoint_layout( |
| 40 | checkpoint_name: str, loaded_name: str, param_name: str |
| 41 | ): |
| 42 | """Real Aria checkpoint weights must reach the right parameter and layout.""" |
| 43 | config = AriaTextConfig( |
| 44 | hidden_size=32, |
| 45 | intermediate_size=64, |
| 46 | moe_num_experts=2, |
| 47 | moe_topk=1, |
| 48 | moe_num_shared_experts=1, |
| 49 | ) |
| 50 | checkpoint_weight = HFMoE(config).state_dict()[checkpoint_name] |
| 51 | checkpoint_weight.copy_( |
| 52 | torch.arange( |
| 53 | checkpoint_weight.numel(), dtype=checkpoint_weight.dtype |
| 54 | ).reshape_as(checkpoint_weight) |
| 55 | ) |
| 56 | expected = checkpoint_weight.clone() |
| 57 | if checkpoint_name.startswith("experts."): |
| 58 | expected = expected.transpose(-1, -2) |
| 59 | |
| 60 | layer = AriaTextMoELayer(config, quant_config=None, prefix="mlp") |
| 61 | loaded = AutoWeightsLoader(layer).load_weights( |
| 62 | [(checkpoint_name, checkpoint_weight)], mapper=AriaTextModel.hf_to_vllm_mapper |
| 63 | ) |
| 64 | |
| 65 | assert loaded == {loaded_name} |
| 66 | torch.testing.assert_close( |
| 67 | dict(layer.named_parameters())[param_name], expected, rtol=0, atol=0 |
| 68 | ) |
| 69 | |
| 70 | |
| 71 | @pytest.mark.cpu_test |
| 72 | def test_aria_quant_config_renames_expert_modules(): |
| 73 | quant_config = CompressedTensorsConfig( |
| 74 | target_scheme_map={}, |
| 75 | ignore=[ |
| 76 | "model.layers.0.mlp.experts.fc1", |
| 77 | "model.layers.0.mlp.experts.fc2", |
| 78 | ], |
| 79 | quant_format="pack-quantized", |
| 80 | ) |
| 81 | AriaTextModel.__new__( |
| 82 | AriaTextModel, vllm_config=VllmConfig(quant_config=quant_config) |
| 83 | ) |
| 84 | |
| 85 | assert quant_config.ignore == [ |
| 86 | "model.layers.0.mlp.experts.gate_up_proj", |
| 87 | "model.layers.0.mlp.experts.down_proj", |
| 88 | ] |
| 89 | |
| 90 | |
| 91 | @pytest.mark.cpu_test |
| 92 | @pytest.mark.parametrize( |
| 93 | "name,expected", [("fc1", "gate_up_proj"), ("fc2", "down_proj")] |
| 94 | ) |
| 95 | def test_aria_expert_scale_names(name: str, expected: str): |
| 96 | suffixes = ["weight_scale", "input_scale"] |
| 97 | names = [f"model.layers.0.mlp.experts.{name}.{suffix}" for suffix in suffixes] |
| 98 | assert AriaTextModel.hf_to_vllm_mapper.apply_list(names) == [ |
| 99 | f"model.layers.0.mlp.experts.{expected}.{suffix}" for suffix in suffixes |
| 100 | ] |