Treinamento distribuído tolerante a falhas no Amazon EKS com NVRx

O problema real de treinar modelos em escala

Treinar modelos de linguagem de grande porte é um processo que dura horas ou dias, distribuído por dezenas de nós. Nessa escala, interrupções não são exceção — são estatisticamente inevitáveis. Partições de rede, erros de memória, exceções de software ou eventos de infraestrutura vão, em algum momento, derrubar ao menos um worker.

O problema é que uma única falha de GPU desencadeia uma cascata: timeouts da Biblioteca de Comunicação Coletiva da NVIDIA (NCCL — NVIDIA Collective Communication Library) se propagam para os workers saudáveis, pods reiniciam fora de sincronia, e o cluster queima horas caras de GPU sem fazer progresso real de treinamento.

Existe ainda um segundo gargalo: o checkpointing síncrono. A cada salvamento, todos os ranks ficam bloqueados aguardando a conclusão da escrita em disco — um overhead que, nos experimentos documentados pela AWS, consumiu até 40% do tempo total de execução.

Para endereçar ambos os problemas, a AWS publicou um guia técnico detalhando como integrar o NVIDIA Resiliency Extension (NVRx) ao treinamento com PyTorch Paralelismo de Dados Totalmente Fragmentado (FSDP — Fully Sharded Data Parallel) no Serviço Gerenciado de Kubernetes da Amazon (Amazon EKS — Amazon Elastic Kubernetes Service). O código completo para reprodução está disponível no repositório.

Visão geral da solução

A arquitetura combina os primitivos de tolerância a falhas do NVRx com um ambiente de treinamento baseado em EKS, projetado para cargas de trabalho de GPU multi-nó de alto desempenho. O NVRx cuida da resiliência na camada de aplicação — checkpointing assíncrono, reinicialização in-process e reinicialização in-job — enquanto o cluster EKS fornece a fundação de infraestrutura: agendamento de GPU, rede de alta largura de banda e armazenamento compartilhado para persistência de checkpoints.

O que é o NVRx

O NVIDIA Resiliency Extension (NVRx) é uma camada Python instalável via pip (pip install nvidia-resiliency-ext) que adiciona primitivos de tolerância a falhas ao PyTorch — sem kernels customizados, sem fork do PyTorch, sem recompilação. Os primitivos se encaixam em um script FSDP existente como importações comuns, sem alterar o modelo ou o código de treinamento. Cada capacidade pode ser adotada de forma independente.

O guia da AWS exercita três funcionalidades principais: checkpointing assíncrono, reinicialização in-process e ft_launcher (reinicialização in-job).

Infraestrutura no Amazon EKS

O Amazon EKS é um serviço gerenciado de Kubernetes que cuida do plano de controle, atualizações e disponibilidade do servidor de API. A configuração utiliza grupos de nós autogerenciados de instâncias p5.48xlarge, cada uma com 8 GPUs NVIDIA H100 de 80 GB e 32 interfaces de Rede de Fabric Elástico (EFA — Elastic Fabric Adapter). Os pods de treinamento rodam como Kubernetes Jobs com Services headless para descoberta de peers via DNS, permitindo que substituições de pods se reconectem sem reconfigurar o job.

Para armazenamento de checkpoints, a solução utiliza o Amazon FSx para Lustre (SCRATCH_2, 1,2 TB) montado em todos os pods de treinamento via driver FSx CSI. Posicionar o FSx na mesma Zona de Disponibilidade dos nós GPU minimiza a latência de leitura durante a recuperação — ponto crítico, pois o carregamento do checkpoint (não o mecanismo de reinicialização em si) domina o tempo de recuperação em escala.

Os principais serviços AWS envolvidos são: Amazon EKS para gerenciamento do ciclo de vida dos pods e Jobs; instâncias EC2 p5.48xlarge com 8x H100 80 GB e 32x adaptadores EFA; EFA fornecendo 3.200 Gbps de largura de banda de rede para operações all-reduce do NCCL; Amazon FSx para Lustre como sistema de arquivos compartilhado POSIX para I/O distribuído de checkpoints; e Amazon Registro de Contêineres Elástico (Amazon ECR — Amazon Elastic Container Registry) para a imagem de treinamento.

As três capacidades do NVRx em detalhe

Checkpointing assíncrono: desacoplando I/O do treinamento

