Carregar dados no ambiente de execução da IA

Importante

Este recurso está no Public Preview.

Os ativos de dados e de modelos são essenciais para cargas de trabalho de aprendizagem profunda e de pós-treino de modelos de linguagem de grande dimensão (LLMs) e modelos de visão e linguagem (VLMs). Com o AI Runtime, todos os dados e ativos do modelo são acedidos através do Unity Catalog:

  • Volumes do Catálogo Unity: usado principalmente para grandes conjuntos de dados e ficheiros não estruturados, incluindo imagens, áudio e texto.
  • Tabelas do Catálogo Unity: usadas para dados estruturados e tabulares, acedidas através do Spark Connect.

Os seus volumes e tabelas devem estar registados no Catálogo Unity e acessíveis ao seu utilizador ou principal de serviço.

Volume do Unity Catalog para dados não estruturados

Os volumes do Unity Catalog fornecem acesso regulado a dados não tabulares em qualquer formato, incluindo dados estruturados, semi-estruturados e não estruturados. No AI Runtime, os volumes são o principal mecanismo para aceder a grandes conjuntos de dados, texto, ativos de modelos e checkpoints de modelo.

Os utilizadores podem listar, ler e escrever ficheiros em volumes do Catálogo Unity usando operações familiares do sistema de ficheiros, semelhante a trabalhar com ficheiros num 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 forma semelhante, as operações de shell funcionam da mesma forma:

%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

Algumas características dos volumes do Unity Catalog tornam-nos particularmente adequados para cargas de trabalho de aprendizagem automática:

  • Armazenamento distribuído: O Unity Catalog é suportado por armazenamento distribuído, permitindo que cargas de trabalho em tempo de execução com IA leiam e escrevam dados e modelem ativos em toda a plataforma, tanto a partir de notebooks como de cargas baseadas em CLI.
  • Otimizado para padrões de acesso ML: Os caminhos de armazenamento e acesso subjacentes estão otimizados para cargas de trabalho comuns de ML, particularmente ficheiros grandes com leituras e escritas sequenciais. Isto torna o Unity Catalog particularmente adequado para o carregamento de dados de treino, o carregamento de recursos do modelo e a gravação de pontos de verificação do modelo.
  • Acesso semelhante ao sistema de ficheiros: Os utilizadores podem listar, ler e escrever ficheiros em volumes do Catálogo Unity usando operações familiares do sistema de ficheiros, semelhante a trabalhar com ficheiros num disco local.

Devido a commits automáticos em segundo plano, os utilizadores podem esperar acesso consistente aos volumes do Unity Catalog:

  • Escritas: o AI Runtime confirma automaticamente as escritas, tornando as alterações visíveis para outras aplicações e cargas de trabalho com acesso ao mesmo volume do Unity Catalog.
  • Leituras: O AI Runtime deteta automaticamente alterações no volume, sem exigir ao utilizador qualquer operação explícita de atualização ou sincronização.

Ajustar o desempenho do volume

Como mencionado, os volumes do Catálogo Unity são suportados por armazenamento distribuído e otimizados para ficheiros grandes com leituras e escritas sequenciais.

Algumas dicas podem ajudar-te a obter o melhor desempenho do AI Runtime:

  • Concatene dados em ficheiros maiores: Sempre que possível, consolide os dados em menos ficheiros maiores, aproximadamente 1 GiB a 10 GiB por ficheiro. Isto permite ao AI Runtime pré-buscar dados de forma agressiva e alcançar automaticamente um desempenho de leitura sequencial quase ótimo.

  • Para cargas de trabalho de ficheiros pequenos, use disco local: Se a sua carga de trabalho envolver muitos ficheiros pequenos, considere copiar os ficheiros para o disco local usando cópias paralelas antes de os processar. Isto pode reduzir a sobrecarga de aceder repetidamente a muitos ficheiros pequenos através do volume.

    # 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

  • Podes usar UCVolumeDataset para as tuas cargas de trabalho de aprendizagem automática. Incorpora as otimizações descritas acima para proporcionar acesso e carregamento eficiente de dados a partir dos volumes do Catálogo Unity. Veja as secções seguintes para obter mais detalhes.

