El cargador de datos de investigación de IBM mejora el entrenamiento de modelos de IA para la comunidad de código abierto.

¿Cómo se superan los cuellos de botella cuando se entrenan modelos de IA en cantidades masivas de datos? En la conferencia PyTorch de este año, IBM Research presentó un innovador cargador de datos para el entrenamiento de LLM a gran escala. La herramienta, ahora disponible para los usuarios de PyTorch, tiene como objetivo simplificar el entrenamiento a gran escala para una audiencia lo más amplia posible.

La idea del cargador de datos de alto rendimiento surgió de los problemas prácticos que los científicos de investigación observaron durante el entrenamiento del modelo, ya que su trabajo requería una herramienta que pudiera procesar grandes cantidades de datos en múltiples dispositivos, todo mientras se mantenía al día con GPU progresivamente eficientes. Como IBM Research señala en su blog sobre el lanzamiento, "todo gracias a un equipo de investigadores que simplemente estaban construyendo las herramientas que necesitaban para hacer un trabajo".

Davis Wertheimer de IBM Investigación explica algunos de los desafíos que pueden surgir durante el entrenamiento a gran escala: “Hay una especie de regla 80/20 cuando se trata de entrenamiento a gran escala. El ochenta por ciento de toda la literatura publicada analiza las compensaciones algorítmicas entre la memoria de GPU y la comunicación y la computación. Pero cuando realmente intentas construir algo, el 80 % de las veces, puedes depender de una cola muy larga de todos estos otros problemas prácticos porque el pipeline se ejecuta a la velocidad del cuello de botella más estrecho”.

A medida que el equipo de IBM desarrollaba su plataforma de capacitación, continuaron encontrando cuellos de botella. "A medida que mejoramos en el uso de nuestras GPU, cada vez más a menudo el cuello de botella es el cargador de datos", observa Wertheimer.

Esta realización llevó a un proceso de desarrollo dual. "Ha habido un proceso paralelo de, por un lado, evolucionar nuestra plataforma de entrenamiento y, por otro lado, evolucionar constantemente nuestro cargador de datos para mantenerse al día con las demandas de velocidad de nuestra plataforma de entrenamiento para evitar cuellos de botella", explica.

Linsong Chu, de IBM Investigación, describe las características esenciales del cargador de datos:

Con estado y con puntos de control: "Cada vez que guarda un modelo, también se guarda el estado de su cargador de datos, y cada vez que se recupera de un punto de control, tanto el estado del modelo como los estados del cargador de datos deben recuperarse al mismo tiempo", dice Chu.

Reescalado automático de puntos de control: el cargador de datos se ajusta automáticamente a los cambios en la carga de trabajo durante sesiones de entrenamiento prolongadas. "La capacitación podría llevar fácilmente semanas o meses, y hay muchas razones por las que podría tener que reescalar su carga de trabajo en el medio", señala Chu.

Transmisión de datos eficiente: el sistema admite la transmisión de datos sin sobrecarga de compilación para la mezcla.

Operación distribuida asíncrona: "Queremos que el cargador de datos no se bloquee", explica Chu. "Al guardar el estado del cargador de datos, queremos que el ahorro se distribuya de una forma en la que no haya comunicación".

Mezcla dinámica de datos: el cargador de datos puede adaptarse a diferentes proporciones de mezcla de datos, lo que es útil para las necesidades de capacitación en evolución.

Barajado global eficiente: La herramienta aborda los cuellos de botella de memoria cuando se manejan grandes conjuntos de datos, lo que hace que el barajado sea eficiente incluso a medida que los datos crecen.

PyTorch nativo, modular y extensivo: Diseñado para la adaptabilidad y escalabilidad, el cargador de datos está preparado para el crecimiento futuro. “¿Y si el próximo año tenemos que lidiar con 30 billones, 50 billones o 100 billones de tokens?” pregunta Chu. "El mundo está cambiando rápidamente, por lo que necesitamos construir el cargador de datos para que no solo pueda sobrevivir hoy, sino también para el mañana".

El equipo de investigación de IBM probó rigurosamente su cargador de datos durante varios meses, ejecutando cientos de trabajos pequeños y grandes. Observaron números de código estables y suaves. Además, todo el cargador de datos funciona de forma asíncrona y no bloquea.

"Aprovechamos muchas capacidades integradas de PyTorch para que todo esto sucediera", dice Wertheimer. “Es por eso que estamos contribuyendo, contribuyendo de nuevo”.