Carga de datos en el entorno de ejecución de IA

Importante

Esta característica está en versión preliminar pública.

Los datos y los activos de modelos son fundamentales para el aprendizaje profundo y las cargas de trabajo post-entrenamiento para grandes modelos de lenguaje (LLM) y modelos de visión y lenguaje (VLM). Con AI Runtime, todos los datos y activos de modelos se acceden a través del Unity Catalog:

Volúmenes del Catálogo Unity: se utiliza principalmente para grandes conjuntos de datos y archivos no estructurados, incluyendo imágenes, audio y texto.
Tablas del Catálogo de Unity: se utilizan para datos estructurados y tabulares, accedidas a través de Spark Connect.

Tus volúmenes y tablas deben estar registrados en el Catálogo de Unity y accesibles para tu usuario o principal de servicio.

Volumen del Catálogo de Unity para datos no estructurados

Los volúmenes del Catálogo Unity proporcionan acceso gobernado a datos no tabulares en cualquier formato, incluyendo datos estructurados, semiestructurados y no estructurados. En AI Runtime, los volúmenes son el mecanismo principal para acceder a grandes conjuntos de datos, texto, recursos de modelos y puntos de control de modelos.

Los usuarios pueden listar, leer y escribir archivos en volúmenes del Catálogo Unity utilizando operaciones familiares del sistema de archivos, similar a trabajar con archivos en un disco local:

import os

dir_path = "/Volumes/<catalog-name>/<schema-name>/<volume-name>/sub-dir"
file_path = os.path.join(dir_path, "test_file")

os.makedirs(dir_path, exist_ok=True)

# Write to the file
with open(file_path, "w") as file:
    file.write("Hello, World!")

De manera similar, las operaciones de shell funcionan igual:

%sh ls -l /Volumes/<catalog-name>/<schema-name>/<volume-name>
%sh mkdir -p /Volumes/<catalog-name>/<schema-name>/<volume-name>/sub-dir
%sh touch /Volumes/<catalog-name>/<schema-name>/<volume-name>/sub-dir/test_file

Algunas características de los volúmenes del Catálogo Unity los hacen muy adecuados para cargas de trabajo de aprendizaje automático:

Almacenamiento distribuido: Unity Catalog cuenta con almacenamiento distribuido, lo que permite que las cargas de trabajo en tiempo de ejecución con IA lean y escriban datos y modelen activos en toda la plataforma, tanto desde portátiles como cargas basadas en CLI.
Optimizado para patrones de acceso ML: Los caminos de almacenamiento y acceso subyacentes están optimizados para cargas de trabajo comunes de ML, especialmente archivos grandes con lecturas y escrituras secuenciales. Esto hace que Unity Catalog sea muy adecuado para entrenar la carga de datos, la carga de activos de modelos y la escritura de puntos de control de modelos.
Acceso similar al sistema de archivos: Los usuarios pueden listar, leer y escribir archivos en volúmenes del Catálogo Unity utilizando operaciones familiares del sistema de archivos, similar a trabajar con archivos en un disco local.

Debido a las confirmaciones automáticas en segundo plano, los usuarios pueden contar con un acceso constante a los datos de volúmenes de Unity Catalog:

Escrituras: AI Runtime confirma automáticamente las escrituras, de modo que los cambios sean visibles para otras aplicaciones y cargas de trabajo que acceden al mismo volumen de Unity Catalog.
Lectura: AI Runtime detecta automáticamente los cambios realizados en el volumen sin requerir ninguna operación explícita de actualización o sincronización por parte del usuario.

Optimizar el rendimiento del volumen

Como se ha mencionado, los volúmenes del Catálogo Unity están respaldados por almacenamiento distribuido y optimizados para archivos grandes con lecturas y escrituras secuenciales.

Algunos consejos pueden ayudarte a obtener el mejor rendimiento de AI Runtime:

Concatena los datos en archivos más grandes: Cuando sea posible, consolida los datos en menos archivos grandes, aproximadamente de 1 GiB a 10 GiB por archivo. Esto permite que AI Runtime pre-receta datos de forma agresiva y logre automáticamente un rendimiento de lectura secuencial casi óptimo.

Para cargas de trabajo de archivos pequeños, usa disco local: si tu carga de trabajo implica muchos archivos pequeños, considera copiar los archivos al disco local usando copias paralelas antes de procesarlos. Esto puede reducir la sobrecarga que supone acceder repetidamente a muchos archivos pequeños a través del volumen de almacenamiento.

