mirror of
https://github.com/carlospolop/hacktricks
synced 2024-11-26 14:40:37 +00:00
667 lines
28 KiB
Markdown
667 lines
28 KiB
Markdown
# 5. Arquitectura LLM
|
|
|
|
## Arquitectura LLM
|
|
|
|
{% hint style="success" %}
|
|
El objetivo de esta quinta fase es muy simple: **Desarrollar la arquitectura del LLM completo**. Juntar todo, aplicar todas las capas y crear todas las funciones para generar texto o transformar texto a IDs y viceversa.
|
|
|
|
Esta arquitectura se utilizará tanto para entrenar como para predecir texto después de haber sido entrenada.
|
|
{% endhint %}
|
|
|
|
Ejemplo de arquitectura LLM de [https://github.com/rasbt/LLMs-from-scratch/blob/main/ch04/01\_main-chapter-code/ch04.ipynb](https://github.com/rasbt/LLMs-from-scratch/blob/main/ch04/01\_main-chapter-code/ch04.ipynb):
|
|
|
|
Se puede observar una representación de alto nivel en:
|
|
|
|
<figure><img src="../../.gitbook/assets/image (3) (1).png" alt="" width="563"><figcaption><p><a href="https://camo.githubusercontent.com/6c8c392f72d5b9e86c94aeb9470beab435b888d24135926f1746eb88e0cc18fb/68747470733a2f2f73656261737469616e72617363686b612e636f6d2f696d616765732f4c4c4d732d66726f6d2d736372617463682d696d616765732f636830345f636f6d707265737365642f31332e776562703f31">https://camo.githubusercontent.com/6c8c392f72d5b9e86c94aeb9470beab435b888d24135926f1746eb88e0cc18fb/68747470733a2f2f73656261737469616e72617363686b612e636f6d2f696d616765732f4c4c4d732d66726f6d2d736372617463682d696d616765732f636830345f636f6d707265737365642f31332e776562703f31</a></p></figcaption></figure>
|
|
|
|
1. **Entrada (Texto Tokenizado)**: El proceso comienza con texto tokenizado, que se convierte en representaciones numéricas.
|
|
2. **Capa de Embedding de Tokens y Capa de Embedding Posicional**: El texto tokenizado pasa a través de una **capa de embedding de tokens** y una **capa de embedding posicional**, que captura la posición de los tokens en una secuencia, crítica para entender el orden de las palabras.
|
|
3. **Bloques de Transformer**: El modelo contiene **12 bloques de transformer**, cada uno con múltiples capas. Estos bloques repiten la siguiente secuencia:
|
|
* **Atención Multi-Cabeza enmascarada**: Permite que el modelo se enfoque en diferentes partes del texto de entrada a la vez.
|
|
* **Normalización de Capa**: Un paso de normalización para estabilizar y mejorar el entrenamiento.
|
|
* **Capa Feed Forward**: Responsable de procesar la información de la capa de atención y hacer predicciones sobre el siguiente token.
|
|
* **Capas de Dropout**: Estas capas previenen el sobreajuste al eliminar aleatoriamente unidades durante el entrenamiento.
|
|
4. **Capa de Salida Final**: El modelo produce un **tensor de 4x50,257 dimensiones**, donde **50,257** representa el tamaño del vocabulario. Cada fila en este tensor corresponde a un vector que el modelo utiliza para predecir la siguiente palabra en la secuencia.
|
|
5. **Objetivo**: El objetivo es tomar estos embeddings y convertirlos de nuevo en texto. Específicamente, la última fila de la salida se utiliza para generar la siguiente palabra, representada como "forward" en este diagrama.
|
|
|
|
### Representación de Código
|
|
```python
|
|
import torch
|
|
import torch.nn as nn
|
|
import tiktoken
|
|
|
|
class GELU(nn.Module):
|
|
def __init__(self):
|
|
super().__init__()
|
|
|
|
def forward(self, x):
|
|
return 0.5 * x * (1 + torch.tanh(
|
|
torch.sqrt(torch.tensor(2.0 / torch.pi)) *
|
|
(x + 0.044715 * torch.pow(x, 3))
|
|
))
|
|
|
|
class FeedForward(nn.Module):
|
|
def __init__(self, cfg):
|
|
super().__init__()
|
|
self.layers = nn.Sequential(
|
|
nn.Linear(cfg["emb_dim"], 4 * cfg["emb_dim"]),
|
|
GELU(),
|
|
nn.Linear(4 * cfg["emb_dim"], cfg["emb_dim"]),
|
|
)
|
|
|
|
def forward(self, x):
|
|
return self.layers(x)
|
|
|
|
class MultiHeadAttention(nn.Module):
|
|
def __init__(self, d_in, d_out, context_length, dropout, num_heads, qkv_bias=False):
|
|
super().__init__()
|
|
assert d_out % num_heads == 0, "d_out must be divisible by num_heads"
|
|
|
|
self.d_out = d_out
|
|
self.num_heads = num_heads
|
|
self.head_dim = d_out // num_heads # Reduce the projection dim to match desired output dim
|
|
|
|
self.W_query = nn.Linear(d_in, d_out, bias=qkv_bias)
|
|
self.W_key = nn.Linear(d_in, d_out, bias=qkv_bias)
|
|
self.W_value = nn.Linear(d_in, d_out, bias=qkv_bias)
|
|
self.out_proj = nn.Linear(d_out, d_out) # Linear layer to combine head outputs
|
|
self.dropout = nn.Dropout(dropout)
|
|
self.register_buffer('mask', torch.triu(torch.ones(context_length, context_length), diagonal=1))
|
|
|
|
def forward(self, x):
|
|
b, num_tokens, d_in = x.shape
|
|
|
|
keys = self.W_key(x) # Shape: (b, num_tokens, d_out)
|
|
queries = self.W_query(x)
|
|
values = self.W_value(x)
|
|
|
|
# We implicitly split the matrix by adding a `num_heads` dimension
|
|
# Unroll last dim: (b, num_tokens, d_out) -> (b, num_tokens, num_heads, head_dim)
|
|
keys = keys.view(b, num_tokens, self.num_heads, self.head_dim)
|
|
values = values.view(b, num_tokens, self.num_heads, self.head_dim)
|
|
queries = queries.view(b, num_tokens, self.num_heads, self.head_dim)
|
|
|
|
# Transpose: (b, num_tokens, num_heads, head_dim) -> (b, num_heads, num_tokens, head_dim)
|
|
keys = keys.transpose(1, 2)
|
|
queries = queries.transpose(1, 2)
|
|
values = values.transpose(1, 2)
|
|
|
|
# Compute scaled dot-product attention (aka self-attention) with a causal mask
|
|
attn_scores = queries @ keys.transpose(2, 3) # Dot product for each head
|
|
|
|
# Original mask truncated to the number of tokens and converted to boolean
|
|
mask_bool = self.mask.bool()[:num_tokens, :num_tokens]
|
|
|
|
# Use the mask to fill attention scores
|
|
attn_scores.masked_fill_(mask_bool, -torch.inf)
|
|
|
|
attn_weights = torch.softmax(attn_scores / keys.shape[-1]**0.5, dim=-1)
|
|
attn_weights = self.dropout(attn_weights)
|
|
|
|
# Shape: (b, num_tokens, num_heads, head_dim)
|
|
context_vec = (attn_weights @ values).transpose(1, 2)
|
|
|
|
# Combine heads, where self.d_out = self.num_heads * self.head_dim
|
|
context_vec = context_vec.contiguous().view(b, num_tokens, self.d_out)
|
|
context_vec = self.out_proj(context_vec) # optional projection
|
|
|
|
return context_vec
|
|
|
|
class LayerNorm(nn.Module):
|
|
def __init__(self, emb_dim):
|
|
super().__init__()
|
|
self.eps = 1e-5
|
|
self.scale = nn.Parameter(torch.ones(emb_dim))
|
|
self.shift = nn.Parameter(torch.zeros(emb_dim))
|
|
|
|
def forward(self, x):
|
|
mean = x.mean(dim=-1, keepdim=True)
|
|
var = x.var(dim=-1, keepdim=True, unbiased=False)
|
|
norm_x = (x - mean) / torch.sqrt(var + self.eps)
|
|
return self.scale * norm_x + self.shift
|
|
|
|
class TransformerBlock(nn.Module):
|
|
def __init__(self, cfg):
|
|
super().__init__()
|
|
self.att = MultiHeadAttention(
|
|
d_in=cfg["emb_dim"],
|
|
d_out=cfg["emb_dim"],
|
|
context_length=cfg["context_length"],
|
|
num_heads=cfg["n_heads"],
|
|
dropout=cfg["drop_rate"],
|
|
qkv_bias=cfg["qkv_bias"])
|
|
self.ff = FeedForward(cfg)
|
|
self.norm1 = LayerNorm(cfg["emb_dim"])
|
|
self.norm2 = LayerNorm(cfg["emb_dim"])
|
|
self.drop_shortcut = nn.Dropout(cfg["drop_rate"])
|
|
|
|
def forward(self, x):
|
|
# Shortcut connection for attention block
|
|
shortcut = x
|
|
x = self.norm1(x)
|
|
x = self.att(x) # Shape [batch_size, num_tokens, emb_size]
|
|
x = self.drop_shortcut(x)
|
|
x = x + shortcut # Add the original input back
|
|
|
|
# Shortcut connection for feed forward block
|
|
shortcut = x
|
|
x = self.norm2(x)
|
|
x = self.ff(x)
|
|
x = self.drop_shortcut(x)
|
|
x = x + shortcut # Add the original input back
|
|
|
|
return x
|
|
|
|
|
|
class GPTModel(nn.Module):
|
|
def __init__(self, cfg):
|
|
super().__init__()
|
|
self.tok_emb = nn.Embedding(cfg["vocab_size"], cfg["emb_dim"])
|
|
self.pos_emb = nn.Embedding(cfg["context_length"], cfg["emb_dim"])
|
|
self.drop_emb = nn.Dropout(cfg["drop_rate"])
|
|
|
|
self.trf_blocks = nn.Sequential(
|
|
*[TransformerBlock(cfg) for _ in range(cfg["n_layers"])])
|
|
|
|
self.final_norm = LayerNorm(cfg["emb_dim"])
|
|
self.out_head = nn.Linear(
|
|
cfg["emb_dim"], cfg["vocab_size"], bias=False
|
|
)
|
|
|
|
def forward(self, in_idx):
|
|
batch_size, seq_len = in_idx.shape
|
|
tok_embeds = self.tok_emb(in_idx)
|
|
pos_embeds = self.pos_emb(torch.arange(seq_len, device=in_idx.device))
|
|
x = tok_embeds + pos_embeds # Shape [batch_size, num_tokens, emb_size]
|
|
x = self.drop_emb(x)
|
|
x = self.trf_blocks(x)
|
|
x = self.final_norm(x)
|
|
logits = self.out_head(x)
|
|
return logits
|
|
|
|
GPT_CONFIG_124M = {
|
|
"vocab_size": 50257, # Vocabulary size
|
|
"context_length": 1024, # Context length
|
|
"emb_dim": 768, # Embedding dimension
|
|
"n_heads": 12, # Number of attention heads
|
|
"n_layers": 12, # Number of layers
|
|
"drop_rate": 0.1, # Dropout rate
|
|
"qkv_bias": False # Query-Key-Value bias
|
|
}
|
|
|
|
torch.manual_seed(123)
|
|
model = GPTModel(GPT_CONFIG_124M)
|
|
out = model(batch)
|
|
print("Input batch:\n", batch)
|
|
print("\nOutput shape:", out.shape)
|
|
print(out)
|
|
```
|
|
### **Función de Activación GELU**
|
|
```python
|
|
# From https://github.com/rasbt/LLMs-from-scratch/tree/main/ch04
|
|
class GELU(nn.Module):
|
|
def __init__(self):
|
|
super().__init__()
|
|
|
|
def forward(self, x):
|
|
return 0.5 * x * (1 + torch.tanh(
|
|
torch.sqrt(torch.tensor(2.0 / torch.pi)) *
|
|
(x + 0.044715 * torch.pow(x, 3))
|
|
))
|
|
```
|
|
#### **Propósito y Funcionalidad**
|
|
|
|
* **GELU (Unidad Lineal de Error Gaussiano):** Una función de activación que introduce no linealidad en el modelo.
|
|
* **Activación Suave:** A diferencia de ReLU, que anula las entradas negativas, GELU mapea suavemente las entradas a salidas, permitiendo valores pequeños y no nulos para entradas negativas.
|
|
* **Definición Matemática:**
|
|
|
|
<figure><img src="../../.gitbook/assets/image (2) (1).png" alt=""><figcaption></figcaption></figure>
|
|
|
|
{% hint style="info" %}
|
|
El objetivo del uso de esta función después de las capas lineales dentro de la capa FeedForward es cambiar los datos lineales a no lineales para permitir que el modelo aprenda relaciones complejas y no lineales.
|
|
{% endhint %}
|
|
|
|
### **Red Neuronal FeedForward**
|
|
|
|
_Las formas se han añadido como comentarios para entender mejor las formas de las matrices:_
|
|
```python
|
|
# From https://github.com/rasbt/LLMs-from-scratch/tree/main/ch04
|
|
class FeedForward(nn.Module):
|
|
def __init__(self, cfg):
|
|
super().__init__()
|
|
self.layers = nn.Sequential(
|
|
nn.Linear(cfg["emb_dim"], 4 * cfg["emb_dim"]),
|
|
GELU(),
|
|
nn.Linear(4 * cfg["emb_dim"], cfg["emb_dim"]),
|
|
)
|
|
|
|
def forward(self, x):
|
|
# x shape: (batch_size, seq_len, emb_dim)
|
|
|
|
x = self.layers[0](x)# x shape: (batch_size, seq_len, 4 * emb_dim)
|
|
x = self.layers[1](x) # x shape remains: (batch_size, seq_len, 4 * emb_dim)
|
|
x = self.layers[2](x) # x shape: (batch_size, seq_len, emb_dim)
|
|
return x # Output shape: (batch_size, seq_len, emb_dim)
|
|
```
|
|
#### **Propósito y Funcionalidad**
|
|
|
|
* **Red FeedForward por Posición:** Aplica una red completamente conectada de dos capas a cada posición de manera separada e idéntica.
|
|
* **Detalles de la Capa:**
|
|
* **Primera Capa Lineal:** Expande la dimensionalidad de `emb_dim` a `4 * emb_dim`.
|
|
* **Activación GELU:** Aplica no linealidad.
|
|
* **Segunda Capa Lineal:** Reduce la dimensionalidad de nuevo a `emb_dim`.
|
|
|
|
{% hint style="info" %}
|
|
Como puedes ver, la red Feed Forward utiliza 3 capas. La primera es una capa lineal que multiplicará las dimensiones por 4 usando pesos lineales (parámetros a entrenar dentro del modelo). Luego, se utiliza la función GELU en todas esas dimensiones para aplicar variaciones no lineales y capturar representaciones más ricas y, finalmente, se utiliza otra capa lineal para volver al tamaño original de las dimensiones.
|
|
{% endhint %}
|
|
|
|
### **Mecanismo de Atención Multi-Cabeza**
|
|
|
|
Esto ya fue explicado en una sección anterior.
|
|
|
|
#### **Propósito y Funcionalidad**
|
|
|
|
* **Autoatención Multi-Cabeza:** Permite que el modelo se enfoque en diferentes posiciones dentro de la secuencia de entrada al codificar un token.
|
|
* **Componentes Clave:**
|
|
* **Consultas, Claves, Valores:** Proyecciones lineales de la entrada, utilizadas para calcular puntajes de atención.
|
|
* **Cabezas:** Múltiples mecanismos de atención que funcionan en paralelo (`num_heads`), cada uno con una dimensión reducida (`head_dim`).
|
|
* **Puntajes de Atención:** Calculados como el producto punto de consultas y claves, escalados y enmascarados.
|
|
* **Enmascaramiento:** Se aplica una máscara causal para evitar que el modelo preste atención a tokens futuros (importante para modelos autorregresivos como GPT).
|
|
* **Pesos de Atención:** Softmax de los puntajes de atención enmascarados y escalados.
|
|
* **Vector de Contexto:** Suma ponderada de los valores, de acuerdo con los pesos de atención.
|
|
* **Proyección de Salida:** Capa lineal para combinar las salidas de todas las cabezas.
|
|
|
|
{% hint style="info" %}
|
|
El objetivo de esta red es encontrar las relaciones entre tokens en el mismo contexto. Además, los tokens se dividen en diferentes cabezas para prevenir el sobreajuste, aunque las relaciones finales encontradas por cabeza se combinan al final de esta red.
|
|
|
|
Además, durante el entrenamiento se aplica una **máscara causal** para que los tokens posteriores no se tengan en cuenta al buscar las relaciones específicas con un token y también se aplica algo de **dropout** para **prevenir el sobreajuste**.
|
|
{% endhint %}
|
|
|
|
### **Normalización de Capa**
|
|
```python
|
|
# From https://github.com/rasbt/LLMs-from-scratch/tree/main/ch04
|
|
class LayerNorm(nn.Module):
|
|
def __init__(self, emb_dim):
|
|
super().__init__()
|
|
self.eps = 1e-5 # Prevent division by zero during normalization.
|
|
self.scale = nn.Parameter(torch.ones(emb_dim))
|
|
self.shift = nn.Parameter(torch.zeros(emb_dim))
|
|
|
|
def forward(self, x):
|
|
mean = x.mean(dim=-1, keepdim=True)
|
|
var = x.var(dim=-1, keepdim=True, unbiased=False)
|
|
norm_x = (x - mean) / torch.sqrt(var + self.eps)
|
|
return self.scale * norm_x + self.shift
|
|
```
|
|
#### **Propósito y Funcionalidad**
|
|
|
|
* **Normalización por Capas:** Una técnica utilizada para normalizar las entradas a través de las características (dimensiones de incrustación) para cada ejemplo individual en un lote.
|
|
* **Componentes:**
|
|
* **`eps`:** Una constante pequeña (`1e-5`) añadida a la varianza para prevenir la división por cero durante la normalización.
|
|
* **`scale` y `shift`:** Parámetros aprendibles (`nn.Parameter`) que permiten al modelo escalar y desplazar la salida normalizada. Se inicializan en uno y cero, respectivamente.
|
|
* **Proceso de Normalización:**
|
|
* **Calcular Media (`mean`):** Calcula la media de la entrada `x` a través de la dimensión de incrustación (`dim=-1`), manteniendo la dimensión para la difusión (`keepdim=True`).
|
|
* **Calcular Varianza (`var`):** Calcula la varianza de `x` a través de la dimensión de incrustación, también manteniendo la dimensión. El parámetro `unbiased=False` asegura que la varianza se calcule utilizando el estimador sesgado (dividiendo por `N` en lugar de `N-1`), lo cual es apropiado al normalizar sobre características en lugar de muestras.
|
|
* **Normalizar (`norm_x`):** Resta la media de `x` y divide por la raíz cuadrada de la varianza más `eps`.
|
|
* **Escalar y Desplazar:** Aplica los parámetros aprendibles `scale` y `shift` a la salida normalizada.
|
|
|
|
{% hint style="info" %}
|
|
El objetivo es asegurar una media de 0 con una varianza de 1 a través de todas las dimensiones del mismo token. El objetivo de esto es **estabilizar el entrenamiento de redes neuronales profundas** al reducir el cambio interno de covariables, que se refiere al cambio en la distribución de las activaciones de la red debido a la actualización de parámetros durante el entrenamiento.
|
|
{% endhint %}
|
|
|
|
### **Bloque Transformer**
|
|
|
|
_Las formas se han añadido como comentarios para entender mejor las formas de las matrices:_
|
|
```python
|
|
# From https://github.com/rasbt/LLMs-from-scratch/tree/main/ch04
|
|
|
|
class TransformerBlock(nn.Module):
|
|
def __init__(self, cfg):
|
|
super().__init__()
|
|
self.att = MultiHeadAttention(
|
|
d_in=cfg["emb_dim"],
|
|
d_out=cfg["emb_dim"],
|
|
context_length=cfg["context_length"],
|
|
num_heads=cfg["n_heads"],
|
|
dropout=cfg["drop_rate"],
|
|
qkv_bias=cfg["qkv_bias"]
|
|
)
|
|
self.ff = FeedForward(cfg)
|
|
self.norm1 = LayerNorm(cfg["emb_dim"])
|
|
self.norm2 = LayerNorm(cfg["emb_dim"])
|
|
self.drop_shortcut = nn.Dropout(cfg["drop_rate"])
|
|
|
|
def forward(self, x):
|
|
# x shape: (batch_size, seq_len, emb_dim)
|
|
|
|
# Shortcut connection for attention block
|
|
shortcut = x # shape: (batch_size, seq_len, emb_dim)
|
|
x = self.norm1(x) # shape remains (batch_size, seq_len, emb_dim)
|
|
x = self.att(x) # shape: (batch_size, seq_len, emb_dim)
|
|
x = self.drop_shortcut(x) # shape remains (batch_size, seq_len, emb_dim)
|
|
x = x + shortcut # shape: (batch_size, seq_len, emb_dim)
|
|
|
|
# Shortcut connection for feedforward block
|
|
shortcut = x # shape: (batch_size, seq_len, emb_dim)
|
|
x = self.norm2(x) # shape remains (batch_size, seq_len, emb_dim)
|
|
x = self.ff(x) # shape: (batch_size, seq_len, emb_dim)
|
|
x = self.drop_shortcut(x) # shape remains (batch_size, seq_len, emb_dim)
|
|
x = x + shortcut # shape: (batch_size, seq_len, emb_dim)
|
|
|
|
return x # Output shape: (batch_size, seq_len, emb_dim)
|
|
|
|
```
|
|
#### **Propósito y Funcionalidad**
|
|
|
|
* **Composición de Capas:** Combina atención multi-cabeza, red de avance, normalización de capas y conexiones residuales.
|
|
* **Normalización de Capas:** Aplicada antes de las capas de atención y avance para un entrenamiento estable.
|
|
* **Conexiones Residuales (Atajos):** Agrega la entrada de una capa a su salida para mejorar el flujo de gradientes y permitir el entrenamiento de redes profundas.
|
|
* **Dropout:** Aplicado después de las capas de atención y avance para regularización.
|
|
|
|
#### **Funcionalidad Paso a Paso**
|
|
|
|
1. **Primer Camino Residual (Autoatención):**
|
|
* **Entrada (`shortcut`):** Guarda la entrada original para la conexión residual.
|
|
* **Norma de Capa (`norm1`):** Normaliza la entrada.
|
|
* **Atención Multi-Cabeza (`att`):** Aplica autoatención.
|
|
* **Dropout (`drop_shortcut`):** Aplica dropout para regularización.
|
|
* **Agregar Residual (`x + shortcut`):** Combina con la entrada original.
|
|
2. **Segundo Camino Residual (FeedForward):**
|
|
* **Entrada (`shortcut`):** Guarda la entrada actualizada para la siguiente conexión residual.
|
|
* **Norma de Capa (`norm2`):** Normaliza la entrada.
|
|
* **Red de Avance (`ff`):** Aplica la transformación de avance.
|
|
* **Dropout (`drop_shortcut`):** Aplica dropout.
|
|
* **Agregar Residual (`x + shortcut`):** Combina con la entrada del primer camino residual.
|
|
|
|
{% hint style="info" %}
|
|
El bloque transformer agrupa todas las redes y aplica alguna **normalización** y **dropouts** para mejorar la estabilidad y los resultados del entrenamiento.\
|
|
Nota cómo se realizan los dropouts después del uso de cada red mientras que la normalización se aplica antes.
|
|
|
|
Además, también utiliza atajos que consisten en **agregar la salida de una red con su entrada**. Esto ayuda a prevenir el problema del gradiente que se desvanece al asegurarse de que las capas iniciales contribuyan "tanto" como las últimas.
|
|
{% endhint %}
|
|
|
|
### **GPTModel**
|
|
|
|
_Las formas se han añadido como comentarios para entender mejor las formas de las matrices:_
|
|
```python
|
|
# From https://github.com/rasbt/LLMs-from-scratch/tree/main/ch04
|
|
class GPTModel(nn.Module):
|
|
def __init__(self, cfg):
|
|
super().__init__()
|
|
self.tok_emb = nn.Embedding(cfg["vocab_size"], cfg["emb_dim"])
|
|
# shape: (vocab_size, emb_dim)
|
|
|
|
self.pos_emb = nn.Embedding(cfg["context_length"], cfg["emb_dim"])
|
|
# shape: (context_length, emb_dim)
|
|
|
|
self.drop_emb = nn.Dropout(cfg["drop_rate"])
|
|
|
|
self.trf_blocks = nn.Sequential(
|
|
*[TransformerBlock(cfg) for _ in range(cfg["n_layers"])]
|
|
)
|
|
# Stack of TransformerBlocks
|
|
|
|
self.final_norm = LayerNorm(cfg["emb_dim"])
|
|
self.out_head = nn.Linear(cfg["emb_dim"], cfg["vocab_size"], bias=False)
|
|
# shape: (emb_dim, vocab_size)
|
|
|
|
def forward(self, in_idx):
|
|
# in_idx shape: (batch_size, seq_len)
|
|
batch_size, seq_len = in_idx.shape
|
|
|
|
# Token embeddings
|
|
tok_embeds = self.tok_emb(in_idx)
|
|
# shape: (batch_size, seq_len, emb_dim)
|
|
|
|
# Positional embeddings
|
|
pos_indices = torch.arange(seq_len, device=in_idx.device)
|
|
# shape: (seq_len,)
|
|
pos_embeds = self.pos_emb(pos_indices)
|
|
# shape: (seq_len, emb_dim)
|
|
|
|
# Add token and positional embeddings
|
|
x = tok_embeds + pos_embeds # Broadcasting over batch dimension
|
|
# x shape: (batch_size, seq_len, emb_dim)
|
|
|
|
x = self.drop_emb(x) # Dropout applied
|
|
# x shape remains: (batch_size, seq_len, emb_dim)
|
|
|
|
x = self.trf_blocks(x) # Pass through Transformer blocks
|
|
# x shape remains: (batch_size, seq_len, emb_dim)
|
|
|
|
x = self.final_norm(x) # Final LayerNorm
|
|
# x shape remains: (batch_size, seq_len, emb_dim)
|
|
|
|
logits = self.out_head(x) # Project to vocabulary size
|
|
# logits shape: (batch_size, seq_len, vocab_size)
|
|
|
|
return logits # Output shape: (batch_size, seq_len, vocab_size)
|
|
```
|
|
#### **Propósito y Funcionalidad**
|
|
|
|
* **Capas de Embedding:**
|
|
* **Embeddings de Tokens (`tok_emb`):** Convierte índices de tokens en embeddings. Como recordatorio, estos son los pesos dados a cada dimensión de cada token en el vocabulario.
|
|
* **Embeddings Posicionales (`pos_emb`):** Agrega información posicional a los embeddings para capturar el orden de los tokens. Como recordatorio, estos son los pesos dados a los tokens según su posición en el texto.
|
|
* **Dropout (`drop_emb`):** Aplicado a los embeddings para regularización.
|
|
* **Bloques de Transformer (`trf_blocks`):** Pila de `n_layers` bloques de transformer para procesar embeddings.
|
|
* **Normalización Final (`final_norm`):** Normalización de capa antes de la capa de salida.
|
|
* **Capa de Salida (`out_head`):** Proyecta los estados ocultos finales al tamaño del vocabulario para producir logits para la predicción.
|
|
|
|
{% hint style="info" %}
|
|
El objetivo de esta clase es usar todas las otras redes mencionadas para **predecir el siguiente token en una secuencia**, lo cual es fundamental para tareas como la generación de texto.
|
|
|
|
Nota cómo **usará tantos bloques de transformer como se indique** y que cada bloque de transformer utiliza una red de atención multi-cabeza, una red de avance y varias normalizaciones. Así que si se utilizan 12 bloques de transformer, multiplica esto por 12.
|
|
|
|
Además, se agrega una capa de **normalización** **antes** de la **salida** y se aplica una capa lineal final al final para obtener los resultados con las dimensiones adecuadas. Nota cómo cada vector final tiene el tamaño del vocabulario utilizado. Esto se debe a que está tratando de obtener una probabilidad por cada token posible dentro del vocabulario.
|
|
{% endhint %}
|
|
|
|
## Número de Parámetros a entrenar
|
|
|
|
Teniendo la estructura de GPT definida, es posible averiguar el número de parámetros a entrenar:
|
|
```python
|
|
GPT_CONFIG_124M = {
|
|
"vocab_size": 50257, # Vocabulary size
|
|
"context_length": 1024, # Context length
|
|
"emb_dim": 768, # Embedding dimension
|
|
"n_heads": 12, # Number of attention heads
|
|
"n_layers": 12, # Number of layers
|
|
"drop_rate": 0.1, # Dropout rate
|
|
"qkv_bias": False # Query-Key-Value bias
|
|
}
|
|
|
|
model = GPTModel(GPT_CONFIG_124M)
|
|
total_params = sum(p.numel() for p in model.parameters())
|
|
print(f"Total number of parameters: {total_params:,}")
|
|
# Total number of parameters: 163,009,536
|
|
```
|
|
### **Cálculo Paso a Paso**
|
|
|
|
#### **1. Capas de Embedding: Embedding de Tokens y Embedding de Posición**
|
|
|
|
* **Capa:** `nn.Embedding(vocab_size, emb_dim)`
|
|
* **Parámetros:** `vocab_size * emb_dim`
|
|
```python
|
|
token_embedding_params = 50257 * 768 = 38,597,376
|
|
```
|
|
* **Capa:** `nn.Embedding(context_length, emb_dim)`
|
|
* **Parámetros:** `context_length * emb_dim`
|
|
```python
|
|
position_embedding_params = 1024 * 768 = 786,432
|
|
```
|
|
**Total de Parámetros de Embedding**
|
|
```python
|
|
embedding_params = token_embedding_params + position_embedding_params
|
|
embedding_params = 38,597,376 + 786,432 = 39,383,808
|
|
```
|
|
#### **2. Bloques de Transformador**
|
|
|
|
Hay 12 bloques de transformador, así que calcularemos los parámetros para un bloque y luego multiplicaremos por 12.
|
|
|
|
**Parámetros por Bloque de Transformador**
|
|
|
|
**a. Atención Multi-Cabeza**
|
|
|
|
* **Componentes:**
|
|
* **Capa Lineal de Consulta (`W_query`):** `nn.Linear(emb_dim, emb_dim, bias=False)`
|
|
* **Capa Lineal de Clave (`W_key`):** `nn.Linear(emb_dim, emb_dim, bias=False)`
|
|
* **Capa Lineal de Valor (`W_value`):** `nn.Linear(emb_dim, emb_dim, bias=False)`
|
|
* **Proyección de Salida (`out_proj`):** `nn.Linear(emb_dim, emb_dim)`
|
|
* **Cálculos:**
|
|
* **Cada uno de `W_query`, `W_key`, `W_value`:**
|
|
|
|
```python
|
|
qkv_params = emb_dim * emb_dim = 768 * 768 = 589,824
|
|
```
|
|
|
|
Dado que hay tres capas de este tipo:
|
|
|
|
```python
|
|
total_qkv_params = 3 * qkv_params = 3 * 589,824 = 1,769,472
|
|
```
|
|
* **Proyección de Salida (`out_proj`):**
|
|
|
|
```python
|
|
out_proj_params = (emb_dim * emb_dim) + emb_dim = (768 * 768) + 768 = 589,824 + 768 = 590,592
|
|
```
|
|
* **Total de Parámetros de Atención Multi-Cabeza:**
|
|
|
|
```python
|
|
mha_params = total_qkv_params + out_proj_params
|
|
mha_params = 1,769,472 + 590,592 = 2,360,064
|
|
```
|
|
|
|
**b. Red FeedForward**
|
|
|
|
* **Componentes:**
|
|
* **Primera Capa Lineal:** `nn.Linear(emb_dim, 4 * emb_dim)`
|
|
* **Segunda Capa Lineal:** `nn.Linear(4 * emb_dim, emb_dim)`
|
|
* **Cálculos:**
|
|
* **Primera Capa Lineal:**
|
|
|
|
```python
|
|
ff_first_layer_params = (emb_dim * 4 * emb_dim) + (4 * emb_dim)
|
|
ff_first_layer_params = (768 * 3072) + 3072 = 2,359,296 + 3,072 = 2,362,368
|
|
```
|
|
* **Segunda Capa Lineal:**
|
|
|
|
```python
|
|
ff_second_layer_params = (4 * emb_dim * emb_dim) + emb_dim
|
|
ff_second_layer_params = (3072 * 768) + 768 = 2,359,296 + 768 = 2,360,064
|
|
```
|
|
* **Total de Parámetros FeedForward:**
|
|
|
|
```python
|
|
ff_params = ff_first_layer_params + ff_second_layer_params
|
|
ff_params = 2,362,368 + 2,360,064 = 4,722,432
|
|
```
|
|
|
|
**c. Normalizaciones de Capa**
|
|
|
|
* **Componentes:**
|
|
* Dos instancias de `LayerNorm` por bloque.
|
|
* Cada `LayerNorm` tiene `2 * emb_dim` parámetros (escala y desplazamiento).
|
|
* **Cálculos:**
|
|
|
|
```python
|
|
layer_norm_params_per_block = 2 * (2 * emb_dim) = 2 * 768 * 2 = 3,072
|
|
```
|
|
|
|
**d. Total de Parámetros por Bloque de Transformador**
|
|
```python
|
|
pythonCopy codeparams_per_block = mha_params + ff_params + layer_norm_params_per_block
|
|
params_per_block = 2,360,064 + 4,722,432 + 3,072 = 7,085,568
|
|
```
|
|
**Total de Parámetros para Todos los Bloques de Transformadores**
|
|
```python
|
|
pythonCopy codetotal_transformer_blocks_params = params_per_block * n_layers
|
|
total_transformer_blocks_params = 7,085,568 * 12 = 85,026,816
|
|
```
|
|
#### **3. Capas Finales**
|
|
|
|
**a. Normalización de la Capa Final**
|
|
|
|
* **Parámetros:** `2 * emb_dim` (escalar y desplazar)
|
|
```python
|
|
pythonCopy codefinal_layer_norm_params = 2 * 768 = 1,536
|
|
```
|
|
**b. Capa de Proyección de Salida (`out_head`)**
|
|
|
|
* **Capa:** `nn.Linear(emb_dim, vocab_size, bias=False)`
|
|
* **Parámetros:** `emb_dim * vocab_size`
|
|
```python
|
|
pythonCopy codeoutput_projection_params = 768 * 50257 = 38,597,376
|
|
```
|
|
#### **4. Resumiendo Todos los Parámetros**
|
|
```python
|
|
pythonCopy codetotal_params = (
|
|
embedding_params +
|
|
total_transformer_blocks_params +
|
|
final_layer_norm_params +
|
|
output_projection_params
|
|
)
|
|
total_params = (
|
|
39,383,808 +
|
|
85,026,816 +
|
|
1,536 +
|
|
38,597,376
|
|
)
|
|
total_params = 163,009,536
|
|
```
|
|
## Generar Texto
|
|
|
|
Teniendo un modelo que predice el siguiente token como el anterior, solo es necesario tomar los últimos valores de token de la salida (ya que serán los del token predicho), que será un **valor por entrada en el vocabulario** y luego usar la función `softmax` para normalizar las dimensiones en probabilidades que sumen 1 y luego obtener el índice de la entrada más grande, que será el índice de la palabra dentro del vocabulario.
|
|
|
|
Código de [https://github.com/rasbt/LLMs-from-scratch/blob/main/ch04/01\_main-chapter-code/ch04.ipynb](https://github.com/rasbt/LLMs-from-scratch/blob/main/ch04/01\_main-chapter-code/ch04.ipynb):
|
|
```python
|
|
def generate_text_simple(model, idx, max_new_tokens, context_size):
|
|
# idx is (batch, n_tokens) array of indices in the current context
|
|
for _ in range(max_new_tokens):
|
|
|
|
# Crop current context if it exceeds the supported context size
|
|
# E.g., if LLM supports only 5 tokens, and the context size is 10
|
|
# then only the last 5 tokens are used as context
|
|
idx_cond = idx[:, -context_size:]
|
|
|
|
# Get the predictions
|
|
with torch.no_grad():
|
|
logits = model(idx_cond)
|
|
|
|
# Focus only on the last time step
|
|
# (batch, n_tokens, vocab_size) becomes (batch, vocab_size)
|
|
logits = logits[:, -1, :]
|
|
|
|
# Apply softmax to get probabilities
|
|
probas = torch.softmax(logits, dim=-1) # (batch, vocab_size)
|
|
|
|
# Get the idx of the vocab entry with the highest probability value
|
|
idx_next = torch.argmax(probas, dim=-1, keepdim=True) # (batch, 1)
|
|
|
|
# Append sampled index to the running sequence
|
|
idx = torch.cat((idx, idx_next), dim=1) # (batch, n_tokens+1)
|
|
|
|
return idx
|
|
|
|
|
|
start_context = "Hello, I am"
|
|
|
|
encoded = tokenizer.encode(start_context)
|
|
print("encoded:", encoded)
|
|
|
|
encoded_tensor = torch.tensor(encoded).unsqueeze(0)
|
|
print("encoded_tensor.shape:", encoded_tensor.shape)
|
|
|
|
model.eval() # disable dropout
|
|
|
|
out = generate_text_simple(
|
|
model=model,
|
|
idx=encoded_tensor,
|
|
max_new_tokens=6,
|
|
context_size=GPT_CONFIG_124M["context_length"]
|
|
)
|
|
|
|
print("Output:", out)
|
|
print("Output length:", len(out[0]))
|
|
```
|
|
## Referencias
|
|
|
|
* [https://www.manning.com/books/build-a-large-language-model-from-scratch](https://www.manning.com/books/build-a-large-language-model-from-scratch)
|