Vaibhaav Tiwari condense DeepSeek-V3 en 190 lignes de JAX pur
TL;DR
- Vaibhaav Tiwari a condensé toute l'architecture DeepSeek-V3 (MLA, DeepSeekMoE, multi-token prediction) en 190 lignes de JAX pur, publiées sous licence MIT.
- Le modèle jouet compte 36,2 M de paramètres dont 13,6 M actifs par token, pour une validation loss de 3,51 après 12 000 pas d'entraînement.
- L'entraînement a tourné environ 2,6 h sur un MacBook Air M5 24 Go, à ~10,7k tokens/s, sur 21,7 M de tokens issus de 207 romans policiers.
Vaibhaav Tiwari a condensé l'architecture DeepSeek-V3 en « 190 lines of pure JAX », sans Flax ni PyTorch, et publié le dépôt sous licence MIT.
Le fichier couvre la Multi-head Latent Attention (compression de 480 à 144 nombres par token et par couche), une DeepSeekMoE avec « 1 shared expert that always runs + 16 routed experts, top-4 per token », un équilibrage sans loss auxiliaire via biais par expert, et la multi-token prediction pondérée à λ = 0,3. Les dépendances sont épinglées à jax 0.4.34 et jax-metal 0.1.1.
La version jouet compte « 36.2M parameters, 13.6M active per token ». Elle a été entraînée sur un MacBook Air M5 doté de 24 Go de mémoire, à environ 10,7k tokens/s, sur 21,7 M de tokens tirés de 207 romans policiers du domaine public. En « 12,000 steps, ~98M tokens (~4.5 epochs), ~2.6 h », la validation loss passe de 5,02 au pas 250 à 3,51.
C'est la troisième lecture légère du code DeepSeek qu'on voit passer cette semaine sur notre radar open source.
Article original publié par github.com
Lire l'article original →Titre original : Un développeur réécrit DeepSeek-V3 en 190 lignes de JAX pur, entraîné sur un MacBook Air M5