初始化项目，由ModelHub XC社区提供模型
Model: flwrlabs/Lizzy-7B Source: Original Platform
This commit is contained in:
ModelHub XC
211
vllm_patches/transformers_lizzy_tp.py
Normal file
211
vllm_patches/transformers_lizzy_tp.py
Normal file
@@ -0,0 +1,211 @@
|
||||
"""Compat patch for Lizzy TP under vLLM's generic Transformers backend."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from torch import nn
|
||||
|
||||
_PATCH_ATTR = "_flwr_transformers_lizzy_tp_patch_applied"
|
||||
|
||||
|
||||
class _TensorParallelSliceNorm(nn.Module):
|
||||
"""Apply a full-width checkpoint norm to a TP-local activation slice."""
|
||||
|
||||
def __init__(self, base_norm: nn.Module, start_idx: int, end_idx: int):
|
||||
super().__init__()
|
||||
self.start_idx = start_idx
|
||||
self.end_idx = end_idx
|
||||
self.weight = base_norm.weight
|
||||
if getattr(base_norm, "bias", None) is not None:
|
||||
self.bias = base_norm.bias
|
||||
else:
|
||||
self.register_parameter("bias", None)
|
||||
self.eps = float(
|
||||
getattr(base_norm, "eps", getattr(base_norm, "variance_epsilon", 1e-6)),
|
||||
)
|
||||
self.norm_kind = (
|
||||
"layernorm" if isinstance(base_norm, nn.LayerNorm) else "rmsnorm"
|
||||
)
|
||||
|
||||
@property
|
||||
def local_size(self) -> int:
|
||||
return self.end_idx - self.start_idx
|
||||
|
||||
def _slice_param(self, param: torch.Tensor | None) -> torch.Tensor | None:
|
||||
if param is None:
|
||||
return None
|
||||
if param.shape[0] == self.local_size:
|
||||
return param
|
||||
return param[self.start_idx : self.end_idx]
|
||||
|
||||
def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
|
||||
weight = self.weight
|
||||
bias = self.bias
|
||||
if hidden_states.shape[-1] != self.weight.shape[0]:
|
||||
if hidden_states.shape[-1] != self.local_size:
|
||||
msg = (
|
||||
"Unexpected hidden size for TP-sliced norm: "
|
||||
f"{hidden_states.shape[-1]} "
|
||||
f"(expected {self.weight.shape[0]} or {self.local_size})"
|
||||
)
|
||||
raise RuntimeError(msg)
|
||||
weight = self._slice_param(weight)
|
||||
bias = self._slice_param(bias)
|
||||
|
||||
if self.norm_kind == "layernorm":
|
||||
return F.layer_norm(
|
||||
hidden_states,
|
||||
(hidden_states.shape[-1],),
|
||||
weight,
|
||||
bias,
|
||||
self.eps,
|
||||
)
|
||||
|
||||
input_dtype = hidden_states.dtype
|
||||
hidden_states_fp32 = hidden_states.to(torch.float32)
|
||||
variance = hidden_states_fp32.pow(2).mean(dim=-1, keepdim=True)
|
||||
hidden_states_norm = hidden_states_fp32 * torch.rsqrt(variance + self.eps)
|
||||
hidden_states_norm = hidden_states_norm.to(input_dtype)
|
||||
output = weight * hidden_states_norm
|
||||
if bias is not None:
|
||||
output = output + bias
|
||||
return output
|
||||
|
||||
|
||||
def _maybe_patch_lizzy_attention_for_tp(
|
||||
*,
|
||||
module: nn.Module,
|
||||
prefix: str,
|
||||
tp_size: int,
|
||||
tp_rank: int,
|
||||
log_replacement: Any, # noqa: ANN401
|
||||
) -> None:
|
||||
if tp_size <= 1 or type(module).__name__ != "LizzyAttention":
|
||||
return
|
||||
|
||||
num_heads = getattr(module, "num_heads", None)
|
||||
num_key_value_heads = getattr(module, "num_key_value_heads", None)
|
||||
head_dim = getattr(module, "head_dim", None)
|
||||
q_norm = getattr(module, "q_norm", None)
|
||||
k_norm = getattr(module, "k_norm", None)
|
||||
if not all(
|
||||
isinstance(value, int)
|
||||
for value in (num_heads, num_key_value_heads, head_dim)
|
||||
):
|
||||
return
|
||||
if num_heads % tp_size != 0 or num_key_value_heads % tp_size != 0:
|
||||
return
|
||||
|
||||
local_num_heads = num_heads // tp_size
|
||||
local_num_key_value_heads = num_key_value_heads // tp_size
|
||||
local_q_dim = local_num_heads * head_dim
|
||||
local_kv_dim = local_num_key_value_heads * head_dim
|
||||
|
||||
module.num_heads = local_num_heads
|
||||
module.num_key_value_heads = local_num_key_value_heads
|
||||
module.num_key_value_groups = local_num_heads // local_num_key_value_heads
|
||||
|
||||
if q_norm is not None and getattr(q_norm, "weight", None) is not None:
|
||||
start = tp_rank * local_q_dim
|
||||
end = start + local_q_dim
|
||||
module.q_norm = _TensorParallelSliceNorm(q_norm, start, end)
|
||||
log_replacement(f"{prefix}.q_norm", q_norm, module.q_norm)
|
||||
|
||||
if k_norm is not None and getattr(k_norm, "weight", None) is not None:
|
||||
start = tp_rank * local_kv_dim
|
||||
end = start + local_kv_dim
|
||||
module.k_norm = _TensorParallelSliceNorm(k_norm, start, end)
|
||||
log_replacement(f"{prefix}.k_norm", k_norm, module.k_norm)
|
||||
|
||||
|
||||
def patch_vllm_transformers_lizzy_tp() -> None:
|
||||
"""Patch the generic vLLM Transformers backend for Lizzy TP norms/heads."""
|
||||
import vllm.model_executor.models.transformers as transformers_mod
|
||||
|
||||
transformers_base = transformers_mod.TransformersBase
|
||||
if getattr(transformers_base, _PATCH_ATTR, False):
|
||||
return
|
||||
|
||||
PreTrainedModel = transformers_mod.PreTrainedModel
|
||||
maybe_prefix = transformers_mod.maybe_prefix
|
||||
replace_linear_class = transformers_mod.replace_linear_class
|
||||
get_feature_request_tip = transformers_mod.get_feature_request_tip
|
||||
re = transformers_mod.re
|
||||
log_replacement = transformers_mod.log_replacement
|
||||
get_tp_rank = getattr(transformers_mod, "get_tensor_model_parallel_rank", None)
|
||||
if get_tp_rank is None:
|
||||
try:
|
||||
from vllm.distributed import ( # noqa: PLC0415
|
||||
get_tensor_model_parallel_rank as get_tp_rank,
|
||||
)
|
||||
except Exception:
|
||||
get_tp_rank = lambda: 0
|
||||
|
||||
def tensor_parallel(self: Any) -> None: # noqa: ANN401
|
||||
"""Apply the model's tensor parallel plan plus Lizzy attention fixes."""
|
||||
is_pretrained_model = lambda m: isinstance(m, PreTrainedModel)
|
||||
supports_tp_plan = lambda m: m.config.base_model_tp_plan is not None
|
||||
pretrained_models = filter(is_pretrained_model, self.model.modules())
|
||||
models_with_tp_plan = filter(supports_tp_plan, pretrained_models)
|
||||
|
||||
if not any(models_with_tp_plan) and self.tp_size > 1:
|
||||
tip = get_feature_request_tip(
|
||||
self.model_config.model,
|
||||
self.model_config.trust_remote_code,
|
||||
)
|
||||
raise ValueError(
|
||||
f"{type(self.model)} does not support tensor parallel. {tip}",
|
||||
)
|
||||
|
||||
tp_rank = get_tp_rank()
|
||||
|
||||
def _tensor_parallel(
|
||||
module: nn.Module,
|
||||
prefix: str = "",
|
||||
tp_plan: dict[str, str] | None = None,
|
||||
) -> None:
|
||||
local_tp_plan = tp_plan or {}
|
||||
|
||||
if isinstance(module, PreTrainedModel):
|
||||
local_tp_plan = module.config.base_model_tp_plan or {}
|
||||
local_tp_plan = {
|
||||
maybe_prefix(prefix, key): value
|
||||
for key, value in local_tp_plan.items()
|
||||
}
|
||||
|
||||
for child_name, child_module in module.named_children():
|
||||
qual_name = maybe_prefix(prefix, child_name)
|
||||
if isinstance(child_module, nn.Linear):
|
||||
generator = (p for p in local_tp_plan if re.match(p, qual_name))
|
||||
pattern = next(generator, None)
|
||||
style = local_tp_plan.get(pattern, "replicate")
|
||||
new_module = replace_linear_class(
|
||||
child_module,
|
||||
style,
|
||||
self.quant_config,
|
||||
prefix=qual_name,
|
||||
)
|
||||
setattr(module, child_name, new_module)
|
||||
log_replacement(qual_name, child_module, new_module)
|
||||
else:
|
||||
_tensor_parallel(
|
||||
child_module,
|
||||
prefix=qual_name,
|
||||
tp_plan=local_tp_plan,
|
||||
)
|
||||
|
||||
_maybe_patch_lizzy_attention_for_tp(
|
||||
module=module,
|
||||
prefix=prefix,
|
||||
tp_size=self.tp_size,
|
||||
tp_rank=tp_rank,
|
||||
log_replacement=log_replacement,
|
||||
)
|
||||
|
||||
_tensor_parallel(self.model)
|
||||
|
||||
transformers_base.tensor_parallel = tensor_parallel
|
||||
setattr(transformers_base, _PATCH_ATTR, True)
|
||||
Reference in New Issue
Block a user