Koda, un LLM entraîné de zéro

KodaLite-1.3B est un modèle de langage que j'ai entraîné de zéro, non pas pour la taille, mais pour comprendre les internals : architecture decoder-only façon LLaMA (24 couches, hidden 2048, GQA 32/8, SwiGLU, RMSNorm pre-norm, RoPE), tokenizer GPT-2 BPE. Le pipeline complet tourne en JAX + Flax NNX sur 2 GPU NVIDIA L40S en bf16 : pré-entraînement sur SlimPajama (~1,6 milliard de tokens, ~25 heures) avec un orchestrateur à reprise sur crash, SFT LoRA sur Dolly et OASST, extension de contexte NTK-aware de 1024 à 2048 tokens. Les checkpoints sont publiés sur Hugging Face avec exports vers Transformers, GGUF (llama.cpp, Ollama, LM Studio) et MLX, plus un benchmark maison de 8 tâches zero-shot.

Défis

• Entraîner un modèle de 1,27 Md de paramètres sur un budget GPU limité (2x L40S, 96 GB VRAM)
• Tenir un run de pré-entraînement d'environ 25 heures sans perdre de progression
• Étendre le contexte de 1024 à 2048 tokens après le pré-entraînement
• Rendre le modèle utilisable en dehors de JAX (Transformers, GGUF, MLX)

Solutions

• Implémentation JAX + Flax NNX en bf16 avec orchestrateur de reprise sur crash
• Pré-entraînement SlimPajama puis SFT LoRA (Dolly, OASST) et fix du token EOS
• Extension de contexte NTK-aware sans ré-entraînement complet
• Pipeline d'export vers Hugging Face Transformers, GGUF et MLX (fp16 et 8 bits)

Résultats

• Modèle KodaLite-1.3B publié sur Hugging Face (YoAbriel/KodaLite-1.3B, variantes GGUF et MLX)
• Pré-entraînement complet : ~1,6 Md de tokens SlimPajama en ~25 h sur 2x L40S
• Benchmark maison de 8 tâches zero-shot pour mesurer ce que le modèle sait vraiment faire
• Code public sur GitHub (Koda-v0.1)

Technologies

JAX · Flax NNX · Python · LoRA · SlimPajama · Hugging Face · GGUF · MLX