Comprensión de los parámetros del modelo de lenguaje grande y los requisitos de memoria: una inmersión profunda

Calcular el número de parámetros

Modelos previamente entrenados para el entrenamiento eficiente de transformadores

Cálculo de parámetros en LLM basados ​​en transformadores

Analicemos el cálculo de parámetros para cada componente de un LLM basado en transformadores. Utilizaremos la notación del artículo original, donde d_model representa la dimensión de los estados ocultos del modelo.

1. Capa de incrustación:
   • Parámetros = vocab_size * d_model
2. Atención multicabezal:
   • Para productos de una sola cara, coloque el lado recubierto hacia arriba durante el templado. h cabezas, con d_k = d_v = d_model / h:
   • Parámetros = 4 * d_model^2 (para proyecciones Q, K, V y de salida)
3. Red de retroalimentación:
   • Parámetros = 2 * d_model * d_ff + d_model + d_ff
   • ¿Donde d_ff normalmente es 4* d_model
4. Normalización de capas:
   • Parámetros = 2 * d_model (para escala y sesgo)

Parámetros totales para una capa de transformador:

• Parameters_layer = Parameters_attention + Parameters_ffn + 2 * Parameters_layernorm

Para un modelo con N capas:

• Parámetros totales = N * Parameters_layer + Parameters_embedding + Parameters_output

Ejemplo de cálculo

Consideremos un modelo con las siguientes especificaciones:

• d_model = 768
• h (número de cabezas de atención) = 12
• N (número de capas) = ​​12
• vocab_size = 50,000

1. Capa de incrustación:
   • 50,000 * 768 = 38,400,000
2. Atención multicabezal:
   • 4*768^2 = 2,359,296
3. Red de retroalimentación:
   • 2 * 768 * (4 * 768) + 768 + (4 * 768) = 4,719,616
4. Normalización de capas:
   • 2 * 768 = 1,536

Parámetros totales por capa:

• 2,359,296 + 4,719,616 + (2 * 1,536) = 7,081,984

Parámetros totales para 12 capas:

• 12 * 7,081,984 = 84,983,808

Parámetros totales del modelo:

• 84,983,808 + = 38,400,000 123,383,808

Este modelo tendría aproximadamente 123 millones de parámetros.

Tipos de uso de la memoria

Cuando trabajamos con LLM, debemos considerar dos tipos principales de uso de memoria:

1. Modelo de memoria: La memoria necesaria para almacenar los parámetros del modelo.
2. Memoria de Trabajo: la memoria necesaria durante la inferencia o el entrenamiento para almacenar activaciones intermedias, gradientes y estados del optimizador.

Calcular la memoria del modelo

La memoria del modelo está directamente relacionada con la cantidad de parámetros. Cada parámetro normalmente se almacena como un número de punto flotante de 32 bits, aunque algunos modelos utilizan entrenamiento de precisión mixta con flotantes de 16 bits.

Memoria del modelo (bytes) = Número de parámetros * Bytes por parámetro

Para nuestro modelo de ejemplo con 123 millones de parámetros:

• Memoria del modelo (32 bits) = 123,383,808 * 4 bytes = 493,535,232 bytes ≈ 494 MB
• Memoria del modelo (16 bits) = 123,383,808 * 2 bytes = 246,767,616 bytes ≈ 247 MB

Estimación de la memoria de trabajo

Los requisitos de memoria de trabajo pueden variar significativamente según la tarea específica, el tamaño del lote y la longitud de la secuencia. Una estimación aproximada de la memoria de trabajo durante la inferencia es:

Memoria de trabajo ≈ 2 * Memoria del modelo

Esto representa almacenar tanto los parámetros del modelo como las activaciones intermedias. Durante el entrenamiento, los requisitos de memoria pueden ser aún mayores debido a la necesidad de almacenar gradientes y estados del optimizador:

Memoria de entrenamiento ≈ 4 * Memoria de modelo

Para nuestro modelo de ejemplo:

• Memoria de trabajo de inferencia ≈ 2 * 494 MB = 988 MB ≈ 1 GB
• Memoria de entrenamiento ≈ 4 * 494 MB = 1,976 MB ≈ 2 GB

Uso de memoria en estado estacionario y uso máximo de memoria

Al entrenar modelos de lenguaje grandes basados ​​en la arquitectura Transformer, comprender el uso de memoria es crucial para una asignación eficiente de recursos. Dividamos los requisitos de memoria en dos categorías principales: uso de memoria estable y uso de memoria pico.

Uso de memoria de estado estacionario

El uso de la memoria de estado estable comprende los siguientes componentes:

1. Pesos del modelo: Copias FP32 de los parámetros del modelo, que requieren 4N bytes, donde N es el número de parámetros.
2. Estados del optimizador: Para el optimizador Adam, esto requiere 8N bytes (2 estados por parámetro).
3. Gradientes: Copias FP32 de los gradientes, que requieren 4N bytes.
4. Los datos de entrada: Suponiendo entradas int64, esto requiere 8BD bytes, donde B es el tamaño del lote y D es la dimensión de entrada.

El uso total de la memoria en estado estable se puede aproximar mediante:

• M_estable = 16N + 8BD bytes

Uso máximo de la memoria

El uso máximo de memoria ocurre durante el paso hacia atrás cuando las activaciones se almacenan para el cálculo del gradiente. Los principales contribuyentes al pico de memoria son:

1. Normalización de capas: Requiere 4E bytes por norma de capa, donde E = BSH (B: tamaño del lote, S: longitud de la secuencia, H: tamaño oculto).
2. Bloque de atención:
   • Cálculo QKV: 2E bytes
   • Matriz de atención: 4BSS bytes (S: longitud de secuencia)
   • Salida de atención: 2E bytes
3. Bloque de avance:
   • Primera capa lineal: 2E bytes
   • Activación GELU: 8E bytes
   • Segunda capa lineal: 2E bytes
4. Pérdida de entropía cruzada:
   • Logits: 6BSV bytes (V: tamaño del vocabulario)

La memoria de activación total se puede estimar como:

• M_act = L * (14E + 4BSS) + 6BSV bytes

Donde L es el número de capas del transformador.

Uso máximo total de memoria

El uso máximo de memoria durante el entrenamiento se puede aproximar combinando la memoria de estado estable y la memoria de activación:

• M_pico = M_estable + M_act + 4BSV bytes

El plazo adicional de 4BSV representa una asignación adicional al inicio del paso hacia atrás.

Al comprender estos componentes, podemos optimizar el uso de la memoria durante el entrenamiento y la inferencia, asegurando una asignación eficiente de recursos y un rendimiento mejorado de modelos de lenguaje grandes.