來(lái)源：Coggle數(shù)據(jù)科學(xué)

在當(dāng)今人工智能領(lǐng)域，大型語(yǔ)言模型（LLM）的開(kāi)發(fā)已經(jīng)成為一個(gè)熱門話題。這些模型通過(guò)學(xué)習(xí)大量的文本數(shù)據(jù)，能夠生成自然語(yǔ)言文本，完成各種復(fù)雜的任務(wù)，如寫(xiě)作、翻譯、問(wèn)答等。

https://github.com/FareedKhan-dev/train-llm-from-scratch

本文將為你提供一個(gè)簡(jiǎn)單直接的方法，從下載數(shù)據(jù)到生成文本，帶你一步步構(gòu)建大院模型。

步驟1：GPU設(shè)備

在開(kāi)始訓(xùn)練語(yǔ)言模型之前，你需要對(duì)面向?qū)ο?a href="http://www.www27dydycom.cn/v/tag/1315/" target="_blank">編程（OOP）、神經(jīng)網(wǎng)絡(luò)（NN）和 PyTorch 有基本的了解。

訓(xùn)練語(yǔ)言模型需要強(qiáng)大的計(jì)算資源，尤其是 GPU。不同的 GPU 在內(nèi)存容量和計(jì)算能力上有所不同，適合不同規(guī)模的模型訓(xùn)練。以下是一個(gè)詳細(xì)的 GPU 對(duì)比表，幫助你選擇合適的硬件。

13M LLM 訓(xùn)練

• 參數(shù)規(guī)模：1300 萬(wàn)參數(shù)。
• 應(yīng)用場(chǎng)景：適合初學(xué)者和小型項(xiàng)目，例如簡(jiǎn)單的文本生成、情感分析或語(yǔ)言理解任務(wù)。
• 硬件需求：相對(duì)較低。大多數(shù)中高端 GPU（如 NVIDIA RTX 3060 或更高）都可以勝任。
• 特點(diǎn)：訓(xùn)練速度快，資源消耗低，適合快速迭代和實(shí)驗(yàn)。

2B LLM 訓(xùn)練

• 參數(shù)規(guī)模：20 億參數(shù)。
• 應(yīng)用場(chǎng)景：適合更復(fù)雜的任務(wù)，如高質(zhì)量的文本生成、多語(yǔ)言翻譯或更高級(jí)的語(yǔ)言理解任務(wù)。
• 硬件需求：較高。需要至少 16GB 內(nèi)存的 GPU，如 NVIDIA RTX 3090 或更高配置。
• 特點(diǎn)：能夠生成更流暢、更自然的文本，但訓(xùn)練時(shí)間長(zhǎng)，資源消耗大。

步驟2：導(dǎo)入環(huán)境

在開(kāi)始之前，我們需要導(dǎo)入一些必要的 Python 庫(kù)。這些庫(kù)將幫助我們處理數(shù)據(jù)、構(gòu)建模型以及訓(xùn)練模型。

# PyTorch for deep learning functions and tensors
import torch
import torch.nn as nn
import torch.nn.functional as F

# Numerical operations and arrays handling
import numpy as np

# Handling HDF5 files
import h5py

# Operating system and file management
import os

# Command-line argument parsing
import argparse

# HTTP requests and interactions
import requests

# Progress bar for loops
from tqdm import tqdm

# JSON handling
import json

# Zstandard compression library
import zstandard as zstd

# Tokenization library for large language models
import tiktoken

# Math operations (used for advanced math functions)
import math

步驟3：加載數(shù)據(jù)集

The Pile 數(shù)據(jù)集是一個(gè)大規(guī)模、多樣化的開(kāi)源數(shù)據(jù)集，專為語(yǔ)言模型訓(xùn)練設(shè)計(jì)。它由 22 個(gè)子數(shù)據(jù)集組成，涵蓋了書(shū)籍、文章、維基百科、代碼、新聞等多種類型的文本。