O ponto de partida é um loop FSDP padrão lançado pelo torchrun com salvamento síncrono de checkpoints. Nesse modelo, o salvamento bloqueia todos os ranks até que os shards de cada rank sejam gravados no armazenamento compartilhado — e uma falha de worker derruba o job inteiro, exigindo reinício completo a partir do último checkpoint.

A substituição é direta: troca-se o torch.save pelo TorchAsyncCheckpoint do NVRx. Instancia-se uma vez com persistent_queue=True, chama-se async_save(state_dict, path) no lugar do torch.save, e chama-se finalize_async_save(blocking=True) uma vez ao encerrar o job. Um processo em background assume a escrita real; a thread principal segue imediatamente para o próximo passo de forward/backward. Com FSDP LOCAL_STATE_DICT, cada rank escreve seu próprio shard diretamente, sem all-gather e sem gargalo no rank 0.

import torch
import torch.distributed as dist
from nvidia_resiliency_ext.checkpointing.async_ckpt.torch_ckpt import TorchAsyncCheckpoint

# Initialize the async checkpoint manager once after model/optimizer setup.
async_ckpt = TorchAsyncCheckpoint(persistent_queue=True)

# Training loop
for step, batch in enumerate(dataloader):
    loss = model(batch)
    loss.backward()
    optimizer.step()

    # Checkpoint every N steps
    if step % checkpoint_interval == 0:
        state_dict = build_state_dict(model, optimizer, step)  # FSDP LOCAL_STATE_DICT, CPU-staged
        # torch.save(state_dict, path)  # was blocking
        async_ckpt.async_save(state_dict, path)  # returns immediately

# Drain any in-flight save at job exit.
async_ckpt.finalize_async_save(blocking=True)

Reinicialização in-process: recuperação sem tocar no ciclo de vida do contêiner

A função de treinamento é envolvida com inprocess.Wrapper. O wrapper assume o loop de reinicialização: uma exceção na função envolvida (ou um travamento detectado pelo watchdog) dispara uma re-entrada em vez de encerrar o processo. O construtor configura quatro aspectos: timeouts (soft_timeout, hard_timeout, barrier_timeout, completion_timeout), verificações de saúde (CudaHealthCheck + FaultCounter), um RetryController que limita o total de reinicializações e define um tamanho mínimo de world size sobrevivente, e atribuição de ranks (ActivateAllRanks + ShiftRanks) que compacta os sobreviventes para manter o world contíguo.

import nvidia_resiliency_ext.inprocess as inprocess
from nvidia_resiliency_ext.inprocess import CallWrapper

def train_with_inprocess_restart(args, restart_metrics, base_store=None, call_wrapper: CallWrapper = None):
    # Re-init dist, rebuild FSDP, load latest checkpoint, train.
    # call_wrapper.ping() each step;
    # call_wrapper.iteration tells you which restart you're on.
    ...

wrapped_train = inprocess.Wrapper(
    store_kwargs = {"host_name": master_addr, "port": master_port + 1},
    soft_timeout = datetime.timedelta(seconds=args.soft_timeout_seconds),
    hard_timeout = datetime.timedelta(seconds=args.hard_timeout_seconds),
    barrier_timeout = datetime.timedelta(seconds=args.barrier_timeout_seconds),
    completion_timeout = datetime.timedelta(seconds=args.barrier_timeout_seconds),
    health_check = inprocess.Compose(
        inprocess.health_check.CudaHealthCheck(),
        inprocess.health_check.FaultCounter(max_rank_faults=20)),
    initialize = inprocess.initialize.RetryController(
        max_iterations=args.max_restarts,
        min_active_world_size=1),
    rank_assignment = inprocess.Compose(
        inprocess.rank_assignment.ActivateAllRanks(),
        inprocess.rank_assignment.ShiftRanks()),
)(train_with_inprocess_restart)

wrapped_train(args, restart_metrics, base_store)

Por baixo dos panos, um ProgressWatchdog usa Py_AddPendingCall para registrar timestamps entre instruções de bytecode, tornando travamentos dentro do NCCL (uma extensão C) detectáveis. Uma MonitorThread monitora o TCPStore do inprocess em busca de uma flag de interrupção e injeta RankShouldRestart na thread principal ao detectá-la. O AbortTorchDistributed então coleta rastreamentos do Flight Recorder, aborta os backends NCCL e destrói o grupo de processos. O processo Python permanece vivo — apenas o grupo de processos distribuídos é reconstruído.

Reinicialização in-job com ft_launcher: cobrindo falhas mais graves

