Nuevo estándar para el almacenamiento de modelos de IA: concepto y ventajas

Cuando navegas por hubs de modelos de IA como Hugging Face, comienzas a ver cada vez más archivos con la extensión .safetensors en lugar de los archivos .bin o .pth que solían ser comunes en el pasado.

En esta publicación, exploraremos en detalle qué es .safetensors, por qué ha surgido y qué ventajas poderosas ofrece con respecto a los métodos existentes desde una perspectiva técnica.

1. ¿Qué es .safetensors?

.safetensors es un nuevo formato de almacenamiento de tensores desarrollado por Hugging Face.

Los modelos de aprendizaje profundo están compuestos por miles de millones de parámetros (pesos), y este formato se utiliza para guardar y recuperar estos enormes conjuntos de números. Se creó para resolver las desventajas críticas (seguridad y velocidad) del método de almacenamiento basado en el módulo pickle de Python, que era el estándar de facto anteriormente.

En términos simples, es un "archivo de modelo más seguro y rápido".

2. ¿Por qué ha surgido?: Problemas del método existente (Pickle)

Los modelos de PyTorch existentes (.bin, .pth) utilizan internamente el módulo pickle de Python para serializar (Serialization) los datos. Sin embargo, pickle tiene un problema crítico.

Vulnerabilidad de seguridad (Ejecución de código arbitrario)

pickle no solo almacena datos, sino que también almacena el objeto en sí de Python. En este proceso, puede incluir código de Python, lo que permite a un hacker malintencionado insertar código en el archivo del modelo que puede dañar el sistema o robar información personal. En el momento en que el usuario carga el modelo sin saberlo, se ejecuta ese código malicioso.

Ejemplo de situación problemática:

El usuario carga un model.bin descargado de Internet -> se ejecuta el código oculto dentro del archivo -> la clave SSH o la contraseña del usuario se envían al servidor del hacker.

.safetensors ha surgido para bloquear estas amenazas de seguridad desde la raíz.

3. Características clave de .safetensors

3.1. Seguridad

.safetensors es, como su nombre indica, seguro. Este formato solo almacena datos de tensor puros y metadatos en formato JSON. No hay espacio para incluir código ejecutable, por lo que puedes cargar archivos incluso de fuentes no confiables sin preocupación.

3.2. Zero-Copy y velocidad

La velocidad al cargar modelos LLM grandes o modelos de Stable Diffusion es notablemente más rápida.

• Método existente: Copiar el archivo a la memoria CPU -> deserializar (Unpickling) -> volver a convertir a tensor -> mover a GPU. (Se produce un proceso de copia innecesario)

• safetensors: Utiliza la tecnología de mapeo de memoria (Memory Mapping, mmap). El sistema operativo mapea directamente el archivo a una dirección de memoria, por lo que se pueden usar los datos del disco sin ningún proceso de copia innecesario. Esto se llama Zero-Copy.

3.3. Carga perezosa (Lazy Loading)

Puedes leer rápidamente solo las partes necesarias sin tener que cargar todo el modelo en la memoria.

Por ejemplo, si deseas verificar solo los pesos de una capa específica en un archivo de modelo de 100GB, con el método existente tendrías que leer los 100GB completos, pero con .safetensors puedes leer solo esa parte específica. Esto es muy ventajoso en entornos de aprendizaje distribuido o optimización de inferencia.

3.4. Compatibilidad con frameworks

No está atado a un framework de aprendizaje profundo específico (como PyTorch).

• PyTorch

• TensorFlow

• JAX

• PaddlePaddle

Está diseñado para ser leído y escrito fácilmente en una variedad de frameworks como los mencionados anteriormente.

4. Estructura del archivo

El archivo .safetensors tiene una estructura muy simple.

1. Encabezado (Header): Se encuentra al principio del archivo y sigue el formato JSON. Contiene información sobre el nombre de cada tensor, tipo de datos (dtype), forma (shape) y la ubicación (offset) donde se almacena el dato.

2. Datos (Data): Es el bloque de datos binarios que sigue al encabezado. Está compuesto solo por valores de tensors.

Gracias a esta estructura, también es posible entender la estructura del modelo solo leyendo el encabezado sin tener que leer el archivo completo.

4-1. Consejo práctico: Verificación de encabezados y metadatos en el terminal

Cargar completamente un modelo de varios gigabytes para verificar si el archivo .safetensors es un modelo cuantificado o qué capas incluye es ineficiente.

La biblioteca safetensors ofrece una función que permite escanear rápidamente solo la información del encabezado sin cargar el archivo completo. Puedes ingresar un Python One-liner en el terminal como el siguiente para verificarlo al instante.

Por supuesto, también puedes verificar haciendo clic en file info en la parte superior derecha de la página donde descargaste el modelo de Hugging Face. Sin embargo, aquí explicaré cómo verificarlo directamente en mi terminal.

Preparativos

Primero, debes tener la biblioteca instalada.

pip install safetensors

Comando (entrada en terminal)

Introduce la ruta real del archivo en lugar de model.safetensors.

python -c "from safetensors import safe_open; \
with safe_open('model.safetensors', framework='pt', device='cpu') as f: \
    print('--- Metadata ---'); \
    print(f.metadata()); \
    print('\n--- Tensor Keys (Layers) ---'); \
    print(list(f.keys())[:5])" # Para no ser muy largo, solo mostrar los primeros 5

Interpretación de los resultados

Al ejecutar este comando, puedes obtener dos informaciones importantes.

1. Metadatos: Información que ha inserido el creador del modelo. Si contiene información como format: gptq o quantization: int4, puedes saber que es un modelo cuantificado, incluso si no está escrito en el nombre del archivo. (Sin embargo, si el creador ha dejado los metadatos en blanco, puede aparecer como None.)

2. Claves: Nombres de las capas que componen el modelo. Esto te ayuda a entender la estructura del modelo.

5. Resumen comparativo: .bin (Pickle) vs .safetensors

6. Ejemplo de uso (Python)

A continuación, un ejemplo simple de cómo guardar y cargar tensores usando la biblioteca safetensors.

import torch
from safetensors.torch import save_file, load_file

# 1. Crear y guardar tensores
tensors = {
    "embedding": torch.zeros((1024, 512)),
    "attention": torch.rand((512, 512))
}

# Guardar el tensor en formato de diccionario en un archivo
save_file(tensors, "model.safetensors")

# 2. Cargar el archivo
loaded = load_file("model.safetensors")
print(loaded["embedding"].shape) 
# Salida: torch.Size([1024, 512])

7. Conclusión

.safetensors no es solo un cambio simple de extensión de archivo, sino una evolución necesaria para la seguridad y eficiencia de los modelos de IA. Las principales comunidades, incluida Hugging Face, ya lo han adoptado como formato básico.

Se recomienda utilizar .safetensors en lugar de .bin siempre que sea posible al descargar o distribuir modelos, ya que esto ofrece una protección contra amenazas de seguridad y reduce drásticamente el tiempo de carga del modelo.