Curso / Unsloth / Arquitectura interna y Triton
Unsloth · 1 de 3

Unsloth · Fine-tuning acelerado y Triton · 1

Arquitectura interna de Unsloth y optimización en Triton

PyTorch y Hugging Face priorizan generalidad sobre eficiencia de silicio: un fine-tuning estándar gasta gigabytes en activaciones intermedias y satura el bus de memoria de la GPU. Unsloth reescribe el backward a mano en OpenAI Triton y fusiona kernels en SRAM. Esta página explica la física del cálculo que lo hace 2× a 5× más rápido con 70% menos VRAM.

Antes de leerHF cap. 3 §5, LoRA · PyTorch 3 §1, autograd y VJP · LLM 2026 · 1, el bloque transformer

🎯 Al terminar vas a poder

  • Explicar con fundamentos de hardware por qué el stack clásico de Hugging Face y bitsandbytes satura el ancho de banda de memoria global (HBM) de la GPU.
  • Demostrar matemáticamente por qué la matriz de logits en la CrossEntropyLoss estándar produce fallas de CUDA Out of Memory y cómo el kernel de Cross Entropy fusionado en Triton lo soluciona.
  • Identificar cómo la fusión de kernels en memoria compartida (SRAM) para RMSNorm, RoPE y SwiGLU reduce round-trips innecesarios a VRAM.
  • Comprender la diferencia entre el gradient checkpointing tradicional de PyTorch y el enfoque selectivo de Unsloth.

0. El cuello de botella de PyTorch + PEFT tradicional

El stack tradicional para hacer fine-tuning eficiente de parámetros (PEFT) combina la librería transformers de Hugging Face, adaptadores peft y cuantización de 4 bits mediante bitsandbytes. Este diseño fue pensado para ser modular: cualquier arquitectura de modelo puede aprovechar cualquier método de adaptación sin cambiar una sola línea de código interno.

Sin embargo, esa modularidad impone un costo severo en hardware. PyTorch opera en modo eager construyendo un grafo computacional dinámico mediante autograd. Para cada operación en la pasada hacia adelante (forward), el framework reserva memoria en VRAM para almacenar tensores intermedios (activations), de modo que puedan ser leídos durante el backward pass para computar los gradientes.

En modelos de lenguaje modernos (como Llama 3, Gemma o Mistral), la gran mayoría de las operaciones no están limitadas por la capacidad aritmética de los Tensor Cores (compute-bound), sino por la velocidad a la que la GPU puede mover datos entre la memoria global (HBM) y los registros de cómputo (memory-bound). Unsloth demostró que no hacía falta inventar un nuevo algoritmo de optimización: derivando analíticamente las derivadas parciales en papel y escribiendo kernels específicos en OpenAI Triton, se eliminan por completo las escrituras y lecturas redundantes en la memoria de la placa de video.

1. La física de la GPU: HBM vs. SRAM y operaciones memory-bound

Para entender por qué un modelo se entrena lento, hay que mirar la jerarquía física de memoria de una GPU:

  • Memoria Global (HBM / VRAM): Es la memoria masiva de la GPU (por ejemplo, 24 GB en una RTX 3090/4090/5090 o 80 GB en una A100/H100). Tiene una gran capacidad, pero un ancho de banda comparativamente bajo (1 a 3 TB/s).
  • Memoria Compartida y Registros (SRAM): Vive directamente dentro de los Streaming Multiprocessors (SMs). Es extremadamente rápida (más de 15 a 30 TB/s), pero diminuta (apenas unos cientos de kilobytes por SM).

Las multiplicaciones matriciales densas (como las proyecciones de atención y los matmuls de las capas feed-forward) tienen una intensidad aritmética alta: realizan O(N3) operaciones sobre O(N2) datos cargados. Por el contrario, operaciones como RMSNorm, RoPE, sumas residuales y funciones de activación (SiLU/GELU) son estrictamente elemento a elemento (O(N)).

En PyTorch estándar, cada una de estas funciones se despacha como un kernel de CUDA independiente. Si se tiene la secuencia:

y=RMSNorm(x)⊙σ(z)

