Guardar, cargar, y el cambio de reglas de PyTorch 2.0

1. state_dict vs. guardar el modelo completo

Un modelo entrenado guarda sus parámetros aprendidos en un diccionario interno, el state_dict — mapea el nombre de cada parámetro (y de cada buffer persistente, como viste en la página 2) a su tensor de valores. Es la forma recomendada de persistir un modelo:

torch.save(model.state_dict(), 'model_weights.pth')

# para cargar: primero instanciar la MISMA clase, después cargar los pesos
model = models.vgg16()
model.load_state_dict(torch.load('model_weights.pth', weights_only=True))
model.eval()   # no lo olvides — ver la página anterior

Notá que cargar el state_dict requiere tener antes una instancia de exactamente la misma clase de modelo — el state_dict son solo números, no sabe nada de la arquitectura que los usa. La alternativa, torch.save(model, 'model.pth'), guarda el objeto completo (arquitectura incluida) usando pickle de Python — más cómodo, pero acopla el archivo guardado a que la definición exacta de la clase siga existiendo y siendo importable cuando se cargue de vuelta. La documentación oficial es clara: state_dict es la práctica recomendada; guardar el objeto completo es "un caso de uso legado".

weights_only y map_location

torch.load usa pickle por dentro, y pickle puede ejecutar código arbitrario al deserializar: un archivo .pth de origen desconocido es, literalmente, un programa. weights_only=True restringe la carga a tensores, diccionarios y tipos primitivos. Desde PyTorch 2.6 es el default — así que el weights_only=True explícito del ejemplo ya es redundante, y lo que cambia es lo contrario: un checkpoint viejo que contenía objetos arbitrarios (un argparse.Namespace, una clase propia) ahora falla con un UnpicklingError que te dice qué tipo rechazó. Las salidas: weights_only=False solo si confiás en el archivo, o registrar el tipo con torch.serialization.add_safe_globals([MiClase]).

map_location decide en qué device aterrizan los tensores. Un state_dict guardado desde GPU recuerda que vivía en cuda:0; cargarlo en una máquina sin GPU (o con menos GPUs) sin map_location falla. La costumbre segura es cargar siempre a CPU y dejar que model.to(device) mueva después:

sd = torch.load("model_weights.pth", map_location="cpu")   # weights_only=True implícito en 2.6+
model.load_state_dict(sd)
model.to(device)

El checkpoint completo para reanudar

Los pesos solos alcanzan para inferencia, pero no para continuar un entrenamiento: el estado de Adam (momento y varianza por parámetro), la posición del scheduler y el epoch en que ibas también son estado. El patrón es un diccionario con todo, guardado de forma atómica (escribir a un archivo temporal y renombrar, para que un corte a mitad de escritura no deje un checkpoint corrupto):

def save_checkpoint(path, model, optimizer, scheduler, epoch, step, scaler=None):
    ckpt = {
        "model_state_dict": model.state_dict(),
        "optimizer_state_dict": optimizer.state_dict(),
        "scheduler_state_dict": scheduler.state_dict(),
        "scaler_state_dict": scaler.state_dict() if scaler else None,
        "epoch": epoch, "step": step,
        "rng": torch.get_rng_state(),
    }
    tmp = path + ".tmp"
    torch.save(ckpt, tmp)
    os.replace(tmp, path)                            # atómico en el mismo filesystem

def load_checkpoint(path, model, optimizer, scheduler, scaler=None):
    ckpt = torch.load(path, map_location="cpu")
    model.load_state_dict(ckpt["model_state_dict"])
    optimizer.load_state_dict(ckpt["optimizer_state_dict"])   # mueve el estado al device de los params
    scheduler.load_state_dict(ckpt["scheduler_state_dict"])
    if scaler and ckpt["scaler_state_dict"]: scaler.load_state_dict(ckpt["scaler_state_dict"])
    torch.set_rng_state(ckpt["rng"])
    return ckpt["epoch"] + 1, ckpt["step"]

