Pular para o conteúdo principal
Última atualização em

Carregar dados no AI Runtime

info

Visualização

Este recurso está em Pré-visualização Pública.

Os ativos de dados e de modelos são críticos para a aprendizagem profunda e para cargas de trabalho de pós-treinamento de grandes modelos de linguagem (LLMs) e modelos de visão-linguagem (VLMs). Com o AI Runtime, todos os ativos de dados e de modelos são acessados por meio do Unity Catalog:

  • Volumes do Unity Catalog: usados principalmente para grandes datasets e arquivos não estruturados, incluindo imagens, áudio e texto.
  • Tabelas do Unity Catalog: usadas para dados estruturados e tabulares, acessadas por meio do Spark Connect.

Seus volumes e tabelas devem estar registrados no Unity Catalog e acessíveis ao seu usuário ou service principal.

Volume do Unity Catalog para dados não estruturados​

Os volumes do Unity Catalog fornecem acesso governado a dados não tabulares em qualquer formato, incluindo dados estruturados, semiestruturados e não estruturados. No AI Runtime, os volumes são o principal mecanismo para acessar grandes datasets, texto, ativos de modelos e pontos de verificação de modelos.

Os usuários podem listar, ler e gravar arquivos em volumes do Unity Catalog usando operações de sistema de arquivos familiares, semelhantes a trabalhar com arquivos em um disco local:

Python
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!")

Da mesma forma, as operações de shell funcionam da mesma maneira:

sh
%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 os tornam altamente adequados para cargas de trabalho de machine learning:

  • Armazenamento distribuído : o Unity Catalog é sustentado por armazenamento distribuído, permitindo que cargas de trabalho do AI Runtime leiam e gravem dados e ativos de modelo em toda a plataforma, tanto de notebooks quanto de cargas de trabalho baseadas em CLI.
  • Otimizado para padrões de acesso de ML : O armazenamento subjacente e os caminhos de acesso são otimizados para cargas de trabalho de ML comuns, particularmente arquivos grandes com leituras e gravações sequenciais. Isso torna o Unity Catalog adequado para o carregamento de dados de treinamento, carregamento de ativos de modelos e gravação de pontos de verificação de modelos.
  • Acesso semelhante ao sistema de arquivos : os usuários podem listar, ler e gravar arquivos em volumes do Unity Catalog usando operações familiares de sistema de arquivos, de forma semelhante a trabalhar com arquivos em um disco local.

Devido aos commits automáticos em segundo plano, os usuários podem esperar acesso consistente aos dados do volume do Unity Catalog:

  • Gravações : O AI Runtime confirma (commit) automaticamente as gravações, tornando as alterações visíveis para outros aplicativos e cargas de trabalho que acessam o mesmo volume do Unity Catalog.
  • Leituras : O AI Runtime detecta automaticamente as alterações no volume sem exigir nenhuma operação explícita de refresh ou sincronização por parte do usuário.

Ajustar o desempenho do volume​

Como mencionado, os volumes do Unity Catalog contam com armazenamento distribuído e são otimizados para arquivos grandes com leituras e gravações sequenciais.

Algumas dicas podem ajudar você a obter o melhor desempenho do AI Runtime:

  • Concatenar dados em arquivos maiores: quando possível, consolide os dados em menos arquivos e maiores, aproximadamente 1 GiB a 10 GiB por arquivo. Isso permite que o AI Runtime faça o prefetch de dados de forma agressiva e atinja automaticamente um desempenho de leitura sequencial quase ideal.

  • Para cargas de trabalho de arquivos pequenos, use o disco local: se a carga de trabalho envolver muitos arquivos pequenos, considere copiar os arquivos para o disco local usando cópias paralelas antes de processá-los. Isso pode reduzir a sobrecarga de acessar repetidamente muitos arquivos pequenos por meio do volume.

    sh
    # 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

  • Você pode usar UCVolumeDataset para suas cargas de trabalho de machine learning. Ele incorpora as otimizações descritas acima para fornecer acesso e carregamento eficientes de dados de volumes do Unity Catalog. Consulte as seções a seguir para obter mais detalhes.

Carregar dados não estruturados com UCVolumeDataset​

Para dados não estruturados, como imagens, áudio e arquivos de texto armazenados em volumes do Unity Catalog, use UCVolumeDataset do módulo databricks.air.data. UCVolumeDataset é um PyTorch IterableDataset que copia cada arquivo do volume para um cache local rápido no primeiro acesso e retorna o caminho do arquivo local em cache. Ele lida com as questões de desempenho e distribuição que você, de outra forma, implementaria manualmente:

  • Armazenamento em cache local. Os arquivos são copiados do ponto de montagem FUSE para um diretório de cache local no primeiro acesso e servidos do cache depois, para que o treinamento multi-época não releia o volume.
  • Particionamento automático. Quando torch.distributed é inicializado, os arquivos são particionados por classificações e, então, divididos entre DataLoader workers, para que cada par (rank, worker) receba uma fatia não sobreposta sem configuração adicional.
nota

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

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

UCVolumeDataset retorna caminhos de arquivo locais RAW. Para decodificar esses arquivos em tensores, envolva-o em um segundo IterableDataset que consome a transmissão de caminho e aplica sua lógica de análise. Isso mantém as preocupações de I/O e de análise separadas.

Python
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 wrapper recebe caminhos locais já em cache, portanto, o passo de parsing nunca acessa o volume. Você pode encadear wrappers adicionais para aumento, tokenização ou filtragem.

Para obter o desempenho ideal, associe UCVolumeDataset a databricks.air.data.DataLoader em vez do PyTorch padrão DataLoader. Ele é otimizado para I/O do AI Runtime e busca e faz cache de arquivos simultaneamente enquanto a GPU executa o compute.

Modelos de ponto de verificação em volumes​

Para fazer o checkpoint do seu modelo para que você possa retomar o treinamento a partir do snapshot mais recente ou se recuperar de uma falha, você pode usar os volumes do Unity Catalog exatamente como um sistema de arquivos local.

A Databricks recomenda o uso de um checkpoint distribuído (DCP) para obter um melhor desempenho em cargas de trabalho de GPU única e de múltiplas GPUs. Consulte Fast, fault-tolerant PyTorch treinamento on AI Runtime no blog de engenharia da Databricks.

Python
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 modelo em GPU única,

    Python
    # 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 o treinamento distribuído iniciado por meio do torchrun,

    Python
    # 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​

Utilize Spark Connect para carregar dados tabulares machine learning a partir de tabelasDelta.

Para treinamento de nó único, você pode converter DataFrames Apache Spark em DataFrames Pandas usando o métodoPySpark toPandas() e, opcionalmente, converter para o formato NumPy usando o métodoPySpark to_numpy().

nota

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

Spark Connect é compatível com a maioria APIs PySpark , incluindo Spark SQL, Pandas API on Spark, transmissão estruturada e MLlib (baseada em DataFrame). Consulte a documentação de referência da API PySpark para obter informações sobre as APIs suportadas mais recentes.

Para outras limitações, consulte Limitações compute 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 Unity Catalog e carregue-os diretamente usando PyTorch ou Hugging Face:

Python
# 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)
Python
# 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")

Essa abordagem evita a sobrecarga Spark durante o treinamento e funciona bem tanto para treinamento com uma única GPU quanto para treinamento distribuído com fluxo de trabalho.

Nesta página