O Google usa tecnologia de IA na tradução de conteúdos para seu idioma de preferência. As traduções com IA podem ter erros.

Treinamento resiliente com programas de aprendizado

O Pathways oferece benefícios de resiliência das seguintes maneiras:

Suspender e retomar: permite que sua carga de trabalho tolere interrupções planejadas, como avisos de preempção, sem exigir um código personalizado de processamento de preempção.
Treinamento elástico: permite que sua carga de trabalho tolere falhas de hardware não planejadas sem causar falhas no cliente usando a lógica de recuperação no código de treinamento.

Antes de começar

Antes de começar, faça o seguinte:

Suspender e retomar

O Google Kubernetes Engine (GKE) normalmente envia um aviso de preempção para um pod acelerador antes que ele seja interrompido. A tolerância de preempção do Pathways é ativada por padrão em todas as implantações na nuvem, e os jobs do acelerador do Pathways detectam esses avisos.

Quando um aviso de preempção chega, o Pathways primeiro determina se pode salvar e restaurar a carga de trabalho de forma transparente. Se sim, ele tenta suspender de forma transparente a carga de trabalho de ML gravando o estado atual dela em um armazenamento permanente, como o Cloud Storage, antes que o GKE remova os jobs do acelerador. Quando o GKE reprogramar os jobs mais tarde, o Pathways vai retomar a carga de trabalho de ML lendo o estado persistente dela.

Se a carga de trabalho não for restaurável, o Pathways vai encerrar o job do acelerador e encaminhar a falha para o job se o treinamento elástico estiver configurado. Se o treinamento elástico não estiver configurado, o GKE vai reiniciar toda a carga de trabalho com base na política de reinicialização do JobSet.

As cargas de trabalho típicas de ML definidas usando o JAX dependem de componentes XLA do Pathways sem estado, que podem ser restaurados usando um snapshot de memória de alta largura de banda (HBM, na sigla em inglês). Algumas cargas de trabalho de ML, como as definidas usando a API Python colocalizada do JAX, dependem de componentes do Pathways com estado. Essas cargas de trabalho não são restauráveis.

Treinamento elástico

O treinamento elástico permite que o job de treinamento continue quando ocorrem falhas de hardware ou preempções de fração, sem reiniciar toda a carga de trabalho. Em um job de treinamento de várias frações, o Pathways mantém o processo do cliente de controlador único e as frações de TPU íntegras ativas quando uma fração individual fica indisponível.

O treinamento elástico combina os recursos do sistema Pathways com a lógica de recuperação no seu código de treinamento:

Detecção de falhas: quando ocorre uma falha de hardware ou preempção em uma fração de TPU, o Pathways cancela a atribuição da fração afetada e gera um jax.errors.JaxRuntimeError na próxima vez que sua carga de trabalho acessar dados nessa fração. Como o processo do cliente permanece ativo, seu código pode detectar esse erro, reconfigurar os recursos e continuar o treinamento.
Modos de recuperação: a classe pathwaysutils.elastic.manager.Manager e o decorador elastic_retry oferecem suporte a duas estratégias de recuperação:
- Pausar e retomar: quando minimum_slice_count for None ou igual ao número total de frações no job, o treinamento será pausado quando uma fração for perdida e aguardará até que todas as frações estejam íntegras novamente antes de ser retomado. Como a contagem de frações não muda entre as tentativas, sua carga de trabalho pode reutilizar a malha de dispositivos e os executáveis compilados.
- Redimensionamento de réplica: quando minimum_slice_count é menor que o número total de frações, o treinamento continua imediatamente nas frações ativas restantes reconstruindo a malha do dispositivo e recompilando para a topologia menor. Uma linha de execução de monitoramento em segundo plano verifica periodicamente as fatias inativas e registra as fatias recém-íntegras em elastic_manager.available_inactive_slices. O loop de treinamento pode recriar a malha do dispositivo com as fatias de substituição, reinicializar para a topologia expandida e colocar o estado atual na nova malha.
Recuperação de estado: para retomar o treinamento após um evento elástico, sua carga de trabalho restaura os parâmetros do modelo, o estado do otimizador e o estado do iterador de dados usando snapshots na memória. A classe Snapshotter em orbax.checkpoint copia de forma assíncrona o estado replicado da HBM da TPU para a memória da CPU pinned_host em cada fração. Durante a recuperação, o Snapshotter reconstrói o estado a partir de réplicas de host sobreviventes e o refragmenta diretamente na malha de TPU ativa sem ler do armazenamento remoto.

Ativar a elasticidade no servidor proxy do Pathways

Para ativar conjuntos de dispositivos elásticos no Pathways no GKE, adicione a flag --num_elastic_instances ao campo args do contêiner pathways-proxy no manifesto do JobSet e especifique o número de frações de TPU elásticas na carga de trabalho:

