El cargador de datos de IBM Research mejora la formación 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 con 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 varios dispositivos, todo ello manteniendo el ritmo de las GPU cada vez más eficientes. Como señala IBM 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, explica algunos de los retos que pueden surgir durante la formación a gran escala: "Existe una especie de regla del 80/20 cuando se trata de formación a gran escala. El ochenta por ciento de toda la literatura publicada analiza las compensaciones algorítmicas entre la memoria de la 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 la canalización funciona a la velocidad del cuello de botella más estrecho".

A medida que el equipo de IBM continuaba desarrollando su plataforma de formación, seguían 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 viaje paralelo de, por un lado, evolucionar nuestra plataforma de formación y, por otro lado, evolucionar constantemente nuestro cargador de datos para mantenerse al día con las demandas de velocidad de nuestra plataforma de formación para evitar cuellos de botella", explica.

Linsong Chu, de IBM Research, describe las características esenciales del cargador de datos:

Con estado y controlable: "Cada vez que guarda un modelo, el estado de su cargador de datos también se guarda, y cada vez que se recupera de un punto de control, tanto el estado del modelo como el estado 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 las sesiones de entrenamiento prolongadas. "El entrenamiento podría llevar fácilmente semanas o meses, y hay un montón de razones por las que podría tener que reescalar su carga de trabajo a mitad de camino", señala Chu.

Transmisión de datos eficiente: el sistema soporta la transmisión de datos con una sobrecarga de construcción nula para el barajado.

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 almacenamiento se distribuya de una forma que no implique ninguna 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 la evolución de las necesidades de formació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 crecen los datos.

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 año que viene 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 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 de las capacidades integradas de PyTorch para hacer todo esto posible", afirma Wertheimer. "Por eso contribuimos, devolviendo lo que nos ha dado".