# Download validation dataset
!wget https://huggingface.co/datasets/monology/pile-uncopyrighted/resolve/main/val.jsonl.zst

# Download the first part of the training dataset
!wget https://huggingface.co/datasets/monology/pile-uncopyrighted/resolve/main/train/00.jsonl.zst

# Download the second part of the training dataset
!wget https://huggingface.co/datasets/monology/pile-uncopyrighted/resolve/main/train/01.jsonl.zst

# Download the third part of the training dataset
!wget https://huggingface.co/datasets/monology/pile-uncopyrighted/resolve/main/train/02.jsonl.zst

最終處理好的數(shù)據(jù)集格式如下：

#### OUTPUT ####
Line: 0
{
"text":"Effect of sleep quality ... epilepsy.",
"meta": {
"pile_set_name":"PubMed Abstracts"
}
}

Line: 1
{
"text":"LLMops a new GitHub Repository ...",
"meta": {
"pile_set_name":"Github"
}
}

步驟4：Transformer 架構(gòu)

Transformer 通過(guò)將文本分解成更小的單元，稱為“標(biāo)記”（token），并預(yù)測(cè)序列中的下一個(gè)標(biāo)記來(lái)工作。Transformer 由多個(gè)層組成，這些層被稱為 Transformer 塊，它們一層疊一層，最后通過(guò)一個(gè)最終層來(lái)進(jìn)行預(yù)測(cè)。

每個(gè) Transformer 塊包含兩個(gè)主要組件：

1. 自注意力頭（Self-Attention Heads）

自注意力頭的作用是確定輸入中哪些部分對(duì)模型來(lái)說(shuō)最為重要。例如，在處理一個(gè)句子時(shí)，自注意力頭可以突出顯示單詞之間的關(guān)系，比如代詞與其所指代的名詞之間的關(guān)系。通過(guò)這種方式，模型能夠更好地理解句子的結(jié)構(gòu)和語(yǔ)義。

2. 多層感知器（MLP，Multi-Layer Perceptron）

多層感知器是一個(gè)簡(jiǎn)單的前饋神經(jīng)網(wǎng)絡(luò)。它接收自注意力頭強(qiáng)調(diào)的信息，并進(jìn)一步處理這些信息。MLP 包含：

• 輸入層：接收來(lái)自自注意力頭的數(shù)據(jù)。
• 隱藏層：為處理過(guò)程增加復(fù)雜性。
• 輸出層：將處理結(jié)果傳遞給下一個(gè) Transformer 塊。

1. 輸入嵌入與位置編碼 ：輸入的文本被分解為標(biāo)記（tokens），然后轉(zhuǎn)換為嵌入向量（embeddings）。同時(shí)，加入位置編碼（position embeddings）以提供標(biāo)記的位置信息。
2. Transformer 塊堆疊 ：模型由多個(gè)（例如 64 個(gè)）相同的 Transformer 塊組成，這些塊依次對(duì)數(shù)據(jù)進(jìn)行處理。
3. 多頭注意力機(jī)制 ：每個(gè) Transformer 塊首先通過(guò)多頭注意力機(jī)制（例如 16 個(gè)頭）分析標(biāo)記之間的關(guān)系，捕捉不同類型的關(guān)聯(lián)。
4. MLP 處理 ：注意力機(jī)制處理后的數(shù)據(jù)通過(guò)一個(gè) MLP（多層感知器）進(jìn)行進(jìn)一步處理，先擴(kuò)展到更大維度（例如 4 倍），再壓縮回原始維度。
5. 殘差連接與層歸一化 ：每個(gè) Transformer 塊中使用殘差連接（residual connections）和層歸一化（layer normalization），以幫助信息流動(dòng)并穩(wěn)定訓(xùn)練。
6. 最終預(yù)測(cè) ：理后的數(shù)據(jù)通過(guò)最終層，轉(zhuǎn)換為詞匯表大小的預(yù)測(cè)結(jié)果，生成下一個(gè)最有可能的標(biāo)記。
7. 文本生成 ：模型通過(guò)反復(fù)預(yù)測(cè)下一個(gè)標(biāo)記，逐步生成完整的文本序列。

