GPUs, memoria y cuantización · IA generativa: de la neurona al chatbot

Qué lograrás

Calcular cuánta memoria ocupa un modelo a partir de sus parámetros y de los bytes por parámetro (FP32, FP16, INT8, INT4), y decidir en qué GPU cabe.
Explicar por qué entrenar con Adam pide unos dieciséis bytes por parámetro y por qué, al escribir, el límite suele ser el ancho de banda de la memoria.
Reconocer las tres formas básicas de repartir un modelo entre varias GPU y qué gana y qué arriesga la cuantización, con mediciones reales.

Ideas clave

Un modelo trabaja desde la memoria de la GPU, la VRAM. Ocupa parámetros × bytes por parámetro: 70 mil millones en FP16 son 140 GB, más que los 80 de una H100.
Cada formato guarda el número con más o menos detalle: FP32 (4 bytes, unas 7 cifras), FP16 (2 bytes, unas 3), INT8 (1 byte) e INT4 (½ byte), que redondean a marcas de una regla.
Entrenar con Adam pide unos 16 bytes por parámetro: el parámetro y su gradiente en FP16, una copia en FP32 y los dos promedios de Adam. Ocho veces lo que pide usarlo, sin contar las activaciones.
Caber no basta: al escribir, por cada token se leen todos los parámetros, y el ancho de banda de la memoria suele marcar la velocidad, no el cálculo.
Si no cabe, se reparte (datos, tubería, tensor; Llama 3.1 405B usó 8 × 16 × 128 = 16 384 H100) o se cuantiza. En Pythia, 8 y 4 bits casi no cambian la pérdida; 2 bits la rompen.

Lo que mucha gente cree

Creencia común Si un modelo cabe en la memoria de mi GPU, también lo puedo entrenar ahí.

Lo correcto Usarlo pide unos 2 bytes por parámetro en FP16; entrenarlo con Adam, unos 16, más las activaciones. Un modelo que apenas cabe para usarse necesita unas ocho veces más memoria para entrenarse. Por eso se reparte entre muchas GPU, o se ajusta solo una pequeña parte de sus parámetros (lección 27).

Compruébalo

Cuatro preguntas. Sin nota ni cuenta: sirven para ver si la idea quedó.

Por dentroPara quien quiere las fórmulas y el código. No hace falta para seguir el curso.

Bytes por parámetro y formatos

La memoria que ocupan los pesos es una multiplicación:

M_{\text{pesos}} = N \cdot b,

con $N$ parámetros y $b$ bytes por parámetro. Para Llama 3.1 70B en FP16: $70\times10^{9}\cdot 2 = 1.4\times10^{11}$ bytes $=140$ GB.

Los formatos con decimales (norma IEEE 754) guardan un signo, un exponente (el alcance) y una mantisa (las cifras):

formato	bytes	signo · exponente · mantisa	cifras decimales	mayor valor
FP32	4	1 · 8 · 23	$\log_{10}2^{24}\approx 7.2$	$\approx 3.4\times10^{38}$
FP16	2	1 · 5 · 10	$\log_{10}2^{11}\approx 3.3$	65 504
BF16	2	1 · 8 · 7	$\log_{10}2^{8}\approx 2.4$	el de FP32

BF16 (Kalamkar et al., 2019) sacrifica cifras para conservar el alcance de FP32; Llama 3 se entrenó en BF16. INT8 e INT4 no tienen exponente: guardan un entero y una escala compartida (ver abajo).

Un detalle de unidades: las fichas dicen «80 GB» de VRAM, pero son $80\cdot 2^{30}$ bytes (GiB), un 7 % más que $80\times10^{9}$ . En la lección, 1 GB = $10^{9}$ bytes, como en los anchos de banda; la diferencia no cambia ninguna conclusión.

Entrenar: dieciséis bytes por parámetro

Con Adam en precisión mixta (Micikevicius et al., 2018), por cada parámetro se guardan el peso y el gradiente en FP16 y, en FP32, una copia del peso y los dos promedios de Adam (Kingma y Ba, 2015). ZeRO (Rajbhandari et al., 2020, §3.1) lo resume así:

