Saturno Un nuevo enfoque para entrenar modelos de lenguaje grandes y otras redes neuronales

Saturno Un enfoque novedoso para entrenar modelos de lenguaje y redes neuronales

Nota del editor: Kabir Nagrecha es un ponente en ODSC West 2023 este otoño. ¡Asegúrese de ver su charla, “Democratizando el ajuste fino de modelos grandes de código abierto con optimización conjunta de sistemas”!

La escala del modelo se ha convertido en un aspecto absolutamente esencial de la práctica moderna de aprendizaje profundo. El éxito de los modelos grandes de lenguaje multimillonario (LLM) ha alentado a los profesionales a empujar los límites del tamaño del modelo. Los investigadores y profesionales ahora están construyendo arquitecturas de modelos cada vez más grandes y que consumen mucha memoria para aumentar la precisión y la calidad del modelo.

Pero estas escalas introducen desafíos infraestructurales que no se pueden ignorar. Las demandas de recursos para el entrenamiento de modelos se han disparado. Y si hay algo que hemos aprendido de nuestras colaboraciones en UCSD y con nuestros socios de la industria, es que los trabajos de aprendizaje profundo nunca se ejecutan de forma aislada. Si estás entrenando un modelo, probablemente estés entrenando una docena: la optimización de hiperparámetros, los clústeres multiusuario y la exploración iterativa motivan el entrenamiento de múltiples modelos, aumentando aún más las demandas de cálculo.

En esta publicación de blog, describiremos nuestro trabajo en Saturn, nuestra herramienta de código abierto recientemente lanzada para entrenar múltiples modelos grandes simultáneamente. Saturn se integra perfectamente con herramientas estándar de la industria para el entrenamiento de modelos grandes, como HuggingFace, FSDP de PyTorch, GPipe de Google, FairScale, XFormers y más. Demostraremos cómo Saturn puede ayudar a optimizar la ejecución en este nuevo y crítico entorno, reduciendo costos y aumentando la eficiencia del entrenamiento.

¿Por qué entrenar modelos grandes? ¿Y por qué entrenar muchos de ellos?

“Aumentar el tamaño de nuestros modelos definitivamente ha mejorado nuestra precisión. Se ha convertido en una regla general: más grande es mejor” (Usuario de Saturn en la industria)

Estudios recientes han demostrado que los modelos más grandes tienden a tener un mejor rendimiento en una variedad de benchmarks.

Los modelos de lenguaje son aprendices de pocos disparos (Brown et al., 2020)

Entonces, si estás buscando construir el modelo más preciso para impulsar tu aplicación, probablemente utilizarás una arquitectura de modelo a gran escala, tal vez un modelo LLaMA, tal vez GPT-J, tal vez incluso GPT-2.

Parece claro que hay amplia motivación para entrenar modelos grandes. Pero la idea principal detrás de Saturn es optimizar múltiples modelos grandes al mismo tiempo.

“Probablemente pasamos por cientos… tal vez miles de iteraciones de modelos diferentes antes de lanzar algo a producción. Y este proceso se repite a diario para el reentrenamiento” (Usuario de Saturn sobre sus canalizaciones de producción)

El aprendizaje profundo en la práctica casi siempre consiste en múltiples trabajos de entrenamiento. Los clústeres de la industria reciben trabajos de cientos de usuarios y canalizaciones. Las canalizaciones de producción automatizadas pueden activar docenas de trabajos de entrenamiento con diferentes configuraciones (por ejemplo, tasas de aprendizaje, tamaños de lote, variaciones de la arquitectura del modelo) para encontrar el más preciso para implementar. Los investigadores individuales pueden enviar múltiples trabajos exploratorios para evaluar diferentes enfoques. Al final del día, es posible que un clúster haya recibido miles de trabajos para gestionar.

Hay un problema sin resolver que aborda los desafíos del entrenamiento de modelos grandes y la ejecución de múltiples modelos que debe resolverse. Con Saturn, optimizamos explícitamente esta configuración.

¿Qué es Saturn? ¿Y cómo puedo usarlo?

Cuando se abordan múltiples modelos grandes simultáneamente, hay tres problemas críticos a tener en cuenta.

Primero, paralelización. Los modelos grandes suelen necesitar múltiples GPUs para la distribución de memoria y un mayor rendimiento de entrenamiento. Sin embargo, diseñar el esquema de ejecución paralela es un desafío: ¡hay docenas (¡o incluso cientos!) de técnicas diferentes! FSDP, paralelismo de canalización, paralelismo 3D y más: la mejor opción dependerá de tus GPUs, interconexiones, arquitectura del modelo e incluso hiperparámetros. Es un espacio difícil de navegar, pero elegir el enfoque incorrecto puede tener serias implicaciones en el rendimiento.

