CodeGen2.5: pequeño, pero poderoso

Contribución igual entre Erik Nijkamp y Hiroaki Hayashi.

Paper
Code
Tweet

Abstract

La familia de modelos CodeGen de Salesforce crece con CodeGen2.5 – ¡un modelo pequeño pero poderoso! Aunque ha habido una tendencia reciente de grandes modelos de lenguaje (LLM) de tamaño cada vez mayor, demostramos que un modelo pequeño puede obtener un rendimiento sorprendentemente bueno, cuando se entrena bien.

Las contribuciones clave hacia la producción de estos modelos son:

• Lanzamiento de CodeGen2.5 LLM con el estado del arte en HumanEval para 7B parámetros.
• CodeGen2.5 con 7B está a la par con >15B modelos de generación de código (CodeGen1-16B, CodeGen2-16B, StarCoder-15B), menos de la mitad del tamaño.
• Con robusto muestreo de relleno, es decir, el modelo puede leer texto tanto del tamaño de la mano izquierda como de la derecha de la posición actual.
• Optimizado para muestreo rápido bajo atención Flash para un servicio optimizado y despliegue local en máquinas personales.
• Licencia permisiva en Apache 2.0.

Motivación

En 2022, Salesforce Research lanzó CodeGen [1,2], uno de los primeros LLM para Síntesis de Programas con parámetros 16B. El modelo CodeGen permite a los usuarios «traducir» lenguaje natural, como el inglés, a lenguajes de programación, como Python. Para este tipo de modelos, desde el descubrimiento de las leyes de escalado, leyes de potencia que relacionan el tamaño de un modelo y el conjunto de datos, la tendencia dominante ha sido escalar los LLM a tamaños mayores.

Cuello de botella de los datos. Aunque el tamaño de los conjuntos de datos se ha escalado de acuerdo con estas leyes, este escalado de los datos se ha limitado a la cantidad de datos disponibles. Es decir, un conjunto de datos puede contener como máximo un millón de documentos. Durante el entrenamiento, una iteración sobre este millón de documentos se denomina «una época». Anteriormente, se creía que un modelo sólo debía observar un documento una vez durante el entrenamiento, lo que implicaba que se debía entrenar como máximo durante «una época». Esta limitación agota rápidamente los datos disponibles en el mundo, ya que son finitos. Nosotros cuestionamos esta creencia y entrenamos durante «más de una época» con una receta especializada. Esto puede permitir entrenar un modelo más pequeño con más datos, en lugar de un modelo más grande, que es costoso de servir y mantener en entornos de producción. La afirmación de «ver una observación una sola vez» para un entrenamiento óptimo puede seguir siendo cierta, pero, en nuestro entorno, creamos variantes (o alteraciones) de estas observaciones, de tal manera que el modelo puede ser entrenado para múltiples épocas. Aunque no afirmamos que este aumento de datos sea estrictamente necesario para eliminar las restricciones de datos, los siguientes resultados empíricos indican que, efectivamente, se puede entrenar un modelo pequeño pero muy potente con esta receta.

Muestreo de relleno. La creación de código fuente es un proceso único que implica editar y mejorar continuamente el código a lo largo del tiempo, corrigiendo errores y mejorando su calidad. Esto significa que el modelo de generación de código debe tener en cuenta no sólo el código anterior a la ubicación de edición, sino también el código posterior, permitiendo el relleno dentro del código existente. El relleno se consigue formateando la secuencia de entrada de manera que prediga el contexto que se va a rellenar utilizando los contextos de prefijo y sufijo [2,3,4]. Sin embargo, hemos observado que el enfoque «fill-in-the-middle» [3] a veces puede tener dificultades para determinar cuándo dejar de rellenar. Esta capacidad es crucial, especialmente en casos prácticos. CodeGen2.5 emplea en su lugar la corrupción de span [2,4], que no sufre esta tendencia.