- name: pathways-proxy
  image: us-docker.pkg.dev/cloud-tpu-v2-images/pathways/proxy_server:latest
  args:
  - --resource_manager_address=$(PATHWAYS_HEAD):29001
  - --server_port=29000
  - --gcs_scratch_location=gs://BUCKET_NAME
  - --num_elastic_instances=NUM_SLICES

Substitua:

BUCKET_NAME: o bucket do Cloud Storage usado para arquivos temporários do Pathways.
NUM_SLICES: o número de frações de TPU na carga de trabalho.

Filtrar dispositivos ativos e criar a malha de dispositivos

Em uma carga de trabalho do Pathways, jax.devices() retorna todos os dispositivos na alocação inicial durante o ciclo de vida do processo do cliente. Ao usar o treinamento elástico, use elastic_manager.active_devices para que a malha do dispositivo e o tamanho do lote usem apenas frações íntegras e ativas:

import jax
import numpy as np
from pathwaysutils.elastic import manager

elastic_manager = manager.Manager()


def create_mesh(
    elastic_manager: manager.Manager, devices_per_slice: int
) -> jax.sharding.Mesh:
  device_array = np.asarray(elastic_manager.active_devices).reshape(
      (elastic_manager.active_slice_count, devices_per_slice)
  )
  return jax.sharding.Mesh(device_array, axis_names=("replica", "model"))

Quando o número de frações ativas mudar no modo de redimensionamento de réplica, use um carregador de dados elástico, como grain.experimental.ElasticIterator. Os iteradores de dados padrão por host salvam o estado por host e não podem ser restaurados em um número diferente de hosts, enquanto ElasticIterator rastreia uma única posição de conjunto de dados global que pode ser restaurada em qualquer número de frações ativas.

Loop de treinamento elástico

Um loop de treinamento elástico coordena o monitoramento de frações, a recuperação de falhas e os snapshots na memória, como mostrado no exemplo de código a seguir:

A função train cria uma única instância manager.Manager() e um Snapshotter do Orbax fora do wrapper de nova tentativa para que o rastreamento de frações e os snapshots de memória do host persistam entre as tentativas.
Decorar a função run_training_loop interna com @elastic_manager.elastic_retry aguarda frações íntegras, define jax.default_device como uma fração ativa, mantém a linha de execução do monitor em segundo plano ativa em todas as etapas e repete o loop quando ocorre um jax.errors.JaxRuntimeError.
Ao se recuperar de uma falha de fração no início de uma tentativa, seu código precisa recriar a malha do dispositivo de elastic_manager.active_devices e chamar snapshotter.load para restaurar o snapshot mais recente nas fragmentações de malha atualizadas, permitindo que os buffers do dispositivo da tentativa com falha saiam do escopo.
Durante cada etapa do loop, o código chama periodicamente snapshotter.save para fazer backup do estado de treinamento e verifica elastic_manager.available_inactive_slices para detectar fatias de substituição íntegras. Quando as frações de substituição estão disponíveis, o código as adiciona a elastic_manager.active_slice_indices, recria a malha, reinicializa o loop de treinamento e coloca o estado na malha expandida.

O exemplo a seguir mostra como implementar um loop de treinamento elástico combinando manager.Manager, elastic_retry e o Snapshotter na memória do Orbax:

import jax
import orbax.checkpoint as ocp
import pathwaysutils
from pathwaysutils.elastic import manager

pathwaysutils.initialize()


def train(
    final_step: int,
    snapshot_interval: int,  # in steps
    devices_per_slice: int,
):
  elastic_manager = manager.Manager()
  snapshotter = ocp.v1.training.pathways.Snapshotter(replica_axis_index=0)

  @elastic_manager.elastic_retry(
      minimum_slice_count=1,
      poll_interval=10,  # in seconds
      timeout=300,  # in seconds
      retry_policy=manager.ElasticRetryLimit(max_attempts=10),
  )
  def run_training_loop():
    mesh = create_mesh(elastic_manager, devices_per_slice)
    state, abstract_state, jitted_train_step, data_iterator = (
        initialize_training_loop(mesh)
    )
    step = 0

    # Restore from the latest in-memory snapshot after a slice failure.
    if snapshotter.latest is not None:
      snapshotter.join()
      restored = snapshotter.load({
          "state": abstract_state,
          "step": step,
          "iterator_state": data_iterator.get_state(),
      })
      state = restored["state"]
      step = int(restored["step"]) + 1
      data_iterator.set_state(restored["iterator_state"])

    while step < final_step:
      batch = next(data_iterator)
      state = jitted_train_step(state, batch)

      if step % snapshot_interval == 0:
        snapshotter.save(
            step,
            {
                "state": state,
                "step": step,
                "iterator_state": data_iterator.get_state(),
            },
        )

      # If replacement slices are healthy, reinitialize and place state onto the new mesh.
      if elastic_manager.available_inactive_slices:
        elastic_manager.active_slice_indices |= (
            elastic_manager.available_inactive_slices
        )
        mesh = create_mesh(elastic_manager, devices_per_slice)
        iterator_state = data_iterator.get_state()
        _, abstract_state, jitted_train_step, data_iterator = (
            initialize_training_loop(mesh)
        )
        data_iterator.set_state(iterator_state)
        state_shardings = jax.tree.map(lambda x: x.sharding, abstract_state)
        state = jax.device_put(state, state_shardings)

      step += 1

    return state

  state = run_training_loop()
  snapshotter.join()
  return state

