File size: 5,580 Bytes
8ad3ba6 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 | # decoder_adapters.py
# Pluggable transfer of backbone weights from HF SeqCLS (Fed fine-tuned) to CausalLM for generate().
from __future__ import annotations
from abc import ABC, abstractmethod
from typing import List, Type
import torch.nn as nn
class DecoderAdapter(ABC):
"""Maps a fine-tuned ``ForSequenceClassification`` inner model into a ``ForCausalLM`` backbone."""
@staticmethod
@abstractmethod
def matches(model_name: str) -> bool:
"""Return True if this adapter handles ``model_name`` (HF id or path)."""
@abstractmethod
def transfer_backbone(
self,
seq_cls_inner: nn.Module,
causal_lm: nn.Module,
) -> None:
"""
Copy shared backbone weights from the classification model into ``causal_lm``.
``seq_cls_inner`` is typically ``GPTNeoXForSequenceClassification`` (possibly merged from PEFT).
``causal_lm`` is ``GPTNeoXForCausalLM``. ``lm_head`` stays as loaded from ``from_pretrained``.
"""
class PythiaNeoXAdapter(DecoderAdapter):
"""EleutherAI Pythia / GPT-NeoX sequence-classification -> causal LM."""
_NEEDLE = "pythia"
@staticmethod
def matches(model_name: str) -> bool:
m = (model_name or "").lower()
return PythiaNeoXAdapter._NEEDLE in m or "gpt-neox" in m
def transfer_backbone(self, seq_cls_inner: nn.Module, causal_lm: nn.Module) -> None:
inner = seq_cls_inner
if hasattr(inner, "merge_and_unload"):
inner = inner.merge_and_unload()
src_sd = inner.state_dict()
dst_sd = causal_lm.state_dict()
to_load = {}
for k, v in src_sd.items():
if not k.startswith("gpt_neox."):
continue
if k not in dst_sd or dst_sd[k].shape != v.shape:
continue
to_load[k] = v.to(device=dst_sd[k].device, dtype=dst_sd[k].dtype)
if not to_load:
raise RuntimeError(
"PythiaNeoXAdapter: no gpt_neox.* keys matched between SeqCLS and CausalLM. "
"Check model_name and transformers versions."
)
causal_lm.load_state_dict(to_load, strict=False)
class Qwen2Adapter(DecoderAdapter):
"""Qwen2 / Qwen2.5 sequence-classification -> causal LM (shared ``model.*`` backbone)."""
@staticmethod
def matches(model_name: str) -> bool:
m = (model_name or "").lower()
return "qwen2" in m
def transfer_backbone(self, seq_cls_inner: nn.Module, causal_lm: nn.Module) -> None:
inner = seq_cls_inner
if hasattr(inner, "merge_and_unload"):
inner = inner.merge_and_unload()
src_sd = inner.state_dict()
dst_sd = causal_lm.state_dict()
to_load = {}
for k, v in src_sd.items():
if not k.startswith("model."):
continue
if k not in dst_sd or dst_sd[k].shape != v.shape:
continue
to_load[k] = v.to(device=dst_sd[k].device, dtype=dst_sd[k].dtype)
if not to_load:
raise RuntimeError(
"Qwen2Adapter: no model.* keys matched between SeqCLS and CausalLM. "
"Check model_name and transformers versions."
)
causal_lm.load_state_dict(to_load, strict=False)
class LlamaAdapter(DecoderAdapter):
"""Llama-family sequence-classification -> causal LM (shared ``model.*`` backbone).
Covers meta-llama/Llama-3.2-*, TinyLlama, Sheared-LLaMA, MobileLLaMA — any HF id
containing "llama". Same state-dict layout as Qwen2 (``model.*`` prefix).
Tied embeddings (e.g. Llama-3.2-1B) are handled implicitly: transferring
``model.embed_tokens.weight`` updates the tied ``lm_head`` as well.
"""
@staticmethod
def matches(model_name: str) -> bool:
m = (model_name or "").lower()
return "llama" in m
def transfer_backbone(self, seq_cls_inner: nn.Module, causal_lm: nn.Module) -> None:
inner = seq_cls_inner
if hasattr(inner, "merge_and_unload"):
inner = inner.merge_and_unload()
src_sd = inner.state_dict()
dst_sd = causal_lm.state_dict()
to_load = {}
for k, v in src_sd.items():
if not k.startswith("model."):
continue
if k not in dst_sd or dst_sd[k].shape != v.shape:
continue
to_load[k] = v.to(device=dst_sd[k].device, dtype=dst_sd[k].dtype)
if not to_load:
raise RuntimeError(
"LlamaAdapter: no model.* keys matched between SeqCLS and CausalLM. "
"Check model_name and transformers versions."
)
causal_lm.load_state_dict(to_load, strict=False)
# Registry: first match wins (order matters for overlapping patterns).
ADAPTER_REGISTRY: List[Type[DecoderAdapter]] = [
Qwen2Adapter,
LlamaAdapter,
PythiaNeoXAdapter,
]
def resolve_adapter(model_name: str) -> DecoderAdapter:
"""
Select an adapter for the given Hugging Face model id.
Raises:
ValueError: If no registered adapter matches.
"""
for cls in ADAPTER_REGISTRY:
if cls.matches(model_name):
return cls()
registered = ", ".join(c.__name__ for c in ADAPTER_REGISTRY)
raise ValueError(
f"No DecoderAdapter registered for model_name={model_name!r}. "
f"Implement a new adapter class, append it to ADAPTER_REGISTRY in decoder_adapters.py, "
f"and implement transfer_backbone for that architecture. "
f"Currently registered: {registered}"
)
|