Muestreo rápido. En productos como los asistentes para la finalización de código, los usuarios finales esperan ver las finalizaciones en, digamos, dos segundos. Internamente, las finalizaciones tienen que ser muestreadas a partir del modelo LLM. Este proceso es lento, ya que cada token (o «palabra») tiene que ser muestreado secuencialmente, es decir, generar una compleción de 32 tokens requiere 32 llamadas al modelo, lo que puede incurrir en una latencia significativa en el muestreo. Para hacer frente a este reto, el modelo CodeGen2.5 se optimizó específicamente para un muestreo rápido bajo autoatención con computación Flash [5] y compatibilidad con NVIDIA Triton [6]. El modelo no sólo es pequeño, lo que reduce la latencia de muestreo, sino que también goza de optimización interna para obtener mayores ganancias.

Despliegue local. En el futuro, estos asistentes se ejecutarán en la máquina local, como los MacBook con chips M2. Este despliegue local no sólo permite una alta personalización de los modelos para cada usuario, sino que también garantiza la privacidad de los datos. Con marcos técnicos recientes como llama.cpp, que está optimizado para el despliegue local de LLM pequeños, CodeGen2.5 disfruta de un rendimiento similar al de LLaMA-7B en un MacBook equipado con un chip M1. ¡Esto hace que el despliegue de asistentes basados en CodeGen en máquinas locales sea factible hoy en día.

Tabla 1: HumanEval pass-rates with n=200, Un punto de referencia que mide si los programas generados a partir de un modelo son funcionalmente correctos. Obsérvese que CodeGen2.5, con sólo 7B parámetros, supera a modelos anteriores de más del doble de su tamaño. Los modelos multilingües se entrenan con diversos lenguajes de programación. Los modelos monolingües sólo se ajustan en Python. ¡

Tabla 2: Índices de aprobación de relleno de línea simple de HumanEval con n=40. La prueba comparativa mide la capacidad de un modelo para «rellenar el centro» de un fragmento de código cuyo «centro» se ha enmascarado. ¡Tenga en cuenta, para la productización, CodeGen2.5 introduce un token centinela especializado para el truncamiento y cuenta con un rendimiento de relleno muy alto.

Tabla 3: Tasas de aprobación de HumanEval con n=200, en modelos ajustados a las instrucciones. Nuestros modelos ajustados a las instrucciones se ajustan en conjuntos de datos de instrucciones específicas para mejorar la capacidad de generar código basado en instrucciones en inglés.

Findings

Calidad del modelo. Como se ha señalado anteriormente, los LLM están ávidos de datos. Aunque entrenar modelos durante una época, es decir, sólo una iteración sobre los datos, puede ser eficiente, este lujo es inviable en entornos con datos limitados, como las síntesis de programas en las que sólo existe una cierta cantidad de código en el mundo. En su lugar, nuestra hipótesis es que el entrenamiento con los siguientes ingredientes puede conducir a un modelo competitivo de menor tamaño: (1) múltiples épocas, (2) aumento de datos, (3) altas tasas de aprendizaje. Para las épocas múltiples, durante el entrenamiento recorremos el iterador de datos varias veces en idéntico orden.

Entrenamos el modelo en StarCoderData, un conjunto de datos de lenguaje de programación desarrollado por BigCode [10]. Como muestra la Figura 1, una epoch constituye unos 300B tokens, mientras que el modelo está preentrenado para 1,4T tokens, alcanzando más de 4 epochs. Además, incorporamos nuestro formato de relleno específico [2] en la función objetivo, que puede servir como forma de aumento de datos en el entrenamiento de varias épocas. Es decir, de forma similar a la corrupción de span, algunos tokens de una secuencia observada se mueven a diferentes posiciones, de forma que la misma secuencia no se observa dos veces. Aunque se trata de una hipótesis, esta forma de corrupción o aumento de los datos puede mitigar las limitaciones de los datos. Por último, la función de decaimiento de la tasa de aprendizaje se programa para un presupuesto de tokens significativamente mayor, lo que hace que el modelo se entrene con tasas de aprendizaje más altas durante un periodo más largo. Aunque se puede argumentar que un modelo funciona mejor cuando se entrena hasta la convergencia, observamos un rendimiento muy competitivo o incluso igual al principio del entrenamiento, por lo que un entrenamiento más largo con presupuestos de fichas mucho mayores parece una receta razonable. En concreto, el índice de aprobados de la Figura 1 parece aumentar de forma constante tanto con índices de aprendizaje elevados como con un régimen de entrenamiento multipunto.