El estado del optimizador pesa: con AdamW son dos tensores fp32 por parámetro entrenable, o sea el doble que los pesos en bf16. Por eso los checkpoints de entrenamiento de un modelo de 7B rondan las decenas de GB, mientras que los pesos para inferencia del mismo modelo ocupan 14. Es también la razón por la que un checkpoint de LoRA es diminuto: solo el adapter y su estado de Adam.

2. safetensors: el formato del ecosistema HF

Cuando bajás un modelo del Hub, el archivo se llama model.safetensors, no pytorch_model.bin. safetensors es el formato que Hugging Face creó para reemplazar a pickle, y resuelve tres problemas de torch.save de una vez:

• Seguridad: el archivo es un header JSON (nombre, dtype, shape y offset de cada tensor) seguido de los bytes crudos. No hay código que ejecutar al cargar — es por diseño imposible que un archivo malicioso haga algo.
• Carga lazy y memory-mapped: se puede leer el header sin tocar los datos, cargar un solo tensor, o mapear el archivo directamente a memoria. Cargar un modelo grande es mucho más rápido que despicklear, y es lo que permite que device_map="auto" reparta tensores entre GPUs y CPU sin cargar todo antes.
• Interoperable: el mismo archivo se lee desde PyTorch, JAX, TensorFlow, Rust (candle) o cualquier lenguaje con un parser JSON.

from safetensors.torch import save_file, load_file, save_model, load_model

# nivel state_dict: un dict {nombre: tensor}
save_file(model.state_dict(), "model.safetensors")
sd = load_file("model.safetensors", device="cpu")      # dict listo para load_state_dict
model.load_state_dict(sd)

# nivel modelo: maneja pesos compartidos (tied embeddings) que save_file rechaza
save_model(model, "model.safetensors")
load_model(model, "model.safetensors")

# inspeccionar sin cargar
from safetensors import safe_open
with safe_open("model.safetensors", framework="pt") as f:
    for k in list(f.keys())[:3]:
        print(k, f.get_slice(k).get_shape(), f.get_slice(k).get_dtype())

Dos restricciones de save_file que aparecen en la práctica: los tensores tienen que ser contiguos (.contiguous() antes de guardar un peso transpuesto) y no puede haber dos claves que compartan memoria — el caso típico es el embedding de entrada y la cabeza de salida atados (tied weights); save_model lo resuelve guardando una sola copia. Y lo que safetensors no guarda: el estado del optimizador ni ningún objeto que no sea un tensor. Para el checkpoint de reanudación de §1, torch.save sigue siendo la herramienta; para publicar pesos, safetensors. model.save_pretrained() de Transformers y PEFT hace safetensors por defecto, más un config.json que reemplaza a "instanciar la misma clase".

3. El problema que PyTorch cargó durante 5 años: eager es lento

La fortaleza histórica de PyTorch —ejecución eager, cada línea corre inmediatamente en Python, sin un grafo estático de por medio— es también su límite de performance. Desde que se lanzó PyTorch en 2017, el hardware (GPUs) se volvió ~15× más rápido en cómputo, pero mantener esa velocidad en modo eager obligó a mover cada vez más lógica interna de Python a C++, lo cual va en contra del valor central de PyTorch: que el código sea hackeable y fácil de extender.

El problema concreto tiene dos caras. Una es el overhead de Python: cada operación pasa por el intérprete, el dispatcher y el lanzamiento de un kernel de CUDA; cuando el kernel dura microsegundos (un modelo chico, un batch chico, la fase de decode de un LLM que procesa un token por paso), la GPU pasa más tiempo esperando órdenes que calculando. La otra es la falta de fusión: x.mul(2).add(1).relu() en eager son tres kernels que leen y escriben el tensor completo tres veces desde la memoria de la GPU, cuando un kernel fusionado lo leería y escribiría una sola vez. En la mayoría de las capas que no son matmuls, el cuello de botella es el ancho de banda de memoria, no el cómputo — y ahí la fusión es la única forma de ir más rápido.