Configurar o gerenciador elástico e o comportamento de nova tentativa

Você instancia manager.Manager(devices=None) uma vez por processo. Se você omitir devices, o gerenciador vai usar jax.devices() e agrupar os dispositivos pelo atributo slice_index.

A instância Manager expõe as seguintes propriedades para inspecionar a topologia do cluster durante o treinamento:

slice_to_devices: um mapeamento de cada slice_index para a sequência de objetos jax.Device que pertencem a essa fração.
all_slice_indices: uma frozenset de todos os índices de fração na alocação.
active_slice_indices: um frozenset de índices de fração que participam da malha de computação ativa.
active_devices: uma lista de objetos jax.Device que pertencem às frações ativas, ordenadas por slice_index.
inactive_slice_indices: um frozenset de índices de fração na alocação que não estão participando da malha de computação.
available_inactive_slices: um frozenset de índices de fração inativos que a linha de execução do monitor em segundo plano detectou como íntegros e prontos para voltar.
total_slice_count e active_slice_count: o número total de frações alocadas e o número de frações ativas.
default_device: dispositivo 0 de uma fração ativa, que o elastic_retry define automaticamente como o dispositivo JAX padrão para cada tentativa.

É possível personalizar o comportamento de nova tentativa e escalonamento transmitindo os seguintes argumentos para @elastic_manager.elastic_retry:

minimum_slice_count: o número mínimo de frações íntegras necessárias antes de iniciar a execução inicial ou uma nova tentativa. Se definido como None (o padrão), elastic_retry aguarda todas as frações na alocação (total_slice_count), operando no modo de pausa e retomada. Se definido como um número inteiro menor (como 1), elastic_retry opera no modo de redimensionamento de réplica.
poll_interval: o intervalo em segundos entre as verificações de integridade de fração realizadas pelo gerenciador de treinamento elástico (Manager) durante o treinamento. O padrão é 10.
timeout: o número máximo de segundos para aguardar até que as frações de minimum_slice_count fiquem ativas antes de gerar um TimeoutError. O padrão é None (sem tempo limite).
pre_callback: um callback opcional que é executado no início de cada tentativa após a atualização de active_slice_indices e antes da execução da função de loop de treinamento encapsulada.
on_elastic_event_callback: um callback opcional que é executado depois que uma redução de fração jax.errors.JaxRuntimeError é detectada, antes do início da próxima tentativa de repetição. Use esse callback para limpar checkpoints incompletos ou redefinir iteradores de dados.
retry_policy: um (attempt: int, error: Exception) -> bool que pode ser chamado e determina se as novas tentativas devem continuar após um evento elástico. Use manager.ElasticRetryLimit(max_attempts) para limitar o número total de tentativas. Se omitido, elastic_retry tenta novamente indefinidamente.

Em um evento elástico, o JAX libera os buffers do dispositivo da tentativa com falha quando o frame de pilha é desfeito e nenhuma variável faz referência a essas matrizes. Evite manter referências a matrizes de dispositivos de uma tentativa com falha dentro de variáveis globais ou callbacks de longa duração para que a memória seja liberada antes do início da próxima tentativa.

Snapshots na memória com o Orbax

A classe ocp.v1.training.pathways.Snapshotter (disponível na versão 0.12.6 e mais recentes do orbax-checkpoint) mantém um backup na memória do estado de treinamento na memória pinned_host em todas as frações:

Snapshotter(replica_axis_index=0): inicializa o snapshotter e especifica qual eixo de malha representa as réplicas de fração paralelas de dados. Todas as matrizes JAX no estado salvo precisam ser replicadas ao longo desse eixo de malha para que as frações sobreviventes mantenham uma cópia completa do modelo e do estado do otimizador quando uma fração falhar. Se o estado tiver matrizes de dispositivo único, replique-as na malha antes de chamar save.
save(step, state): transfere de forma assíncrona state para a memória pinned_host e o divide em árvores por réplica ao longo de replica_axis_index. Se uma transferência de snapshot anterior ainda estiver em andamento, save vai pular a nova solicitação para evitar que o loop de treinamento seja interrompido.
load(abstract_state, reset_snapshot_state=True): sonda as réplicas de host salvas para identificar fatias sobreviventes, concatena as réplicas íntegras e refragmenta o estado nas fragmentações de destino especificadas por abstract_state.
join(): bloqueia até que qualquer operação de snapshot em fila seja concluída.
latest: retorna metadados (incluindo step) para o snapshot concluído mais recentemente ou None se nenhum snapshot tiver sido concluído.

