Treinamento distribuído tolerante a falhas no EKS usa NVRx

Treinamento distribuído tolerante a falhas no EKS usa NVRx

A Amazon Web Services (AWS) detalhou em 16 de setembro de 2026, em seu blog de machine learning, uma arquitetura para tornar treinamentos distribuídos de larga escala tolerantes a falhas no Amazon Elastic Kubernetes Service (Amazon EKS) com a NVIDIA Resiliency Extension (NVRx). A proposta combina três primitivas de resiliência aplicadas a scripts PyTorch Fully Sharded Data Parallel (FSDP) já existentes.

O material relata recuperação de falhas de GPU em segundos, com eficiência de treinamento acima de 99% em benchmarks com GPUs H100, segundo a publicação da AWS.

O problema que motiva a solução é estatístico. Jobs de treinamento nessa escala rodam por horas ou dias em dezenas de nós.

Nesse horizonte, interrupções são inevitáveis: partições de rede, erros de memória, exceções de software ou eventos de infraestrutura acabam derrubando ao menos um worker.

Uma única falha de GPU gera efeito cascata. Timeouts da NVIDIA Collective Communication Library (NCCL) se propagam para workers saudáveis, pods caem e reiniciam fora de sincronia.

O cluster consome horas caras de GPU sem produzir progresso.

A isso soma-se o checkpointing síncrono, em que cada salvamento bloqueia todos os ranks em operações de I/O. Nos tamanhos de cluster descritos no material, esse bloqueio chegou a responder por até 40% do tempo total de parede.

Três camadas independentes de recuperação

Para atacar esse cenário, a AWS recorreu ao NVRx. A biblioteca Python é instalável via pip install NVIDIA-resiliency-ext e adiciona primitivas de tolerância a falhas ao PyTorch sem exigir kernels customizados, fork do framework ou recompilação.

As primitivas entram em um script FSDP como importações comuns, deixando modelo e código de treinamento intactos. Cada uma pode ser adotada de forma isolada.

A primeira é o checkpointing assíncrono, exposto pelo TorchAsyncCheckpoint. Ele substitui o torch.save por uma chamada async_save(), que entrega o state dict a um processo em segundo plano e retorna imediatamente.

Um finalize_async_save() antes do salvamento seguinte confirma a escrita anterior. Combinado com o LOCAL_STATE_DICT do FSDP, cada rank grava seu próprio shard diretamente, sem all-gather e sem gargalo no rank 0.

A segunda é o reinício in-process, via inprocess.Wrapper, que envolve a função de treino para que uma falha transitória não mate o processo Python. Isso vale para exceção não tratada ou travamento de NCCL.

O NVRx aborta o process group ativo, executa verificações de saúde por rank (GPU, NVLink e NIC), re-rendezvousa os sobreviventes e reentra na função a partir do último checkpoint, preservando interpretador, alocador CUDA e objetos de escopo externo.

A terceira é o ft_launcher, que cobre o que o reinício in-process não alcança: SIGKILL, mortes por out-of-memory (OOM) e travamentos em nível de sistema operacional. Cada rank executa um RankMonitorClient.

O launcher compara heartbeats com timeouts definidos por linha de comando e, diante de parada ou morte, encerra os sobreviventes, recupera a memória de GPU e respawna workers novos no mesmo job, que recarregam o último checkpoint.

Segundo o texto, as camadas são independentes e cobrem classes distintas de falha: in-process para faltas leves, ft_launcher para faltas duras e o orquestrador de cluster para perda de nó.

A base de infraestrutura no amazon eks

Do lado de infraestrutura, o Amazon EKS sustenta node groups autogerenciados de instâncias p5.48xlarge, cada uma com 8 GPUs NVIDIA H100 de 80 GB e 32 interfaces Elastic Fabric Adapter (EFA). O EKS é um serviço gerenciado de Kubernetes que cuida do control plane, das atualizações e da disponibilidade do API server.

Os pods de treinamento rodam como Kubernetes Jobs com headless Services para descoberta de pares. Os workers se encontram por DNS em vez de IPs fixos, o que permite que substitutos reentrem no job sem reconfiguração.

As GPUs e os adaptadores EFA aparecem como recursos estendidos por meio dos device plugins da NVIDIA e do EFA. O scheduler posiciona os pods com node affinity e tolerations para obter a alocação completa de 8 GPUs por nó.

Para checkpoints, a arquitetura usa Amazon FSx for Lustre (SCRATCH_2, 1,2 TB) montado em todos os pods pelo driver CSI da FSx. Manter o FSx na mesma Zona de Disponibilidade dos nós de GPU reduz a latência de leitura durante a recuperação.

Esse ponto é crítico, já que o carregamento do checkpoint, e não o mecanismo de reinício, domina o tempo de recuperação em escala.

O que os números indicam

Os benchmarks apresentados no post da AWS foram executados em GPUs H100, de 2 a 8 nós. Reportaram eficiência de treinamento acima de 99% com recuperação na casa dos segundos.

O material afirma ainda que todo o código está disponível para reprodução, o que permite a times de engenharia avaliar a adoção das primitivas de forma incremental, sem reescrever o pipeline de treinamento existente.

A relevância prática está no custo. Em treinamentos longos, cada minuto de GPU parada é gasto sem aprendizado.

Mecanismos que encurtam a recuperação e sobrepõem o I/O de checkpoint ao cálculo atacam diretamente essa ineficiência.

A estratégia de camadas independentes, conforme descrito na publicação da AWS, permite que cada equipe escolha a proteção compatível com seus modos de falha mais frequentes, em vez de adotar um pacote único de resiliência para todos os cenários.

Fontes e links

Matéria publicada originalmente em AWS

Últimas Notícias

Gusto, Insight e Leland debatem agentes de IA em startups
OpenAI e AARP levam oficina gratuita de ChatGPT a 1.000 idosos
Suleyman alerta contra 'bem-estar de modelos' e direitos à IA
OpenAI lança framework para reportar desalinhamento de modelos