Segundo, asignación de recursos. Dado un clúster con 100 GPUs y 30 trabajos presentados, ¿cómo se deben distribuir las GPUs entre los trabajos? Optimizar el rendimiento requiere encontrar automáticamente un plan de distribución de recursos optimizado.

Tercero, programación. ¿Cómo se debe ordenar la ejecución para minimizar los tiempos de ejecución de principio a fin? ¿Debería ejecutarse el Modelo A antes del Modelo B? ¿O el Modelo C antes del Modelo A?

Ten en cuenta que cada uno de estos problemas están conectados. El plan de asignación de recursos que utilices limitará el cronograma e influirá en las selecciones de paralelismo óptimas, y viceversa. Por lo tanto, debemos resolver esto como un problema conjunto.

Y eso es exactamente lo que Saturno hace. Con Saturno, puedes simplemente enviar un lote de trabajos de entrenamiento y observar cómo se genera automáticamente un plan que resuelve todos estos problemas. Saturno realiza un escaneo de perfilado rápido de todos tus trabajos y luego utiliza un solucionador de programación lineal entera mixta para producir un plan optimizado para la ejecución. Es fácil de ampliar y registrar nuevas técnicas de paralelización, por lo que no hay riesgo de quedarse atrás a medida que avanza la investigación. Combinado con algunos mecanismos como la introspección para una programación eficiente, es capaz de reducir drásticamente los tiempos de ejecución de los trabajos de múltiples modelos en lotes. En la práctica, encontramos que puede reducir los tiempos de ejecución y los costos hasta en un factor de 2X en comparación con la práctica actual, al tiempo que reduce la carga de trabajo del desarrollador.

Utilizar Saturno para ejecutar tus trabajos de manera más eficiente es un proceso sencillo. Nuestro repositorio de GitHub contiene un flujo de trabajo de ejemplo, pero aquí hay un breve recorrido.

Primero, envuelve tu función de carga/inicialización del modelo en la estructura “Task” de Saturno.

from saturn import Task
t1 = Task(load_model, load_dataloader, loss_function, hyperparameters)

A continuación, registra las técnicas de paralelización que deseas utilizar con la “Library” de Saturno.

from saturn.library import register
from saturn.core.executors.Technique import BaseTechnique 

class MyExecutor(BaseTechnique):
   def execute(task, gpus, task_id, batch_count):
      # entrenar el modelo
      return
   def search(task, gpus, task_id):
      # optimizar cualquier parámetro interno
      # estos se asocian automáticamente con la tarea enviada
      return parameters

register("my_parallelism", MyExecutor)

Finalmente, envía una lista de tareas para el perfilado y la ejecución.

from saturn.trial_runner.PerformanceEvaluator import search
from saturn import orchestrate

search([t1, t2, t3, t4]) # perfila y evalúa cada tarea
orchestrate([t1, t2, t3, t4]) # orquesta y gestiona la ejecución

Conclusiones y puntos clave

Saturno permite aceleraciones dramáticas de >2X para cargas de trabajo que entrenan múltiples modelos grandes simultáneamente. Es perfectamente compatible con herramientas estándar como HuggingFace. La próxima vez que intentes ajustar finamente un modelo LLaMA o construir tu propio LLM, ¡considera usar Saturno para reducir costos y acelerar los tiempos de ejecución! Si estás interesado en obtener más información, puedes encontrar nuestro artículo sobre Saturno aquí.

Actualmente estamos buscando nuevos colaboradores y usuarios, por lo que si estás interesado en Saturno, ¡considera visitar nuestro repositorio recientemente publicado en GitHub en https://github.com/knagrecha/saturn!

Acerca del autor/ponente ODSC West 2023:

Kabir Nagrecha es candidato a doctorado en UC San Diego, trabajando con los profesores Arun Kumar y Hao Zhang. Es beneficiario de la Beca de Investigación Meta, así como de becas del Instituto de Ciencia de Datos Halicioglu y la Escuela de Ingeniería Jacobs de UCSD.

Kabir es el estudiante de doctorado más joven en la historia de UCSD, habiendo comenzado su doctorado a los 17 años. Previamente ha trabajado con empresas como Apple, Meta y Netflix para construir la infraestructura principal que soporta servicios ampliamente utilizados como Siri y los algoritmos de recomendación de Netflix. 

Linkedin: https://www.linkedin.com/in/kabir-nagrecha-952591152/

Website: https://kabirnagrecha.com/