Presentamos Phisch

Hace unas semanas lanzamos Phisch, nuestra primera arquitectura de modelo de lenguaje, desarrollada íntegramente en Rust utilizando el framework de aprendizaje automático Burn. En aquel momento no había forma de verificar su funcionamiento, ya que no existían modelos preentrenados disponibles para cargar y probar; por ello, trabajamos discretamente hasta tener algo listo.

Hoy nos complace explicar qué es Phisch y, al mismo tiempo, lanzar el primer modelo Phisch totalmente entrenado: Poetry Phisch.

¿Por qué crear Phisch?

Crab Lab surgió a raíz de dos problemas coincidentes en los modelos de IA actuales: son demasiado grandes para el hardware de la mayoría de los usuarios y son de código cerrado, lo que significa que nadie ajeno a las empresas que los desarrollan puede entrenar uno de estos modelos de forma independiente. Incluso los modelos de pesos abiertos han empezado a aumentar de tamaño para competir con las alternativas de código cerrado más avanzadas. Esto fue lo que nos inspiró a crear y lanzar nuestro propio modelo.

Al crear Phisch, nos centramos en dos aspectos: garantizar que todas las herramientas desarrolladas en torno al modelo estuvieran bien documentadas y fueran fáciles de seguir para cualquier entusiasta del aprendizaje automático, y facilitar el ajuste de los hiperparámetros para adaptarlos al hardware de cada usuario.

¿Qué es Phisch?

Phisch es un modelo de lenguaje para la generación de texto. Se basa en una arquitectura Transformer de tipo decoder-only (solo decodificador), modernizada con mejoras introducidas por Llama (Touvron et al., 2023): RMSNorm, RoPE y SwiGLU.

Elegimos Rust debido a las herramientas y capacidades que ofrece este lenguaje. Gracias a estas características, compilar Phisch en diversos entornos y sistemas resulta sencillo; además, un único binario se encarga del entrenamiento, la generación de texto, el despliegue del modelo y la exportación de modelos en formato safetensors.

Desarrollado sobre Burn, Phisch utiliza su backend wgpu, lo que permite su ejecución en cualquier GPU compatible con Vulkan, Metal o DirectX 12, sin depender de CUDA ni de ningún fabricante de hardware específico. Esto, sumado a hiperparámetros configurables —como el número de capas, las dimensiones de los embeddings, las cabezas de atención y la longitud del contexto—, permite adaptar el modelo a los recursos disponibles, ya sea una GPU dedicada o integrada.

Conoce a Poetry Phisch

Para demostrar las capacidades de Phisch, entrenamos Poetry Phisch, un modelo de 3 millones de parámetros, utilizando un conjunto de datos personalizado compuesto por 2.230 libros de poesía clásica en inglés provenientes de Project Gutenberg. Solo incluimos libros que se encuentran en el dominio público en los Estados Unidos y cuyos colaboradores acreditados fallecieron hace al menos 70 años. En total, recopilamos aproximadamente 87 millones de tokens: cerca de 83 millones para el entrenamiento y 4,4 millones para la validación.

El entrenamiento se llevó a cabo en una GPU de consumo (una NVIDIA GTX 1660 con 6 GB de VRAM) y duró unas cuatro horas y media, divididas en dos ciclos de tres épocas cada uno. En el primer ciclo, la pérdida de entrenamiento se estabilizó en torno a 3,49 al finalizar la tercera época. El segundo ciclo se reanudó a partir del mejor punto de control (checkpoint) del primero, utilizando una tasa de aprendizaje más baja para confirmar la convergencia; la pérdida de validación mejoró de forma insignificante, situándose finalmente en 3,69.

El resultado es un modelo que imita la estructura y el estilo de los textos utilizados para su entrenamiento, aunque todavía no logra generar oraciones con significado ni mantener una idea coherente a lo largo de varias líneas. En nuestras pruebas, tampoco encontramos casos en los que el modelo reprodujera fragmentos literales del conjunto de datos. Algunos ejemplos:

Consigna: The bird

The bird is singing, and the stars are white.
And I am sick of all things long ago--
I cannot tell you why I love you.

Consigna: Love

Love is a bird of the moon.
And I have heard my voice, and can say:
'The world is in the night,' we said.
He looked at me from my eyes.

Primeros pasos

Para probar Poetry Phisch, descárgalo desde Hugging Face y sigue las instrucciones de la ficha del modelo (model card).

Para instalar Phisch y entrenar tu propio modelo, consulta el archivo README, donde se explica paso a paso un flujo de trabajo de entrenamiento simplificado. Para obtener una explicación detallada de cada comando y parámetro, revisa USAGE.md.

¿Qué sigue?

Puedes probar Poetry Phisch ahora mismo en nuestro entorno de pruebas, donde también estarán disponibles futuros modelos de Phisch. Y si entrenas tu propio modelo, ¡nos encantaría saberlo! No dudes en abrir una issue en GitHub para compartir comentarios, preguntas o resultados.