(self, tllm_key, weights, **kwargs)
| 369 | return x |
| 370 | |
| 371 | def postprocess(self, tllm_key, weights, **kwargs): |
| 372 | using_head_as_leading_dim = kwargs.get("using_head_as_leading_dim", |
| 373 | False) |
| 374 | config = kwargs.get("config", None) |
| 375 | if self.is_qkv: |
| 376 | if isinstance(weights, list): |
| 377 | head_size = config.hidden_size // config.num_attention_heads if config.head_size is None else config.head_size |
| 378 | if getattr(config, "remove_duplicated_kv_heads", False): |
| 379 | if config.remove_duplicated_kv_heads: |
| 380 | k, v = weights[1:] |
| 381 | k = k.reshape([ |
| 382 | k.shape[0] // head_size // 2, 2, head_size, |
| 383 | self.in_features |
| 384 | ]) |
| 385 | v = v.reshape([ |
| 386 | v.shape[0] // head_size // 2, 2, head_size, |
| 387 | self.in_features |
| 388 | ]) |
| 389 | assert (k[:, 0] == k[:, 1]).all() |
| 390 | assert (v[:, 0] == v[:, 1]).all() |
| 391 | k = k[:, 0].reshape([-1, self.in_features]) |
| 392 | v = v[:, 0].reshape([-1, self.in_features]) |
| 393 | weights[1] = k |
| 394 | weights[2] = v |
| 395 | # Duplicate kv heads in case of invalid TP size |
| 396 | tp_size = config.mapping.tp_size |
| 397 | num_kv_heads = config.num_key_value_heads |
| 398 | if num_kv_heads < tp_size: |
| 399 | for qkv_idx in range(3): |
| 400 | v = weights[qkv_idx] |
| 401 | if qkv_idx > 0: |
| 402 | assert tp_size % num_kv_heads == 0 |
| 403 | reps = tp_size // num_kv_heads |
| 404 | if tllm_key.endswith("bias"): |
| 405 | v = v.reshape(num_kv_heads, |
| 406 | head_size)[:, None, :].expand( |
| 407 | num_kv_heads, reps, head_size) |
| 408 | v = v.reshape(num_kv_heads * reps * head_size) |
| 409 | else: |
| 410 | v = v.reshape(num_kv_heads, head_size, |
| 411 | -1)[:, None, :, :].expand( |
| 412 | num_kv_heads, reps, head_size, |
| 413 | v.shape[1]) |
| 414 | v = v.reshape(num_kv_heads * reps * head_size, |
| 415 | -1) |
| 416 | weights[qkv_idx] = v.chunk( |
| 417 | tp_size, self.tp_dim)[config.mapping.tp_rank] |
| 418 | |
| 419 | weights = torch.cat(weights) |
| 420 | if using_head_as_leading_dim: |
| 421 | # Reorder [n_head, 3, head_dim, ...] into [3, n_head, head_dim, ...] |
| 422 | assert config.num_attention_heads == config.num_key_value_heads, "using_head_as_leading_dim require head_size to be multiple of 3." |
| 423 | num_heads = config.num_attention_heads |
| 424 | head_dim = self.out_features // (3 * num_heads) |
| 425 | w = weights.reshape(num_heads, 3, head_dim, -1) |
| 426 | w = w.transpose(0, 1) |
| 427 | if w.shape[-1] > 1: |
| 428 | weights = w.reshape(-1, self.in_features) # Weight |
no test coverage detected