# Recommended using parallel copy (256 concurrency in this example, you can tune)
#
# This takes only 22 seconds to copy 15,375 150KiB small image files.
%sh cd /Volumes/<catalog-name>/<schema-name>/<volume-name>/sub-dir/ && find . -type f -print0 | xargs -0 -P 256 -I {} cp --parents "{}" /tmp/


# !!! Avoid doing this !!!
#
# Because the files are copied in serial, this copies the same 15,375 150KiB small image files much more slowly.
# %sh cp -r /Volumes/<catalog-name>/<schema-name>/<volume-name>/sub-dir/* /tmp

Puedes usarlo UCVolumeDataset para tus cargas de trabajo de aprendizaje automático. Incorpora las optimizaciones descritas anteriormente para proporcionar un acceso y carga eficiente de datos desde los volúmenes del Catálogo de Unity. Consulte las secciones siguientes para obtener más información.

Cargar datos no estructurados con `UCVolumeDataset`

Para datos no estructurados como imágenes, audio y archivos de texto almacenados en volúmenes del Catálogo de Unity, utiliza UCVolumeDataset desde el databricks.air.data módulo. UCVolumeDataset es un IterableDataset de PyTorch que copia cada archivo del volumen en una caché local rápida la primera vez que se accede a él y devuelve la ruta local del archivo almacenado en caché. Controla los problemas de rendimiento y distribución que, de lo contrario, implementaría a mano:

Almacenamiento en caché local. Los archivos se copian del montaje FUSE a un directorio de caché local la primera vez que se accede a ellos y, a partir de entonces, se sirven desde la caché, por lo que el entrenamiento durante varias épocas no vuelve a leer el volumen.
Creación automática de particiones. Cuando se inicializa torch.distributed, los archivos se reparten entre los rangos y luego se subdividen entre los DataLoader trabajadores, de modo que cada par de (rank, worker) recibe una porción no superpuesta sin necesidad de configuración adicional.

Nota:

UCVolumeDataset y databricks.air.data.DataLoader vienen del paquete databricks-sdk-air. Instálalo con el data extra, que también instala un compatible torch:

%pip install "databricks-sdk-air[data]"

UCVolumeDataset produce rutas de acceso a archivos locales sin procesar. Para descodificar esos archivos en tensores, envuélvelo en un segundo IterableDataset que consume el flujo de rutas y aplica tu lógica de análisis. Esto mantiene separadas las cuestiones de E/S y análisis.

from databricks.air.data import UCVolumeDataset
from torch.utils.data import IterableDataset
from PIL import Image
import torchvision.transforms.functional as TF

class ImageDataset(IterableDataset):
    """Decodes each cached file path from UCVolumeDataset into a tensor."""

    def __init__(self, path_dataset: UCVolumeDataset):
        self._path_dataset = path_dataset

    def __iter__(self):
        for local_path in self._path_dataset:
            image = Image.open(local_path).convert("RGB")
            yield TF.to_tensor(image)

path_dataset = UCVolumeDataset("/Volumes/catalog/schema/my_volume/images")
dataset = ImageDataset(path_dataset)

El envoltorio recibe rutas locales ya almacenadas en caché, por lo que el paso de análisis nunca toca el volumen. Puede encadenar contenedores adicionales para aumentar, tokenizar o filtrar.

Para obtener un rendimiento óptimo, combine UCVolumeDataset con databricks.air.data.DataLoader en lugar del DataLoader estándar de PyTorch. Está ajustado para E/S de IA en tiempo de ejecución y recoge y almacena archivos en caché simultáneamente mientras la GPU calcula.

Modelos de punto de control en volúmenes

Para guardar puntos de control de tu modelo y poder reanudar el entrenamiento desde la instantánea más reciente o recuperarte tras un fallo, puedes usar volúmenes de Unity Catalog como si fuera un sistema de archivos local.

Databricks recomienda usar un punto de control distribuido (DCP) para un mejor rendimiento tanto en cargas de trabajo de una sola GPU como de múltiples GPU. Consulta Entrenamiento rápido y tolerante a fallos de PyTorch en tiempo de ejecución de IA en el blog de ingeniería de Databricks.

import torch.distributed.checkpoint as dcp
from torch.distributed.checkpoint.state_dict import get_state_dict, set_state_dict

