Cómo entrenar modelos de lenguaje grande (LLM) usando Databricks, Hugging Face y MosaicML
introducir
Los grandes modelos de lenguaje, como GPT-4 de OpenAI o PaLM de Google, han conquistado el mundo de la IA. Sin embargo, la mayoría de las empresas actualmente no tienen la capacidad de entrenar estos modelos y dependen completamente de un puñado de grandes empresas tecnológicas como proveedores de tecnología.
En Replit, hemos invertido mucho en la infraestructura necesaria para entrenar nuestros propios modelos de lenguaje a gran escala desde cero. En esta publicación de blog, describimos cómo capacitamos a los LLM, desde los datos sin procesar hasta la implementación en entornos de producción orientados al usuario. Analizaremos los desafíos de ingeniería que enfrentamos en el camino y cómo aprovechamos a los proveedores que creemos que conforman la pila moderna de LLM: Databricks, Hugging Face y MosaicML.
Si bien nuestro modelo está diseñado principalmente para el caso de uso de la generación de código, las técnicas y lecciones discutidas son aplicables a todos los tipos de LLM, incluidos los modelos de lenguaje general. Planeamos sumergirnos en los aspectos prácticos de nuestro proceso en una serie de publicaciones de blog en las próximas semanas y meses.
¿Por qué formarse para un LLM?
Una de las preguntas más comunes que hace el equipo de IA de Replit es "¿Por qué capacitar su propio modelo?" Hay muchas razones por las que una empresa podría decidir capacitar a su propio LL.M., desde la privacidad y seguridad de los datos hasta un mayor control sobre las actualizaciones y mejoras. .
En Replit, nos preocupamos principalmente por la personalización, la reducción de dependencias y la rentabilidad.
- personalizado _ Entrenar un modelo personalizado nos permite adaptarlo a nuestras necesidades y requisitos específicos, incluidas las funciones, la terminología y el contexto específicos de la plataforma, que se encuentran en modelos de propósito general como GPT-4 o incluso modelos específicos de código como Codex Not bien cubierto Por ejemplo, nuestros modelos fueron entrenados para funcionar mejor con lenguajes web específicos populares en Replit, incluidos Javascript React (JSX) y Typescript React (TSX).
- Reducir las dependencias . Si bien siempre usaremos el modelo correcto para la tarea en cuestión, creemos que depender menos de algunos proveedores de IA tiene beneficios. Esto se aplica no solo a Replit, sino también a la comunidad de desarrolladores en general. Es por eso que planeamos abrir el código fuente de algunos de nuestros modelos, lo cual no podríamos hacer sin una forma de entrenarlos.
- rentable _ Aunque los costos seguirán disminuyendo, LLM sigue siendo prohibitivamente caro de usar en la comunidad global de desarrolladores. En Replit, nuestra misión es poner en línea a los próximos mil millones de creadores de software. Creemos que los estudiantes que programan teléfonos móviles en India deberían tener acceso a la misma IA que los desarrolladores profesionales en Silicon Valley. Para que esto sea posible, entrenamos modelos personalizados que son más pequeños, más eficientes y se pueden alojar a un costo drásticamente reducido.
tubería de datos
LLM requiere una gran cantidad de datos para la formación. Capacitarlos requiere construir canales de datos robustos que estén altamente optimizados pero lo suficientemente flexibles para incluir fácilmente nuevas fuentes de datos públicos y propietarios.
la pila
Comenzamos con The Stack disponible en Hugging Face como nuestra principal fuente de datos. Hugging Face es un gran recurso para conjuntos de datos y modelos previamente entrenados. También proporcionan varias herramientas útiles como parte de la biblioteca de Transformers, incluidas herramientas para tokenización, inferencia de modelos y evaluación de código.
Stack es proporcionado por el proyecto BigCode. Los detalles de la construcción del conjunto de datos son proporcionados por Kocetkov et al . (2022) . Después de la deduplicación, la versión 1.2 del conjunto de datos contiene aproximadamente 2,7 TB de código fuente con licencia escrito en más de 350 lenguajes de programación.
La biblioteca de Transformers hace un buen trabajo al abstraer muchos de los desafíos asociados con el entrenamiento de modelos, incluido el manejo de datos a gran escala. Sin embargo, lo encontramos insuficiente para nuestro proceso, ya que necesitábamos un control adicional sobre los datos y la capacidad de procesarlos de forma distribuida.
procesamiento de datos
Cuando se requiere un procesamiento de datos más avanzado, usamos Databricks para construir nuestras canalizaciones. Este enfoque también nos permite traer fácilmente otras fuentes de datos, como Replit o Stack Overflow, a nuestro proceso, lo que planeamos hacer en futuras iteraciones.
El primer paso es descargar los datos sin procesar de Hugging Face. Usamos Apache Spark para paralelizar el proceso de creación de conjuntos de datos en cada lenguaje de programación. Luego volvemos a particionar los datos y los reescribimos en formato parquet utilizando configuraciones optimizadas para el procesamiento posterior.
A continuación, pasamos a limpiar y preprocesar nuestros datos. Por lo general, es importante deduplicar los datos y solucionar varios problemas de codificación, pero The Stack lo ha hecho por nosotros utilizando la técnica de deduplicación aproximada descrita por Kocetkov et al. (2022). Sin embargo, una vez que comencemos a traer datos de Replit a nuestra canalización, tendremos que volver a ejecutar el proceso de deduplicación. Esa es la belleza de tener una herramienta como Databricks donde podemos tratar los datos Stack, Stackoverflow y Replit como tres fuentes en un lago de datos más grande y usarlos en nuestros procesos posteriores según sea necesario.
Otro beneficio de usar Databricks es que podemos ejecutar análisis escalables y manejables en los datos subyacentes. Ejecutamos todo tipo de estadísticas resumidas en nuestras fuentes de datos, examinamos distribuciones de cola larga y diagnosticamos cualquier problema o inconsistencia en el proceso. Todo esto se hace en cuadernos Databricks, que también se pueden integrar con MLFlow para rastrear y reproducir todo nuestro análisis durante todo el proceso. Este paso equivale a una radiografía regular de nuestros datos y también ayuda a informar los diversos pasos que tomamos para el preprocesamiento.
Para el preprocesamiento, seguimos los siguientes pasos:
- Anonimizamos los datos eliminando cualquier información de identificación personal (PII), incluido el correo electrónico, las direcciones IP y las claves de cifrado.
- Usamos varias heurísticas para detectar y eliminar el código generado automáticamente.
- Para algunos idiomas, eliminamos el código que no se compilaba ni analizaba con analizadores estándar.
- Filtramos los archivos según la longitud de línea promedio, la longitud de línea máxima y el porcentaje de caracteres alfanuméricos.
Tokenización y Entrenamiento de Vocabulario
Antes de la tokenización, entrenamos nuestro propio vocabulario personalizado usando una submuestra aleatoria de los mismos datos que usamos para el entrenamiento del modelo. Los vocabularios personalizados permiten que nuestro modelo comprenda mejor y genere contenido de código. Esto mejora el rendimiento del modelo y acelera el entrenamiento y la inferencia del modelo.
Este paso es uno de los más importantes del proceso, ya que se usa en las tres etapas de nuestro proceso (canalización de datos, entrenamiento del modelo, inferencia). Destaca la importancia de proporcionar una infraestructura robusta y totalmente integrada para el proceso de formación del modelo.
Planeamos explorar la tokenización con más profundidad en una futura publicación de blog. En un nivel alto, algunas de las cosas importantes que debemos considerar son el tamaño del vocabulario, los tokens especiales y el espacio reservado para los tokens de marcado.
Una vez que hayamos entrenado nuestro vocabulario personalizado, etiquetaremos nuestros datos. Finalmente, construimos nuestro conjunto de datos de entrenamiento y lo escribimos en un formato fragmentado optimizado para el proceso de entrenamiento del modelo.
entrenamiento modelo
Entrenamos nuestro modelo usando MosaicML . Habiendo implementado previamente nuestro propio grupo de capacitación, descubrimos que la plataforma MosaicML nos brindaba algunas ventajas clave.
- Múltiples proveedores de nube . Mosaic nos permite utilizar GPU de diferentes proveedores de la nube sin los gastos generales de configurar una cuenta y todas las integraciones necesarias.
- Configuración de la formación LL.M. La biblioteca Composer tiene muchas configuraciones bien ajustadas para entrenar varios modelos y diferentes tipos de objetivos de entrenamiento.
- Infraestructura de alojamiento . Su infraestructura de alojamiento nos brinda orquestación, optimización de la eficiencia y tolerancia a fallas (es decir, recuperación de fallas de nodos).
Al determinar los parámetros de nuestro modelo, consideramos varias compensaciones entre el tamaño del modelo, la ventana de contexto, el tiempo de inferencia, la huella de memoria, etc. Los modelos más grandes generalmente brindan un mejor rendimiento y son más capaces de transferir el aprendizaje. Sin embargo, estos modelos tienen requisitos computacionales más altos tanto para el entrenamiento como para la inferencia. Esto último es especialmente importante para nosotros. Replit es un IDE nativo de la nube que funciona como una aplicación nativa de escritorio, por lo que nuestro modelo de finalización de código debía ser ultrarrápido. Por esta razón, normalmente elegimos modelos más pequeños con huellas de memoria más pequeñas e inferencia de baja latencia.
Además de los parámetros del modelo, elegimos entre una variedad de objetivos de entrenamiento, cada uno con sus propias fortalezas y debilidades únicas. El objetivo de entrenamiento más común es la predicción del siguiente marcador. Esto suele funcionar bien para completar el código, pero no tiene en cuenta el contexto posterior de la documentación. Esto se puede mitigar mediante el uso de un objetivo de "relleno en el medio", en el que se oculta un rango de tokens en el documento y el modelo debe usar el contexto circundante para predecirlos. Otro enfoque es UL2 (Aprendizaje de lenguaje latente no supervisado), que enmarca una función objetivo diferente para entrenar un modelo de lenguaje como una tarea de eliminación de ruido, donde el modelo tiene que recuperar las subsecuencias que faltan dada una entrada.
Una vez que decidimos la configuración de nuestro modelo y los objetivos de capacitación, lanzamos nuestra ejecución de capacitación en un clúster de GPU de varios nodos. Podemos ajustar la cantidad de nodos asignados para cada ejecución según el tamaño del modelo que estamos entrenando y la rapidez con la que queremos que se complete el proceso de entrenamiento. La ejecución de grandes clústeres de GPU es costosa, por lo que es importante utilizarlos de la manera más eficiente posible. Supervisamos de cerca la utilización de la GPU y la memoria para asegurarnos de obtener el máximo uso posible de nuestros recursos informáticos.
Usamos ponderaciones y sesgos para monitorear el proceso de capacitación, incluida la utilización de recursos y el progreso de la capacitación. Monitoreamos nuestras curvas de pérdida para garantizar que el modelo aprenda de manera eficiente en cada paso del proceso de capacitación. También nos centramos en los picos de pérdidas. Estos son aumentos repentinos en los valores de pérdida que generalmente indican un problema con los datos de entrenamiento subyacentes o la arquitectura del modelo. Debido a que estos eventos a menudo requieren una mayor investigación y posibles ajustes, aplicamos la certeza de los datos en nuestros procesos para que podamos reproducir, diagnosticar y abordar más fácilmente la fuente potencial de tales picos de pérdida.
Evaluar
Para probar nuestro modelo, usamos una variante del marco HumanEval descrito en Chen et al . (2021) . Dada una firma de función y una cadena de documentos, usamos el modelo para generar una pieza de código de Python. Luego ejecutamos casos de prueba contra las funciones generadas para determinar si los bloques de código generados funcionan como se esperaba. Ejecutamos varias muestras y analizamos los números Pass@K correspondientes.
Este enfoque funciona mejor con Python, donde se pueden usar evaluadores y casos de prueba. Pero dado que Replit admite múltiples lenguajes de programación, debemos evaluar el rendimiento del modelo en varios otros lenguajes. Hemos encontrado esto difícil de hacer, y no existen herramientas o marcos ampliamente adoptados que brinden una solución integral. Dos desafíos específicos incluyen la creación de entornos de tiempo de ejecución reproducibles en cualquier lenguaje de programación y la ambigüedad de los lenguajes de programación en ausencia de estándares ampliamente utilizados para casos de prueba (por ejemplo, HTML, CSS, etc.). Afortunadamente, "Entornos de tiempo de ejecución reproducibles para cualquier lenguaje de programación" es la especialidad de Replit. Actualmente estamos construyendo un marco de evaluación que permite a cualquier investigador conectarse y probar sus puntos de referencia multilingües. a nosotros'
implementar en producción
Una vez que hemos entrenado y evaluado nuestro modelo, es hora de implementarlo en producción. Como mencionamos antes, nuestro modelo de finalización de código debería sentirse rápido, con una latencia muy baja entre solicitudes. Usamos FasterTransformer y Triton Server de NVIDIA para acelerar nuestro proceso de inferencia. FasterTransformer es una biblioteca que implementa un motor de aceleración para la inferencia de redes neuronales basadas en transformadores, y Triton es un servidor de inferencia estable y rápido que es fácil de configurar. Esta combinación nos brinda una capa altamente optimizada entre el modelo del convertidor y el hardware GPU subyacente, y permite una inferencia distribuida ultrarrápida en modelos grandes.
Después de implementar nuestro modelo en producción, pudimos usar nuestra infraestructura de Kubernetes para escalarlo automáticamente para satisfacer la demanda. Aunque hemos discutido el ajuste de escala automático en publicaciones de blog anteriores, vale la pena mencionar que alojar un servidor de inferencia presenta un conjunto único de desafíos. Estos incluyen artefactos grandes (es decir, pesos de modelos) y requisitos especiales de hardware (es decir, diferentes tamaños/números de GPU). Diseñamos nuestras implementaciones y configuraciones de clúster para que podamos entregar de manera rápida y confiable. Por ejemplo, nuestros clústeres están diseñados para abordar la escasez de GPU en regiones individuales y buscar los nodos disponibles más baratos.
Nos gusta probarlo nosotros mismos y tener una idea de la "vibración" del modelo antes de ponerlo frente a los usuarios reales. Los resultados de la prueba HumanEval que calculamos anteriormente son útiles, pero no hay nada como usar el modelo para tener una idea, incluida su latencia, la consistencia de las recomendaciones y la ayuda general. Poner un modelo frente a un equipo Replit es tan fácil como encender un interruptor. Una vez que estemos satisfechos con él, activaremos otro interruptor y lo implementaremos para nuestros otros usuarios.
Continuaremos monitoreando el rendimiento del modelo y las métricas de uso. Para el rendimiento del modelo, monitoreamos métricas como la latencia de solicitud y la utilización de GPU. Para el uso, hacemos un seguimiento de las tasas de aceptación de las sugerencias de código y las desglosamos en múltiples dimensiones, incluido el lenguaje de programación. Esto también nos permite realizar pruebas A/B de diferentes modelos y obtener una medida cuantitativa de cómo se compara un modelo con otro.
Retroalimentación e iteración
Nuestra plataforma de entrenamiento de modelos nos permite convertir datos sin procesar en modelos implementados en producción en menos de un día. Pero lo que es más importante, nos permite entrenar e implementar modelos, recopilar comentarios y luego iterar rápidamente en función de esos comentarios.
También es importante que nuestro proceso se mantenga sólido ante cualquier cambio en la fuente de datos subyacente, el objetivo de entrenamiento del modelo o la arquitectura del servidor. Esto nos permite aprovechar los nuevos avances y características en un campo en rápida evolución donde parece haber anuncios nuevos y emocionantes todos los días.
A continuación, ampliaremos nuestra plataforma para permitirnos usar Replit para mejorar nuestros modelos. Esto incluye técnicas como el aprendizaje de refuerzo basado en comentarios humanos (RLHF) y el ajuste de instrucciones utilizando datos recopilados de Replit Bounties.
Próximo paso
Si bien hemos recorrido un largo camino, todavía estamos en las primeras etapas de capacitación de un LL.M. Tenemos muchas mejoras que hacer y muchos acertijos que resolver. Esta tendencia solo se acelerará a medida que los modelos lingüísticos continúen mejorando. Habrá un nuevo conjunto de desafíos relacionados con la evaluación de datos, algoritmos y modelos.
Si está entusiasmado con los muchos desafíos de ingeniería de capacitar a un LLM, nos encantaría hablar con usted. Nos encantan los comentarios y nos encantaría saber de usted sobre lo que nos falta y lo que podría hacer de manera diferente.