Treinamento elástico em TPUs: como o MaxText recupera uma falha em segundos

Quem já treinou modelos grandes distribuídos em várias máquinas conhece a dor: um nó morre no meio do treino, a operação de all-reduce fica esperando dados que nunca chegam, um timeout dispara e todos os workers caem. A saída padrão é o scheduler (Slurm, Kubernetes, Ray) reagendar tudo e relançar o job a partir do último checkpoint, pagando o custo cheio de restart e perdendo cada passo desde o último salvamento.

Um artigo do Google Developers Blog propõe uma alternativa usando o stack JAX (MaxText e PathwayS) em Cloud TPUs. Os autores mataram um worker de propósito no meio do treino e viram a recuperação acontecer no mesmo processo, mesmo PID, sem relançar nada. O downtime total, da morte até o próximo passo, ficou em menos de dois minutos, dominado pela espera do Kubernetes agendar um pod substituto.

Por que o modelo de controlador único muda o jogo

O detalhe que sustenta tudo é a arquitetura do PathwayS. Em launchers tradicionais de treino distribuído, roda-se um processo Python por nó, todos iguais, coordenando como pares (o clássico SPMD). Com PathwayS existe exatamente um processo Python, num nó CPU comum, que enxerga todos os chips da TPU como se fossem locais. Um jax.devices() devolve todos os chips; as máquinas TPU rodam apenas um binário worker fino que recebe programas XLA compilados e executa.

A consequência é direta: quando uma máquina TPU morre, ainda existe um processo Python vivo e saudável no nó CPU capaz de fazer algo a respeito. Esse "algo" é o que o texto chama de treinamento elástico.

O que 'elástico' significa aqui

Na prática, a falha de hardware vira uma jax.errors.JaxRuntimeError capturável, em vez de terminar o processo. Como você continua dentro de um processo vivo, com config e imports já carregados e as fatias sobreviventes ainda de pé, surgem duas estratégias:

• Pause and resume: captura a exceção, espera a fatia falha ser substituída, recarrega o último checkpoint válido e continua na malha completa.
• Replica resize: recarrega o checkpoint imediatamente nas fatias sobreviventes e segue treinando com throughput reduzido, voltando ao tamanho cheio quando o substituto sobe.

Ambas já existem no MaxText. A biblioteca pathways-utils oferece um decorator, elastic_retry, que envolve a função de treino inteira, captura a exceção, limpa estado parcial, restaura o último checkpoint viável e chama a função de novo, tudo no mesmo processo.

O que realmente economiza (e o que não)

Vale o realismo do próprio artigo: a recuperação elástica não pula tanto quanto parece. Chamar a função de treino de novo refaz setup do modelo, dataloader e restore do checkpoint, coisas que você pagaria num restart completo também. O agendamento do pod substituto continua dominando o relógio. O que a elasticidade economiza é o teardown de todo o workload: um restart completo derruba e reagenda o pod head, todos os workers saudáveis e o processo Python controlador. A recuperação elástica deixa tudo isso rodando e troca só a fatia que morreu. Compilação também não é economia: o PathwayS mantém cache persistente de XLA no Cloud Storage. A diferença entre os dois caminhos, resumindo, é o teardown, e no experimento isso foi a distância entre centenas de segundos e vários minutos.

Outra confusão que o texto desfaz: suspend-resume não é a mesma coisa. Para TPUs Spot com preempções planejadas, o PathwayS ouve o aviso de preempção, salva estado no Cloud Storage e retoma sozinho. O treinamento elástico é o caminho para as falhas sem aviso nenhum.

Como se configura

O comportamento elástico no MaxText é ligado por flags:

• enable_single_controller=True roteia o JAX pelo proxy do PathwayS (requisito duro).
• elastic_enabled=true aplica o decorator elastic_retry.
• elastic_timeout_seconds=300 limita a espera pela fatia substituta.
• elastic_max_retries=10 é o orçamento de falhas do run inteiro.

Há uma quinta flag implícita, elastic_min_slice_count, com default -1 (todas as fatias), que caracteriza o pause and resume. Um valor entre 1 e numSlices - 1 habilita o replica resize.

Um cuidado prático: se uma fatia falha durante uma escrita ativa de checkpoint, a versão atual do MaxText sai em vez de tentar de novo. Um checkpoint_period frequente cria janelas seguras entre escritas, ou use enable_continuous_checkpointing=True para o Orbax começar o próximo save assim que o anterior termina.

No Kubernetes, a submissão via xpk empacota o comando num JobSet. A linha decisiva é o backoffLimit no Job dos workers: ele permite reiniciar os pods de uma fatia falha no nível do Job, mantendo a falha local em vez de escalar para um restart do JobSet inteiro.

Por que importa para quem constrói no Brasil

Para times que escalam treino em cloud com orçamento apertado, cada restart completo custa dinheiro e tempo. A abordagem tem arestas (dependência do ecossistema JAX/TPU, o caso da falha durante checkpoint, o custo de agendamento de pod que não some), mas muda a matemática de confiabilidade em runs longos. O experimento reproduzível do artigo custou cerca de US$ 30 em meia hora com 48 chips v5e, um valor acessível para testar antes de comprometer um treino de dias.

Fonte: Google Developers Blog