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 el entrenamiento como para predecir texto después de haber sido entrenada.
Se puede observar una representación de alto nivel en:
Entrada (Texto Tokenizado): El proceso comienza con texto tokenizado, que se convierte en representaciones numéricas.
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.
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.
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.
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
import torchimport torch.nn as nnimport tiktokenclassGELU(nn.Module):def__init__(self):super().__init__()defforward(self,x):return0.5* x * (1+ torch.tanh(torch.sqrt(torch.tensor(2.0/ torch.pi)) *(x +0.044715* torch.pow(x, 3))))classFeedForward(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"]),)defforward(self,x):return self.layers(x)classMultiHeadAttention(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_outself.num_heads = num_headsself.head_dim = d_out // num_heads # Reduce the projection dim to match desired output dimself.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 outputsself.dropout = nn.Dropout(dropout)self.register_buffer('mask', torch.triu(torch.ones(context_length, context_length), diagonal=1))defforward(self,x):b, num_tokens, d_in = x.shapekeys = 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 maskattn_scores = queries @ keys.transpose(2, 3)# Dot product for each head# Original mask truncated to the number of tokens and converted to booleanmask_bool = self.mask.bool()[:num_tokens,:num_tokens]# Use the mask to fill attention scoresattn_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_dimcontext_vec = context_vec.contiguous().view(b, num_tokens, self.d_out)context_vec = self.out_proj(context_vec)# optional projectionreturn context_vecclassLayerNorm(nn.Module):def__init__(self,emb_dim):super().__init__()self.eps =1e-5self.scale = nn.Parameter(torch.ones(emb_dim))self.shift = nn.Parameter(torch.zeros(emb_dim))defforward(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.shiftclassTransformerBlock(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"])defforward(self,x):# Shortcut connection for attention blockshortcut = xx = 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 blockshortcut = xx = self.norm2(x)x = self.ff(x)x = self.drop_shortcut(x)x = x + shortcut # Add the original input backreturn xclassGPTModel(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 _ inrange(cfg["n_layers"])])self.final_norm =LayerNorm(cfg["emb_dim"])self.out_head = nn.Linear(cfg["emb_dim"], cfg["vocab_size"], bias=False)defforward(self,in_idx):batch_size, seq_len = in_idx.shapetok_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 logitsGPT_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
# From https://github.com/rasbt/LLMs-from-scratch/tree/main/ch04classGELU(nn.Module):def__init__(self):super().__init__()defforward(self,x):return0.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 pequeños valores no cero para entradas negativas.
Definición Matemática:
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.
Red Neuronal FeedForward
Las formas se han añadido como comentarios para entender mejor las formas de las matrices:
# From https://github.com/rasbt/LLMs-from-scratch/tree/main/ch04classFeedForward(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"]),)defforward(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.
Como puedes ver, la red Feed Forward utiliza 3 capas. La primera es una capa lineal que multiplicará las dimensiones por 4 utilizando 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.
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.
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 un dropout para prevenir el sobreajuste.
Normalización de Capa
# From https://github.com/rasbt/LLMs-from-scratch/tree/main/ch04classLayerNorm(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))defforward(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
Layer Normalization: Una técnica utilizada para normalizar las entradas a través de las características (dimensiones de embedding) 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 unos y ceros, respectivamente.
Proceso de Normalización:
Calcular Media (mean): Calcula la media de la entrada x a través de la dimensión de embedding (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 embedding, 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.
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 de covariables internas, 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.
Bloque Transformer
Las formas se han añadido como comentarios para entender mejor las formas de las matrices:
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
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.
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.
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 los dropouts se realizan 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 asegurar que las capas iniciales contribuyan "tanto" como las últimas.
GPTModel
Las formas se han añadido como comentarios para entender mejor las formas de las matrices:
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.
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 está utilizando 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ónantes 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.
Número de Parámetros a entrenar
Teniendo la estructura de GPT definida, es posible averiguar el número de parámetros a entrenar:
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
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.
defgenerate_text_simple(model,idx,max_new_tokens,context_size):# idx is (batch, n_tokens) array of indices in the current contextfor _ inrange(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 contextidx_cond = idx[:,-context_size:]# Get the predictionswith 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 probabilitiesprobas = torch.softmax(logits, dim=-1)# (batch, vocab_size)# Get the idx of the vocab entry with the highest probability valueidx_next = torch.argmax(probas, dim=-1, keepdim=True)# (batch, 1)# Append sampled index to the running sequenceidx = torch.cat((idx, idx_next), dim=1)# (batch, n_tokens+1)return idxstart_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 dropoutout =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]))