M_{\text{entrenar}} \approx (2 + 2 + \underbrace{4 + 4 + 4}_{K\,=\,12})\,N = 16N \text{ bytes}.

La copia en FP32 existe porque el ajuste de cada paso (gradiente por tasa de aprendizaje) suele ser miles de veces menor que el peso: sumado en FP16, que tiene 10 bits de mantisa, se perdería al redondear. Lo medimos en una RTX 3060 (measure-memory.py, arquitectura de Pythia-160M con pesos al azar): 16.19 bytes por parámetro con FP32 completo y 16.05 con precisión mixta, tras el primer paso de AdamW.

Aparte quedan las activaciones, los resultados intermedios que la retropropagación necesita. Por capa de un transformer en 16 bits, con atención de memoria lineal, son unos $34\,s\,b\,h$ bytes ( $s$ = tokens por secuencia, $b$ = secuencias por lote, $h$ = dimensión; Korthikanti et al., 2022). En la misma RTX 3060 medimos 0.949 veces esa fórmula. Se reducen recalculándolas en el regreso, a cambio de más cómputo.

ZeRO y su versión en PyTorch, FSDP (la que usó Meta), reparten esos 16 bytes entre las copias de datos, para que cada GPU guarde solo una parte de los estados.

Ancho de banda: el techo al escribir

Al generar un token a la vez, cada paso lee todos los pesos:

\text{tokens por segundo} \le \frac{\text{ancho de banda}}{N\cdot b}.

Una H100 SXM mueve 3 350 GB/s; un modelo de 8 mil millones en FP16 (16 GB) queda por debajo de unos 209 tokens por segundo, por rápido que calcule. Su pico de cálculo en FP16 sin dispersión es de unos 989 billones de operaciones por segundo (la ficha da 1 979 «con dispersión»): $989\times10^{12}/3.35\times10^{12}\approx 295$ operaciones por byte, unas 590 por cada parámetro de 2 bytes, que en un producto matriz–vector se usa en solo 2 (multiplicar y sumar).

Medido en una RTX 3060 (measure-memory.py): copiar 1 GiB, 333 GB/s; multiplicar una matriz de $16384\times16384$ en FP16 por un vector, 1.57 ms, 341 GB/s y apenas 0.34 billones de operaciones por segundo; por 4 096 vectores a la vez, 27.7 billones: 81 veces más con el mismo chip, porque cada peso leído se reutiliza. Agrupar peticiones en lotes es la salida que veremos en la lección 26.

Paralelismo

El número de GPU es el producto de los tres repartos: $G = T\cdot P\cdot D$ (tensor, tubería, datos).

Datos: cada copia ve otro trozo del lote; al final del paso se promedian los gradientes (una suma entre todas, all-reduce).
Tubería (GPipe, Huang et al., 2019): $K$ tramos de capas; el lote se parte en $M$ micro-lotes para que los tramos trabajen a la vez. La fracción de tiempo ociosa (la «burbuja») es del orden de $\frac{K-1}{M+K-1}$ ; GPipe la encontró despreciable con $M \ge 4K$ .
Tensor (Megatron-LM, Shoeybi et al., 2019): en la red de cada bloque, la primera tabla se parte por columnas y la segunda por filas, de modo que basta una suma entre GPU por bloque en cada sentido. Como ocurre en cada capa, se mantiene dentro del servidor, donde las GPU se unen por NVLink (900 GB/s en la H100).

Llama 3.1 405B (Llama Team, 2024, tabla 4): 8 de tensor × 16 de tubería × 128 de datos = 16 384 H100, con un 41 % de aprovechamiento del cálculo; en la etapa de contexto largo añadió un cuarto reparto, por contexto (8 × 16 de contexto × 16 × 8). Y para usarlo: en BF16 no cabe en un servidor de 8 H100 (§6.1), por eso Meta lo sirve repartido en dos servidores o cuantizado a FP8.

Cuantización por grupos