Reduzir a fragmentação em caso de falha de fração

Quando uma fração de TPU falha ou é interrompida durante a execução, o Pathways gera um jax.errors.JaxRuntimeError (como DATA_LOSS, DEADLINE_EXCEEDED, NOT_FOUND ou INTERNAL). O decorador elastic_retry captura essa exceção e verifica elastic.is_error_due_to_slice_down(error):

Se o erro não for causado por um evento de redução, elastic_retry vai gerar novamente a exceção imediatamente.
Se o erro for causado por um evento de redução de fração, elastic_retry interrompe qualquer trace ativo do criador de perfil do JAX, executa on_elastic_event_callback e avalia retry_policy(attempt, error).
Antes de iniciar a próxima tentativa, elastic_retry chama elastic.wait_for_slices uma vez, que pesquisa a integridade da fração até que pelo menos minimum_slice_count frações estejam ativas, atualiza elastic_manager.active_slice_indices e executa a função de loop de treinamento nas frações sobreviventes.

Aumentar a fragmentação quando as frações se juntarem novamente

Quando a carga de trabalho é executada em um número reduzido de frações no modo de redimensionamento de réplica, a linha de execução do monitor em segundo plano executa periodicamente verificações de integridade em elastic_manager.inactive_slice_indices. Quando as frações reparadas ou reprogramadas passam na verificação de integridade, a linha de execução do monitor adiciona os índices de fração a elastic_manager.available_inactive_slices.

Quando elastic_manager.available_inactive_slices não está vazio no final de uma etapa de treinamento, a carga de trabalho é escalonada verticalmente sem perder o progresso do treinamento:

Adicionar os índices de fração em elastic_manager.available_inactive_slices a elastic_manager.active_slice_indices incorpora as frações reunidas em elastic_manager.active_devices e as limpa de available_inactive_slices.
A recriação da malha do dispositivo de elastic_manager.active_devices e a reinicialização da etapa de treinamento e do iterador de dados adaptam o cálculo à topologia expandida.
Chamar jax.device_put coloca o estado de treinamento atual nas novas fragmentações de malha sem restaurar de um snapshot.

Hot-swap

A troca rápida se refere a um recurso da API JobSet do GKE em que um job de maior prioridade pode assumir rapidamente os recursos de um job de menor prioridade, minimizando o tempo de inatividade e garantindo uma recuperação mais rápida.

Quando um JobSet é criado, o GKE programa a carga de trabalho em várias frações, conforme especificado na configuração do JobSet. Se ocorrer uma falha de hardware em uma ou mais frações, os pods afetados serão marcados como com falha. Se o cluster do GKE tiver uma fração extra executando um job de prioridade mais baixa, o sistema do JobSet vai remapear a carga de trabalho da fração com falha do job de prioridade mais alta para essa fração extra. Esse remapeamento geralmente leva menos de um minuto.

Quando um JobSet é reiniciado, a troca automática pode ocorrer nas seguintes situações:

Modo padrão: se houver frações de TPU ociosas e de reserva disponíveis no mesmo cluster, o programador do Kubernetes prioriza a programação dos jobs reiniciados nessas frações, em vez de esperar que as frações com falha sejam reparadas. Esse comportamento de programação oferece uma recuperação mais rápida.
Cargas de trabalho heterogêneas: em clusters que executam várias cargas de trabalho com um recurso PriorityClass do Kubernetes configurado, um JobSet reiniciado pode acionar uma troca a quente. Se a afinidade do job reiniciado corresponder aos recursos de um job de prioridade mais baixa, o Kubernetes vai interromper o job de prioridade mais baixa para que o job de prioridade mais alta possa ser iniciado sem esperar pelo reparo da fração. Por exemplo, é possível configurar os pods de worker do Pathways com prioridades diferentes usando um recurso PriorityClass.

Para usar prioridades no cluster, defina um manifesto PriorityClass:

kind: PriorityClass
metadata:
  name: high-prior-job
value: 2000
globalDefault: false
description: "This priority class should be used for high priority job."

Aplique esse YAML ao cluster do GKE:

kubectl apply -f high-prior-job.yaml

Em seguida, anexe a nova classe de prioridade ao job de worker do Pathways adicionando o seguinte texto à podspec do pod pathways-worker.

priorityClassName: high-prior-job

Treinamento resiliente com programas de aprendizado Mantenha tudo organizado com as coleções Salve e categorize o conteúdo com base nas suas preferências.