El equipo de PyTorch probó varios enfoques de compilación a lo largo de 5 años — torch.jit.trace, TorchScript, FX tracing, Lazy Tensors— y ninguno resolvía el problema completo: algunos eran flexibles pero lentos, otros rápidos pero exigían reescribir el modelo de forma no trivial.

4. torch.compile: un decorador, no una reescritura

torch.compile (PyTorch 2.0, 2023) resolvió esto sin pedirle al usuario que cambie una sola línea de su modelo:

model = torch.compile(model)   # eso es todo — el resto del código no cambia

Por dentro corren cuatro tecnologías nuevas, cada una resolviendo una parte distinta del problema de compilación:

• TorchDynamo: captura el grafo de cómputo de forma segura usando un hook del intérprete de CPython (PEP-0523) — sobre 7.000+ proyectos reales de GitHub usados como validación, acertó a capturar el grafo correctamente el 99% de las veces, contra menos del 50% de TorchScript.
• AOTAutograd: genera de antemano ("ahead-of-time") el grafo del backward pass, no solo el forward.
• PrimTorch: reduce los más de 2000 operadores de PyTorch a un conjunto cerrado de ~250 operaciones primitivas — simplifica muchísimo escribir un backend nuevo.
• TorchInductor: el compilador final, que genera código Triton para GPU y C++/OpenMP para CPU.

La pieza conceptualmente nueva es Dynamo: en vez de trazar el modelo ejecutándolo con tensores de mentira (lo que rompía con cualquier if sobre valores) o de parsear un subconjunto de Python (TorchScript), intercepta el bytecode de cada función justo antes de ejecutarse, extrae las operaciones sobre tensores a un grafo FX y deja el resto del Python como estaba. Cuando encuentra algo que no puede capturar (un print, una llamada a una librería externa, un if que depende del valor de un tensor), corta el grafo ahí — un graph break — compila el pedazo anterior, ejecuta esa línea en eager y sigue capturando después. Por eso "funciona" en el 99 % de los casos: el peor escenario no es un error, es un modelo con más pedazos compilados de los que te gustaría.

5. Los números reales, y el matiz importante

Sobre un benchmark de 163 modelos open-source (46 de Hugging Face Transformers, 61 de TIMM, 56 de TorchBench), sin modificar el código de esos modelos más allá de envolverlos en torch.compile:

Sylvain Gugger, mantenedor principal de Hugging Face Transformers en ese momento, lo resumió así: "con solo una línea de código, PyTorch 2.0 da un speedup de entre 1.5x y 2x entrenando modelos Transformer — lo más emocionante desde que se introdujo mixed precision training". El matiz honesto que la propia documentación marca: los speedups son menores en GPUs de consumo (una RTX 3090, por ejemplo) que en GPUs de servidor como A100, y TorchInductor no soporta todavía todas las arquitecturas de GPU por igual.

Otro matiz que los números de arriba no muestran: el speedup depende de dónde estaba el cuello de botella. Un modelo grande con batches grandes ya pasa casi todo su tiempo en matmuls que cuBLAS ejecuta igual de rápido compilado o no — ahí la ganancia viene solo de fusionar lo que rodea a las matmuls (norms, activaciones, residuales) y suele ser del orden del 10-30 %. Un modelo chico, o la fase de generación token a token de un LLM, está dominado por overhead de lanzamiento, y ahí mode="reduce-overhead" (ver §6) puede dar 2× o más. Medir antes y después con el mismo input, y después de calentar, es la única forma de saber.

6. torch.compile en la práctica: modos, shapes dinámicos, recompilaciones

"Una línea" es cierto, pero hay cuatro cosas que conviene saber antes de agregarla a un entrenamiento o a un servidor.

El costo de la primera iteración

