SafeTensors Reader en Zig 0.15

Objet du projet

Exercice candidature spontanée au sein de ZML

Ce projet implémente un loader SafeTensors en Zig 0.15 basé sur std.Io.Reader, qui expose lui-même un std.Io.Reader capable de déquantifier à la volée les formats FP4 + block vers f32 (ou bf16). Bonus :Implémentation SIMD pour optimiser les performances de déquantification.

Structure du projet

├── src/
│   ├── main.zig                    # Point d'entrée principal
│   ├── safetensor_reader.zig       # Implémentation du loader SafeTensors
│   ├── utils.zig                   # Utilitaires (décodeurs FP4, structures de données)
│   ├── experimental_safetensor_reader.zig  # Versions expérimentales (traces du raisonnement)
│   └── root.zig                    # Module racine
├── build.zig                       # Configuration de build
├── build.zig.zon                   # Dépendances du projet
├── *.safetensors                   # Fichiers de test (modèles quantifiés)
└── dequantized_output_*/           # Dossiers de sortie des déquantifications

Présentation des fichiers

src/main.zig

Point d'entrée du programme qui :

• Ouvre un fichier SafeTensors FP4
• Initialise un reader buffered avec std.Io.Reader
• Lance la déquantification SIMD
• Affiche les résultats et les premiers poids déquantifiés

src/safetensor_reader.zig

Cœur du projet - Implémente SafeTensorReader avec :

Fonctionnalités principales :

• readHeader() : Parse l'en-tête JSON des SafeTensors
• streamDequantizeNVFP4() : Déquantification NVFP4 streaming
• streamDequantizeMXFP4() : Déquantification MXFP4 streaming
• streamSIMDDequantizeMXFP4() : Version SIMD optimisée

Optimisations implémentées :

• Buffering intelligent : Buffer pour réduire les appels système
• SIMD vectoriel : Traitement de 8 valeurs FP4 simultanément
• Streaming : Pas de chargement complet en mémoire

src/utils.zig

Utilitaires et structures de données :

• decodeFp4E2M1() : Décodeur FP4 (exposant 2 bits, mantisse 1 bit)
• DequantizedLayers : Structure pour gérer les couches déquantifiées
• Fonctions de debug

Démarche technique

• Phase 1 : Apprentissage de Zig
• Phase 2 : Compréhension du format Safetensor
• Phase 3 : Lecture séquentielle simple avec Io.Reader
• Phase 4 : Récupération des poids en fp4 (Je me suis appliqué la contrainte d'utiliser seulement std.Io.Reader et pas std.fs.File)
• Phase 5 : Compréhension de la déquantification (NVFP4 puis rapidement MXFP4 par manque de doc sur le premier)
• Phase 6 : Déquantification à la volée (+2h30 pour traiter safetensor de 5GB)
• Phase 7 : Implémentation du SIMD (1h pour traiter safetensor de 5GB)
• Phase 8 : Bufferisation de l'écriture (moins de 10 minutes pour traiter safetensor de 5GB)

Rendu final

Fonctionnalités implémentées

• Loader SafeTensors complet avec std.Io.Reader
• Déquantification MXFP4 streaming (NVFP4 reste expérimentale)
• Optimisations SIMD avec vectorisation
• Sortie au format fp32

Comment exécuter le projet

Prérequis

• Zig 0.15.1
• Fichiers SafeTensors : Placez vos fichiers .safetensors dans le répertoire racine

Compilation et exécution

# Compilation
zig build

# Exécution directe
zig run src/main.zig

Possibles amélioration

• Rendre flexible pour tout type de safetensor MXFP4 (actuellement fonctionne bien pour openai/gpt-oss-20b)
• Multithreading pour traiter simultanèment plusieurs couches