Acelera el entrenamiento de GPU en tablas de Unity Catalog con UCDataset

Las GPU suelen ser el recurso más costoso en cualquier proceso de entrenamiento, y la mayoría de los equipos no logran mantenerlas ocupadas. El cuello de botella normalmente no es el modelo. Es llevar los datos desde donde se encuentran hasta la GPU con la suficiente rapidez; y cada segundo que esos chips pasan esperando, estás desperdiciando tiempo y dinero.

Hoy lanzamos Unity Catalog Dataset (UCDataset) en versión beta en AI Runtime. UCDataset es un cargador de datos nativo de PyTorch que transmite tablas de Unity Catalog directamente a cada GPU en tu tarea de entrenamiento, sin exportaciones, sin copias y sin código de carga personalizado.

En cargas de trabajo de producción simuladas, UCDataset redujo el tiempo de entrenamiento en 10x en comparación con la lectura de una copia en Parquet exportada de los mismos datos.

El problema: Tus GPU están esperando tus datos

La mayor parte de los datos de entrenamiento para tus modelos ya se encuentra en las tablas de Unity Catalog. Llevarlos a una tarea de entrenamiento distribuido implica elegir una de tres alternativas, y cada una te cuesta tiempo y dinero:

• Exportaciones lentas y tediosas antes de cada ejecución: El patrón más común es exportar la tabla a Parquet (u otro formato fragmentado) en un volumen antes de cada ejecución, y luego apuntar un conjunto de datos a los archivos. Pagas por almacenar una segunda copia, queda desactualizada en cuanto llegan nuevos datos y queda fuera del linaje de Unity Catalog.
• Cargar toda la tabla en la memoria no escala: Leer la tabla en pandas con toPandas() funciona para tablas pequeñas y falla con las grandes. Además, concentra todo en un solo nodo, lo cual no ayuda a una tarea que se entrena en 32 GPU.
• Transmitir datos a través de un almacén es lento: Los almacenes no fueron diseñados originalmente para alimentar GPU y, sin optimización, se ralentizarán cuando muchas GPU lean al mismo tiempo.

Luego está el cargador en sí. Cada GPU y cada trabajador necesitan su propia porción equilibrada de los datos. Los equipos pierden tiempo creando esto manualmente, y un error sutil puede hacer que el modelo se entrene con filas duplicadas o omitidas sin que nadie se dé cuenta.

El resultado: GPU costosas inactivas, un entrenamiento que toma más tiempo del debido y ingenieros de ML que mantienen la infraestructura de datos en lugar de mejorar los modelos.

Presentación de UCDataset

UCDataset convierte las tablas de Unity Catalog en una fuente de datos de primer nivel para el entrenamiento distribuido de PyTorch. Describes qué leer (una tabla, las columnas que necesitas y un filtro de filas) y UCDataset transmite lotes listos para el entrenamiento a cada GPU. Los beneficios incluyen:

• Diseñado para la velocidad: UCDataset lee los archivos de la tabla directamente desde el almacenamiento en la nube, por lo que nada se interpone entre tus datos y tus GPU. Divide las filas coincidentes de manera uniforme entre cada GPU y cada trabajador del cargador de datos.
• Entrena en el lake, no en una copia: Apunta a la tabla en vivo. Los filtros y los vectores de eliminación se aplican a medida que se leen los datos, y la siguiente ejecución toma automáticamente las filas recién añadidas.
• Gobernanza de extremo a extremo: Las tareas de entrenamiento leen los datos con la identidad que las ejecuta, y los permisos de Unity Catalog se aplican exactamente igual que en cualquier otro lugar de Databricks.

Prueba comparativa: Entrenamiento más rápido en una carga de trabajo de un modelo de recomendación

Probamos UCDataset en una carga de trabajo de producción simulada, transmitiendo una tabla de características de 300 GB en 8 GPU NVIDIA H100 para entrenar un modelo de recomendación, y observamos en promedio una mejora de 10x en el tiempo de entrenamiento en comparación con cargar los datos directamente desde Parquet.

Cómo usarlo:

