"""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)