Una implementación de 190 líneas en JAX puro reproduce toda la arquitectura de DeepSeek V3: atención latente multi-cabeza, mezcla de expertos sin pérdida auxiliar y predicción multi-token incluidos. El desarrollador Vaibhaav Tiwari la entrenó de punta a punta en una MacBook Air M5, usando como único corpus novelas policiales de dominio público, y publicó el código completo en GitHub con el último push registrado el 14 de noviembre de 2026.

Qué hay dentro de las 190 líneas de DeepSeek V3 en JAX

El archivo no usa Flax, ni Equinox, ni PyTorch: los parámetros son diccionarios planos de arrays, y cada pieza del modelo es una función que se lee de arriba a abajo en deepseek_v3.py. Eso lo distingue de la mayoría de reimplementaciones educativas, que suelen apoyarse en un framework de alto nivel para esconder la mecánica interna. Tiwari también publicó un escrito técnico con diagramas y la explicación función por función en su sitio personal.

La tabla que publica el propio repositorio mapea cada idea del informe técnico de DeepSeek V3 a una función concreta:

  • Atención latente multi-cabeza (MLA): en vez de guardar las llaves y valores completos de cada token, cachea un único vector latente más chico (480 a 144 números por token y por capa en esta implementación), en las funciones mla e init_cache. Es como guardar un resumen comprimido de la conversación en lugar del historial completo, y descomprimirlo recién cuando el modelo necesita atender a ese token.- RoPE desacoplado: una sola llave rotatoria compartida entre todas las cabezas, mantenida fuera del latente comprimido, en rope y apply_rope.- DeepSeekMoE: un experto compartido que siempre corre, más 16 expertos enrutados de los que se activan 4 por token, en route y moe. Funciona parecido a un médico de guardia que atiende siempre, más hasta cuatro especialistas que se llaman según el caso.- Balanceo sin pérdida auxiliar: un sesgo por experto que empuja el enrutamiento hacia los expertos menos usados, sin tocar sus pesos, en update_bias y route.- Predicción multi-token (MTP): un módulo extra que además predice el token siguiente al próximo, con un peso de 0,3 en la función de pérdida, en forward y loss_fn. Diagrama simplificado de la mezcla de expertos (DeepSeekMoE). ## A quién le sirve leer este código

El público natural no es alguien que busca entrenar un modelo de producción, sino un desarrollador o estudiante que ya leyó sobre DeepSeek V3 y quiere ver, línea por línea, cómo se traduce cada idea (MLA, el enrutamiento de expertos, el balanceo) a código que corre de verdad. Sirve, por ejemplo, para un ingeniero de backend que quiere entender qué hace realmente un router de MoE antes de decidir si justifica la complejidad en un proyecto propio, o para alguien que prepara una charla interna y necesita un ejemplo que se pueda proyectar completo en una pantalla.

Para un ingeniero de ML en LATAM sin acceso a un clúster de GPU, es también la forma más barata de tocar estos conceptos con las manos: todo el entrenamiento del ejemplo corrió en una laptop de consumo, no en infraestructura de nube. No sirve, en cambio, para quien busca una librería lista para producción: el propio README aclara que se trata de una demostración de arquitectura. El autor lo dice así tras mostrar una muestra de texto generado: no apunta a competir como modelo de lenguaje.

Lo que probamos: la instalación falla fuera de macOS

Clonamos el repositorio en un sandbox aislado con Linux y seguimos los pasos que indica el README. El clon tardó 1,208 s y el repositorio pesó 16,4 MB; crear el entorno virtual tomó 1,812 s. La instalación de dependencias falló a los 0,716 s:

$ .venv/bin/python -m pip install --no-input -r requirements.txt Collecting jax==0.4.34 (from -r requirements.txt (line 1)) Downloading jax-0.4.34-py3-none-any.whl.metadata (22 kB) ERROR: Could not find a version that satisfies the requirement jax-metal==0.1.1 (from versions: none) ERROR: No matching distribution found for jax-metal==0.1.1

El motivo no es un bug: jax-metal solo se publica para macOS con Apple Silicon, y el requirements.txt fija las versiones exactas jax==0.4.34 y jax-metal==0.1.1, sin alternativa. El README documenta la instalación así:

git clone https://github.com/Vaibhaav-Tiwari/deepseek-v3-jax.git && cd deepseek-v3-jax uv venv -p 3.12 .venv uv pip install --python .venv/bin/python -r requirements.txt