torch.compile es lazy: no compila nada hasta la primera llamada con inputs reales. Esa primera iteración tarda desde unos segundos (un MLP) hasta varios minutos (un LLM grande con max-autotune), porque Dynamo captura, Inductor genera código Triton y lo compila. Todo benchmark que incluya la primera llamada está midiendo compilación, no ejecución: calentar con dos o tres iteraciones y recién después cronometrar. Para servidores, la caché en disco (habilitada por defecto en versiones recientes) amortiza el costo entre reinicios.

Modos

Shapes dinámicos y recompilaciones

El código compilado está especializado en los shapes que vio. Si la próxima llamada trae un shape distinto (un batch más chico al final del epoch, una secuencia más larga), Dynamo tiene que recompilar. El comportamiento por defecto (dynamic=None) es pragmático: la primera vez asume shapes estáticos; al ver un segundo shape, marca esa dimensión como dinámica y compila una versión que sirve para cualquier valor. Si sabés de antemano que el shape varía (largo de secuencia en texto), dynamic=True lo hace desde el principio y te ahorra una recompilación; si sabés que no varía nunca, dynamic=False genera código algo más rápido. Hay un límite de recompilaciones por función (8 por defecto) tras el cual Dynamo se rinde y esa función vuelve a eager sin avisar en la salida normal — por eso el primer diagnóstico cuando "compilé y no fue más rápido" es mirar los logs:

model = torch.compile(model, mode="reduce-overhead", dynamic=True)   # o fullgraph=True para que un graph break sea un error

# en la shell, antes de correr:
# TORCH_LOGS="graph_breaks,recompiles" python train.py

# o desde Python, para ver cuántos grafos y por qué se cortaron:
explicacion = torch._dynamo.explain(model)(x)
print(explicacion.graph_count, explicacion.graph_break_count)
print(explicacion.break_reasons)

De torch.compile al deployment: torch.export y ExecuTorch

torch.compile necesita Python en runtime — es un acelerador del programa que ya tenés. Cuando el destino no tiene Python (un servidor en C++, un teléfono, un microcontrolador), torch.export captura el grafo completo (sin graph breaks: si hay uno, es un error) en un ExportedProgram serializable; AOTInductor lo compila a una librería que se carga desde C++, y ExecuTorch es el runtime liviano de PyTorch para móvil y edge que consume ese mismo artefacto. Es el sucesor de lo que TorchScript prometía, construido sobre las mismas piezas de 2.0.

7. Errores comunes

Los de carga vienen con mensajes largos pero precisos; los de compilación son más silenciosos.

8. Resumen

1. state_dict es la forma recomendada de guardar un modelo — solo números, requiere reinstanciar la clase para cargar; guardar el objeto completo con pickle es un caso de uso legado. map_location="cpu" siempre, y weights_only=True ya es el default desde 2.6.
2. Para reanudar un entrenamiento hace falta un checkpoint completo (modelo, optimizador, scheduler, epoch, scaler); para publicar pesos, safetensors: seguro por construcción, lazy, memory-mapped, y el formato de todo el ecosistema HF.
3. La ejecución eager de PyTorch es flexible pero deja rendimiento sobre la mesa — overhead de Python por operación y kernels sin fusionar — el problema que PyTorch 2.0 resolvió después de 5 años de intentos previos.
4. torch.compile no requiere reescribir el modelo — TorchDynamo captura el grafo con 99% de éxito sin tocar el código original, y donde no puede, corta (graph break) en vez de fallar.
5. Los números reales: 93% de modelos compilan sin problema, con speedups de 21% (FP32) a 51% (mixed precision) en GPUs de servidor; menos en GPUs de consumo, más en cargas dominadas por overhead con reduce-overhead.
6. En la práctica: la primera llamada paga la compilación; dynamic=True para shapes variables; los graph breaks y recompilaciones se ven con TORCH_LOGS; torch.export + ExecuTorch es el camino cuando el destino no tiene Python.
7. Compilación y cuantización no compiten — atacan ejes distintos del mismo problema (velocidad vs. memoria) y se combinan en producción real.

