初始化项目，由ModelHub XC社区提供模型
Model: pathcosmos/frankenstallm Source: Original Platform
This commit is contained in:
ModelHub XC
370
source/model/transformer.py
Normal file
370
source/model/transformer.py
Normal file
@@ -0,0 +1,370 @@
|
||||
"""
|
||||
Full transformer: TransformerBlock and top-level LLM model.
|
||||
Supports pure Transformer and hybrid Mamba-2 + Transformer architectures.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from pathlib import Path
|
||||
from typing import Optional
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
from .config import LMConfig
|
||||
from .layers import RMSNorm, RotaryEmbedding, SwiGLU
|
||||
from .attention import MultiHeadAttention
|
||||
from .mamba_block import Mamba2Block
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Optional TransformerEngine import (FP8 support)
|
||||
# ---------------------------------------------------------------------------
|
||||
try:
|
||||
import transformer_engine.pytorch as te # type: ignore[import]
|
||||
HAS_TE = True
|
||||
except ImportError:
|
||||
te = None # type: ignore[assignment]
|
||||
HAS_TE = False
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# HuggingFace ↔ Custom weight conversion helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def _load_hf_state_dict(path: Path) -> dict[str, torch.Tensor]:
|
||||
"""Load weights from HF safetensors (or pytorch_model.bin fallback)."""
|
||||
safetensors_path = path / "model.safetensors"
|
||||
if safetensors_path.exists():
|
||||
from safetensors.torch import load_file
|
||||
return load_file(str(safetensors_path), device="cpu")
|
||||
bin_path = path / "pytorch_model.bin"
|
||||
if bin_path.exists():
|
||||
return torch.load(bin_path, map_location="cpu", weights_only=True)
|
||||
raise FileNotFoundError(f"No model.safetensors or pytorch_model.bin in {path}")
|
||||
|
||||
|
||||
def _convert_hf_to_custom(hf_sd: dict[str, torch.Tensor], config: LMConfig) -> dict[str, torch.Tensor]:
|
||||
"""Convert HuggingFace LlamaForCausalLM state dict to our custom format.
|
||||
|
||||
Key mapping:
|
||||
HF: model.embed_tokens.weight → embedding.weight
|
||||
HF: model.layers.{i}.self_attn.q/k/v_proj.weight → layers.{i}.attn.qkv_proj.weight (fused)
|
||||
HF: model.layers.{i}.self_attn.o_proj.weight → layers.{i}.attn.out_proj.weight
|
||||
HF: model.layers.{i}.input_layernorm.weight → layers.{i}.attn_norm.weight
|
||||
HF: model.layers.{i}.mlp.gate_proj.weight → layers.{i}.ffn.gate_proj.weight
|
||||
HF: model.layers.{i}.mlp.up_proj.weight → layers.{i}.ffn.up_proj.weight
|
||||
HF: model.layers.{i}.mlp.down_proj.weight → layers.{i}.ffn.down_proj.weight
|
||||
HF: model.layers.{i}.post_attention_layernorm.weight → layers.{i}.ffn_norm.weight
|
||||
HF: model.norm.weight → norm.weight
|
||||
HF: lm_head.weight → lm_head.weight
|
||||
"""
|
||||
sd: dict[str, torch.Tensor] = {}
|
||||
|
||||
sd["embedding.weight"] = hf_sd["model.embed_tokens.weight"]
|
||||
sd["norm.weight"] = hf_sd["model.norm.weight"]
|
||||
sd["lm_head.weight"] = hf_sd["lm_head.weight"]
|
||||
|
||||
for i in range(config.n_layers):
|
||||
pfx = f"model.layers.{i}"
|
||||
out = f"layers.{i}"
|
||||
|
||||
# Fuse Q, K, V into single qkv_proj
|
||||
q = hf_sd[f"{pfx}.self_attn.q_proj.weight"]
|
||||
k = hf_sd[f"{pfx}.self_attn.k_proj.weight"]
|
||||
v = hf_sd[f"{pfx}.self_attn.v_proj.weight"]
|
||||
sd[f"{out}.attn.qkv_proj.weight"] = torch.cat([q, k, v], dim=0)
|
||||
|
||||
sd[f"{out}.attn.out_proj.weight"] = hf_sd[f"{pfx}.self_attn.o_proj.weight"]
|
||||
sd[f"{out}.attn_norm.weight"] = hf_sd[f"{pfx}.input_layernorm.weight"]
|
||||
|
||||
sd[f"{out}.ffn.gate_proj.weight"] = hf_sd[f"{pfx}.mlp.gate_proj.weight"]
|
||||
sd[f"{out}.ffn.up_proj.weight"] = hf_sd[f"{pfx}.mlp.up_proj.weight"]
|
||||
sd[f"{out}.ffn.down_proj.weight"] = hf_sd[f"{pfx}.mlp.down_proj.weight"]
|
||||
sd[f"{out}.ffn_norm.weight"] = hf_sd[f"{pfx}.post_attention_layernorm.weight"]
|
||||
|
||||
return sd
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Transformer Block
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class TransformerBlock(nn.Module):
|
||||
"""Single pre-norm transformer decoder block.
|
||||
|
||||
Layout:
|
||||
x = x + Attention( RMSNorm(x) )
|
||||
x = x + FFN( RMSNorm(x) )
|
||||
"""
|
||||
|
||||
def __init__(self, config: LMConfig) -> None:
|
||||
super().__init__()
|
||||
self.attn_norm = RMSNorm(config.d_model)
|
||||
self.attn = MultiHeadAttention(config)
|
||||
self._use_fp8 = config.use_fp8 and HAS_TE
|
||||
|
||||
if self._use_fp8:
|
||||
# te.LayerNormMLP fuses RMSNorm + gate/up/down projections into one kernel.
|
||||
# It applies normalisation internally, so ffn_norm is not needed.
|
||||
self.ffn_norm = None
|
||||
self.ffn = te.LayerNormMLP(
|
||||
hidden_size=config.d_model,
|
||||
ffn_hidden_size=config.d_ffn,
|
||||
bias=config.bias,
|
||||
activation="swiglu",
|
||||
normalization="RMSNorm",
|
||||
)
|
||||
else:
|
||||
self.ffn_norm = RMSNorm(config.d_model)
|
||||
self.ffn = SwiGLU(config.d_model, config.d_ffn, bias=config.bias)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
x: torch.Tensor,
|
||||
cos: torch.Tensor,
|
||||
sin: torch.Tensor,
|
||||
) -> torch.Tensor:
|
||||
"""
|
||||
Args:
|
||||
x: (B, T, C)
|
||||
cos: (T, head_dim // 2)
|
||||
sin: (T, head_dim // 2)
|
||||
|
||||
Returns:
|
||||
(B, T, C)
|
||||
"""
|
||||
# Pre-norm attention with residual
|
||||
x = x + self.attn(self.attn_norm(x), cos, sin)
|
||||
# FFN with residual — te.LayerNormMLP applies norm internally
|
||||
if self._use_fp8:
|
||||
x = x + self.ffn(x)
|
||||
else:
|
||||
x = x + self.ffn(self.ffn_norm(x))
|
||||
return x
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Full Language Model
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class LLM(nn.Module):
|
||||
"""Decoder-only transformer language model.
|
||||
|
||||
Features:
|
||||
- Learned token embeddings with weight tying to the LM head
|
||||
- Rotary positional embeddings (no learned position embeddings)
|
||||
- Stack of pre-norm TransformerBlocks
|
||||
- Final RMSNorm before the LM head
|
||||
- Optional cross-entropy loss computation (for training)
|
||||
"""
|
||||
|
||||
def __init__(self, config: LMConfig) -> None:
|
||||
super().__init__()
|
||||
self.config = config
|
||||
|
||||
# --- Embedding -------------------------------------------------------
|
||||
self.embedding = nn.Embedding(config.vocab_size, config.d_model)
|
||||
|
||||
# --- Layers (pure Transformer or hybrid Mamba-Transformer) -----------
|
||||
if config.use_hybrid and config.hybrid_pattern:
|
||||
pattern = config.hybrid_pattern.strip().split()
|
||||
if len(pattern) != config.n_layers:
|
||||
raise ValueError(
|
||||
f"hybrid_pattern has {len(pattern)} entries but "
|
||||
f"n_layers={config.n_layers}"
|
||||
)
|
||||
layers: list[nn.Module] = []
|
||||
# Track which layers are Mamba vs Attention for forward dispatch
|
||||
self._layer_types: list[str] = pattern
|
||||
for layer_type in pattern:
|
||||
if layer_type == "M":
|
||||
layers.append(Mamba2Block(
|
||||
d_model=config.d_model,
|
||||
d_state=config.mamba_d_state,
|
||||
head_dim=config.mamba_head_dim,
|
||||
expand=config.mamba_expand,
|
||||
conv_kernel=config.mamba_conv_kernel,
|
||||
n_groups=config.mamba_n_groups,
|
||||
chunk_size=config.mamba_chunk_size,
|
||||
))
|
||||
elif layer_type == "A":
|
||||
layers.append(TransformerBlock(config))
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Unknown layer type '{layer_type}' in hybrid_pattern. "
|
||||
f"Use 'M' (Mamba) or 'A' (Attention)."
|
||||
)
|
||||
self.layers = nn.ModuleList(layers)
|
||||
else:
|
||||
self._layer_types = ["A"] * config.n_layers
|
||||
self.layers = nn.ModuleList(
|
||||
[TransformerBlock(config) for _ in range(config.n_layers)]
|
||||
)
|
||||
|
||||
# --- Final normalisation and LM head ---------------------------------
|
||||
self.norm = RMSNorm(config.d_model)
|
||||
# NOTE: lm_head는 nn.Linear 유지 — embedding weight tying + TE FP8 호환성
|
||||
self.lm_head = nn.Linear(config.d_model, config.vocab_size, bias=False)
|
||||
|
||||
# Weight tying: share embedding and LM-head weight matrices
|
||||
self.lm_head.weight = self.embedding.weight
|
||||
|
||||
# --- Rotary embeddings -----------------------------------------------
|
||||
self.rope = RotaryEmbedding(
|
||||
dim=config.head_dim,
|
||||
max_seq_len=config.max_seq_len,
|
||||
theta=config.rope_theta,
|
||||
)
|
||||
|
||||
# --- Initialise weights ----------------------------------------------
|
||||
self.apply(self._init_weights)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Weight initialisation
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
@staticmethod
|
||||
def _init_weights(module: nn.Module) -> None:
|
||||
"""Apply standard initialisation:
|
||||
- Linear / Embedding weights: N(0, 0.02)
|
||||
- Bias parameters: zeros
|
||||
- te.Linear / te.LayerNormMLP: skipped (TE manages its own init)
|
||||
- Mamba2Block: skipped (manages its own init)
|
||||
"""
|
||||
# TE modules handle their own weight initialisation.
|
||||
if HAS_TE and isinstance(module, (te.Linear, te.LayerNormMLP)):
|
||||
return
|
||||
# Mamba2Block handles its own parameter init (A_log, D, dt_bias, etc.)
|
||||
if isinstance(module, Mamba2Block):
|
||||
return
|
||||
if isinstance(module, (nn.Linear, nn.Embedding)):
|
||||
nn.init.normal_(module.weight, mean=0.0, std=0.02)
|
||||
if isinstance(module, nn.Linear) and module.bias is not None:
|
||||
nn.init.zeros_(module.bias)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Forward pass
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def forward(
|
||||
self,
|
||||
input_ids: torch.Tensor,
|
||||
targets: Optional[torch.Tensor] = None,
|
||||
) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
|
||||
"""
|
||||
Args:
|
||||
input_ids: (B, T) long tensor of token indices
|
||||
targets: (B, T) long tensor of target token indices, or None.
|
||||
Use -1 (ignore_index) to mask positions.
|
||||
|
||||
Returns:
|
||||
logits: (B, T, vocab_size)
|
||||
loss: scalar cross-entropy loss, or None if targets is None
|
||||
"""
|
||||
B, T = input_ids.shape
|
||||
device = input_ids.device
|
||||
|
||||
# Token embeddings: (B, T, C)
|
||||
x = self.embedding(input_ids)
|
||||
|
||||
# Rotary cos/sin for this sequence length: (T, head_dim // 2)
|
||||
# Only needed for Attention layers, but precomputed once for all.
|
||||
cos, sin = self.rope(T, device)
|
||||
|
||||
# Run through blocks — Mamba blocks ignore cos/sin
|
||||
for layer, ltype in zip(self.layers, self._layer_types):
|
||||
if ltype == "M":
|
||||
x = layer(x)
|
||||
else:
|
||||
x = layer(x, cos, sin)
|
||||
|
||||
# Final normalisation
|
||||
x = self.norm(x)
|
||||
|
||||
# LM head: (B, T, vocab_size)
|
||||
logits = self.lm_head(x)
|
||||
|
||||
# Compute loss if targets are provided
|
||||
loss: Optional[torch.Tensor] = None
|
||||
if targets is not None:
|
||||
loss = F.cross_entropy(
|
||||
logits.view(-1, logits.size(-1)),
|
||||
targets.view(-1),
|
||||
ignore_index=-1,
|
||||
)
|
||||
|
||||
return logits, loss
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Properties
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
@property
|
||||
def num_params(self) -> int:
|
||||
"""Number of trainable parameters."""
|
||||
return sum(p.numel() for p in self.parameters() if p.requires_grad)
|
||||
|
||||
def get_input_embeddings(self) -> nn.Embedding:
|
||||
"""HuggingFace-compatible accessor for the token embedding layer."""
|
||||
return self.embedding
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Constructors
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
@classmethod
|
||||
def from_config(cls, config: LMConfig) -> "LLM":
|
||||
"""Construct an LLM from an LMConfig instance."""
|
||||
return cls(config)
|
||||
|
||||
@classmethod
|
||||
def from_pretrained(cls, path: str | Path) -> "LLM":
|
||||
"""Load model from a checkpoint directory.
|
||||
|
||||
Supports two formats (auto-detected):
|
||||
1. Custom: config.yaml + model.pt
|
||||
2. HuggingFace: config.json + model.safetensors (LlamaForCausalLM)
|
||||
"""
|
||||
path = Path(path)
|
||||
|
||||
# --- Custom format ---
|
||||
if (path / "config.yaml").exists():
|
||||
config = LMConfig.from_yaml(path / "config.yaml")
|
||||
model = cls(config)
|
||||
state_dict = torch.load(
|
||||
path / "model.pt",
|
||||
map_location="cpu",
|
||||
weights_only=True,
|
||||
)
|
||||
model.load_state_dict(state_dict)
|
||||
return model
|
||||
|
||||
# --- HuggingFace format ---
|
||||
if (path / "config.json").exists():
|
||||
config = LMConfig.from_hf_config(path / "config.json")
|
||||
model = cls(config)
|
||||
hf_sd = _load_hf_state_dict(path)
|
||||
our_sd = _convert_hf_to_custom(hf_sd, config)
|
||||
model.load_state_dict(our_sd)
|
||||
return model
|
||||
|
||||
raise FileNotFoundError(
|
||||
f"No config.yaml or config.json found in {path}"
|
||||
)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Persistence
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def save_pretrained(self, path: str | Path) -> None:
|
||||
"""Save config and model weights to a directory.
|
||||
|
||||
Creates:
|
||||
<path>/config.yaml
|
||||
<path>/model.pt
|
||||
"""
|
||||
path = Path(path)
|
||||
path.mkdir(parents=True, exist_ok=True)
|
||||
self.config.to_yaml(path / "config.yaml")
|
||||
torch.save(self.state_dict(), path / "model.pt")
|
||||
Reference in New Issue
Block a user