El pico de memoria de toda la prueba (clon, entorno e instalación fallida) fue de 86,1 MB. El README menciona una ruta alternativa, "No Mac? See Running on CPU", pero seguir los pasos por defecto rompe antes de llegar a esa opción: hay que editar el requirements.txt a mano para sacar la dependencia de Metal, y el extracto que tuvimos a mano no detalla cuánto tarda esa vía en CPU.

Los números que reporta el propio entrenamiento

Esto no lo probamos nosotros: son las cifras que publica el autor sobre su propia corrida, en una MacBook Air M5 de 24 GB. El modelo tiene 36,2 millones de parámetros totales, de los cuales solo 13,6 millones se activan por token gracias al enrutamiento de expertos. Entrenó con 207 libros (21,7 millones de tokens de entrenamiento, con las novelas de Dashiell Hammett repetidas 3 veces) durante 12.000 pasos, unos 98 millones de tokens vistos en total (~4,5 épocas), a 10.700 tokens por segundo de mediana, en 2,6 horas.

La pérdida de validación bajó de 5,02 en el paso 250 a 3,51 al final del entrenamiento. El balanceo sin pérdida auxiliar deja el dato más concreto en el paso 2.000. Con el sesgo activado, la carga por experto quedó entre 5,7% y 6,9% (lo parejo sería 6,25%); sin el sesgo, osciló entre 0,09% y 18,6%. Pese a esa diferencia enorme en el reparto, la pérdida de validación en ese mismo paso fue prácticamente igual con balanceo y sin él (3,921 frente a 3,923), lo que confirma el argumento del paper original: el sesgo reparte trabajo entre expertos sin costarle nada al modelo (fuente).

El propio modelo, generado a partir del prompt "The fog came in off the": "The fog came in off the bay, and the great trees were spreading and the stars were shining" (fuente).

flowchart TD A["Token de entrada"] --> B["Experto compartido (siempre activo)"] A --> C["Router: top-4 de 16 expertos"] C --> D["Experto 1"] C --> E["Experto 2"] C --> F["Experto N"] B --> G["Suma ponderada"] D --> G E --> G F --> G G --> H["Salida de la capa MoE"]

DeepSeek V3 en JAX frente a otras reimplementaciones

No es la única forma de leer una arquitectura de LLM moderna sin depender de un framework pesado. Vale la pena ubicarla frente a dos referencias conocidas del mismo espacio:
ProyectoFramework y escalaPara qué sirvedeepseek-v3-jaxJAX puro, sin Flax ni Equinox; pensado para una sola laptopLeer de punta a punta cómo se implementan juntos MLA, MoE y MTPDeepSeek-V3 oficial (deepseek-ai)PyTorch, pensado para inferencia distribuida en producciónCorrer el modelo real, no estudiar la arquitectura paso a pasoMaxText (Google)JAX, orientado a entrenar a escala en TPU y GPUPunto de partida para entrenar modelos grandes, no para leerlo como tutorial
La diferencia de fondo es de propósito, no de calidad: el repositorio oficial y MaxText están pensados para escalar, este para entenderse. Ninguno de los tres reemplaza a los otros dos.

Limitaciones: qué no hay que esperar de este código

  • Es una demostración de arquitectura. El propio autor aclara en el README que no apunta a ser un modelo de lenguaje competitivo.- Solo corre de forma directa en macOS con Apple Silicon, por la dependencia fijada a jax-metal==0.1.1; en Linux la instalación por defecto falla, como confirmamos en nuestra prueba.- El checkpoint publicado es un pickle de Python, un formato que puede ejecutar código arbitrario al cargarse; el propio repositorio recomienda cargar solo checkpoints de fuentes de confianza.- Los hashes de los datos de entrenamiento pueden no coincidir entre corridas distintas, porque algunos libros vienen de Project Gutenberg y ese archivo cambia con el tiempo (los de Hammett, en cambio, están fijados a un commit exacto de Standard Ebooks y sí son reproducibles).

Qué significa para vos si programás con IA en LATAM

Si trabajás con modelos de lenguaje y nunca leíste código fuente de una arquitectura MoE moderna, este repositorio es más rápido de recorrer que un informe técnico y más transparente que una librería que te esconde los detalles. Para correrlo de verdad necesitás una Mac con Apple Silicon: nuestra prueba confirma que en un entorno Linux estándar, el típico de un pipeline de CI o una instancia cloud, la instalación por defecto se rompe por jax-metal, y el README no deja un atajo documentado para saltarse esa dependencia sin editar requirements.txt a mano.

Si tu objetivo es otro (correr un modelo DeepSeek real o entrenar algo a escala), este código no es el punto de partida; para eso están el repositorio oficial de DeepSeek-V3 o un framework como MaxText.

Referencias

Seguimos el tema en @programacion.