PyTorch lee x de la VRAM, calcula la normalización, escribe el resultado en VRAM; luego otro kernel lee z, calcula el sigmoide, escribe en VRAM; y finalmente un tercer kernel lee ambos tensores temporales para multiplicarlos. Los multiprocesadores de la GPU pasan la mayor parte del tiempo esperando a que los datos viajen por el bus de memoria global, con los Tensor Cores inactivos.

Comparativa de ejecución: PyTorch estándar vs. Kernel Fusionado en Triton de Unsloth A la izquierda el flujo estándar de PyTorch donde cada operación viaja de VRAM a SRAM y vuelve a VRAM repetidas veces. A la derecha, Unsloth carga los datos una sola vez a SRAM, realiza RMSNorm, activación y producto en registros y escribe el resultado final una sola vez. PYTORCH ESTÁNDAR · MÚLTIPLES VIAJES A VRAM (MEMORY-BOUND) VRAM: Tensores en Memoria Global x, z guardados en DRAM externa Kernel RMSNorm Kernel SiLU Act Kernel Multiplicación UNSLOTH · UN SOLO KERNEL FUSIONADO EN TRITON VRAM: Memoria Global Una sola lectura inicial · Una sola escritura final SRAM de los SMs (Memoria Ultrarrápida) RMSNorm + SiLU + Hadamard en Registros 0 escrituras intermedias a memoria global Al fusionar operaciones en SRAM, el ancho de banda efectivo se multiplica y la GPU pasa a ejecutar a velocidad de cómputo real.
Comparación entre la ejecución estándar de operaciones compuestas en PyTorch y la ejecución fusionada en Triton de Unsloth. Al evitar escribir activaciones intermedias en la memoria global (VRAM), la latencia cae drásticamente.

2. Derivación analítica del backward y OpenAI Triton

PyTorch implementa autodiferenciación automática mediante grafos dinámicos. Esto es sumamente versátil para investigación porque permite que cualquier expresión en Python sea derivable sin esfuerzo manual. Sin embargo, la autodiferenciación no optimiza automáticamente la memoria: cada nodo del grafo retiene sus entradas para poder computar las reglas locales de la cadena.

OpenAI Triton es un lenguaje de programación y compilador intermedio que permite escribir kernels de GPU de altísimo rendimiento en una sintaxis semejante a Python. A diferencia de CUDA en C++ (donde hay que gestionar manualmente la sincronización de hilos, los accesos coalescentes a memoria y las primitivas de sincronización como __syncthreads()), Triton opera sobre bloques de datos (tiles) y delega la optimización de registros y jerarquías al compilador de LLVM.

La innovación central de Daniel y Michael Han en Unsloth fue tomar las operaciones estándar de los LLMs (atención, normalizaciones, funciones de pérdida y adaptadores de LoRA) y derivar analíticamente en papel los productos vector-Jacobiano (VJP) del backward pass. Al implementar tanto el forward como el backward en kernels artesanales de Triton, lograron eliminar por completo la necesidad de que PyTorch mantenga el grafo de autograd para esas capas.

3. Fused Cross Entropy: eliminando la matriz de logits en VRAM

La causa número uno de colapso de memoria (OOM) en secuencias medianas o largas durante el fine-tuning no es el modelo en sí, sino la función de pérdida CrossEntropyLoss.

Sea un lote con B secuencias, longitud de contexto T y dimensión oculta d. La proyección del modelo hacia el vocabulario de tamaño V calcula:

Z=XWvocab⊤∈ℝ(B·T)×V

Analicemos la huella de memoria para un caso estándar en 2026:

  • Tokens totales por lote: B·T=4096.
  • Tamaño del vocabulario (Gemma 2/4 o Llama 3): V=256.000 tokens.
  • Tipo de dato en forward: FP16 o BF16 (2 bytes por elemento).
Memoria Logits=4096×256.000×2 bytes≈2,097 GB

Cómo se lee. Solo almacenar la matriz entera de logits de salida requiere más de 2 GB de memoria global. Si la secuencia sube a 8192 o 16384 tokens, esta sola matriz supera los 8 GB a 16 GB, superando la capacidad de la mayoría de las placas de video.

