Carregar dados no AI Runtime

Os ativos de dados e modelos são críticos para cargas de aprendizado profundo e pós-treinamento para grandes modelos de linguagem (LLMs) e modelos de visão-linguagem (VLMs). Com o AI Runtime, todos os dados e ativos de modelos são acessados através do Unity Catalog:

• Volumes do Unity Catalog: usado principalmente para grandes conjuntos de dados e arquivos não estruturados, incluindo imagens, áudio e texto.
• Tabelas do Catálogo Unity: usadas para dados estruturados e tabulares, acessadas via Spark Connect.

Seus volumes e tabelas devem estar registrados no Catálogo Unity e acessíveis ao seu usuário ou principal de serviço.

Volume do Catálogo Unity 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 ambiente de execução de IA, os volumes são o principal mecanismo para acessar grandes conjuntos de dados, texto, ativos de modelo e checkpoints de modelo.

Os usuários podem listar, ler e gravar arquivos em volumes do Catálogo Unity usando operações familiares do sistema de arquivos, semelhante ao trabalho com arquivos em um 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 Catálogo Unity os tornam bem adequados para cargas de trabalho de aprendizado de máquina:

• Armazenamento distribuído: O Unity Catalog é respaldado 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 de notebooks quanto de cargas de trabalho baseadas em CLI.
• Otimizado para padrões de acesso de ML: O armazenamento e os caminhos de acesso subjacentes são otimizados para cargas de trabalho comuns de ML, especialmente arquivos grandes com leituras e gravações sequenciais. Isso torna o Unity Catalog muito adequado para o carregamento de dados de treinamento, o carregamento de artefatos do modelo e a gravação de checkpoints do modelo.
• Acesso semelhante ao sistema de arquivos: Usuários podem listar, ler e gravar arquivos em volumes do Unity Catalog usando operações familiares do sistema de arquivos, semelhante ao trabalho com arquivos em um disco local.

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

• Gravações: o AI Runtime confirma 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 atualização ou sincronização por parte do usuário.

Ajustar o desempenho do volume

Como mencionado, os volumes do Catálogo Unity são suportados por armazenamento distribuído e otimizados para arquivos grandes com leituras e gravações sequenciais.

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

• Concatene dados em arquivos maiores: Quando possível, consolide os dados em menos arquivos maiores, aproximadamente de 1 GiB a 10 GiB por arquivo. Isso permite que o AI Runtime pré-busque dados de forma agressiva e alcance automaticamente um desempenho de leitura sequencial quase ideal.

• Para cargas de trabalho de arquivos pequenos, use disco local: Se sua 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 causada pelo acesso repetido a muitos arquivos pequenos via o 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
  

• Você pode usar UCVolumeDataset para suas cargas de trabalho de aprendizado de máquina. Ele incorpora as otimizações descritas acima para fornecer 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 arquivos de texto armazenados em volumes do Unity Catalog, use UCVolumeDataset do módulo databricks.air.data. UCVolumeDataset é um IterableDataset PyTorch que copia cada arquivo do volume para um cache local rápido no primeiro acesso e retorna o caminho local do arquivo em cache. Ele lida com as preocupações de desempenho e distribuição que você implementaria manualmente:

• Armazenamento em cache local. Os arquivos são copiados da montagem do FUSE para um diretório de cache local no primeiro acesso e servidos do cache posteriormente, para que o treinamento de várias épocas não leia novamente o volume.
• Particionamento automático. Quando torch.distributed é inicializado, os arquivos são particionados entre as classificações e, em seguida, divididos ainda mais entre trabalhadores DataLoader, de modo que cada par (rank, worker) recebe uma fatia sem sobreposição e sem configuração adicional.

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

UCVolumeDataset gera caminhos de arquivo locais brutos. Para decodificar esses arquivos em tensores, encapsule-os em um segundo IterableDataset que consome o fluxo de caminhos e aplica sua lógica de análise. Isso mantém separadas as questões de E/S e de parsing.

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 encapsulador recebe caminhos locais já armazenados em cache, então a etapa de análise nunca acessa o volume. Você pode encadear wrappers adicionais para aumento, tokenização ou filtragem.

Para obter o melhor desempenho, use UCVolumeDataset com databricks.air.data.DataLoader em vez do PyTorch DataLoader padrão. Ele é otimizado para a entrada e saída (I/O) do runtime de IA e busca e armazena arquivos em cache simultaneamente enquanto a GPU realiza os cálculos.

Modelos de ponto de verificação em volumes

Para fazer checkpoint no seu modelo e retomar o treinamento a partir do snapshot mais recente ou se recuperar de um crash, você pode usar volumes do Unity Catalog assim como um sistema de arquivos local.

A Databricks recomenda o uso de um checkpoint distribuído (DCP) para melhor desempenho tanto em cargas de trabalho de GPU única quanto de múltiplas GPUs. Veja Treinamento rápido e tolerante a falhas com PyTorch em tempo de execução de IA no blog de engenharia 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 checkpointing de modelo 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 treinamento distribuído lançado 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 de tabela

Use o Spark Connect para carregar dados de machine learning tabulares de tabelas Delta.

Para treinamento de nó único, você pode converter DataFrames do Apache Spark em DataFrames do Pandas usando o método PySparktoPandas() e, opcionalmente, convertê-los no formato do NumPy usando o método PySparkto_numpy().

O Spark Connect dá suporte à maioria das APIs do PySpark, incluindo o Spark SQL, a API do Pandas no Spark, o Streaming Estruturado e o MLlib (baseado em DataFrame). Consulte a documentação de referência da API do PySpark para obter as APIs mais recentes com suporte.

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

Carregue grandes tabelas Delta usando volumes do Unity Catalog

Para tabelas Delta que são grandes demais para serem convertidas com toPandas(), exporte os dados para um volume do Catálogo do Unity 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")

Essa abordagem evita a sobrecarga associada ao Spark durante o treinamento e funciona bem tanto para fluxos de trabalho de treinamento com um único GPU quanto para distribuídos.