File size: 5,277 Bytes
c675b8d | 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 | """Architecture adapters: tensor names and structure for the supported Hugging Face causal LMs.
Supported model_type: llama, mistral, qwen2 (incl. Llama-2/3/3.1/3.3-70B, Qwen2.5-72B), opt (pre-LN), gpt2.
A linear module is described by (module_name, input_key, role):
input_key in {"ln1", "attn_out", "ln2", "act"} - which activation it reads
role in {"q","k","v","qkv","o","gate","up","down"}
Weights are returned in torch Linear layout (out, in); GPT-2 Conv1D weights are transposed on read.
"""
from dataclasses import dataclass, field
from typing import List, Optional, Tuple
@dataclass
class Linear:
name: str # module path, e.g. "model.layers.3.self_attn.q_proj"
input_key: str
role: str
conv1d: bool = False # stored (in, out)
@dataclass
class Arch:
model_type: str
n_layer: int
d: int
n_head: int
n_kv: int
head_dim: int
inter: int
norm: str # "rms" | "ln"
norm_eps: float
act: str # "swiglu" | "relu" | "gelu_new" | "gelu"
embed: str
pos_embed: Optional[str]
pos_offset: int
final_norm: str
lm_head: Optional[str]
vocab: int
max_pos: int
layers: List[List[Linear]] = field(default_factory=list)
norms: List[Tuple[str, str]] = field(default_factory=list) # per layer: (ln1 prefix, ln2 prefix)
def norm_names(self, L, which):
p = self.norms[L][0 if which == "ln1" else 1]
return p + ".weight", (p + ".bias" if self.norm == "ln" else None)
def from_config(cfg: dict) -> Arch:
mt = cfg.get("model_type")
if mt in ("llama", "mistral", "qwen2"):
d, H = cfg["hidden_size"], cfg["num_attention_heads"]
act = cfg.get("hidden_act", "silu")
if act != "silu":
raise NotImplementedError(f"{mt} with hidden_act={act}")
a = Arch(mt, cfg["num_hidden_layers"], d, H, cfg.get("num_key_value_heads", H), cfg.get("head_dim") or d // H,
cfg["intermediate_size"], "rms", cfg.get("rms_norm_eps", 1e-6), "swiglu",
"model.embed_tokens.weight", None, 0, "model.norm.weight",
None if cfg.get("tie_word_embeddings", False) else "lm_head.weight",
cfg["vocab_size"], cfg.get("max_position_embeddings", 4096))
for L in range(a.n_layer):
p = f"model.layers.{L}."
a.layers.append([Linear(p + "self_attn.q_proj", "ln1", "q"), Linear(p + "self_attn.k_proj", "ln1", "k"),
Linear(p + "self_attn.v_proj", "ln1", "v"), Linear(p + "self_attn.o_proj", "attn_out", "o"),
Linear(p + "mlp.gate_proj", "ln2", "gate"), Linear(p + "mlp.up_proj", "ln2", "up"),
Linear(p + "mlp.down_proj", "act", "down")])
a.norms.append((p + "input_layernorm", p + "post_attention_layernorm"))
return a
if mt == "opt":
if not cfg.get("do_layer_norm_before", True):
raise NotImplementedError("post-LN OPT (opt-350m) is not supported")
if cfg.get("word_embed_proj_dim", cfg["hidden_size"]) != cfg["hidden_size"]:
raise NotImplementedError("OPT with project_in/out is not supported")
d, H = cfg["hidden_size"], cfg["num_attention_heads"]
a = Arch(mt, cfg["num_hidden_layers"], d, H, H, d // H, cfg["ffn_dim"], "ln", 1e-5,
{"relu": "relu"}.get(cfg.get("activation_function", "relu"), cfg.get("activation_function")),
"model.decoder.embed_tokens.weight", "model.decoder.embed_positions.weight", 2,
"model.decoder.final_layer_norm.weight", None if cfg.get("tie_word_embeddings", True) else "lm_head.weight",
cfg["vocab_size"], cfg.get("max_position_embeddings", 2048))
for L in range(a.n_layer):
p = f"model.decoder.layers.{L}."
a.layers.append([Linear(p + "self_attn.q_proj", "ln1", "q"), Linear(p + "self_attn.k_proj", "ln1", "k"),
Linear(p + "self_attn.v_proj", "ln1", "v"), Linear(p + "self_attn.out_proj", "attn_out", "o"),
Linear(p + "fc1", "ln2", "up"), Linear(p + "fc2", "act", "down")])
a.norms.append((p + "self_attn_layer_norm", p + "final_layer_norm"))
return a
if mt == "gpt2":
d, H = cfg["n_embd"], cfg["n_head"]
a = Arch(mt, cfg["n_layer"], d, H, H, d // H, cfg.get("n_inner") or 4 * d, "ln", cfg.get("layer_norm_epsilon", 1e-5),
{"gelu_new": "gelu_new"}.get(cfg.get("activation_function", "gelu_new"), cfg.get("activation_function")),
"transformer.wte.weight", "transformer.wpe.weight", 0, "transformer.ln_f.weight", None,
cfg["vocab_size"], cfg.get("n_positions", 1024))
for L in range(a.n_layer):
p = f"transformer.h.{L}."
a.layers.append([Linear(p + "attn.c_attn", "ln1", "qkv", True), Linear(p + "attn.c_proj", "attn_out", "o", True),
Linear(p + "mlp.c_fc", "ln2", "up", True), Linear(p + "mlp.c_proj", "act", "down", True)])
a.norms.append((p + "ln_1", p + "ln_2"))
return a
raise NotImplementedError(f"model_type={mt!r} is not supported (supported: llama, mistral, qwen2, opt, gpt2)")
|