80 lines
2.8 KiB
Python
80 lines
2.8 KiB
Python
|
import logging
|
||
|
|
from collections.abc import Iterable
|
||
|
|
from pathlib import Path
|
||
|
|
|
||
|
|
import torch
|
||
|
|
from safetensors.torch import load_file
|
||
|
|
from vllm.model_executor.models.llama_eagle3 import Eagle3LlamaForCausalLM
|
||
|
|
from vllm.model_executor.models.utils import (
|
||
|
|
AutoWeightsLoader,
|
||
|
|
process_eagle_weight,
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
def patch_load_weights(target_config):
|
||
|
|
Eagle3LlamaForCausalLM.load_weights = make_load_weights(target_config)
|
||
|
|
|
||
|
|
|
||
|
|
def make_load_weights(target_config):
|
||
|
|
logger = logging.getLogger(__name__)
|
||
|
|
quant_cfg = target_config.quant_config
|
||
|
|
rotation_matrix3 = None
|
||
|
|
|
||
|
|
model_path = target_config.model_config.model
|
||
|
|
try:
|
||
|
|
rotation_rel_path = quant_cfg.quant_description["optional"]["quarot"]["rotation_map"]["global_rotation"]
|
||
|
|
except KeyError as e:
|
||
|
|
logger.error(
|
||
|
|
"Invalid quant_config: missing key "
|
||
|
|
"quant_description['optional']['quarot']['rotation_map']['global_rotation']. "
|
||
|
|
"If you don't use quarot model, please ignore it. "
|
||
|
|
f"Error: {e}"
|
||
|
|
)
|
||
|
|
else:
|
||
|
|
rotation_path = Path(model_path) / rotation_rel_path
|
||
|
|
try:
|
||
|
|
safetensor_data = load_file(rotation_path)
|
||
|
|
Q = safetensor_data["global_rotation"]
|
||
|
|
rotation_matrix3 = torch.block_diag(Q, Q, Q)
|
||
|
|
except Exception as e:
|
||
|
|
logger.error(
|
||
|
|
f"Failed to load rotation weight from '{rotation_path}'. "
|
||
|
|
"If you don't use quarot model, please ignore it. "
|
||
|
|
f"Error: {e}"
|
||
|
|
)
|
||
|
|
|
||
|
|
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):
|
||
|
|
model_weights = {}
|
||
|
|
includes_draft_id_mapping = False
|
||
|
|
includes_embed_tokens = False
|
||
|
|
for name, loaded_weight in weights:
|
||
|
|
if "t2d" in name:
|
||
|
|
continue
|
||
|
|
if "d2t" in name:
|
||
|
|
name = name.replace("d2t", "draft_id_to_target_id")
|
||
|
|
includes_draft_id_mapping = True
|
||
|
|
elif "lm_head" not in name:
|
||
|
|
name = "model." + name
|
||
|
|
if "fc." in name and rotation_matrix3 is not None:
|
||
|
|
loaded_weight = loaded_weight @ rotation_matrix3.to(loaded_weight.dtype)
|
||
|
|
if "embed_tokens" in name:
|
||
|
|
includes_embed_tokens = True
|
||
|
|
model_weights[name] = loaded_weight
|
||
|
|
process_eagle_weight(self, name)
|
||
|
|
|
||
|
|
skip_substrs = []
|
||
|
|
if not includes_draft_id_mapping:
|
||
|
|
skip_substrs.append("draft_id_to_target_id")
|
||
|
|
if not includes_embed_tokens:
|
||
|
|
skip_substrs.append("embed_tokens")
|
||
|
|
if not self.model.use_aux_hidden_state:
|
||
|
|
skip_substrs.append("fc.")
|
||
|
|
loader = AutoWeightsLoader(
|
||
|
|
self,
|
||
|
|
skip_prefixes=None,
|
||
|
|
skip_substrs=skip_substrs,
|
||
|
|
)
|
||
|
|
loader.load_weights(model_weights.items())
|
||
|
|
|
||
|
|
return load_weights
|