Nieuwe standaard voor het opslaan van AI-modellen: concepten en voordelen

Als je door AI-model hubs zoals Hugging Face bladert, zie je dat er steeds meer bestanden met de extensie .safetensors beschikbaar zijn, in plaats van de eerder gebruikelijke .bin of .pth bestanden.

In deze post zullen we technisch diepgaand verkennen wat .safetensors is, waarom het is ontstaan en welke krachtige voordelen het heeft ten opzichte van de bestaande methoden.

1. Wat is .safetensors?

.safetensors is een nieuwe tensor opslagformat ontwikkeld door Hugging Face.

Deep learning modellen bestaan uit miljarden parameters (gewichten) en zijn verantwoordelijk voor het opslaan en ophalen van deze enorme hoeveelheid cijfers. Het is ontwikkeld om de fatale tekortkomingen (veiligheid en snelheid) van de opslagmethode op basis van de Python pickle module, die feitelijk de standaard was, op te lossen.

Kortom, het is een "veiliger en snellere model opslaan bestand".

2. Waarom is het ontstaan: de problemen van de bestaande methode (Pickle)

Bestaande PyTorch-modellen (.bin, .pth) gebruiken intern de Python pickle module om gegevens te serialiseren. Echter, pickle heeft fatale problemen.

Beveiligingskwetsbaarheid (Arbitrary Code Execution)

pickle slaat niet alleen gegevens op, maar slaat ook Python-objecten zelf op. In dit proces kan Python code worden opgenomen, waardoor kwaadwillende hackers code in het interne modelbestand kunnen verbergen die het systeem kan beschadigen of persoonlijke gegevens kan stelen. Op het moment dat de gebruiker onbewust het model load(), wordt deze kwaadaardige code uitgevoerd.

Voorbeeld van een probleem situatie:

De gebruiker laadt model.bin gedownload van het internet -> De verborgen code in het bestand wordt uitgevoerd -> De SSH-sleutels of wachtwoorden van de gebruiker worden naar de hackerserver verzonden.

.safetensors is ontstaan om dergelijke beveiligingsbedreigingen te voorkomen.

3. Kernfeatures van .safetensors

3.1. Veiligheid (Safety)

.safetensors is, zoals de naam al zegt, veilig. Dit format slaat alleen pure tensorgegevens en metadata (in JSON-vorm) op. Er is geen ruimte voor uitvoerbare code, waardoor je zelfs bestanden die van een onbetrouwbare bron zijn gedownload met vertrouwen kunt laden.

3.2. Zero-Copy en snelheid

Bij het laden van grote LLM's (Large Language Models) of Stable Diffusion-modellen is de snelheid revolutionair snel.

• Bestaande methode: bestand kopiëren naar CPU-geheugen -> deserialisatie (Unpickling) -> opnieuw converteren naar tensorvorm -> verplaatsen naar GPU. (Onnodige kopieerstappen ontstaan)

• safetensors: maakt gebruik van geheugenmapping (Memory Mapping, mmap) technologie. Het besturingssysteem mappt het bestand zelf direct naar een geheugenadres, waardoor gegevens onmiddellijk van de schijf kunnen worden gebruikt zonder onnodige kopieën. Dit wordt Zero-Copy genoemd.

3.3. Lazy Loading

Het is mogelijk om niet het volledige model in het geheugen te laden, maar alleen de noodzakelijke delen snel te lezen.

Bijvoorbeeld, als je alleen de gewichten van een specifieke laag in een modelbestand van 100GB wilt bekijken, moest je met de bestaande methode 100GB helemaal lezen, terwijl .safetensors alleen dat specifieke deel kan ophalen. Dit is zeer voordelig in een gedistribueerde leeromgeving of bij optimalisatie voor inferentie.

3.4. Framework-compatibiliteit

Het is niet afhankelijk van specifieke deep learning frameworks (zoals PyTorch).

• PyTorch

• TensorFlow

• JAX

• PaddlePaddle

Het is ontworpen om gemakkelijk lees- en schrijfbaar te zijn in een verscheidenheid aan frameworks zoals hierboven genoemd.

4. Bestandstructuur

.safetensors bestanden hebben een zeer eenvoudige structuur.

1. Header: bevindt zich aan het begin van het bestand en volgt het JSON-formaat. Het bevat informatie over de naam, datatypes (dtype), vormen (shapes), en de opslaglocatie (offset) van elke tensor.

2. Data: de binaire datablok die volgt na de header. Het bevat puur tensorwaarden die dicht op elkaar zijn gepakt.

Door deze structuur is het mogelijk om het modelstructuur te begrijpen door alleen de header te lezen zonder het hele bestand te lezen.

4-1. Praktische tip: Header- en metadata controleren in de terminal

Het is inefficiënt om een model van meerdere gigabytes te laden om te controleren of een gedownload .safetensors bestand gequantiseerd is of welke lagen zijn inbegrepen.

De safetensors bibliotheek biedt de mogelijkheid om snel de headerinformatie te scannen zonder het hele bestand te lezen. Door de volgende Python One-liner in de terminal in te voeren, kan dit onmiddellijk worden gecontroleerd.

Natuurlijk kun je ook naar de downloadpagina van het Hugging Face-model gaan en op rechtsboven op bestandinfo klikken. Maar hier leg ik uit hoe je het direct in mijn terminal kunt controleren.

Voorbereiding

De bibliotheek moet eerst geïnstalleerd zijn.

pip install safetensors

Opdracht (terminal invoeren)

Vervang model.safetensors door het werkelijke pad naar het bestand.

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])" # limited to the top 5 entries

Uitleg van de output

Wanneer je dit commando uitvoert, krijg je twee belangrijke informatie:

1. Metadata: Informatie die door de modelmaker is ingebed. Als hier informatie zoals format: gptq of quantization: int4 staat, kun je stellen dat het een gequantiseerd model is, ook al staat het niet in de bestandsnaam. (Let op: als de maker de metadata leeg heeft gelaten, kan None verschijnen.)

2. Keys: Namen van de lagen die het model vormen. Dit helpt je om de structuur van het model te begrijpen.

5. Vergelijkingsoverzicht: .bin (Pickle) vs .safetensors

6. Voorbeeldgebruik (Python)

Een eenvoudig voorbeeld van het opslaan en laden van tensors met de safetensors bibliotheek.

import torch
from safetensors.torch import save_file, load_file

# 1. Tensor maken en opslaan
tensors = {
    "embedding": torch.zeros((1024, 512)),
    "attention": torch.rand((512, 512))
}

# Opslaan van tensor in dictionaryvorm naar bestand
save_file(tensors, "model.safetensors")

# 2. Bestand laden
loaded = load_file("model.safetensors")
print(loaded["embedding"].shape) 
# Output: torch.Size([1024, 512])

7. Conclusie

.safetensors is niet alleen een verandering van bestandsextensie, maar een noodzakelijke evolutie voor de veiligheid en efficiëntie van AI-modellen. Het wordt al de standaard formaat in belangrijke gemeenschappen zoals Hugging Face.

Bij het downloaden of distribueren van modellen wordt aanbevolen om .safetensors te gebruiken in plaats van .bin. Het biedt bescherming tegen beveiligingsbedreigingen en vermindert de laadtijd van modellen aanzienlijk.