212 lines
7.8 KiB
Python
212 lines
7.8 KiB
Python
"""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)
|