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.yamlekubectl get jobs -n nvrx-training - Liberar o volume claim com
kubectl delete -f fsx-storage.yaml(o PersistentVolume usapersistentVolumeReclaimPolicy: 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
- Amazon Elastic Kubernetes Service (Amazon EKS)
- Amazon FSx para Lustre
- Amazon Elastic Container Registry (Amazon ECR)
- Elastic Fabric Adapter (EFA)
- Treinamento distribuído com Amazon EKS e Torch Distributed Elastic
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/)
Leave a Reply