371 lines
14 KiB
Python
371 lines
14 KiB
Python
|
"""
|
||
|
|
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")
|