Los modelos de razonamiento en lenguajes grandes (LLM) están diseñados para resolver problemas complejos dividiéndolos en una serie de pasos más pequeños. Estos potentes modelos son particularmente buenos en tareas desafiantes como la programación avanzada y la planificación de varios pasos.
Pero desarrollar modelos de razonamiento requiere una enorme cantidad de cálculo y energía debido a ineficiencias en el proceso de formación. Mientras algunos de los procesadores de alta potencia trabajan continuamente en consultas complicadas, otros del grupo permanecen inactivos.
Investigadores del MIT y de otros lugares han encontrado una manera de utilizar este tiempo de inactividad computacional para acelerar de manera eficiente el entrenamiento de modelos de razonamiento.
Su nuevo método entrena automáticamente un modelo más pequeño y más rápido para predecir los resultados del LLM de razonamiento más amplio, que el modelo más grande verifica. Esto reduce la cantidad de trabajo que debe realizar el modelo de razonamiento, acelerando el proceso de formación.
La clave de este sistema es su capacidad para entrenar e implementar el modelo más pequeño de forma adaptativa para que se active sólo cuando algunos procesadores estén inactivos. Al aprovechar los recursos computacionales que de otro modo se desperdiciarían, acelera la capacitación sin incurrir en gastos generales adicionales.
Cuando se probó en múltiples LLM de razonamiento, el método duplicó la velocidad de entrenamiento preservando la precisión. Esto podría reducir el coste y aumentar la eficiencia energética del desarrollo de LLM avanzados para aplicaciones como la previsión de tendencias financieras o la detección de riesgos en las redes eléctricas.
“La gente quiere modelos que puedan manejar tareas más complejas. Pero si ese es el objetivo del desarrollo de modelos, entonces debemos priorizar la eficiencia. Encontramos una solución sin pérdidas para este problema y luego desarrollamos un sistema completo que puede proporcionar aceleraciones bastante dramáticas en la práctica”, dice Qinghao Hu, postdoctorado en el MIT y coautor principal de un artículo sobre esta técnica.
En el artículo lo acompañan el coautor principal Shang Yang, un estudiante de posgrado en ingeniería eléctrica e informática (EECS); Junxian Guo, estudiante de posgrado de EECS; el autor principal Song Han, profesor asociado de EECS, miembro del Laboratorio de Investigación Electrónica y científico distinguido de NVIDIA; así como otros en NVIDIA, ETH Zurich, MIT-IBM Watson AI Lab y la Universidad de Massachusetts en Amherst. La investigación se presentará en la Conferencia Internacional ACM sobre Soporte Arquitectónico para Lenguajes de Programación y Sistemas Operativos.
Cuello de botella en el entrenamiento
Los desarrolladores quieren que los LLM de razonamiento identifiquen y corrijan errores en su proceso de pensamiento crítico. Esta característica les permite responder consultas complicadas que atascarían un LLM estándar.
Para enseñarles esta habilidad, los desarrolladores entrenan a los LLM de razonamiento utilizando una técnica llamada aprendizaje por refuerzo (RL). El modelo genera múltiples respuestas potenciales a una consulta, recibe una recompensa para el mejor candidato y se actualiza en función de la mejor respuesta. Estos pasos se repiten miles de veces a medida que el modelo aprende.
Pero los investigadores han descubierto que el proceso de generar múltiples respuestas, llamado implementación, puede consumir hasta el 85% del tiempo de ejecución requerido para el entrenamiento de RL.
“En comparación, actualizar el modelo, que es la verdadera parte de ‘formación’, lleva muy poco tiempo”, afirma Hu.
Este cuello de botella ocurre en los algoritmos RL estándar porque todos los procesadores del grupo de entrenamiento deben finalizar sus respuestas antes de pasar al siguiente paso. Debido a que algunos procesadores pueden estar trabajando en respuestas muy largas, otros que generaron respuestas más cortas están esperando a que terminen.
“Nuestro objetivo era convertir este tiempo de inactividad en aceleración sin desperdiciar costes”, añade Hu.
Intentaron utilizar una técnica existente llamada decodificación especulativa para acelerar las cosas. La decodificación especulativa implica entrenar un modelo más pequeño llamado redactor para adivinar rápidamente los resultados futuros del modelo más grande.
El modelo más grande verifica las suposiciones del escritor y las respuestas que acepta se utilizan para el entrenamiento.
Debido a que el modelo más grande puede verificar todas las suposiciones del escritor a la vez, en lugar de generar cada resultado secuencialmente, acelera el proceso.
Una solución adaptativa
Pero en la decodificación especulativa, el modelo del escritor normalmente se entrena sólo una vez y permanece estático. Esto hace inviable la técnica de aprendizaje por refuerzo, ya que el modelo de razonamiento se actualiza miles de veces durante el entrenamiento.
Un escritor estático rápidamente quedaría obsoleto e inútil después de unos pocos pasos.
Para superar este problema, los investigadores crearon un sistema flexible conocido como “Domar la cola larga” o TLT.
La primera parte de TLT es un entrenador de bocetos adaptativo, que utiliza tiempo libre en procesadores inactivos para entrenar el modelo de boceto en tiempo real, manteniéndolo bien alineado con el modelo de destino sin utilizar recursos computacionales adicionales.
El segundo componente, un motor de implementación adaptativo, gestiona la decodificación especulativa para seleccionar automáticamente la estrategia óptima para cada nuevo lote de entradas. Este mecanismo cambia la configuración de decodificación especulativa según las capacidades de la carga de trabajo de entrenamiento, como la cantidad de entradas procesadas por el modelo preliminar y la cantidad de entradas aceptadas por el modelo objetivo durante la verificación.
Además, los investigadores diseñaron el modelo preliminar para que fuera liviano y pudiera entrenarse rápidamente. TLT reutiliza algunos componentes del proceso de entrenamiento del modelo de razonamiento para entrenar al escritor, lo que genera ganancias adicionales de velocidad.
“Tan pronto como algunos procesadores finalizan sus consultas breves y quedan inactivos, inmediatamente los cambiamos para realizar un entrenamiento preliminar del modelo utilizando los mismos datos que están usando para el proceso de implementación. El motor central es nuestra decodificación especulativa adaptativa; estas ganancias no serían posibles sin él”, dice Hu.
Probaron TLT en varios LLM de razonamiento que fueron capacitados utilizando conjuntos de datos del mundo real. El sistema aceleró el entrenamiento entre un 70 y un 210 por ciento preservando al mismo tiempo la precisión de cada modelo.
Como beneficio adicional, la plantilla de borrador pequeño podría usarse fácilmente para una implementación eficiente como un subproducto gratuito.
En el futuro, los investigadores quieren integrar TLT en más tipos de marcos de entrenamiento e inferencia y encontrar nuevas aplicaciones de aprendizaje por refuerzo que puedan acelerarse utilizando este enfoque.
“A medida que el razonamiento continúa convirtiéndose en la principal carga de trabajo que impulsa la demanda de inferencia, el TLT de Qinghao hace un excelente trabajo al abordar el cuello de botella computacional del entrenamiento de estos modelos de razonamiento. Creo que este método será muy útil en el contexto de la informática de IA eficiente”, afirma Han.
Este trabajo cuenta con el apoyo del MIT-IBM Watson AI Lab, el MIT AI Hardware Program, el MIT Amazon Science Hub, Hyundai Motor Company y la National Science Foundation.