9. Ponete a prueba

Respondé antes de abrir cada respuesta. Si fallás una, volvé a la sección indicada.

10. Para practicar

import json, struct, torch, torch.nn as nn
from safetensors.torch import save_file

m = nn.Sequential(nn.Linear(16, 32), nn.ReLU(), nn.Linear(32, 4)).bfloat16()
save_file(m.state_dict(), "mini.safetensors")

def header(path):
    with open(path, "rb") as f:
        (n,) = struct.unpack("<Q", f.read(8))          # largo del header en bytes
        return json.loads(f.read(n))

for k, v in header("mini.safetensors").items():
    if k != "__metadata__":
        nbytes = v["data_offsets"][1] - v["data_offsets"][0]
        print(f"{k:12s} {v['dtype']:5s} {v['shape']}  {nbytes} bytes")

import os, torch, torch.nn as nn
def setup(seed=0):
    torch.manual_seed(seed)
    m = nn.Linear(8, 1); opt = torch.optim.AdamW(m.parameters(), lr=1e-2)
    sch = torch.optim.lr_scheduler.CosineAnnealingLR(opt, T_max=200)
    return m, opt, sch
X, y = torch.randn(64, 8), torch.randn(64, 1)
def step(m, opt, sch):
    loss = ((m(X) - y) ** 2).mean(); opt.zero_grad(); loss.backward(); opt.step(); sch.step()
    return loss.item()

# corrida de referencia, sin corte
m, opt, sch = setup()
for i in range(200):
    if i == 100:
        torch.save({"model": m.state_dict(), "opt": opt.state_dict(), "sch": sch.state_dict()}, "ckpt.pt")
    ref = step(m, opt, sch)
print("referencia:", ref, sch.get_last_lr())

# reanudar con checkpoint completo
m, opt, sch = setup(seed=1)                     # pesos iniciales distintos a propósito
ck = torch.load("ckpt.pt", map_location="cpu")
m.load_state_dict(ck["model"]); opt.load_state_dict(ck["opt"]); sch.load_state_dict(ck["sch"])
for i in range(100, 200): full = step(m, opt, sch)
print("completo:  ", full, sch.get_last_lr())

# reanudar solo con pesos
m, opt, sch = setup(seed=1)
m.load_state_dict(ck["model"])
for i in range(100, 200): solo = step(m, opt, sch)
print("solo pesos:", solo, sch.get_last_lr())

import torch

def con_if(x):
    if x.sum() > 0:          # depende del VALOR del tensor
        return x * 2
    return x - 1

def sin_if(x):
    return torch.where(x.sum() > 0, x * 2, x - 1)

x = torch.randn(8)
for fn in (con_if, sin_if):
    e = torch._dynamo.explain(fn)(x)
    print(fn.__name__, "grafos:", e.graph_count, "breaks:", e.graph_break_count)

torch.compile(sin_if, backend="eager", fullgraph=True)(x)   # ok
torch.compile(con_if, backend="eager", fullgraph=True)(x)   # error: no puede capturar el grafo completo

11. Fuentes y para seguir

• Save and Load the Model — Learn the Basics — el tutorial original de §1, con el ejemplo de VGG16.
• Serialization semantics — la nota oficial sobre torch.save/torch.load: tensores que comparten memoria, weights_only, add_safe_globals y compatibilidad entre versiones.
• safetensors — documentación — API de Python, especificación del formato (header + datos) y las comparaciones de velocidad de carga contra pickle.
• PyTorch 2.0 — anuncio y overview técnico — la fuente de §3 a §5: los benchmarks, la explicación de Dynamo/AOTAutograd/PrimTorch/Inductor y la cita de Gugger.
• Introduction to torch.compile — el tutorial oficial: modos, dynamic, cómo medir el speedup correctamente y cómo leer graph breaks.
• torch.export — la referencia del camino a deployment sin Python; y ExecuTorch, el runtime para móvil y edge que consume ese artefacto.