| """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 |
| input_key: str |
| role: str |
| conv1d: bool = False |
|
|
|
|
| @dataclass |
| class Arch: |
| model_type: str |
| n_layer: int |
| d: int |
| n_head: int |
| n_kv: int |
| head_dim: int |
| inter: int |
| norm: str |
| norm_eps: float |
| act: str |
| 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) |
|
|
| 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)") |
|
|
