Présentation

Le dépôt deepseek-v3-jax reproduit l’architecture décrite dans le rapport technique DeepSeek‑V3 en seulement 190 lignes de code JAX pur. Aucun framework supplémentaire (Flax, Equinox, PyTorch) n’est utilisé ; les paramètres sont stockés dans des dictionnaires d’arrays et chaque composant du modèle est une fonction lisible de haut en bas. La version 1.0 fournit le checkpoint (138 Mo) et le tokenizer, tous deux sous licence MIT.

Architecture et implémentation

Le modèle repose sur un MoE (Mixture‑of‑Experts) dense : chaque expert reçoit chaque token, les sorties non sélectionnées sont simplement mises à zéro. Le routage est calculé exactement, sans dispatch sparses ni parallélisme d’experts. Cette approche simplifie le code mais augmente le coût de calcul, car aucune spécialisation matérielle n’est exploitée. Le modèle ne réalise pas d’absorption de poids dans le module MLA ; le latent est décompressé explicitement, ce qui facilite la lecture du code au prix d’une latence d’inférence supérieure.

tail -n +5 deepseek_v3.py | grep -cvE '^\s*(#|$)'    # 190

Le script deepseek_v3.py contient l’ensemble de l’architecture ; la fonction de génération lit le checkpoint pickle, ce qui impose une vigilance quant à la provenance du fichier, car le format peut exécuter du code Python au chargement.

Entraînement et performances

Le modèle a été entraîné de zéro sur un MacBook Air M5 disposant de 24 Go de RAM, uniquement avec du texte du domaine public (207 romans policiers, soit environ 120 Mo). Le jeu d’entraînement compte 211 documents, 83,4 M de caractères et 21,69 M de tokens, tandis que le jeu de validation comprend 4 documents, 1,1 M de caractères et 0,30 M de tokens. Les hachages SHA‑256 fournis permettent de vérifier l’intégrité des fichiers.

Un entraînement complet de 12 000 étapes dure environ 2,6 h. Une ablation du biais de routage a été réalisée : à l’étape 2000, la perte de validation passe de 3.921 (biais activé) à 3.923 (biais désactivé), confirmant l’affirmation du papier selon laquelle le biais équilibre les experts sans coût supplémentaire. Les métriques sont enregistrées dans metrics.jsonl et les checkpoints sont sauvegardés à chaque évaluation.

La génération de texte utilise une température de 0.8, un top‑k de 50 et produit 300 nouveaux tokens par échantillon. Exemple fourni : « The fog came in off the bay, and the great trees were spreading and the stars were shining. »

Limitations et perspectives

Le modèle reste une démonstration : avec 36,2 M de paramètres, il est loin des modèles de plusieurs centaines de milliards de paramètres (671 B) décrits dans le rapport original. Le contexte maximal est de 256 tokens contre 128 k dans la version complète, et plusieurs optimisations avancées (FP8, DualPipe, YaRN, post‑training) sont absentes. L’entraînement sur CPU est possible mais nettement plus lent, et la version « small » (4 couches, 8 experts, top‑2, contexte 128 tokens) cible les machines disposant de 8‑16 GB de RAM, bien que seules quelques étapes aient été validées.

Le code repose sur jax==0.4.34 et jax-metal==0.1.1, qui ne supportent pas les gradients à travers une contraction à deux axes ; le modèle évite cette opération. Les tests unitaires (39) passent sous pytest, garantissant la reproductibilité du pipeline sur les configurations supportées.