Carregar dados no AI Runtime

Dados e ativos de modelo são críticos para a aprendizagem profunda e cargas de trabalho pós-treinamento para grandes modelos de linguagem (LLMs) e modelos de linguagem de visão (VLMs). Com o AI Runtime, todos os dados e ativos de modelo 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 modelo e pontos de verificação de modelo.

Os usuários podem listar, ler e gravar arquivos em volumes do Unity Catalog usando operações de sistema de arquivos familiares, semelhantes ao trabalho 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 adequados para cargas de trabalho de machine learning:

• Armazenamento distribuído : O Unity Catalog é suportado pelo armazenamento distribuído, permitindo que as cargas de trabalho do AI Runtime leiam e gravem ativos de dados e modelos em toda a plataforma, tanto de notebooks quanto de cargas de trabalho baseadas em CLI.
• Otimizado para padrões de acesso 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 bem adequado para carregamento de dados de treinamento, carregamento de ativos de modelo e gravação de checkpoint de modelo.
• 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, semelhantes a trabalhar com arquivos em um disco local.

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

• Escritas : O AI Runtime commit as escritas automaticamente, 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 seleciona automaticamente as alterações no volume sem exigir que o usuário faça qualquer operação explícita de refresh ou sincronização.

Ajustar desempenho de volume

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

Algumas dicas podem ajudá-lo a obter o melhor desempenho do AI Runtime:

• Concatenar dados em arquivos maiores: sempre que possível, consolide os dados em arquivos menores e maiores, aproximadamente 1 GiB a 10 GiB por arquivo. Isso permite que o AI Runtime pré-busque agressivamente os dados e obtenha automaticamente um desempenho de leitura sequencial quase ideal.

• Para cargas de trabalho com muitos arquivos pequenos, use o 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 de acessar repetidamente muitos arquivos pequenos através 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 o UCVolumeDataset para suas cargas de trabalho de machine learning. Ele incorpora as otimizações descritas acima para fornecer acesso e carregamento de dados eficientes a partir 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 gera o caminho do arquivo local armazenado em cache. Ele lida com as preocupações de desempenho e distribuição que você teria que implementar manualmente:

• Cache local. Os arquivos são copiados da montagem FUSE para um diretório de cache local no primeiro acesso e servidos a partir do cache posteriormente, de modo que o treinamento de múltiplas épocas não leia o volume novamente.
• Particionamento automático. Quando torch.distributed é inicializado, os arquivos são particionados entre ranks e depois divididos ainda mais entre DataLoader workers, de modo que cada par (rank, worker) recebe uma fatia que não se sobrepõe sem configuração extra.

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 produz caminhos de arquivo locais brutos. Para decodificar esses arquivos em tensores, envolva-os em um segundo IterableDataset que consuma a transmissão de caminho e aplique sua lógica de análise. Isso mantém as preocupações de E/S e 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á armazenados em cache, por isso o passo de análise nunca toca no volume. O usuário pode encadear wrappers adicionais para aumento, tokenização ou filtragem.

Para obter o melhor desempenho, emparelhe UCVolumeDataset com databricks.air.data.DataLoader em vez do PyTorch DataLoader stock. Ele é ajustado para I/O do AI Runtime e busca e faz cache de arquivos simultaneamente enquanto a GPU computa.

Modelos de ponto de verificação em volumes

Para verificar seu modelo e retomar o treinamento a partir do último snapshot ou recuperar de uma falha, você pode usar volumes do Unity Catalog da mesma forma que um sistema de arquivos local.

O Databricks recomenda o uso de um ponto de verificação distribuído (DCP) para obter melhor desempenho em cargas de trabalho de GPU única e multi-GPU. Consulte Treinamento PyTorch rápido e tolerante a falhas no AI Runtime no blog de engenharia do 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 verificação de ponto de verificação de modelo de 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 treinamento distribuído iniciado via 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

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

Para treinamento de nó único, você pode converter DataFrames do Apache Spark em DataFrames do Pandas usando o método PySpark toPandas() e, opcionalmente, converter para o formato NumPy usando o método PySpark 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. Consulte Comparar o Spark Connect com o Spark Classic.

O Spark Connect é compatível com a maioria das APIs do PySpark, incluindo Spark SQL, API do Pandas no Spark, Structured Streaming e MLlib (baseado em DataFrame). Consulte a documentação de referência da API do PySpark para obter as APIs compatíveis mais recentes.

Para outras limitações, consulte Limitações de compute serverless.

Carregar tabelas Delta grandes usando volumes do Unity Catalog

Para tabelas Delta grandes que são muito grandes para converter com toPandas(), exporte os dados para um volume do 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")

Esta abordagem evita o overhead do Spark durante o treinamento e funciona bem para fluxos de trabalho de treinamento em GPU única e distribuída.