Supongamos que un equipo mantiene características de usuarios y elementos con una etiqueta de clic almacenada en una tabla main.recsys.ctr_features. Nuevas impresiones llegan a la tabla todos los días. Quieren entrenar un modelo de tasa de clics en 16 GPU en dos nodos, directamente desde la tabla.

1. Describe qué leer y crea el cargador. En train.py, selecciona la tabla, conserva las filas de entrenamiento etiquetadas y selecciona solo las columnas que el modelo necesita. El filtro se aplica a medida que se leen los archivos de la tabla, por lo que las filas que no coinciden nunca llegan a tu código de entrenamiento. UCDataset ya genera lotes completos, así que envuélvelo en un DataLoader estándar de PyTorch con batch_size=None.

2. Entrena e inicia en dos nodos. Entrena con un bucle DDP estándar de PyTorch. UCDataset asigna a cada una de las 16 GPU su propia parte de las filas coincidentes, por lo que no hay código de muestreo o fragmentación que escribir. Describe la tarea en workload.yaml, que inicia train.py con torchrun en cada nodo, y ejecútala con la CLI de AI Runtime:

Para entrenar en más GPU, cambia num_accelerators en workload.yaml. El código de entrenamiento sigue siendo el mismo. El train.py y el workload.yaml completos, con una tabla de ejemplo para probarlos, se encuentran en el libro de recetas de AI Runtime: Modelo de tasa de clics en una tabla de Unity Catalog.

Primeros pasos

Instala el paquete desde PyPI. En un notebook, ejecuta %pip install databricks-uc-dataset.

UCDataset está disponible en versión beta en AI Runtime. Para probarlo:

• Lee la documentación: Carga tablas de Unity Catalog en PyTorch con UCDataset
• Comienza con un ejemplo funcional: receta del libro de recetas
• Habla con tu equipo de cuenta, nos encantaría conocer tus comentarios.

Para datos no estructurados almacenados en volúmenes de Unity Catalog, utiliza UCVolumeDataset para transmitir archivos con almacenamiento en caché local. Si bien UCDataset maneja tablas estructuradas, también lo estamos expandiendo para admitir entrenamiento multimodal a través del tipo FILE, lo que te permite entrenar con imágenes, audio y documentos directamente desde una tabla. La versión beta lee tablas Delta en Unity Catalog.

Preguntas frecuentes

¿Qué es UCDataset?

UCDataset (Unity Catalog Dataset) es un cargador de datos nativo que transmite tablas de Unity Catalog directamente a las GPU en una tarea de entrenamiento. Está disponible en versión beta en Databricks AI Runtime. Seleccionas una tabla y las columnas que necesitas, y UCDataset entrega lotes listos para el entrenamiento sin necesidad de un paso de exportación o código de carga personalizado.

¿Cómo acelera UCDataset el entrenamiento de GPU?

UCDataset mantiene las GPU ocupadas al leer los archivos de datos de una tabla directamente desde el almacenamiento en la nube, sin una copia exportada de por medio. Divide las filas coincidentes de manera uniforme entre cada GPU y cada trabajador del cargador de datos, por lo que cada GPU lee solo su propia parte.

¿Necesito exportar las tablas de Unity Catalog a Parquet para entrenar con ellas?

No. UCDataset lee la tabla en vivo en Unity Catalog, por lo que no hay un paso de exportación ni una segunda copia que almacenar o mantener sincronizada. Aplica tu filtro y los vectores de eliminación de la tabla a medida que lee, y cada nueva ejecución toma las filas añadidas desde la última vez.

¿UCDataset aplica los permisos de Unity Catalog?

Sí. En AI Runtime, las tareas de entrenamiento leen los datos con la identidad que las ejecuta, y se aplican los permisos de Unity Catalog de la tabla. Como la tarea lee la tabla gobernada directamente, no hay una copia exportada que asegurar por separado. En la versión beta, UCDataset no puede leer tablas que tengan filtros de filas o máscaras de columnas.

¿Qué tablas puede leer UCDataset?

UCDataset lee tablas Delta registradas en Unity Catalog, incluidas las tablas con vectores de eliminación.

(Esta entrada del blog ha sido traducida utilizando herramientas basadas en inteligencia artificial) Publicación original