Carregar dados não estruturados com UCVolumeDataset

Para dados não estruturados, como imagens, áudio e ficheiros de texto armazenados em volumes do Catálogo Unity, use UCVolumeDataset do módulo databricks.air.data. UCVolumeDataset é um PyTorch IterableDataset que copia cada ficheiro do volume para uma cache local rápida no primeiro acesso e gera o caminho do ficheiro local em cache. Trata das questões de desempenho e distribuição que, de outra forma, teria de implementar manualmente:

  • Armazenamento em cache local. Os ficheiros são copiados da montagem do FUSE para um diretório de cache local no primeiro acesso e servidos a partir da cache depois, pelo que o treino multi-época não relê o volume.
  • Particionamento automático. Quando torch.distributed é inicializado, os ficheiros são particionados entre fileiras e depois ainda divididos entre DataLoader trabalhadores, de modo que cada (rank, worker) par recebe uma fatia não sobreposta sem configuração adicional.

Observação

UCVolumeDataset E databricks.air.data.DataLoader vêm do pacote databricks-sdk-air. Instale-o com o data extra, que também puxa um compatível torch:

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

UCVolumeDataset gera caminhos de ficheiros locais brutos. Para descodificar esses ficheiros em tensores, envolve-os num segundo IterableDataset que consome o fluxo de caminho e aplique a tua lógica de análise. Isto mantém as preocupações de I/O e de análise separadas.

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)

O módulo recebe caminhos locais já armazenados em cache, pelo que a etapa de análise sintática nunca acede ao volume. Pode encadear wrappers adicionais para aumentação, tokenização ou filtragem.

Para um desempenho ideal, combine UCVolumeDataset com databricks.air.data.DataLoader em vez do PyTorch DataLoader predefinido. Está ajustado para I/O em tempo de execução de IA e recolhe e armazena ficheiros em cache simultaneamente enquanto a GPU faz o cálculo.

Modelos de pontos de controlo em volumes

Para guardar pontos de verificação do seu modelo, de modo a poder retomar o processo de treino a partir do instantâneo mais recente ou recuperar após uma falha, pode usar volumes do Unity Catalog como se fossem um sistema de ficheiros local.

A Databricks recomenda o uso de um checkpoint distribuído (DCP) para melhor desempenho tanto em cargas de trabalho com GPU simples como multi-GPU. Veja Treino rápido e tolerante a falhas em PyTorch no AI Runtime no blogue de engenharia da 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"],
)

A abordagem monolítica torch.save também funciona.

  • Para a criação de pontos de verificação de modelos de GPU única,

    # 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 treino distribuído iniciado via 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"])

Carregar dados tabulares

Use o Spark Connect para carregar dados de aprendizagem automática tabulares de tabelas Delta.

Para o treino em nó único, pode converter DataFrames Apache Spark em DataFrames pandas usando o método PySparktoPandas(), e depois, opcionalmente, converter para formato NumPy usando o método PySparkto_numpy().

Observação

O Spark Connect adia a análise e a resolução de nomes para o tempo de execução, o que pode alterar o comportamento do seu código. Consulte Comparar o Spark Connect com o Spark Classic.

O Spark Connect suporta a maioria das APIs PySpark, incluindo Spark SQL, Pandas API no Spark, Structured Streaming e MLlib (baseado em DataFrame). Consulte a documentação de referência da API PySpark para as APIs mais recentes suportadas.

Para outras limitações, consulte as limitações da computação sem servidor.

Carregar grandes tabelas Delta usando volumes do Unity Catalog

Para tabelas Delta grandes demais para serem convertidas com toPandas(), exporte os dados para um volume do Unity Catalog e carregue-os diretamente usando PyTorch ou 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")

Esta abordagem evita a sobrecarga do Spark durante o treino e funciona bem tanto para fluxos de trabalho de treino com GPU única como distribuídos.

  • Última atualização a