步驟5：多層感知器（MLP）

多層感知器（MLP）是 Transformer 架構(gòu)中前饋神經(jīng)網(wǎng)絡(luò)（Feed-Forward Network, FFN）的核心組成部分。它的主要作用是引入非線性特性，并學(xué)習(xí)嵌入表示中的復(fù)雜關(guān)系。在定義 MLP 模塊時(shí)，一個(gè)重要的參數(shù)是n_embed，它定義了輸入嵌入的維度。

MLP 的整個(gè)序列轉(zhuǎn)換過(guò)程使得它能夠?qū)ψ⒁饬C(jī)制學(xué)習(xí)到的表示進(jìn)行進(jìn)一步的精細(xì)化處理。具體來(lái)說(shuō)：

• 引入非線性：通過(guò) ReLU 激活函數(shù)，MLP 能夠捕捉到嵌入表示中的復(fù)雜非線性關(guān)系。
• 特征增強(qiáng)：隱藏層的擴(kuò)展操作為模型提供了更大的空間來(lái)學(xué)習(xí)更豐富的特征。
• 維度一致性：投影線性層確保 MLP 的輸出維度與輸入維度一致，便于后續(xù)的 Transformer 塊繼續(xù)處理。

# --- MLP (Multi-Layer Perceptron) Class ---

class MLP(nn.Module):
"""
A simple Multi-Layer Perceptron with one hidden layer.

This module is used within the Transformer block for feed-forward processing.
It expands the input embedding size, applies a ReLU activation, and then projects it back
to the original embedding size.
"""
def __init__(self, n_embed):
super().__init__()
self.hidden = nn.Linear(n_embed, 4 * n_embed) # Linear layer to expand embedding size
self.relu = nn.ReLU() # ReLU activation function
self.proj = nn.Linear(4 * n_embed, n_embed) # Linear layer to project back to original size

def forward(self, x):
"""
Forward pass through the MLP.

Args:
x (torch.Tensor): Input tensor of shape (B, T, C), where B is batch size,
T is sequence length, and C is embedding size.

Returns:
torch.Tensor: Output tensor of the same shape as the input.
"""
x = self.forward_embedding(x)
x = self.project_embedding(x)
returnx

def forward_embedding(self, x):
"""
Applies the hidden linear layer followed by ReLU activation.

Args:
x (torch.Tensor): Input tensor.

Returns:
torch.Tensor: Output after the hidden layer and ReLU.
"""
x = self.relu(self.hidden(x))
returnx

def project_embedding(self, x):
"""
Applies the projection linear layer.

Args:
x (torch.Tensor): Input tensor.

Returns:
torch.Tensor: Output after the projection layer.
"""
x = self.proj(x)
returnx

步驟6：Single Head Attention

注意力頭（Attention Head）是 Transformer 模型的核心部分，其主要作用是讓模型能夠?qū)Ｗ⒂谳斎胄蛄兄信c當(dāng)前任務(wù)最相關(guān)的部分。在定義注意力頭模塊時(shí)，有幾個(gè)重要的參數(shù)：

• head_size：決定了鍵（key）、查詢（query）和值（value）投影的維度，從而影響注意力機(jī)制的表達(dá)能力。
• n_embed：輸入嵌入的維度，定義了這些投影層的輸入大小。
• context_length：用于創(chuàng)建因果掩碼（causal mask），確保模型只能關(guān)注前面的標(biāo)記，從而實(shí)現(xiàn)自回歸（autoregressive）特性。

在注意力頭內(nèi)部，我們初始化了三個(gè)無(wú)偏置的線性層（nn.Linear），分別用于鍵、查詢和值的投影。此外，我們注冊(cè)了一個(gè)大小為context_length x context_length的下三角矩陣（tril）作為緩沖區(qū)（buffer），以實(shí)現(xiàn)因果掩碼，防止注意力機(jī)制關(guān)注未來(lái)的標(biāo)記。