El rendimiento final en la prueba de referencia HumanEval, que refleja la capacidad de un modelo para generar un programa funcionalmente correcto, se resume en las Tablas 1. Mientras que CodeGen2.5 sólo tiene un modelo de referencia, CodeGen2.5 no tiene un modelo de referencia. Aunque CodeGen2.5 sólo contiene 7B parámetros de modelo, supera significativamente a CodeGen1.0 y 2.0 con modelos de 16B parámetros, más del doble de su tamaño, y está a la par con el reciente modelo StarCoder 15.5B. CodeGen2.5 es pequeño, ¡pero puede!

Figure 1: HumanEval pass@1 con n=40 en miles de millones de tokens de entrenamiento. Esta prueba mide la capacidad de un modelo para generar programas o fragmentos de código funcionalmente correctos. Una epoch constituye unos 300B tokens, de modo que el modelo se entrenó durante más de 4 epochs. Obsérvese que el rendimiento aumenta de forma constante hasta 1,2T tokens a pesar de entrenar con altas tasas de aprendizaje y múltiples épocas.

Muestreo de relleno. CodeGen2.5 adopta el mismo objetivo de entrenamiento que CodeGen2, lo que permite el muestreo de relleno, es decir, el muestreo de texto teniendo en cuenta el contexto pasado y futuro. También observamos una mejora en las tasas de aprobación de la tarea HumanEval-SingleLineInfill, una tarea de completado de código python de una sola línea. La Tabla 2 resume la comparación con otros modelos con capacidad de relleno.

Ajuste de instrucciones. Además, utilizamos CodeGen2.5-7B-mono y lo ajustamos en conjuntos de datos de instrucciones públicas para mejorar la capacidad de generación de código basado en instrucciones en inglés. Los resultados se resumen en la Tabla 3. Observamos una mejora continua con respecto al modelo monolingüe.

Latencia de muestreo. CodeGen2.5 adopta el forward pass del modelo LLaMA. Esto nos permite disfrutar de las ventajas del servidor de inferencia NVIDIA Triton [6], un rápido marco de inferencia en GPUs NVIDIA. ¡Con FlashAttention activado, conseguimos una inferencia el doble de rápida y un mejor rendimiento en comparación con CodeGen2.0-16B, como se muestra en la Tabla 4.

CodeGen2.0-16B

Tabla 4: Latencia de muestreo en milisegundos con varios marcos de inferencia que admiten Flash attention bajo NVIDIA Triton. La longitud del contexto es de 2.000 tokens y el tamaño del lote está fijado en 2. El número variado de tokens representa una configuración realista para los productos de asistente de código. CodeGen2.5 presenta una latencia más baja, lo que permite mejorar la experiencia del usuario.

Conclusión

La familia de modelos CodeGen da la bienvenida a un nuevo miembro: CodeGen2.5, pequeño pero poderoso. Demostramos que el entrenamiento multipunto puede mitigar las limitaciones de los datos y dar lugar a modelos pequeños pero potentes. Además de su tamaño relativamente pequeño, CodeGen2.5 presenta un robusto muestreo de relleno y un muestreo rápido optimizado para servir con atención Flash, que permiten la productización de estos modelos para asistentes de codificación. En el futuro, ampliaremos los límites de estos modelos.

Explore More

Salesforce AI Research le invita a profundizar en los conceptos tratados en esta entrada de blog (enlaces más abajo). Conéctese con nosotros en las redes sociales y en nuestro sitio web para obtener actualizaciones periódicas sobre éste y otros proyectos de investigación.

Referencias

[1] CodeGen1 (https://arxiv.org/abs/2203.13474)
[2] CodeGen2 (https://arxiv.org/abs/2305.02309)
[3] FIM OpenAI (https://arxiv.org/abs/2207.14255)
[4] InCoder (https://arxiv.org/abs/2204.05999)
[5] Flash Attention (https://arxiv.org/abs/2205.14135)
[6] NVIDIA Triton (https://developer.nvidia.com/triton-inference-server)
[7] abacaj/code-eval (https://github.com/abacaj/code-eval)
[8] StarCoder (https://arxiv.org/abs/2305.06161)
[9] MPT-30B (https://www.mosaicml.com/blog/mpt-30b)
[10] StarCoderData (https://huggingface.co/datasets/bigcode/starcoderdata)