Checkpointing estilo Unsloth¶

Resumen corto: úsalo cuando el job casi cabe, y prueba FFN-only primero si el modelo lo soporta.

El backend unsloth descarga activaciones guardadas a CPU. El backend torch las descarta y las recalcula en backward. Unsloth puede comprar los últimos GiB para subir batch, resolución o frames. No es velocidad gratis. Si el run ya cabe con torch, normalmente torch sigue siendo el mejor default.

Controles¶

{
  "gradient_checkpointing": true,
  "gradient_checkpointing_backend": "torch-ffn"
}

gradient_checkpointing_backend tiene cuatro valores útiles:

En familias compatibles también puedes checkpointar menos bloques:

{
  "gradient_checkpointing": true,
  "gradient_checkpointing_backend": "torch-ffn",
  "gradient_checkpointing_interval": 2
}

gradient_checkpointing_interval: 2 checkpointa chunks contiguos de dos bloques en rutas de bloque completo compatibles. Valores más altos recomputan menos y dejan más activaciones en VRAM.

En esas rutas segmentadas, gradient_checkpointing_segment_stride también funciona con unsloth. Trátalo como una palanca para caber, no para acelerar: los bloques saltados se quedan en GPU, mientras los bloques checkpointed siguen usando CPU offload para tensores guardados. Para el resumen con torch y benchmarks por modelo, consulta Segmented Checkpointing.

gradient_checkpointing_offload_attention es independiente del backend. En bloques compatibles con separación attention/FFN, descarga las activaciones guardadas del lado attention. Puede ejecutarse solo o combinarse con torch, torch-ffn, unsloth o unsloth-ffn cuando el modelo soporte ese backend.

gradient_checkpointing_offload_pin_memory_max_buckets controla el pooling de CPU pinned para tensores guardados offloaded. El valor predeterminado es 12 buckets de tensor distintos; usa 0 para usar solo memoria CPU normal.

torch-ffn y unsloth-ffn soportan actualmente Chroma, Flux, Krea 2, LTXVideo2, MageFlow, Wan y Z-Image. Otras familias fallan claramente hasta que sus bloques expongan el mismo límite seguro.

Qué intercambia¶

• torch: descarta activaciones intermedias y las recalcula en backward.
• unsloth: guarda parte de esos tensores en CPU y los copia de vuelta para backward.
• *-ffn: checkpointa solo el lado feed-forward en modelos con un límite FFN limpio.
• Flash Attention ya evita materializar la matriz grande de atención. Esa idea de "checkpointing gratis" aplica sobre todo a atención, no a todo el bloque transformer.
• CPU offload ayuda más cuando las activaciones son grandes y el pico no viene de parámetros u optimizer.

Requiere CUDA y suficiente RAM de CPU. El ancho de banda PCIe importa. Si las copias CPU-GPU quedan expuestas, el step se vuelve más lento.

Nuestro sweep¶

Bloque transformer sintético, bf16, flash SDPA, pesos base congelados, batch 1. No son garantías de modelo; enseñan la forma del tradeoff.

Latents de imagen empaquetados¶

Con empaquetado 2x2, 64x64, 128x128 y 256x256 pasan a 1024, 4096 y 16384 tokens.

En 1024 tokens, el offload extra casi no importa salvo que ya estés al límite. En 16384 tokens, torch-ffn es el paso barato y whole-layer checkpointing es la palanca grande para caber. unsloth compra alrededor de 1.3 GiB más que torch layer checkpointing.

Transformer más grande¶

32 capas congeladas, ancho 4096, 3072 tokens:

Con pesos completos entrenables, gradientes y optimizer dominaron el pico, así que unsloth no ahorró más que torch en ese run sintético. PEFT se parece más al caso de pesos congelados.

Regla práctica¶

1. Si cabe sin checkpointing, déjalo apagado.
2. Si no cabe, prueba gradient_checkpointing_backend: torch-ffn.
3. Si sigue muy justo, prueba torch.
4. Si torch layer checkpointing aún no cabe, prueba unsloth-ffn y luego unsloth.
5. Si el modelo soporta gradient_checkpointing_interval, usa 2 o más solo después de que el run ya quepa y quieras recuperar velocidad.

Vale la pena cuando permite usar el batch, resolución, frames o rank que querías. No vale mucho para tokens pequeños ni cuando el pico viene de pesos entrenables, gradientes, optimizer, caché VAE o validación.

Notas¶

• Con FSDP activation checkpointing, SimpleTuner desactiva el checkpointing de modelo para evitar conflictos.
• torch-ffn y unsloth-ffn requieren soporte del modelo. SimpleTuner falla de forma explícita en vez de ejecutar otro scope en silencio.
• gradient_checkpointing_interval: 1 equivale al checkpointing normal de cada bloque.
• Algunas familias no tienen interval checkpointing. SimpleTuner avisa e ignora el intervalo.
• torch.compile no rescató el camino de offload en nuestro sweep.