Para falhas que o in-process não consegue capturar — SIGKILL, término por falta de memória (OOM — Out of Memory) e travamentos em nível de sistema operacional — o NVRx oferece o ft_launcher. A mudança é trocar o torchrun pelo ft_launcher na invocação do job, adicionando alguns parâmetros específicos de controle de reinicialização:

# --- before ---
torchrun --nnodes=$NNODES --nproc_per_node=$GPU_PER_NODE \
  --rdzv-backend=c10d --rdzv-endpoint=$MASTER_ADDR:29500 \
  train.py [args]

# --- after ---
ft_launcher --nnodes=$NNODES --nproc_per_node=$GPU_PER_NODE \
  --rdzv_backend=c10d --rdzv_endpoint=$MASTER_ADDR:29500 \
  --max-restarts=20 --ft-restart-policy=any-failed \
  --ft-rank-heartbeat-timeout=900 \
  --ft-initial-rank-heartbeat-timeout=1200 \
  --monitor-interval=5 \
  train_ft_launcher.py [args]

No script de treinamento, instancia-se um RankMonitorClient uma vez após a inicialização distribuída e envia-se um heartbeat a cada step:

import nvidia_resiliency_ext.fault_tolerance as fault_tolerance

ft_client = fault_tolerance.RankMonitorClient()
ft_client.init_workload_monitoring()  # once, after dist init

for step in range(1, args.max_steps + 1):
    loss = train_step(model, next(data_iter), optimizer)
    ft_client.send_heartbeat()  # liveness signal to RankMonitorServer

Um RankMonitorServer por rank rastreia os intervalos de heartbeat contra os timeouts configurados. Ao detectar timeout, o launcher envia SIGTERM para os sobreviventes (SIGKILL para os que não respondem), libera a memória GPU, faz o re-rendezvous e spawna workers frescos. Os workers recarregam do último checkpoint na inicialização.

Como as camadas se complementam

As duas camadas de recuperação cobrem classes de falha distintas: o in-process captura o que cabe dentro de um processo Python (exceções transientes, travamentos NCCL detectáveis pelo watchdog); o ft_launcher captura o que mata o processo ou trava em nível de SO (SIGKILL, OOM, deadlocks abaixo do Python); e o orquestrador de cluster cuida da perda de nó. O checkpointing assíncrono é ortogonal a ambos: pode ser pareado com qualquer camada de recuperação (ou nenhuma) e limita a quantidade de trabalho perdido em caso de falha.

Resultados dos benchmarks

Checkpointing assíncrono vs síncrono

A AWS realizou benchmarks comparando checkpointing assíncrono (NVRx) versus síncrono (torch.save) no treinamento FSDP do LLaMA-3.1-8B, variando de 2 a 8 nós (16 a 64 GPUs H100) com intervalo de checkpoint a cada 1.000 steps.

O checkpointing assíncrono manteve eficiência de treinamento acima de 99% em todas as escalas (99,2% com 2 nós, 99,8% com 8 nós), enquanto o síncrono ficou estagnado entre 57% e 61%, perdendo ~40% do tempo de execução em I/O bloqueante — independentemente do número de nós. A lacuna persiste porque o overhead de checkpoint é limitado pelo armazenamento: o tempo de escrita no FSx Lustre (~275s) é constante, seja com 16 ou 64 GPUs. O modo assíncrono esconde essa latência completamente ao sobrepor o I/O com o próximo segmento de treinamento.

Ao variar a frequência de checkpoint em escala de 8 nós, a vantagem assíncrona fica ainda mais evidente. A cada 1.000 steps, o assíncrono atinge 99,8% de eficiência contra 60,3% do síncrono. A cada 100 steps, o síncrono despenca para 14,7% enquanto o assíncrono degrada graciosamente para 29,6% — ainda 2x mais eficiente. O ponto de cruzamento ocorre quando o intervalo de treinamento entre checkpoints (~280s a 100 steps) se aproxima do tempo de I/O do FSx (~275s).

A implicação prática é direta: equipes que treinam em escala frequentemente checkpointam com pouca frequência para evitar a penalidade de bloqueio, aceitando horas de progresso perdido em caso de falha. O checkpointing assíncrono elimina esse tradeoff — é possível checkpointar agressivamente (a cada poucos minutos) com overhead próximo de zero.

Recuperação de falhas

