mirror of
https://github.com/elder-plinius/OBLITERATUS.git
synced 2026-07-11 22:56:32 +02:00
pliny
235 lines
7.2 KiB
Python
235 lines
7.2 KiB
Python
"""Utilities for navigating different HF model architectures."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import torch.nn as nn
|
|
|
|
from obliteratus.models.loader import ModelHandle
|
|
|
|
# Mapping from model_type -> attribute path to the list of transformer layers.
|
|
_LAYER_ATTR_PATHS: dict[str, list[str]] = {
|
|
"gpt2": ["transformer", "h"],
|
|
"gpt_neo": ["transformer", "h"],
|
|
"gpt_neox": ["gpt_neox", "layers"],
|
|
"llama": ["model", "layers"],
|
|
"mistral": ["model", "layers"],
|
|
"gemma": ["model", "layers"],
|
|
"gemma2": ["model", "layers"],
|
|
"phi": ["model", "layers"],
|
|
"phi3": ["model", "layers"],
|
|
"qwen2": ["model", "layers"],
|
|
"qwen3": ["model", "layers"],
|
|
"qwen3_moe": ["model", "layers"],
|
|
"qwen3_5": ["model", "layers"],
|
|
"qwen3_5_text": ["model", "layers"],
|
|
"minimax_m2": ["model", "layers"],
|
|
"glm_moe_dsa": ["model", "layers"],
|
|
"deepseek_v3": ["model", "layers"],
|
|
"glm4": ["model", "layers"],
|
|
"glm4_moe": ["model", "layers"],
|
|
"glm4_moe_lite": ["model", "layers"],
|
|
"minicpm3": ["model", "layers"],
|
|
"internlm3": ["model", "layers"],
|
|
"falcon": ["transformer", "h"],
|
|
"opt": ["model", "decoder", "layers"],
|
|
"bloom": ["transformer", "h"],
|
|
"mpt": ["transformer", "blocks"],
|
|
"stablelm": ["model", "layers"],
|
|
"chatglm": ["transformer", "encoder", "layers"],
|
|
"glm": ["model", "layers"],
|
|
"gpt_oss": ["model", "layers"],
|
|
"smollm3": ["model", "layers"],
|
|
"cohere": ["model", "layers"],
|
|
"cohere2": ["model", "layers"],
|
|
"olmo": ["model", "layers"],
|
|
"olmo2": ["model", "layers"],
|
|
"internlm2": ["model", "layers"],
|
|
"granite": ["model", "layers"],
|
|
"gemma3": ["model", "layers"],
|
|
}
|
|
|
|
_ATTENTION_ATTR: dict[str, str] = {
|
|
"gpt2": "attn",
|
|
"gpt_neo": "attn.attention",
|
|
"gpt_neox": "attention",
|
|
"llama": "self_attn",
|
|
"mistral": "self_attn",
|
|
"gemma": "self_attn",
|
|
"gemma2": "self_attn",
|
|
"phi": "self_attn",
|
|
"phi3": "self_attn",
|
|
"qwen2": "self_attn",
|
|
"qwen3": "self_attn",
|
|
"qwen3_moe": "self_attn",
|
|
"qwen3_5": "self_attn",
|
|
"qwen3_5_text": "self_attn",
|
|
"minimax_m2": "self_attn",
|
|
"glm_moe_dsa": "self_attn",
|
|
"deepseek_v3": "self_attn",
|
|
"glm4": "self_attn",
|
|
"glm4_moe": "self_attn",
|
|
"glm4_moe_lite": "self_attn",
|
|
"minicpm3": "self_attn",
|
|
"internlm3": "self_attn",
|
|
"falcon": "self_attention",
|
|
"opt": "self_attn",
|
|
"bloom": "self_attention",
|
|
"mpt": "attn",
|
|
"stablelm": "self_attn",
|
|
"chatglm": "self_attention",
|
|
"glm": "self_attn",
|
|
"gpt_oss": "self_attn",
|
|
"smollm3": "self_attn",
|
|
"cohere": "self_attn",
|
|
"cohere2": "self_attn",
|
|
"olmo": "self_attn",
|
|
"olmo2": "self_attn",
|
|
"internlm2": "attention",
|
|
"granite": "self_attn",
|
|
"gemma3": "self_attn",
|
|
}
|
|
|
|
_FFN_ATTR: dict[str, str] = {
|
|
"gpt2": "mlp",
|
|
"gpt_neo": "mlp",
|
|
"gpt_neox": "mlp",
|
|
"llama": "mlp",
|
|
"mistral": "mlp",
|
|
"gemma": "mlp",
|
|
"gemma2": "mlp",
|
|
"phi": "mlp",
|
|
"phi3": "mlp",
|
|
"qwen2": "mlp",
|
|
"qwen3": "mlp",
|
|
"qwen3_moe": "mlp",
|
|
"qwen3_5": "mlp",
|
|
"qwen3_5_text": "mlp",
|
|
"minimax_m2": "mlp",
|
|
"glm_moe_dsa": "mlp",
|
|
"deepseek_v3": "mlp",
|
|
"glm4": "mlp",
|
|
"glm4_moe": "mlp",
|
|
"glm4_moe_lite": "mlp",
|
|
"minicpm3": "mlp",
|
|
"internlm3": "mlp",
|
|
"falcon": "mlp",
|
|
# OPT: fc1/fc2 live directly on the layer — handled by _FLAT_FFN_ARCHS
|
|
"bloom": "mlp",
|
|
"mpt": "ffn",
|
|
"stablelm": "mlp",
|
|
"chatglm": "mlp",
|
|
"glm": "mlp",
|
|
"gpt_oss": "mlp",
|
|
"smollm3": "mlp",
|
|
"cohere": "mlp",
|
|
"cohere2": "mlp",
|
|
"olmo": "mlp",
|
|
"olmo2": "mlp",
|
|
"internlm2": "feed_forward",
|
|
"granite": "mlp",
|
|
"gemma3": "mlp",
|
|
}
|
|
|
|
|
|
# Architectures with hybrid attention (e.g. Qwen3.5 mixes standard multi-head
|
|
# attention with GatedDeltaNet). On layers that lack the primary attribute,
|
|
# try the fallbacks in order.
|
|
_ATTENTION_ATTR_FALLBACKS: dict[str, list[str]] = {
|
|
"qwen3_5": ["linear_attn"],
|
|
"qwen3_5_text": ["linear_attn"],
|
|
}
|
|
|
|
# Architectures where FFN weights (fc1/fc2) live directly on the layer module
|
|
# instead of inside a dedicated MLP submodule. For these, get_ffn_module
|
|
# returns the layer itself so _project_out_advanced can find fc1/fc2.
|
|
_FLAT_FFN_ARCHS: set[str] = {"opt"}
|
|
|
|
|
|
def _resolve_attr(obj, dotted_path: str):
|
|
"""Resolve a dotted attribute path like 'model.layers'."""
|
|
for attr in dotted_path.split("."):
|
|
obj = getattr(obj, attr)
|
|
return obj
|
|
|
|
|
|
def get_layer_modules(handle: ModelHandle) -> nn.ModuleList:
|
|
"""Return the nn.ModuleList of transformer layers for this model."""
|
|
arch = handle.architecture
|
|
if arch in _LAYER_ATTR_PATHS:
|
|
obj = handle.model
|
|
for attr in _LAYER_ATTR_PATHS[arch]:
|
|
obj = getattr(obj, attr)
|
|
return obj
|
|
|
|
# Fallback: walk the model looking for a ModuleList with the right length.
|
|
# If num_layers is known, match exactly; otherwise find the largest ModuleList
|
|
# (which is almost always the transformer layer stack).
|
|
best = None
|
|
for module in handle.model.modules():
|
|
if isinstance(module, nn.ModuleList) and len(module) > 1:
|
|
if handle.num_layers and len(module) == handle.num_layers:
|
|
return module
|
|
if best is None or len(module) > len(best):
|
|
best = module
|
|
if best is not None:
|
|
return best
|
|
raise RuntimeError(
|
|
f"Cannot locate transformer layers for architecture {arch!r}. "
|
|
f"Supported: {sorted(_LAYER_ATTR_PATHS)}"
|
|
)
|
|
|
|
|
|
def get_attention_module(layer_module: nn.Module, architecture: str) -> nn.Module:
|
|
"""Return the attention sub-module inside a single transformer layer.
|
|
|
|
For hybrid architectures (e.g. Qwen3.5 with GatedDeltaNet), tries
|
|
fallback attribute names when the primary one is missing.
|
|
"""
|
|
attr_path = _ATTENTION_ATTR.get(architecture, "self_attn")
|
|
try:
|
|
return _resolve_attr(layer_module, attr_path)
|
|
except AttributeError:
|
|
for fallback in _ATTENTION_ATTR_FALLBACKS.get(architecture, []):
|
|
try:
|
|
return _resolve_attr(layer_module, fallback)
|
|
except AttributeError:
|
|
continue
|
|
raise
|
|
|
|
|
|
def get_ffn_module(layer_module: nn.Module, architecture: str) -> nn.Module:
|
|
"""Return the FFN/MLP sub-module inside a single transformer layer.
|
|
|
|
For architectures with flat FFN layout (e.g. OPT where fc1/fc2 are
|
|
direct attributes of the layer), returns the layer itself.
|
|
"""
|
|
if architecture in _FLAT_FFN_ARCHS:
|
|
return layer_module
|
|
attr_path = _FFN_ATTR.get(architecture, "mlp")
|
|
return _resolve_attr(layer_module, attr_path)
|
|
|
|
|
|
def get_embedding_module(handle: ModelHandle) -> nn.Embedding:
|
|
"""Return the token embedding module."""
|
|
model = handle.model
|
|
# Try common paths
|
|
for path in [
|
|
"transformer.wte",
|
|
"model.embed_tokens",
|
|
"gpt_neox.embed_in",
|
|
"model.decoder.embed_tokens",
|
|
"transformer.word_embeddings",
|
|
]:
|
|
try:
|
|
emb = _resolve_attr(model, path)
|
|
if isinstance(emb, nn.Embedding):
|
|
return emb
|
|
except AttributeError:
|
|
continue
|
|
|
|
# Fallback: find first Embedding
|
|
for module in model.modules():
|
|
if isinstance(module, nn.Embedding):
|
|
return module
|
|
raise RuntimeError("Cannot locate embedding module.")
|