En PyTorch tradicional, la proyección lineal calcula la matriz entera Z y la guarda en VRAM. Luego se llama a torch.nn.CrossEntropyLoss, que toma esa matriz, resta el valor máximo para estabilidad numérica, computa el log-sum-exp, y durante el backward calcula:

∂L∂zi=pi-yi=softmax(z)i-𝟙[i=target]

La solución de Unsloth (Fused Cross Entropy): En lugar de materializar la matriz gigante en VRAM, Unsloth divide los tokens en bloques pequeños (por ejemplo, 64 tokens por bloque) en Triton. Para cada bloque de tokens, carga una porción de las representaciones ocultas X, calcula los logits directamente en la SRAM del procesador, evalúa el log-sum-exp online y el softmax, computa el gradiente analítico pi-yi, y retropropaga directamente el gradiente hacia X y Wvocab.

Resultado: La matriz completa de logits de 2 GB a 16 GB jamás se escribe en la VRAM de la GPU. El consumo de memoria de la pérdida se reduce a prácticamente cero, liberando VRAM para aumentar el tamaño de lote o la ventana de contexto.

4. Fusión de RoPE, RMSNorm y SwiGLU en un solo kernel

Además de la función de pérdida, Unsloth reescribe en Triton los componentes principales de cada bloque del transformer:

Fused RoPE (Rotary Positional Embeddings)

Como se vio en el bloque transformer de 2026, RoPE rota pares de dimensiones de Query y Key:

fm(x2i,x2i+1)=(x2icos(mθi)-x2i+1sin(mθi), x2isin(mθi)+x2i+1cos(θi))

PyTorch suele implementar esto desglosando los tensores en mitades, creando máscaras complejas y duplicando memoria. El kernel de Unsloth rota los vectores in-place directamente en SRAM en una única pasada, sin alojar tensores temporales para senos y cosenos.

Fused SwiGLU

La compuerta feed-forward moderna calcula SwiGLU(x)=silu(xWgate)⊙(xWup). En Unsloth, las salidas de ambas proyecciones lineales se mantienen en memoria compartida, se evalúa la función de activación suave y el producto de Hadamard elemento a elemento, y sólo se vuelca el vector final a VRAM para la proyección descendente.

5. Aceleración de QLoRA y gradient checkpointing selectivo

En QLoRA (Dettmers et al., 2023), los pesos de la red base se cuantizan a 4 bits utilizando el formato NormalFloat4 (NF4), mientras que los adaptadores de bajo rango A y B se entrenan en BF16 o FP16.

En la librería bitsandbytes tradicional, cada operación de multiplicación involucra kernels de CUDA que descomprimen los bloques NF4 a FP16 antes de ejecutar el matmul. Esos kernels originales sufren de sobrecosto de sincronización del driver y desalineación de memoria. Unsloth reescribió la descompresión y multiplicación directamente en Triton: los bloques de 4 bits se cargan empaquetados en registros, se descomprimen al vuelo en los Tensor Cores y se multiplican directamente, aumentando la velocidad de entrenamiento en un 30% a 50% frente al stack convencional de bitsandbytes.

📐 Gradient Checkpointing Quirúrgico

En PyTorch convencional, activar gradient_checkpointing=True descarta todas las activaciones intermedias del bloque transformer durante el forward, obligando a reejecutar el bloque completo durante el backward (un costo adicional del 33% en tiempo de cómputo).

Unsloth implementa un esquema de checkpointing selectivo: retiene en memoria únicamente los tensores pesados cuya recomputación exigiría operaciones matriciales caras (O(d2)), y recalcula al vuelo sobre SRAM las operaciones lineales ligeras (O(d)) como normalizaciones RMSNorm y sumas residuales. De esta forma, se obtiene el ahorro masivo de memoria del checkpointing sin pagar la penalización completa de latencia.

6. Errores comunes