import databricks.air.data

checkpoint_path = "/Volumes/my-catalog/my-schema/my-volume/checkpoints/step_1000"

# Save
model_sd, optim_sd = get_state_dict(model, optimizer)
state_dict = {"model": model_sd, "optim": optim_sd, "step": 1000}
dcp.async_save(
    state_dict,
    storage_writer=databricks.air.data.UCVolumeWriter(checkpoint_path))

# Load
model_sd, optim_sd = get_state_dict(model, optimizer)
state_dict = {"model": model_sd, "optim": optim_sd}
dcp.load(
    state_dict,
    storage_reader=databricks.air.data.UCVolumeReader(checkpoint_path))

set_state_dict(
    model,
    optimizer,
    model_state_dict=state_dict["model"],
    optim_state_dict=state_dict["optim"],
)

El enfoque monolítico torch.save también funciona.

Para el guardado de puntos de control de modelos de una sola GPU,

# The monolithic torch.save approach for single GPU chip

# Save
torch.save({"model": model.state_dict(), "opt": optimizer.state_dict()},
           "/Volumes/<catalog-name>/<schema-name>/<volume-name>/sub-dir/ckpt.pt")


# Load
ckpt = torch.load(
    "/Volumes/<catalog-name>/<schema-name>/<volume-name>/sub-dir/ckpt.pt",
    weights_only=True)
model.load_state_dict(ckpt["model"])
optimizer.load_state_dict(ckpt["opt"])

Para formación distribuida lanzada vía Torchrun,

# The monolithic torch.save approach for multi-GPU distributed training.
# This snippet assumes your launcher has already called
# dist.init_process_group(...).

import os
import torch.distributed as dist

# Save only on rank 0.
if dist.get_rank() == 0:
    torch.save({"model": model.state_dict(), "opt": optimizer.state_dict()},
               "/Volumes/<catalog-name>/<schema-name>/<volume-name>/sub-dir/ckpt.pt")

# Wait for rank 0 to finish writing before any rank reads.
dist.barrier()

# Load on ALL ranks (map to current rank's local GPU).
local_rank = int(os.environ["LOCAL_RANK"])
ckpt = torch.load(
    "/Volumes/<catalog-name>/<schema-name>/<volume-name>/sub-dir/ckpt.pt",
    map_location=f"cuda:{local_rank}",
    weights_only=True)
model.load_state_dict(ckpt["model"])
optimizer.load_state_dict(ckpt["opt"])

Carga de datos tabulares

Use Spark Connect para cargar datos de aprendizaje automático tabulares de tablas Delta.

Para el entrenamiento de un solo nodo, puede convertir DataFrames de Apache Spark en DataFrames de pandas mediante el método PySparktoPandas(), y luego opcionalmente convertir a formato NumPy utilizando el método PySparkto_numpy().

Nota:

Spark Connect aplaza el análisis y la resolución de nombres al tiempo de ejecución, lo que puede cambiar el comportamiento del código. Consulte Comparación de Spark Connect con Spark clásico.

Spark Connect admite la mayoría de las API de PySpark, como Spark SQL, Pandas API en Spark, Structured Streaming y MLlib (basado en dataframe). Consulte la documentación de referencia de la API de PySpark para obtener las API compatibles más recientes.

Para ver otras limitaciones, consulte Limitaciones de proceso sin servidor.

Carga grandes tablas Delta usando volúmenes del Catálogo de Unity

Para tablas Delta grandes que son demasiado extensas para convertir con toPandas(), exporte los datos a un volumen de Unity Catalog y cárguelos directamente usando PyTorch o Hugging Face.

# Step 1: Export the Delta table to Parquet files in a UC volume
output_path = "/Volumes/catalog/schema/my_volume/training_data"
spark.table("catalog.schema.my_table").write.mode("overwrite").parquet(output_path)

# Step 2: Load the exported data directly using Hugging Face datasets
from datasets import load_dataset

dataset = load_dataset("parquet", data_files="/Volumes/catalog/schema/my_volume/training_data/*.parquet")

Este enfoque evita la sobrecarga de Spark durante el entrenamiento y funciona bien tanto para flujos de trabajo de entrenamiento distribuido como para GPU única.

Comentarios

¿Le ha resultado útil esta página?

Se actualizó por última vez el 2026-10-05