Os experimentos de recuperação de falhas foram conduzidos em 2 nós p5.48xlarge (16 GPUs H100 no total) treinando LLaMA-3.1-8B com FSDP, checkpointando a cada 500 steps em uma execução de 2.000 steps. Cinco falhas determinísticas foram injetadas em cada execução (mesma semente, mesmo cronograma de falhas) e três mecanismos de recuperação foram comparados.

O NVRx in-process restart alcançou 31% de goodput de treinamento com 87% de goodput de infraestrutura, recuperando em ~10 segundos com zero reinicializações de contêiner. O ft_launcher alcançou 25,5% de goodput de treinamento (85,9% de infra goodput) com 17s de recuperação por falha. A reinicialização padrão do Kubernetes foi catastrófica: 11,5% de goodput de treinamento, 35,8% de infra goodput e 270s de recuperação por falha.

O modo de falha da baseline é uma cascata: um rank cai, o rank sobrevivente atinge timeout do NCCL (60s), ambos os pods reiniciam fora de sincronia, disparando CrashLoopBackOff e ciclos adicionais de timeout antes que o treinamento possa ser retomado. Em clusters com centenas a milhares de GPUs, onde falhas de hardware ocorrem a cada poucas horas, esse padrão significa que o cluster pode gastar mais tempo se recuperando do que treinando.

Pré-requisitos e como começar

Para reproduzir a solução, são necessários: conta AWS com cota de serviço para instâncias p5.48xlarge (ou p4de.24xlarge); cluster Amazon EKS (v1.28+) com grupos de nós GPU autogerenciados habilitados para EFA e o plugin de dispositivo NVIDIA instalado; sistema de arquivos Amazon FSx para Lustre (SCRATCH_2) na mesma Zona de Disponibilidade dos nós GPU; imagem de contêiner com PyTorch 2.9+ e NVRx 0.4.1 para reproduzir os benchmarks (ou v0.6.0 para implantação atual); kubectl configurado para o cluster; conta HuggingFace com acesso ao meta-llama/Llama-3.1-8B; e dataset de treinamento pré-baixado no armazenamento compartilhado.

Para criação do cluster EKS com nós GPU e rede EFA, os guias de infraestrutura estão disponíveis em awsome-distributed-ai/1.architectures. Para a configuração específica do NVRx incluindo módulos Terraform, build do contêiner e preparação do dataset, consulte o README do caso de teste NVRx.

Limpeza de recursos

Para evitar cobranças contínuas, a AWS orienta remover os recursos na seguinte ordem — objetos Kubernetes que montam FSx devem ser removidos antes do próprio sistema de arquivos, ou a exclusão ficará presa em finalizadores de volume:

Deletar os jobs de treinamento com ./deploy.sh --delete training-job-inprocess.yaml e kubectl get jobs -n nvrx-training
Liberar o volume claim com kubectl delete -f fsx-storage.yaml (o PersistentVolume usa persistentVolumeReclaimPolicy: Retain, então isso desanexa o sistema de arquivos sem deletá-lo — o FSx continua sendo cobrado até ser removido na etapa 4)
Escalar os grupos de nós GPU para zero via AWS CLI
Deletar o sistema de arquivos, repositório de imagens e cluster usando as ferramentas de provisionamento originais (terraform destroy, eksctl delete cluster ou CloudFormation)

Conclusão

A combinação do NVRx com o Amazon EKS resolve dois gargalos clássicos do treinamento distribuído em escala: o bloqueio por checkpoints síncronos e o tempo excessivo de recuperação após falhas. Com checkpointing assíncrono, a eficiência de treinamento se mantém acima de 99% de 16 a 64 GPUs H100. Com reinicialização in-process, a recuperação ocorre em ~10 segundos sem reinicializações de contêiner; com ft_launcher, falhas graves (SIGKILL, OOM) são tratadas sem interromper o job. Ambas representam uma mudança de patamar frente à recuperação padrão do Kubernetes, onde cascatas de timeout NCCL e tempestades de CrashLoopBackOff consomem mais de 4 minutos por falha.

O exemplo completo reproduzível está disponível no repositório awsome-distributed-ai. Para mais contexto sobre o NVRx, consulte a extensão de resiliência da NVIDIA e a palestra da GTC: Build Fault-Tolerant Distributed AI Training at Scale.

Referências

Fonte

Fault tolerant distributed training on Amazon EKS using NVRx (https://aws.amazon.com/blogs/machine-learning/fault-tolerant-distributed-training-on-amazon-eks-using-nvrx/)