# --- Attention Head Class ---

class Head(nn.Module):
def __init__(self, head_size, n_embed, context_length):
super().__init__()
self.key = nn.Linear(n_embed, head_size, bias=False) # Key projection
self.query = nn.Linear(n_embed, head_size, bias=False)# Query projection
self.value = nn.Linear(n_embed, head_size, bias=False)# Value projection
# Lower triangular matrix for causal masking
self.register_buffer('tril', torch.tril(torch.ones(context_length, context_length)))

def forward(self, x):
B, T, C = x.shape
k = self.key(x) # (B, T, head_size)
q = self.query(x) # (B, T, head_size)
scale_factor = 1 / math.sqrt(C)
# Calculate attention weights: (B, T, head_size) @ (B, head_size, T) -> (B, T, T)
attn_weights = q @ k.transpose(-2, -1) * scale_factor
# Apply causal masking
attn_weights = attn_weights.masked_fill(self.tril[:T, :T] == 0,float('-inf'))
attn_weights = F.softmax(attn_weights, dim=-1)
v = self.value(x) # (B, T, head_size)
# Apply attention weights to values
out = attn_weights @ v# (B, T, T) @ (B, T, head_size) -> (B, T, head_size)
returnout

步驟7：Multi Head Attention

多頭注意力（Multi-Head Attention）是 Transformer 架構(gòu)中的關(guān)鍵機(jī)制，用于捕捉輸入序列中多樣化的關(guān)聯(lián)關(guān)系。通過(guò)將多個(gè)獨(dú)立的注意力頭（attention heads）并行運(yùn)行，模型能夠同時(shí)關(guān)注輸入的不同方面，從而更全面地理解序列信息。

• n_head：決定了并行運(yùn)行的注意力頭的數(shù)量。每個(gè)注意力頭獨(dú)立處理輸入數(shù)據(jù)，從而讓模型能夠從多個(gè)角度捕捉輸入序列中的關(guān)系。
• 上下文長(zhǎng)度context_length定義了輸入序列的長(zhǎng)度，用于創(chuàng)建因果掩碼（causal mask），確保模型只能關(guān)注前面的標(biāo)記，從而實(shí)現(xiàn)自回歸特性。

