Treinamento PyTorch rápido e tolerante a falhas no AI Runtime

Em escala, a eficiência do seu treinamento é determinada por uma única métrica: "goodput", a proporção de tempo que suas GPUs passam em computação produtiva em vez de esperando ou se recuperando de falhas. Como as falhas de GPU são o caso esperado em escala, a capacidade de se recuperar de forma rápida e automática de uma falha é a única maneira de manter um alto goodput e gerenciar seus gastos totais com GPU.

Dois subsistemas definem o sucesso ou o fracasso dessa recuperação, mas ambos são rotineiramente tratados como secundários: o pipeline de dados que alimenta seus aceleradores e o mecanismo de checkpointing que salva o estado (snapshot) para que um job possa ser retomado. Se errar em qualquer um deles, cada falha custará muito mais tempo de GPU ociosa do que deveria. Mesmo fora de cenários de falha, um pipeline de dados que não consegue acompanhar o ritmo de seus aceleradores deixará suas GPUs ociosas silenciosamente e reduzirá o goodput com a mesma certeza que uma falha de sistema faria. Vamos analisar os mecanismos e as compensações (trade-offs) de ambos, e como cada um molda seu goodput e seus gastos totais com GPU. Consulte o Guia de desempenho e resiliência de treinamento complementar para ver indicações de código e exemplos.

Para o lado da infraestrutura do mesmo problema, como uma frota detecta e isola GPUs problemáticas antes que elas derrubem um job, consulte o post complementar, Como mantemos as GPUs confiáveis no Databricks AI.

Por que as falhas são o caso esperado em escala

À medida que o número de GPUs em um job cresce, a probabilidade de que ele sobreviva a toda a sua duração sem interrupções cai rapidamente. Um modelo simples de estimativa rápida do post complementar da Databricks assume que cada GPU apresenta uma taxa de falha anualizada de aproximadamente 1%. Sob essa premissa, o post observa que "um job de 256 GPUs executado por 30 dias tem cerca de 19% de chance de sofrer uma falha. Com 1.024 GPUs, esse número sobe para 57%." e esses são apenas problemas no nível de infraestrutura.

Para fundamentar essa estimativa na realidade, o supercomputador delta de 608 GPUs H100 apresentou falhas a cada 1,9 horas, o que significa que, para um job de 32 GPUs, o tempo médio para falha seria de 36 horas. A principal conclusão é que seu job de treinamento provavelmente falhará em algum momento, e tomar as decisões corretas pode tornar seu modelo resiliente e reduzir o tempo total perdido quando isso acontecer.

Impacto 1: O formato do checkpoint decide com que frequência você pode se dar ao luxo de salvar

O checkpointing é onde a resiliência é ganha ou perdida, e o mecanismo que você escolhe tem um efeito de primeira ordem sobre a frequência com que você pode salvar. Este é o principal fator de influência no seu goodput: se você fizer o checkpoint uma vez por dia, uma falha exigirá a reexecução de, em média, 12 horas de trabalho duplicado para trazer você de volta ao estado em que estava quando a falha ocorreu.

O gargalo monolítico do torch.save

O primeiro checkpoint que a maioria das equipes cria é um simples torch.save no rank 0. Dependendo de como seu modelo é treinado, existem potencialmente dois problemas:

1. Para treinamento distribuído, ele reúne todos os estados no rank 0 e grava um único arquivo.
2. Um único processo grava todo o checkpoint de forma síncrona. Isso pode ser bloqueado por coisas como transferências de rede ao salvar em armazenamentos de objetos remotos, como o Unity Catalog (UC).

Esse comportamento de bloqueio deixa suas GPUs ociosas, reduzindo seu goodput. Mas há uma maneira de reduzir a quantidade de tempo que sua GPU gasta fazendo checkpointing: a API de checkpoint distribuído do Torch.

Checkpoint distribuído (DCP): cada rank grava seu próprio shard

O checkpoint distribuído do PyTorch inverte o design. Cada rank grava seu próprio shard distinto em paralelo, junto com um pequeno arquivo .metadata que descreve como os shards se compõem nos tensores completos.

O tempo de salvamento diminui aproximadamente como 1/N com o número de ranks e, como o arquivo .metadata registra o layout global, o mesmo checkpoint pode ser recarregado em um número diferente de GPUs. DCP replaneja quais bytes cada novo rank precisa, de modo que a recuperação em um cluster de capacidade reduzida após a perda de nós simplesmente funciona.

O DCP vale a pena mesmo para jobs simples de dados paralelos

Uma suposição comum é que o DCP serve apenas para modelos fragmentados (sharded), e que um job de paralelismo de dados (DDP), onde cada rank mantém uma réplica idêntica dos pesos, não tem nada a ganhar. Não é verdade: o DCP fragmenta (shards) o estado do modelo e o grava em paralelo em cada worker, mesmo para tarefas de treinamento DDP.

