Acelere o treinamento de GPU em tabelas do Unity Catalog com o UCDataset

As GPUs costumam ser a parte mais cara de qualquer processo de treinamento, e a maioria das equipes não consegue mantê-las ocupadas. O gargalo geralmente não é o modelo. É obter os dados de onde eles estão para a GPU rápido o suficiente - e a cada segundo que esses chips ficam esperando, você está gastando tempo e dinheiro.

Hoje estamos lançando o Unity Catalog Dataset (UCDataset) em Beta no AI Runtime. O UCDataset é um carregador de dados nativo do PyTorch que transmite tabelas do Unity Catalog diretamente para cada GPU no seu job de treinamento, sem exportações, sem cópias e sem código de carregador personalizado.

Em cargas de trabalho de produção simuladas, o UCDataset reduziu o tempo de treinamento em 10x em comparação com a leitura de uma cópia Parquet exportada dos mesmos dados.

O problema: Suas GPUs estão esperando pelos seus dados

A maior parte dos dados de treinamento para seus modelos já está em tabelas do Unity Catalog. Colocá-los em um job de treinamento distribuído significa escolher uma entre três soluções alternativas, e cada uma delas custa tempo e dinheiro:

• Exportações lentas e entediantes antes de cada execução: O padrão mais comum é exportar a tabela para Parquet (ou outro formato de fragmento) em um volume antes de cada execução, e depois apontar um dataset para os arquivos. Você paga para armazenar uma segunda cópia, ela fica desatualizada no momento em que novos dados chegam e fica fora da linhagem do Unity Catalog.
• Puxar a tabela inteira para a memória não escala: Ler a tabela para pandas com toPandas() funciona para tabelas pequenas e falha em tabelas grandes. Isso também coloca tudo em um único nó, o que não ajuda em um job que treina em 32 GPUs.
• Transmitir por um data warehouse é lento: Os warehouses não foram criados nativamente para alimentar GPUs e, sem otimização, serão limitados quando muitas GPUs fizerem leituras ao mesmo tempo.

Depois, há o próprio carregador. Cada GPU e cada worker precisam da sua própria fatia balanceada de dados. As equipes gastam tempo construindo isso manualmente, e um bug sutil pode fazer com que o modelo treine com linhas duplicadas ou ignoradas sem que ninguém perceba.

O resultado: GPUs caras ficam ociosas, o treinamento demora mais do que deveria e os engenheiros de ML ficam fazendo manutenção de dados em vez de melhorar os modelos.

Apresentando o UCDataset

O UCDataset torna as tabelas do Unity Catalog uma fonte de dados de primeira classe para o treinamento distribuído do PyTorch. Você descreve o que ler (uma tabela, as colunas necessárias e um filtro de linha) e o UCDataset transmite lotes prontos para treinamento para cada GPU. Os benefícios incluem:

• Criado para velocidade: O UCDataset lê os arquivos da tabela diretamente do armazenamento em nuvem, então nada fica entre seus dados e suas GPUs. Ele divide as linhas correspondentes uniformemente entre cada GPU e worker do carregador de dados.
• Treine no lake, não em uma cópia: Aponte para a tabela ativa. Filtros e vetores de exclusão são aplicados conforme os dados são lidos, e a próxima execução coleta as novas linhas anexadas automaticamente.
• Governado de ponta a ponta: Os jobs de treinamento leem os dados com a identidade que os executa, e as permissões do Unity Catalog se aplicam exatamente como em qualquer outro lugar no Databricks.

Benchmark: Treinamento mais rápido em uma carga de trabalho de modelo de recomendação

Testamos o UCDataset em uma carga de trabalho de produção simulada, transmitindo uma tabela de recursos de 300 GB em 8 GPUs NVIDIA H100 para treinar um modelo de recomendação e vimos em média uma melhoria de 10x no tempo de treinamento em comparação com o carregamento de dados diretamente do Parquet.

Como usar:

Digamos que uma equipe mantenha recursos de usuários e itens com um rótulo de clique armazenados em uma tabela main.recsys.ctr_features. Novas impressões chegam à tabela todos os dias. Eles querem treinar um modelo de taxa de cliques em 16 GPUs em dois nós, diretamente da tabela.

1. Descreva o que ler e crie o carregador. No train.py, escolha a tabela, mantenha as linhas de treinamento rotuladas e selecione apenas as colunas que o modelo precisa. O filtro é aplicado conforme os arquivos da tabela são lidos, para que as linhas que não correspondem nunca cheguem ao seu código de treinamento. O UCDataset já gera lotes inteiros, portanto, envolva-o em um DataLoader padrão do PyTorch com batch_size=None.

2. Treine e inicie em dois nós. Treine com um loop DDP padrão do PyTorch. O UCDataset dá a cada uma das 16 GPUs sua própria parte das linhas correspondentes, então não há código de sampler ou sharding para escrever. Descreva o job em workload.yaml, que inicia o train.py com o torchrun em cada nó, e execute-o com a CLI do AI Runtime:

Para treinar em mais GPUs, altere num_accelerators em workload.yaml. O código de treinamento permanece o mesmo. O train.py e o workload.yaml completos, com uma tabela de exemplo para testá-los, estão no cookbook do AI Runtime: Modelo de taxa de cliques em uma tabela do Unity Catalog.

Comece a usar

Instale o pacote do PyPI. Em um notebook, execute %pip install databricks-uc-dataset.

O UCDataset está disponível em Beta no AI Runtime. Para testá-lo:

• Leia a documentação: Carregar tabelas do Unity Catalog no PyTorch com o UCDataset
• Comece a partir de um exemplo funcional: receita do cookbook
• Fale com a sua equipe de atendimento, adoraríamos ouvir seu feedback.

Para dados não estruturados armazenados em volumes do Unity Catalog, use UCVolumeDataset para transmitir arquivos com cache local. Embora o UCDataset lidere com tabelas estruturadas, também estamos expandindo-o para dar suporte ao treinamento multimodal por meio do tipo FILE, permitindo que você treine em imagens, áudio e documentos diretamente de uma tabela. A versão Beta lê tabelas Delta no Unity Catalog.

FAQs

O que é o UCDataset?

O UCDataset (Unity Catalog Dataset) é um carregador de dados nativo que transmite tabelas do Unity Catalog diretamente para as GPUs em um job de treinamento. Ele está disponível em Beta no Databricks AI Runtime. Você escolhe uma tabela, as colunas necessárias e o UCDataset entrega lotes prontos para treinamento sem etapa de exportação ou código de carregador personalizado.

Como o UCDataset acelera o treinamento de GPU?

O UCDataset mantém as GPUs ocupadas lendo os arquivos de dados de uma tabela diretamente do armazenamento em nuvem, sem uma cópia exportada no meio. Ele divide as linhas correspondentes uniformemente entre cada GPU e worker do carregador de dados, para que cada GPU leia apenas sua própria parte.

Preciso exportar tabelas do Unity Catalog para Parquet para treiná-las?

Não. O UCDataset lê a tabela ativa no Unity Catalog, portanto, não há etapa de exportação nem segunda cópia para armazenar ou manter sincronizada. Ele aplica seu filtro e os vetores de exclusão da tabela à medida que lê, e cada nova execução coleta as linhas anexadas desde a última.

O UCDataset aplica as permissões do Unity Catalog?

Sim. No AI Runtime, os jobs de treinamento leem os dados como a identidade que os executa, e as permissões do Unity Catalog na tabela são aplicadas. Como o job lê a tabela governada diretamente, não há cópia exportada para proteger separadamente. No Beta, o UCDataset não pode ler tabelas que possuem filtros de linha ou máscaras de coluna.

Quais tabelas o UCDataset pode ler?

O UCDataset lê tabelas Delta registradas no Unity Catalog, incluindo tabelas com vetores de exclusão.

(Esta publicação no blog foi traduzida utilizando ferramentas baseadas em inteligência artificial) Publicação original