class MultiHeadAttention(nn.Module):
"""
Multi-Head Attention module.

This module combines multiple attention heads in parallel. The outputs of each head
are concatenated to form the final output.
"""
def __init__(self, n_head, n_embed, context_length):
super().__init__()
self.heads = nn.ModuleList([Head(n_embed // n_head, n_embed, context_length)for_inrange(n_head)])

def forward(self, x):
"""
Forward pass through the multi-head attention.

Args:
x (torch.Tensor): Input tensor of shape (B, T, C).

Returns:
torch.Tensor: Output tensor after concatenating the outputs of all heads.
"""
# Concatenate the output of each head along the last dimension (C)
x = torch.cat([h(x)forhinself.heads], dim=-1)
returnx

步驟8：Transformer 塊

Transformer 塊是 Transformer 架構(gòu)的核心單元，它通過(guò)組合多頭注意力機(jī)制和前饋網(wǎng)絡(luò)（MLP），并應(yīng)用層歸一化（Layer Normalization）以及殘差連接（Residual Connections），來(lái)處理輸入并學(xué)習(xí)復(fù)雜的模式。

• n_head：多頭注意力中并行注意力頭的數(shù)量。
• n_embed：輸入嵌入的維度，也是層歸一化的參數(shù)維度。
• context_length：上下文長(zhǎng)度，用于定義序列的長(zhǎng)度，并創(chuàng)建因果掩碼。

每個(gè) Transformer 塊包含以下部分：

1. 多頭注意力層（Multi-Head Attention）：負(fù)責(zé)捕捉輸入序列中不同位置之間的關(guān)系。
2. 前饋網(wǎng)絡(luò)（MLP）：用于進(jìn)一步處理多頭注意力層的輸出，引入非線性特性。
3. 層歸一化（Layer Normalization）：在每個(gè)子層之前應(yīng)用，有助于穩(wěn)定訓(xùn)練。
4. 殘差連接（Residual Connections）：在每個(gè)子層之后應(yīng)用，幫助信息流動(dòng)并緩解深層網(wǎng)絡(luò)訓(xùn)練中的梯度消失問(wèn)題。

class Block(nn.Module):
def __init__(self, n_head, n_embed, context_length):
super().__init__()
self.ln1 = nn.LayerNorm(n_embed)
self.attn = MultiHeadAttention(n_head, n_embed, context_length)
self.ln2 = nn.LayerNorm(n_embed)
self.mlp = MLP(n_embed)

def forward(self, x):
# Apply multi-head attention with residual connection
x = x + self.attn(self.ln1(x))
# Apply MLP with residual connection
x = x + self.mlp(self.ln2(x))
returnx

def forward_embedding(self, x):
res = x + self.attn(self.ln1(x))
x = self.mlp.forward_embedding(self.ln2(res))
returnx, res

步驟9：完整模型結(jié)構(gòu)

到目前為止，我們已經(jīng)編寫(xiě)了 Transformer 模型的一些小部件，如多頭注意力（Multi-Head Attention）和 MLP（多層感知器）。接下來(lái)，我們需要將這些部件整合起來(lái)，構(gòu)建一個(gè)完整的 Transformer 模型，用于執(zhí)行序列到序列的任務(wù)。為此，我們需要定義幾個(gè)關(guān)鍵參數(shù)：n_head、n_embed、context_length、vocab_size和N_BLOCKS。

# --- Transformer Model Class ---

class Transformer(nn.Module):
"""
The main Transformer model.

This class combines token and position embeddings with a sequence of Transformer blocks
and a final linear layer for language modeling.
"""
def __init__(self, n_head, n_embed, context_length, vocab_size, N_BLOCKS):
super().__init__()
self.context_length = context_length
self.N_BLOCKS = N_BLOCKS
self.token_embed = nn.Embedding(vocab_size, n_embed)
self.position_embed = nn.Embedding(context_length, n_embed)
self.attn_blocks = nn.ModuleList([Block(n_head, n_embed, context_length)for_inrange(N_BLOCKS)])
self.layer_norm = nn.LayerNorm(n_embed)
self.lm_head = nn.Linear(n_embed, vocab_size)
self.register_buffer('pos_idxs', torch.arange(context_length))

def _pre_attn_pass(self, idx):
B, T = idx.shape
tok_embedding = self.token_embed(idx)
pos_embedding = self.position_embed(self.pos_idxs[:T])
returntok_embedding + pos_embedding

def forward(self, idx, targets=None):
x = self._pre_attn_pass(idx)
forblockinself.attn_blocks:
x = block(x)
x = self.layer_norm(x)
logits = self.lm_head(x)
loss = None
iftargets is not None:
B, T, C = logits.shape
flat_logits = logits.view(B * T, C)
targets = targets.view(B * T).long()
loss = F.cross_entropy(flat_logits, targets)
returnlogits, loss

def forward_embedding(self, idx):
x = self._pre_attn_pass(idx)
residual = x
forblockinself.attn_blocks:
x, residual = block.forward_embedding(x)
returnx, residual

def generate(self, idx, max_new_tokens):
for_inrange(max_new_tokens):
idx_cond = idx[:, -self.context_length:]
logits, _ = self(idx_cond)
logits = logits[:, -1, :]
probs = F.softmax(logits, dim=-1)
idx_next = torch.multinomial(probs, num_samples=1)
idx = torch.cat((idx, idx_next), dim=1)
returnidx

步驟10：訓(xùn)練參數(shù)配置

現(xiàn)在我們已經(jīng)完成了模型的編碼工作，接下來(lái)需要定義訓(xùn)練參數(shù)，包括注意力頭的數(shù)量、Transformer 塊的數(shù)量等，以及數(shù)據(jù)路徑等相關(guān)配置。

• VOCAB_SIZE：詞匯表大小，表示詞匯表中唯一標(biāo)記的數(shù)量。
• CONTEXT_LENGTH：模型能夠處理的最大序列長(zhǎng)度。
• N_EMBED：嵌入空間的維度，決定了標(biāo)記嵌入和位置嵌入的大小。
• N_HEAD：每個(gè) Transformer 塊中的注意力頭數(shù)量。
• N_BLOCKS：模型中 Transformer 塊的數(shù)量，決定了模型的深度。
• T_BATCH_SIZE：訓(xùn)練時(shí)每個(gè)批次的樣本數(shù)量。
• T_CONTEXT_LENGTH：訓(xùn)練批次的上下文長(zhǎng)度。
• T_TRAIN_STEPS：總訓(xùn)練步數(shù)。

步驟11：模型訓(xùn)練

我們使用 AdamW 優(yōu)化器，這是一種改進(jìn)版的 Adam 優(yōu)化器，適用于深度學(xué)習(xí)任務(wù)。

• 高初始損失：20 億參數(shù)模型在訓(xùn)練初期的損失值通常非常高。這是因?yàn)槟Ｐ偷膮?shù)量巨大，初始權(quán)重隨機(jī)分布，導(dǎo)致模型在開(kāi)始時(shí)對(duì)數(shù)據(jù)的理解非常有限。
• 劇烈波動(dòng)：在訓(xùn)練的前幾輪，損失值可能會(huì)出現(xiàn)劇烈波動(dòng)。這是因?yàn)槟Ｐ托枰邶嫶蟮膮?shù)空間中尋找合適的權(quán)重組合，而初始的學(xué)習(xí)率可能過(guò)高，導(dǎo)致優(yōu)化過(guò)程不穩(wěn)定。
• 快速下降：盡管初始損失很高，但隨著訓(xùn)練的進(jìn)行，損失值通常會(huì)迅速下降。這是因?yàn)槟Ｐ椭饾u開(kāi)始學(xué)習(xí)數(shù)據(jù)中的模式和結(jié)構(gòu)。
• 波動(dòng)調(diào)整：在快速下降之后，損失值可能會(huì)出現(xiàn)一些波動(dòng)。這是因?yàn)槟Ｐ驮谡{(diào)整權(quán)重時(shí)，可能會(huì)在不同的局部最優(yōu)解之間徘徊。這種波動(dòng)表明模型在尋找更穩(wěn)定的全局最優(yōu)解。

步驟12：生成文本

接下來(lái)，我們將創(chuàng)建一個(gè)函數(shù)generate_text，用于從保存的模型中生成文本。該函數(shù)接受保存的模型路徑和輸入文本作為輸入，并返回生成的文本。我們還將比較數(shù)百萬(wàn)參數(shù)模型和數(shù)十億參數(shù)模型在生成文本時(shí)的表現(xiàn)。

def generate_text(model_path, input_text, max_length=512, device="gpu"):
# Load the model checkpoint
checkpoint = torch.load(model_path)

# Initialize the model (you should ensure that the Transformer class is defined elsewhere)
model = Transformer().to(device)

# Load the model's state dictionary
model.load_state_dict(checkpoint['model_state_dict'])

# Load the tokenizer for the GPT model (we use 'r50k_base' for GPT models)
enc = tiktoken.get_encoding('r50k_base')

# Encode the input text along with the end-of-text token
input_ids = torch.tensor(
enc.encode(input_text, allowed_special={'<|endoftext|>'}),
dtype=torch.long
)[None, :].to(device) # Add batch dimension and move to the specified device

# Generate text with the model using the encoded input
with torch.no_grad():
# Generate up to 'max_length' tokens of text
generated_output = model.generate(input_ids, max_length)

# Decode the generated tokens back into text
generated_text = enc.decode(generated_output[0].tolist())

returngenerated_text