Também é a mesma API que você precisará no dia em que migrar para FSDP ou paralelismo de tensor, portanto, adotá-la cedo significa que você nunca precisará reescrever o código de resiliência no pior momento possível.

Salvamentos assíncronos tornam a frequência quase gratuita

Mesmo com gravações paralelas, um salvamento síncrono bloqueia o treinamento até que os bytes estejam protegidos no armazenamento — para um checkpoint grande em um volume remoto, isso significa dezenas de segundos de tempo ocioso do acelerador. O async_save divide a operação: uma cópia rápida para um buffer de staging e, em seguida, um upload em segundo plano que ocorre simultaneamente ao treinamento contínuo.

O loop de treinamento paga apenas pela cópia de staging, não pelo upload. Um checkpoint que costumava custar dezenas de segundos de tempo ocioso agora não custa quase nada, o que é exatamente o que torna viável o checkpointing frequente na próxima seção.

No AI Runtime, o UCVolumeWriter e o UCVolumeReader implementam o DCP em volumes do UC, realizando o staging de I/O por meio de NVMe local e marcando um checkpoint como concluído apenas quando seus dados tiverem sido totalmente gravados. Consulte o guia de desempenho e resiliência para obter detalhes completos e exemplos de código.

O exemplo acima exclui o tempo de armazenamento em rede para o torch.save.

Impacto 2: A frequência do checkpoint decide seu custo de recuperação

É aqui que as peças se somam. Quando um job falha, ele perde tudo desde o último checkpoint válido e deve recomputá-lo. Portanto, o trabalho desperdiçado esperado por falha é de cerca de metade do intervalo do checkpoint, e salvamentos assíncronos baratos permitem que você reduza esse intervalo.

Reduzir o intervalo por um fator de 10 reduz o tempo esperado de recuperação por um fator de 10. Lembre-se do número do Llama 3 de ~8,6 interrupções por dia: com essa taxa de falha, fazer o checkpoint a cada 2 horas significa que você espera desperdiçar 8,6 horas por dia em retreinamento, um goodput de 64%. Fazendo o checkpoint a cada 30 minutos, você gasta apenas 2,15 horas, um goodput de 91%.

A recuperação também deve ser automática. Na reinicialização, o job deve encontrar o checkpoint mais recente que terminou de ser gravado, ignorando qualquer um que tenha ficado gravado pela metade devido à falha, e retomar a partir dele sem intervenção humana. O DCP torna isso confiável: o arquivo .metadata é gravado apenas depois que todos os shards são salvos, portanto, sua presença é um marcador confiável de "este salvamento está concluído" para seleção.

Impacto 3: O carregamento de dados (dataloading) decide se suas GPUs ficarão ociosas

Um job de treinamento avança na velocidade de sua entrada mais lenta. Quando os aceleradores esperam pelo próximo lote (batch), seu goodput é reduzido, pois suas GPUs ficam simplesmente ociosas. A única maneira de corrigir esse problema é garantir que seu pipeline de entrada sobreponha a preparação de dados para a próxima etapa com a computação da etapa atual, conforme mostrado na figura abaixo:

Frequentemente vemos clientes que passam a sobrepor o carregamento de dados (dataloading) com a computação perceberem uma redução de 20% a 50% no tempo total de execução (wall-clock time).

O custo de ler diretamente do armazenamento remoto

Em uma plataforma governada, os dados de treinamento residem em armazenamento de objetos remoto. No AI Runtime, os volumes do Unity Catalog (UC) são apresentados como montagens de rede.

Ler arquivos diretamente dessa montagem a cada acesso vincula o tempo da sua etapa à latência da rede e baixa novamente os mesmos arquivos a cada época (epoch).

A solução é um dataloader que copia cada arquivo para um armazenamento local rápido no primeiro acesso, atende às leituras subsequentes a partir desse cache local e busca os próximos arquivos em paralelo enquanto a GPU computa.

Com o AI Runtime, o UCVolumeDataset e o DataLoader fazem exatamente isso (consulte o guia para ver exemplos de código). O UCVolumeDataset faz o streaming de arquivos de um volume do UC, armazenando cada um em cache no NVMe local no primeiro acesso, e particiona os arquivos entre ranks e workers para que cada acelerador receba uma fatia disjunta e sem sobreposição. Nosso DataLoader é uma subclasse direta do DataLoader do PyTorch cujos padrões são ajustados para esse caminho, de modo que os arquivos são buscados e armazenados em cache simultaneamente enquanto a GPU computa, em vez de um por um na thread de treinamento.

Exemplo: treinando um modelo de imagem a partir de arquivos do UC