SíntomaCausa probableQué hacer
ImportError: Unsloth must be imported before torch or transformersUnsloth parchea dinámicamente las clases internas de Hugging Face y PyTorch durante el inicio.Colocar from unsloth import FastLanguageModel en la primera línea absoluta del script, antes de cualquier otro import de ML.
Incompatibilidad de kernels Triton en arquitecturas viejasGPUs anteriores a Turing (Compute Capability < 7.5, como Pascal o GTX 1080) carecen de soporte para los tipos de datos e instrucciones requeridas por los kernels.Utilizar GPUs modernas (Ampere, Ada Lovelace, Hopper, Blackwell) con arquitectura Turing o superior (RTX 2000 en adelante, T4, A100, H100).
CUDA Out of Memory al usar optimizadores estándarUsar AdamW estándar de PyTorch reserva 8 bytes por parámetro para momentum y varianza.Utilizar optimizadores fusionados en 8 bits como paged_adamw_8bit, que además descargan estados a la CPU si la VRAM se satura.
Pérdida de velocidad por batch size desalineadoLos kernels de Triton operan óptimamente con bloques potencia de dos (16, 32, 64, 128).Asegurarse de que el tamaño de lote efectivo y la dimensión de corte de secuencia sean múltiplos de 16 o 32.

7. Resumen

  1. Unsloth no inventa un nuevo algoritmo de optimización: optimiza la ejecución física del cálculo de LoRA/QLoRA en el silicio de la GPU.
  2. La memoria global (HBM) es el cuello de botella en operaciones elemento a elemento; Unsloth fusiona operaciones en la memoria compartida ultrarrápida (SRAM).
  3. La derivación manual del Vector-Jacobian Product (VJP) en OpenAI Triton elimina la necesidad de retener el grafo completo de autograd de PyTorch.
  4. El kernel de Fused Cross Entropy evita materializar la matriz de logits en VRAM, ahorrando entre 2 GB y 16 GB por paso de entrenamiento.
  5. Los kernels acelerados de QLoRA descomprimen los bloques NF4 directamente en registros, superando a la implementación original de bitsandbytes.

8. Ponete a prueba

1. ¿Por qué la matriz de logits en CrossEntropyLoss satura la VRAM en secuencias largas?

Porque su tamaño es de (B·T)×V. Con un contexto de 4096 tokens y un vocabulario moderno de 256.000 palabras (como Llama 3 o Gemma), ocupa más de 2 GB en FP16. En PyTorch tradicional, esa matriz completa se guarda en VRAM para el backward; Unsloth la calcula por bloques en SRAM y retropropaga sin escribirla en memoria global. (§3)

2. ¿Qué significa que una operación sea "memory-bound" frente a "compute-bound"?

Una operación memory-bound (como RMSNorm o SiLU) realiza pocas operaciones matemáticas por cada byte transferido desde la VRAM (O(N)), saturando el bus de memoria. Una operación compute-bound (como un matmul denso) realiza muchas operaciones por byte cargado (O(N3)), aprovechando plenamente los núcleos de cálculo. (§1)

3. ¿Cómo logra Unsloth evitar que PyTorch construya grafos de autograd masivos?

Sus autores derivaron en papel las derivadas parciales y programaron manualmente tanto el forward como el backward pass en kernels de OpenAI Triton, prescindiendo del seguimiento automático de nodos de PyTorch. (§2)

4. ¿Qué ventaja ofrece el gradient checkpointing selectivo frente al clásico?

El checkpointing clásico descarta todo el bloque transformer y recalcula el 100% en el backward, perdiendo un 33% de velocidad. El selectivo guarda las activaciones de multiplicaciones matriciales caras y sólo recalcula operaciones lineales muy baratas (como sumas residuales o normas). (§5)

9. Para practicar

✏️ Ejercicio 1 — Comparar consumo de memoria de CrossEntropy

Escribí un script corto en PyTorch midiendo el pico de memoria VRAM (torch.cuda.max_memory_allocated()) al calcular una proyección lineal hacia un vocabulario de 128.000 tokens con y sin retener el grafo de backward. Observá el salto en gigabytes al materializar los logits.

Pista / qué esperar

Con batch_size=4 y seq_len=2048, la sola llamada a logits = linear(x) consumirá ~2.1 GB de VRAM en float16. En cambio, si se calcula la pérdida por trozos de 128 tokens acumulando el escalar de pérdida, el pico de memoria desciende a menos de 200 MB.

10. Fuentes y para seguir