La regla de la lección (redondeo simétrico al más cercano, sin calibración), para un grupo $g$ de 32 pesos y $b$ bits:

q_{\max} = 2^{b-1}-1,\qquad s_g = \frac{\max_{i\in g}|w_i|}{q_{\max}},\qquad q_i = \operatorname{clip}\!\Big(\operatorname{round}\frac{w_i}{s_g},\,-q_{\max},\,q_{\max}\Big),\qquad \hat w_i = q_i\, s_g.

Con 8 bits hay 255 marcas (de −127 a 127); con 4, 15. Cada grupo guarda además su escala en FP16: $16/32 = 0.5$ bits más por peso, así que «INT4» son en realidad 4.5 bits (el modelo de 70 mil millones ocupa unos 39 GB, no 35).

Medido en Pythia-1.4B (build-data.py: las 96 matrices de sus 24 capas, 1 208 millones de pesos; pérdida en nats por token sobre 287 744 tokens de WikiText-2):

bits	FP16	8	6	5	4	3	2
una escala por grupo de 32	2.5805	2.5806	2.5858	2.5975	2.6565	3.1589	11.045
una escala por matriz	2.5805	2.6002	3.2172	5.3017	10.692	10.445	21.347

Con una sola escala por matriz, el modelo se rompe ya a 4 bits: unos pocos pesos atípicos, hasta 13 veces la desviación típica en la matriz mediana, estiran la regla, y en esa matriz dos de cada tres pesos caen en la marca cero a 4 bits. Por eso se usan grupos pequeños. Los métodos de producción van más lejos: LLM.int8() (Dettmers et al., 2022) separa las dimensiones con valores atípicos y las deja en 16 bits; GPTQ (Frantar et al., 2023) corrige el redondeo con información de segundo orden y lleva modelos de 175 mil millones a 3 o 4 bits con una pérdida de precisión despreciable. Advertencia: Pythia vio Wikipedia al entrenar, así que el nivel absoluto de la pérdida es optimista; lo que se compara es la diferencia, medida con el mismo texto.

En código

import numpy as np

def gb(params, bytes_per_param):
    return params * bytes_per_param / 1e9

for fmt, b in [("FP32", 4), ("FP16", 2), ("INT8", 1), ("INT4", 0.5)]:
    print(f"70 mil M en {fmt}: {gb(70e9, b):5.0f} GB")
print(f"entrenar con Adam (16 B): {gb(70e9, 16):.0f} GB")
print(f"techo al escribir, 8 mil M en FP16 en una H100: {3350 / gb(8e9, 2):.0f} tokens/s")

# cuantizar un grupo de 32 pesos: una regla con 2^(b-1) - 1 marcas por lado
rng = np.random.default_rng(0)
w = rng.normal(0, 0.02, 32)
for bits in (8, 4, 2):
    qmax = 2 ** (bits - 1) - 1
    s = np.abs(w).max() / qmax                     # escala del grupo
    q = np.clip(np.round(w / s), -qmax, qmax)      # los enteros que se guardan
    err = np.linalg.norm(w - q * s) / np.linalg.norm(w)
    print(f"{bits} bits: error relativo {err:.3f}, pesos en cero {np.mean(q == 0):.0%}")

Salida:

70 mil M en FP32:   280 GB
70 mil M en FP16:   140 GB
70 mil M en INT8:    70 GB
70 mil M en INT4:    35 GB
entrenar con Adam (16 B): 1120 GB
techo al escribir, 8 mil M en FP16 en una H100: 209 tokens/s
8 bits: error relativo 0.006, pesos en cero 0%
4 bits: error relativo 0.100, pesos en cero 19%
2 bits: error relativo 0.792, pesos en cero 84%

Qué sigue

Al usar el modelo aparece otra memoria, la caché de la atención, que crece con cada token de la conversación, y la forma de aprovechar el ancho de banda: atender a muchas personas a la vez. Es la lección 26, la inferencia.

Lente de decisión

Qué significa esta idea si tu trabajo es decidir, no programar.