Considere uma carga de trabalho simples de classificação de imagens: decodificar JPEGs de um volume do UC, aplicar técnicas de aumento de dados e treinar um modelo de visão computacional. Vamos analisar duas maneiras de fazer isso na mesma GPU, modelo e tamanho de lote (batch size): o Dataset padrão do PyTorch lendo de um volume do UC versus o UCVolumeDataset combinado com os padrões do DataLoader do Databricks.

Você não precisa adivinhar para onde vai o tempo

Como parte do desenvolvimento do DataLoader, garantimos que ele registre suas métricas no MLflow, facilitando identificar rapidamente se o seu pipeline de dados está bloqueando o treinamento.

A métrica fetch_seconds mede explicitamente quanto tempo o dataloader leva para produzir um lote e, durante esse tempo, sua GPU fica ociosa.

Impacto 4: Esquecer o pipeline de dados corrompe silenciosamente seu modelo

Existe um último bug de resiliência que não gera nenhuma mensagem de erro, nenhuma falha e nenhum job com falha, apenas um modelo sutilmente pior do que deveria ser. Isso acontece quando você salva o checkpoint do modelo, do otimizador e da etapa, mas não a posição do seu pipeline de dados dentro do dataset.

Considere um job interrompido no meio de uma época. Ele restaura o modelo corretamente e retoma o loop de treinamento, mas o dataloader recomeça do início do dataset.

O job retomado treina novamente em exemplos que já viu nesta época e, potencialmente, pula aqueles que ainda não havia alcançado. Com as muitas reinicializações que a escala torna rotina, isso enviesa silenciosamente a distribuição dos seus dados. O modelo ainda treina; ele apenas treina na amostragem errada dos seus dados, precisamente o tipo de falha silenciosa que é a mais cara, porque o job é concluído e ninguém percebe um problema até que as métricas sejam decepcionantes.

A solução é tratar a posição dos dados como parte do checkpoint. Dependendo do seu pipeline, isso significa rastrear um deslocamento (offset) de amostra ou shard e avançar ao retomar, fazer com que um dataset personalizado serialize sua própria posição ou criar checkpoints nos limites das épocas. Tudo isso depende de um pré-requisito: determinismo. O embaralhamento (shuffling) e o aumento de dados dependem de geradores de números aleatórios, portanto, essas sementes (seeds) e estados de RNG também devem fazer parte do checkpoint; caso contrário, a ordem dos dados após a reinicialização não corresponderá à ordem anterior, e uma posição salva apontará para as amostras erradas.

Semente, ordem reproduzível e pipeline de dados retomável são três expressões de uma única ideia. O guia aborda cada estratégia com código.

Resumo

Um treinamento rápido e tolerante a falhas é o resultado de algumas decisões que se somam:

1. Use o Distributed Checkpoint em vez do torch.save, mesmo para DDP, para que os salvamentos sejam paralelos e de baixo custo, em vez de um gargalo serial.
2. Salve de forma assíncrona para que os checkpoints sejam quase gratuitos, o que permite salvar com frequência.
3. Recupere automaticamente para o checkpoint válido mais recente, de modo que uma falha custe minutos de recomputação, não horas.
4. Sobreponha o carregamento de dados à computação fazendo cache e prefetch do armazenamento remoto para que os aceleradores nunca fiquem ociosos esperando por entrada. Isso representa horas de GPU economizadas de forma recorrente em cada etapa.
5. Salve o checkpoint do pipeline de dados e do estado do RNG, para que um job retomado continue com os dados corretos em vez de corromper silenciosamente o seu modelo.

O princípio unificador: checkpoints frequentes, de baixo custo e completos transformam uma falha de hardware de um evento que encerra o job em um mero detalhe, e um pipeline de entrada sobreposto mantém os aceleradores ocupados nesse intervalo. Salvamentos baratos (assíncronos) tornam a frequência viável; salvamentos completos (modelo, dados e RNG) garantem a correção da recuperação. Com ambos implementados e uma frota que detecta e isola hardwares com falha, seu tempo de treinamento efetivo se aproxima do limite máximo que o hardware permite, independentemente de quão instável seja o cluster subjacente.

Referências

• Meta, The Llama 3 Herd of Models (2024): §3.3.1, confiabilidade do treinamento e detalhamento de interrupções em até 16.384 GPUs H100.
• Characterizing the Resilience of Hopper H100 and Ampere A100 GPUs (2025): estudo de campo de 2,5 anos do sistema Delta; MTBE e probabilidade de falha por tamanho de job.
• Databricks, How we keep GPUs reliable across Databricks AI: verificação de integridade de GPU em nível de frota e o modelo probabilístico de falha.

Pronto para testar? Consulte o Guia de desempenho e resiliência de treinamento na documentação do Databricks AI Runtime para obter o código completo e leia How we keep GPUs reliable across Databricks AI para entender o lado da infraestrutura.

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