| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | from types import SimpleNamespace |
| 4 | |
| 5 | import pytest |
| 6 | import torch |
| 7 | from torch import nn |
| 8 | |
| 9 | import vllm.model_executor.models.k2_horizon as k2_horizon |
| 10 | from tests.quantization.utils import is_quant_method_supported |
| 11 | from vllm.config import VllmConfig, set_current_vllm_config |
| 12 | from vllm.distributed import ( |
| 13 | cleanup_dist_env_and_memory, |
| 14 | init_distributed_environment, |
| 15 | initialize_model_parallel, |
| 16 | ) |
| 17 | from vllm.model_executor.layers.linear import ( |
| 18 | LinearBase, |
| 19 | UnquantizedLinearMethod, |
| 20 | ) |
| 21 | from vllm.model_executor.layers.quantization.fp8 import Fp8Config, Fp8LinearMethod |
| 22 | from vllm.model_executor.layers.rotary_embedding import get_rope |
| 23 | from vllm.model_executor.models.k2_horizon import ( |
| 24 | K2HorizonDecoderLayer, |
| 25 | K2HorizonModel, |
| 26 | K2HorizonRMSNorm, |
| 27 | _rope_weight_perm, |
| 28 | apply_partial_rope, |
| 29 | is_rope_weights_folding_supported, |
| 30 | ) |
| 31 | from vllm.platforms import current_platform |
| 32 | from vllm.utils.network_utils import get_open_port |
| 33 | from vllm.utils.torch_utils import set_default_torch_dtype, set_random_seed |
| 34 | |
| 35 | DEVICE = current_platform.device_type |
| 36 | DTYPE = torch.bfloat16 |
| 37 | SHAPES = [ |
| 38 | (128, 64), |
| 39 | (128, 128), |
| 40 | (192, 128), |
| 41 | (96, 32), |
| 42 | ] |
| 43 | |
| 44 | |
| 45 | def get_gptj_rope( |
| 46 | positions: torch.Tensor, |
| 47 | q: torch.Tensor, |
| 48 | k: torch.Tensor, |
| 49 | num_heads: int, |
| 50 | num_kv_heads: int, |
| 51 | head_dim: int, |
| 52 | rope_head_dim: int, |
| 53 | max_position: int, |
| 54 | base: float, |
| 55 | dtype: torch.dtype, |
| 56 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 57 | rotary_emb = get_rope( |
| 58 | rope_head_dim, |
| 59 | max_position=max_position, |
| 60 | is_neox_style=True, |
| 61 | rope_parameters={"rope_theta": base}, |
| 62 | dtype=dtype, |
| 63 | ) |
| 64 | return apply_partial_rope( |
| 65 | rotary_emb, |
| 66 | positions, |
| 67 | q, |
| 68 | k, |
| 69 | num_heads, |
| 70 | num_kv_heads, |
| 71 | head_dim, |
| 72 | rope_head_dim, |
| 73 | fold_rope_weights=False, |
| 74 | ) |
| 75 | |
| 76 | |
| 77 | def _apply_head_perm( |
| 78 | x: torch.Tensor, num: int, head_dim: int, idx: torch.Tensor |
| 79 | ) -> torch.Tensor: |
| 80 | return ( |
| 81 | x.reshape(*x.shape[:-1], num, head_dim)[..., idx].reshape(*x.shape).contiguous() |
| 82 | ) |
| 83 | |
| 84 | |
| 85 | def get_neox_rope( |
| 86 | positions: torch.Tensor, |
| 87 | q: torch.Tensor, |
| 88 | k: torch.Tensor, |
| 89 | num_heads: int, |
| 90 | num_kv_heads: int, |
| 91 | head_dim: int, |
| 92 | rope_head_dim: int, |
| 93 | max_position: int, |
| 94 | base: float, |
| 95 | dtype: torch.dtype, |
| 96 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 97 | rotary_emb = get_rope( |
| 98 | head_dim, |
| 99 | max_position=max_position, |
| 100 | is_neox_style=True, |
| 101 | rope_parameters={"rope_theta": base, "rope_dim": rope_head_dim}, |
| 102 | dtype=dtype, |
| 103 | ) |
| 104 | return apply_partial_rope( |
| 105 | rotary_emb, |
| 106 | positions, |
| 107 | q, |
| 108 | k, |
| 109 | num_heads, |
| 110 | num_kv_heads, |
| 111 | head_dim, |
| 112 | rope_head_dim, |
| 113 | fold_rope_weights=True, |
| 114 | ) |
| 115 | |
| 116 | |
| 117 | def fold_qk_proj_weight( |
| 118 | weight: torch.Tensor, head_dim: int, idx: torch.Tensor |
| 119 | ) -> torch.Tensor: |
| 120 | hidden = weight.shape[-1] |
| 121 | return weight.view(-1, head_dim, hidden)[:, idx, :].reshape(-1, hidden).contiguous() |
| 122 | |
| 123 | |
| 124 | def fold_qk_channels( |
| 125 | vec: torch.Tensor, head_dim: int, idx: torch.Tensor |
| 126 | ) -> torch.Tensor: |
| 127 | return vec.view(-1, head_dim)[:, idx].reshape(-1).contiguous() |
| 128 | |
| 129 | |
| 130 | @pytest.mark.parametrize("head_dim,rope_head_dim", SHAPES) |
| 131 | def test_rope_fold_matches_permute_then_rope( |
| 132 | head_dim, rope_head_dim, default_vllm_config |
| 133 | ): |
| 134 | dtype = DTYPE |
| 135 | set_random_seed(0) |
| 136 | num_tokens = 17 |
| 137 | num_heads = 8 |
| 138 | num_kv_heads = 2 |
| 139 | hidden_size = 512 |
| 140 | max_position = 4096 |
| 141 | base = 10000.0 |
| 142 | |
| 143 | torch.set_default_dtype(dtype) |
| 144 | |
| 145 | positions = torch.randint(0, max_position, (num_tokens,), device=DEVICE) |
| 146 | |
| 147 | scale = hidden_size**-0.5 |
| 148 | x = torch.randn(num_tokens, hidden_size, dtype=dtype, device=DEVICE) |
| 149 | w_q = ( |
| 150 | torch.randn(num_heads * head_dim, hidden_size, dtype=dtype, device=DEVICE) |
| 151 | * scale |
| 152 | ) |
| 153 | w_k = ( |
| 154 | torch.randn(num_kv_heads * head_dim, hidden_size, dtype=dtype, device=DEVICE) |
| 155 | * scale |
| 156 | ) |
| 157 | |
| 158 | idx = _rope_weight_perm(head_dim, rope_head_dim).to(DEVICE) |
| 159 | |
| 160 | q = torch.nn.functional.linear(x, w_q) |
| 161 | k = torch.nn.functional.linear(x, w_k) |
| 162 | q_gptj, k_gptj = get_gptj_rope( |
| 163 | positions, |
| 164 | q, |
| 165 | k, |
| 166 | num_heads, |
| 167 | num_kv_heads, |
| 168 | head_dim, |
| 169 | rope_head_dim, |
| 170 | max_position, |
| 171 | base, |
| 172 | dtype, |
| 173 | ) |
| 174 | |
| 175 | q_folded = torch.nn.functional.linear(x, fold_qk_proj_weight(w_q, head_dim, idx)) |
| 176 | k_folded = torch.nn.functional.linear(x, fold_qk_proj_weight(w_k, head_dim, idx)) |
| 177 | q_neox, k_neox = get_neox_rope( |
| 178 | positions, |
| 179 | q_folded, |
| 180 | k_folded, |
| 181 | num_heads, |
| 182 | num_kv_heads, |
| 183 | head_dim, |
| 184 | rope_head_dim, |
| 185 | max_position, |
| 186 | base, |
| 187 | dtype, |
| 188 | ) |
| 189 | |
| 190 | q_gptj_perm = _apply_head_perm(q_gptj, num_heads, head_dim, idx) |
| 191 | k_gptj_perm = _apply_head_perm(k_gptj, num_kv_heads, head_dim, idx) |
| 192 | |
| 193 | atol, rtol = 1e-2, 1e-2 |
| 194 | |
| 195 | # Two complementary checks. First, an element-wise check that the folded |
| 196 | # NeoX path equals the GPT-J path up to the channel permutation P |
| 197 | # (neox_rope(Pq, Pk) == P @ gptj_rope(q, k)), hence permuting q_gptj/k_gptj. |
| 198 | # Second, the scores check below: P cancels in q @ k.T. |
| 199 | torch.testing.assert_close(q_neox, q_gptj_perm, atol=atol, rtol=rtol) |
| 200 | torch.testing.assert_close(k_neox, k_gptj_perm, atol=atol, rtol=rtol) |
| 201 | |
| 202 | def scores(qt, kt): |
| 203 | qh = qt.reshape(num_tokens, num_heads, head_dim) |
| 204 | kh = kt.reshape(num_tokens, num_kv_heads, head_dim) |
| 205 | g = num_heads // num_kv_heads |
| 206 | qh = qh.reshape(num_tokens, num_kv_heads, g, head_dim).float() |
| 207 | kh = kh.float() |
| 208 | return torch.einsum("tkgd,skd->tksg", qh, kh) |
| 209 | |
| 210 | torch.testing.assert_close( |
| 211 | scores(q_neox, k_neox), scores(q_gptj, k_gptj), atol=atol, rtol=rtol |
| 212 | ) |
| 213 | |
| 214 | |
| 215 | @pytest.mark.parametrize("head_dim,rope_head_dim", SHAPES) |
| 216 | @pytest.mark.parametrize("num_heads", [8, 2]) |
| 217 | def test_qk_bias_fold_matches_permute( |
| 218 | head_dim, rope_head_dim, num_heads, default_vllm_config |
| 219 | ): |
| 220 | dtype = DTYPE |
| 221 | set_random_seed(0) |
| 222 | num_tokens = 17 |
| 223 | hidden_size = 512 |
| 224 | torch.set_default_dtype(dtype) |
| 225 | |
| 226 | scale = hidden_size**-0.5 |
| 227 | x = torch.randn(num_tokens, hidden_size, dtype=dtype, device=DEVICE) |
| 228 | w = ( |
| 229 | torch.randn(num_heads * head_dim, hidden_size, dtype=dtype, device=DEVICE) |
| 230 | * scale |
| 231 | ) |
| 232 | b = torch.randn(num_heads * head_dim, dtype=dtype, device=DEVICE) |
| 233 | |
| 234 | idx = _rope_weight_perm(head_dim, rope_head_dim).to(DEVICE) |
| 235 | |
| 236 | ref = _apply_head_perm( |
| 237 | torch.nn.functional.linear(x, w, b), num_heads, head_dim, idx |
| 238 | ) |
| 239 | folded = torch.nn.functional.linear( |
| 240 | x, |
| 241 | fold_qk_proj_weight(w, head_dim, idx), |
| 242 | fold_qk_channels(b, head_dim, idx), |
| 243 | ) |
| 244 | |
| 245 | atol, rtol = 1e-2, 1e-2 |
| 246 | torch.testing.assert_close(folded, ref, atol=atol, rtol=rtol) |
| 247 | |
| 248 | |
| 249 | @pytest.mark.skipif( |
| 250 | not current_platform.is_cuda_alike(), reason="requires a GPU (CUDA/ROCm)" |
| 251 | ) |
| 252 | @pytest.mark.parametrize("head_dim,rope_head_dim", SHAPES) |
| 253 | @pytest.mark.parametrize("num_heads", [8, 2]) |
| 254 | def test_qk_norm_scale_fold_matches_permute( |
| 255 | head_dim, rope_head_dim, num_heads, default_vllm_config |
| 256 | ): |
| 257 | dtype = DTYPE |
| 258 | set_random_seed(0) |
| 259 | num_tokens = 13 |
| 260 | torch.set_default_dtype(dtype) |
| 261 | |
| 262 | hidden = num_heads * head_dim |
| 263 | idx = _rope_weight_perm(head_dim, rope_head_dim).to(DEVICE) |
| 264 | |
| 265 | x = torch.randn(num_tokens, hidden, dtype=dtype, device=DEVICE) |
| 266 | |
| 267 | norm = K2HorizonRMSNorm(hidden_size=hidden, n_groups=num_heads).to(DEVICE) |
| 268 | with torch.no_grad(): |
| 269 | norm.weight.copy_(torch.randn(hidden, dtype=dtype, device=DEVICE)) |
| 270 | |
| 271 | y_ref = _apply_head_perm(norm(x.clone()), num_heads, head_dim, idx) |
| 272 | |
| 273 | folded_norm = K2HorizonRMSNorm(hidden_size=hidden, n_groups=num_heads).to(DEVICE) |
| 274 | with torch.no_grad(): |
| 275 | folded_norm.weight.copy_(fold_qk_channels(norm.weight, head_dim, idx)) |
| 276 | # Permute x to emulate the folded q_proj/k_proj: at runtime the norm sees the |
| 277 | # already-permuted projection output (P @ x), not the raw activations. |
| 278 | y_folded = folded_norm(_apply_head_perm(x.clone(), num_heads, head_dim, idx)) |
| 279 | |
| 280 | atol, rtol = 1e-2, 1e-2 |
| 281 | torch.testing.assert_close(y_folded, y_ref, atol=atol, rtol=rtol) |
| 282 | |
| 283 | |
| 284 | class _Fp8QKVOnlyConfig(Fp8Config): |
| 285 | def get_quant_method(self, layer, prefix): |
| 286 | if isinstance(layer, LinearBase) and not ( |
| 287 | prefix.endswith("qkv_proj") or prefix.endswith("qk_proj") |
| 288 | ): |
| 289 | return UnquantizedLinearMethod() |
| 290 | return super().get_quant_method(layer, prefix) |
| 291 | |
| 292 | |
| 293 | class _DummyAttention(nn.Module): |
| 294 | def __init__(self, *args, **kwargs) -> None: |
| 295 | super().__init__() |
| 296 | |
| 297 | def forward(self, q, k, v): |
| 298 | return q |
| 299 | |
| 300 | |
| 301 | @pytest.fixture |
| 302 | def dist_init(): |
| 303 | with set_current_vllm_config(VllmConfig()): |
| 304 | init_distributed_environment( |
| 305 | world_size=1, |
| 306 | rank=0, |
| 307 | distributed_init_method=f"tcp://127.0.0.1:{get_open_port()}", |
| 308 | local_rank=0, |
| 309 | backend="gloo", |
| 310 | ) |
| 311 | initialize_model_parallel(1, 1) |
| 312 | yield |
| 313 | cleanup_dist_env_and_memory() |
| 314 | |
| 315 | |
| 316 | def _make_hf_config( |
| 317 | *, |
| 318 | hidden_size: int, |
| 319 | num_heads: int, |
| 320 | num_kv_heads: int, |
| 321 | head_dim: int, |
| 322 | rope_head_dim: int, |
| 323 | attention_bias: bool, |
| 324 | query_key_norm: bool, |
| 325 | ) -> SimpleNamespace: |
| 326 | return SimpleNamespace( |
| 327 | hidden_size=hidden_size, |
| 328 | max_position_embeddings=4096, |
| 329 | dual_chunk_attention_config=None, |
| 330 | mlp_only_layers=[], |
| 331 | num_experts=0, # dense layer: avoids the MoE path |
| 332 | decoder_sparse_step=1, |
| 333 | mova_num_experts=0, |
| 334 | num_attention_heads=num_heads, |
| 335 | num_key_value_heads=num_kv_heads, |
| 336 | rope_parameters={"rope_theta": 10000.0}, |
| 337 | rope_head_dim=rope_head_dim, |
| 338 | query_key_norm=query_key_norm, |
| 339 | rms_norm_eps=1e-6, |
| 340 | attention_bias=attention_bias, |
| 341 | head_dim=head_dim, |
| 342 | attention_gate_func=None, |
| 343 | intermediate_size=2 * hidden_size, |
| 344 | hidden_act="silu", |
| 345 | layernorm_num_groups=1, |
| 346 | vocab_size=1000, |
| 347 | num_hidden_layers=1, |
| 348 | tie_word_embeddings=False, |
| 349 | ) |
| 350 | |
| 351 | |
| 352 | def _build_single_layer_model( |
| 353 | monkeypatch, hf_config: SimpleNamespace, *, quant_config=None |
| 354 | ) -> K2HorizonModel: |
| 355 | monkeypatch.setattr(k2_horizon, "Attention", _DummyAttention) |
| 356 | |
| 357 | vllm_config = VllmConfig() |
| 358 | vllm_config.model_config = SimpleNamespace( |
| 359 | hf_text_config=hf_config, dtype=torch.get_default_dtype() |
| 360 | ) |
| 361 | vllm_config.quant_config = quant_config |
| 362 | |
| 363 | with set_current_vllm_config(vllm_config): |
| 364 | layer = K2HorizonDecoderLayer(vllm_config=vllm_config, prefix="model.layers.0") |
| 365 | |
| 366 | model = K2HorizonModel.__new__(K2HorizonModel) |
| 367 | nn.Module.__init__(model) |
| 368 | model.config = hf_config |
| 369 | model.quant_config = quant_config |
| 370 | model.num_redundant_experts = 0 |
| 371 | model.layers = nn.ModuleList([layer]) |
| 372 | return model |
| 373 | |
| 374 | |
| 375 | def _fold_rows(weight: torch.Tensor, head_dim: int, idx: torch.Tensor) -> torch.Tensor: |
| 376 | hidden = weight.shape[-1] |
| 377 | return weight.view(-1, head_dim, hidden)[:, idx, :].reshape(-1, hidden).contiguous() |
| 378 | |
| 379 | |
| 380 | def _fold_channels(vec: torch.Tensor, head_dim: int, idx: torch.Tensor) -> torch.Tensor: |
| 381 | return vec.view(-1, head_dim)[:, idx].reshape(-1).contiguous() |
| 382 | |
| 383 | |
| 384 | def _non_qk_weights(prefix, *, hidden_size, intermediate_size, q_size, device=None): |
| 385 | return [ |
| 386 | ( |
| 387 | f"{prefix}.self_attn.o_proj.weight", |
| 388 | torch.randn(hidden_size, q_size, device=device), |
| 389 | ), |
| 390 | ( |
| 391 | f"{prefix}.mlp.gate_proj.weight", |
| 392 | torch.randn(intermediate_size, hidden_size, device=device), |
| 393 | ), |
| 394 | ( |
| 395 | f"{prefix}.mlp.up_proj.weight", |
| 396 | torch.randn(intermediate_size, hidden_size, device=device), |
| 397 | ), |
| 398 | ( |
| 399 | f"{prefix}.mlp.down_proj.weight", |
| 400 | torch.randn(hidden_size, intermediate_size, device=device), |
| 401 | ), |
| 402 | ( |
| 403 | f"{prefix}.input_layernorm.weight", |
| 404 | torch.randn(hidden_size, device=device), |
| 405 | ), |
| 406 | ( |
| 407 | f"{prefix}.post_attention_layernorm.weight", |
| 408 | torch.randn(hidden_size, device=device), |
| 409 | ), |
| 410 | ] |
| 411 | |
| 412 | |
| 413 | @pytest.mark.usefixtures("dist_init") |
| 414 | @pytest.mark.parametrize("head_dim,rope_head_dim", SHAPES) |
| 415 | @pytest.mark.parametrize("attention_bias", [False, True]) |
| 416 | def test_load_weights_dense_folds_qk( |
| 417 | head_dim, |
| 418 | rope_head_dim, |
| 419 | attention_bias, |
| 420 | monkeypatch, |
| 421 | ): |
| 422 | dtype = DTYPE |
| 423 | set_random_seed(0) |
| 424 | num_heads = 8 |
| 425 | num_kv_heads = 2 |
| 426 | hidden_size = 128 |
| 427 | intermediate_size = 2 * hidden_size |
| 428 | q_size = num_heads * head_dim |
| 429 | kv_size = num_kv_heads * head_dim |
| 430 | prefix = "layers.0" |
| 431 | |
| 432 | with set_default_torch_dtype(dtype): |
| 433 | hf_config = _make_hf_config( |
| 434 | hidden_size=hidden_size, |
| 435 | num_heads=num_heads, |
| 436 | num_kv_heads=num_kv_heads, |
| 437 | head_dim=head_dim, |
| 438 | rope_head_dim=rope_head_dim, |
| 439 | attention_bias=attention_bias, |
| 440 | query_key_norm=True, |
| 441 | ) |
| 442 | model = _build_single_layer_model(monkeypatch, hf_config, quant_config=None) |
| 443 | |
| 444 | assert model.layers[0].self_attn.fold_rope_weights is True |
| 445 | |
| 446 | raw = { |
| 447 | "q_proj.weight": torch.randn(q_size, hidden_size), |
| 448 | "k_proj.weight": torch.randn(kv_size, hidden_size), |
| 449 | "v_proj.weight": torch.randn(kv_size, hidden_size), |
| 450 | "q_norm.weight": torch.randn(q_size), |
| 451 | "k_norm.weight": torch.randn(kv_size), |
| 452 | } |
| 453 | if attention_bias: |
| 454 | raw.update( |
| 455 | { |
| 456 | "q_proj.bias": torch.randn(q_size), |
| 457 | "k_proj.bias": torch.randn(kv_size), |
| 458 | "v_proj.bias": torch.randn(kv_size), |
| 459 | "o_proj.bias": torch.randn(hidden_size), |
| 460 | } |
| 461 | ) |
| 462 | |
| 463 | weights: list[tuple[str, torch.Tensor]] = [ |
| 464 | (f"{prefix}.self_attn.{name}", tensor) for name, tensor in raw.items() |
| 465 | ] |
| 466 | weights += _non_qk_weights( |
| 467 | prefix, |
| 468 | hidden_size=hidden_size, |
| 469 | intermediate_size=intermediate_size, |
| 470 | q_size=q_size, |
| 471 | ) |
| 472 | |
| 473 | loaded = model.load_weights(iter(weights)) |
| 474 | |
| 475 | params = dict(model.named_parameters()) |
| 476 | assert loaded == set(params) |
| 477 | |
| 478 | idx = _rope_weight_perm(head_dim, rope_head_dim) |
| 479 | |
| 480 | qkv = params[f"{prefix}.self_attn.qkv_proj.weight"] |
| 481 | q_part = qkv[:q_size] |
| 482 | k_part = qkv[q_size : q_size + kv_size] |
| 483 | v_part = qkv[q_size + kv_size :] |
| 484 | |
| 485 | q_norm = params[f"{prefix}.self_attn.q_norm.weight"] |
| 486 | k_norm = params[f"{prefix}.self_attn.k_norm.weight"] |
| 487 | |
| 488 | assert torch.equal(q_part, _fold_rows(raw["q_proj.weight"], head_dim, idx)) |
| 489 | assert torch.equal(k_part, _fold_rows(raw["k_proj.weight"], head_dim, idx)) |
| 490 | assert torch.equal(v_part, raw["v_proj.weight"]) # v never folded |
| 491 | assert torch.equal(q_norm, _fold_channels(raw["q_norm.weight"], head_dim, idx)) |
| 492 | assert torch.equal(k_norm, _fold_channels(raw["k_norm.weight"], head_dim, idx)) |
| 493 | assert qkv.dtype == dtype |
| 494 | |
| 495 | if attention_bias: |
| 496 | qkv_bias = params[f"{prefix}.self_attn.qkv_proj.bias"] |
| 497 | q_bias = qkv_bias[:q_size] |
| 498 | k_bias = qkv_bias[q_size : q_size + kv_size] |
| 499 | v_bias = qkv_bias[q_size + kv_size :] |
| 500 | assert torch.equal(q_bias, _fold_channels(raw["q_proj.bias"], head_dim, idx)) |
| 501 | assert torch.equal(k_bias, _fold_channels(raw["k_proj.bias"], head_dim, idx)) |
| 502 | assert torch.equal(v_bias, raw["v_proj.bias"]) # v bias never folded |
| 503 | |
| 504 | |
| 505 | @pytest.mark.skipif( |
| 506 | not is_quant_method_supported("fp8"), |
| 507 | reason="fp8 is not supported.", |
| 508 | ) |
| 509 | @pytest.mark.usefixtures("dist_init") |
| 510 | @pytest.mark.parametrize("head_dim,rope_head_dim", SHAPES) |
| 511 | @pytest.mark.parametrize("attention_bias", [False, True]) |
| 512 | def test_load_weights_fp8_qk_no_fold( |
| 513 | head_dim, |
| 514 | rope_head_dim, |
| 515 | attention_bias, |
| 516 | monkeypatch, |
| 517 | ): |
| 518 | set_random_seed(0) |
| 519 | num_heads = 8 |
| 520 | num_kv_heads = 2 |
| 521 | hidden_size = 128 |
| 522 | intermediate_size = 2 * hidden_size |
| 523 | q_size = num_heads * head_dim |
| 524 | kv_size = num_kv_heads * head_dim |
| 525 | prefix = "layers.0" |
| 526 | |
| 527 | def rand_fp8(*shape): |
| 528 | return (torch.randn(*shape, device=DEVICE) * 32.0).to(torch.float8_e4m3fn) |
| 529 | |
| 530 | def rand_scale(): |
| 531 | return torch.rand(1, device=DEVICE, dtype=torch.float32) + 0.1 |
| 532 | |
| 533 | with set_default_torch_dtype(torch.bfloat16), torch.device(DEVICE): |
| 534 | hf_config = _make_hf_config( |
| 535 | hidden_size=hidden_size, |
| 536 | num_heads=num_heads, |
| 537 | num_kv_heads=num_kv_heads, |
| 538 | head_dim=head_dim, |
| 539 | rope_head_dim=rope_head_dim, |
| 540 | attention_bias=attention_bias, |
| 541 | query_key_norm=True, |
| 542 | ) |
| 543 | quant_config = _Fp8QKVOnlyConfig( |
| 544 | is_checkpoint_fp8_serialized=True, activation_scheme="dynamic" |
| 545 | ) |
| 546 | model = _build_single_layer_model( |
| 547 | monkeypatch, hf_config, quant_config=quant_config |
| 548 | ) |
| 549 | |
| 550 | qkv_proj = model.layers[0].self_attn.qkv_proj |
| 551 | assert isinstance(qkv_proj.quant_method, Fp8LinearMethod) |
| 552 | assert model.layers[0].self_attn.fold_rope_weights is False |
| 553 | |
| 554 | raw_w = { |
| 555 | "q_proj.weight": rand_fp8(q_size, hidden_size), |
| 556 | "k_proj.weight": rand_fp8(kv_size, hidden_size), |
| 557 | "v_proj.weight": rand_fp8(kv_size, hidden_size), |
| 558 | } |
| 559 | raw_s = { |
| 560 | "q_proj.weight_scale": rand_scale(), |
| 561 | "k_proj.weight_scale": rand_scale(), |
| 562 | "v_proj.weight_scale": rand_scale(), |
| 563 | } |
| 564 | q_norm = torch.randn(q_size, device=DEVICE) |
| 565 | k_norm = torch.randn(kv_size, device=DEVICE) |
| 566 | |
| 567 | weights: list[tuple[str, torch.Tensor]] = [ |
| 568 | (f"{prefix}.self_attn.{name}", tensor) |
| 569 | for name, tensor in {**raw_w, **raw_s}.items() |
| 570 | ] |
| 571 | weights += [ |
| 572 | (f"{prefix}.self_attn.q_norm.weight", q_norm), |
| 573 | (f"{prefix}.self_attn.k_norm.weight", k_norm), |
| 574 | ] |
| 575 | biases = {} |
| 576 | if attention_bias: |
| 577 | biases = { |
| 578 | "q_proj.bias": torch.randn(q_size, device=DEVICE), |
| 579 | "k_proj.bias": torch.randn(kv_size, device=DEVICE), |
| 580 | "v_proj.bias": torch.randn(kv_size, device=DEVICE), |
| 581 | "o_proj.bias": torch.randn(hidden_size, device=DEVICE), |
| 582 | } |
| 583 | weights += [ |
| 584 | (f"{prefix}.self_attn.{name}", tensor) |
| 585 | for name, tensor in biases.items() |
| 586 | ] |
| 587 | weights += _non_qk_weights( |
| 588 | prefix, |
| 589 | hidden_size=hidden_size, |
| 590 | intermediate_size=intermediate_size, |
| 591 | q_size=q_size, |
| 592 | device=DEVICE, |
| 593 | ) |
| 594 | |
| 595 | loaded = model.load_weights(iter(weights)) |
| 596 | |
| 597 | params = dict(model.named_parameters()) |
| 598 | assert loaded == set(params) |
| 599 | |
| 600 | qkv = params[f"{prefix}.self_attn.qkv_proj.weight"] |
| 601 | weight_scale = params[f"{prefix}.self_attn.qkv_proj.weight_scale"] |
| 602 | |
| 603 | assert qkv.dtype == torch.float8_e4m3fn |
| 604 | assert torch.equal(qkv[:q_size], raw_w["q_proj.weight"]) |
| 605 | assert torch.equal(qkv[q_size : q_size + kv_size], raw_w["k_proj.weight"]) |
| 606 | assert torch.equal(qkv[q_size + kv_size :], raw_w["v_proj.weight"]) |
| 607 | |
| 608 | assert torch.equal( |
| 609 | weight_scale, |
| 610 | torch.cat( |
| 611 | [ |
| 612 | raw_s["q_proj.weight_scale"], |
| 613 | raw_s["k_proj.weight_scale"], |
| 614 | raw_s["v_proj.weight_scale"], |
| 615 | ] |
| 616 | ), |
| 617 | ) |
| 618 | |
| 619 | assert torch.equal(params[f"{prefix}.self_attn.q_norm.weight"], q_norm) |
| 620 | assert torch.equal(params[f"{prefix}.self_attn.k_norm.weight"], k_norm) |
| 621 | |
| 622 | if attention_bias: |
| 623 | qkv_bias = params[f"{prefix}.self_attn.qkv_proj.bias"] |
| 624 | assert torch.equal(qkv_bias[:q_size], biases["q_proj.bias"]) |
| 625 | assert torch.equal(qkv_bias[q_size : q_size + kv_size], biases["k_proj.bias"]) |
| 626 | assert torch.equal(qkv_bias[q_size + kv_size :], biases["v_proj.bias"]) |
| 627 | |
| 628 | |
| 629 | def test_is_rope_weights_folding_supported(default_vllm_config): |
| 630 | dense = SimpleNamespace(quant_method=UnquantizedLinearMethod()) |
| 631 | quantized = SimpleNamespace(quant_method=object()) |
| 632 | |
| 633 | assert is_rope_weights_folding_supported(dense, None) is True |
| 634 | assert is_rope_weights_folding_supported(quantized, None) is False |
| 635 | assert is_rope_weights_folding_supported(dense, {"chunk_size": 8}) is False |
| 636 | assert is_rope_weights_folding_supported(SimpleNamespace(), None) is False |