Para estimar el hardware: memoria ≈ parámetros × bytes por parámetro, más un margen para el cálculo (la caché de la lección 26). Un modelo de 8 mil millones en INT8 cabe en una tarjeta de 12 GB; uno de 70 mil millones en FP16 necesita al menos dos H100.
Pregunta en qué formato se sirve un modelo: el mismo nombre puede ofrecerse en FP16 o cuantizado a 4 bits, con costos, velocidad y calidad distintos.
Si te proponen entrenar o ajustar un modelo completo, cuenta unos 16 bytes por parámetro antes de comprar: entrenar pide unas ocho veces más memoria que usar.
Una versión cuantizada puede pasar las pruebas estándar y aun así fallar en casos raros (Meta lo vio con Llama 3). Evalúala con tus propios datos antes de cambiar.

Términos de la lección

Transcripción completa

ComentariosConversación pública de esta lección: dudas, datos por revisar, ideas.

Los comentarios viven en GitHub Discussions, en el repositorio público del curso; para escribir se necesita una cuenta de GitHub. Se cargan desde giscus.app al abrir esta sección. Las personas que revisan el curso los leen, y si un cambio se acepta, se aplica en los dos idiomas. Abrir el hilo en GitHub

Fuentes

NVIDIA (2022). NVIDIA H100 Tensor Core GPU: especificaciones (H100 SXM: 80 GB, 3.35 TB/s, NVLink 900 GB/s)
Andersch, M., Palmer, G., Krashinsky, R., Stam, N., Mehta, V., Brito, G., Ramaswamy, S. (2022). NVIDIA Hopper Architecture In-Depth (16 896 núcleos FP32, 80 GB HBM3 en 5 pilas)
NVIDIA (2021). GeForce RTX 3060 family: especificaciones (12 GB GDDR6, bus de 192 bits)
Llama Team, AI @ Meta (2024). The Llama 3 Herd of Models
Rajbhandari, S., Rasley, J., Ruwase, O., He, Y. (2020). ZeRO: Memory Optimizations Toward Training Trillion Parameter Models
Micikevicius, P., Narang, S., Alben, J., Diamos, G., Elsen, E., Garcia, D., et al. (2018). Mixed Precision Training
Kingma, D. P., Ba, J. (2015). Adam: A Method for Stochastic Optimization
IEEE (2019). IEEE Standard for Floating-Point Arithmetic (IEEE 754-2019)
Kalamkar, D., Mudigere, D., Mellempudi, N., Das, D., Banerjee, K., Avancha, S., et al. (2019). A Study of BFLOAT16 for Deep Learning Training
Korthikanti, V., Casper, J., Lym, S., McAfee, L., Andersch, M., Shoeybi, M., Catanzaro, B. (2022). Reducing Activation Recomputation in Large Transformer Models
Huang, Y., Cheng, Y., Bapna, A., Firat, O., Chen, M. X., Chen, D., et al. (2019). GPipe: Efficient Training of Giant Neural Networks using Pipeline Parallelism
Shoeybi, M., Patwary, M., Puri, R., LeGresley, P., Casper, J., Catanzaro, B. (2019). Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism
Narayanan, D., Shoeybi, M., Casper, J., LeGresley, P., Patwary, M., Korthikanti, V., et al. (2021). Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM
Dettmers, T., Lewis, M., Belkada, Y., Zettlemoyer, L. (2022). LLM.int8(): 8-bit Matrix Multiplication for Transformers at Scale
Frantar, E., Ashkboos, S., Hoefler, T., Alistarh, D. (2023). GPTQ: Accurate Post-Training Quantization for Generative Pre-trained Transformers
Biderman, S., Schoelkopf, H., Anthony, Q., Bradley, H., O’Brien, K., Hallahan, E., et al. (2023). Pythia: A Suite for Analyzing Large Language Models Across Training and Scaling
EleutherAI (2023). EleutherAI/pythia-1.4b (ficha del modelo, licencia Apache-2.0)
Merity, S., Xiong, C., Bradbury, J., Socher, R. (2016). Pointer Sentinel Mixture Models (presenta WikiText-2)