Treinamento resiliente com programas de aprendizado

Nesta página, explicamos como configurar recursos de resiliência para cargas de trabalho do Pathways no Google Kubernetes Engine (GKE) para que os jobs de treinamento em grande escala possam se recuperar de interrupções planejadas e falhas de hardware não planejadas.

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.
  • Checkpoints em vários níveis: permite que sua carga de trabalho salve checkpoints locais e replicados por pares de alta frequência em workers da TPU usando o Python colocalizado, com backups assíncronos do Cloud Storage para reduzir o tempo de recuperação.
  • Troca em funcionamento: permite que uma carga de trabalho do JobSet do GKE de maior prioridade assuma rapidamente as frações de TPU de menor prioridade ou de reserva quando ocorre uma falha de hardware.

Antes de começar

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

Se você usa o treinamento elástico, instale os seguintes pacotes na imagem do contêiner da carga de trabalho:

  • pathwaysutils
  • Orbax (orbax-checkpoint): >= 0.12.6 (necessário para snapshots na memória com ocp.v1.training.pathways.Snapshotter)
  • Grão (grain): obrigatório se você usar grain.experimental.ElasticIterator

Se você usa a verificação em várias camadas, também precisa atender aos seguintes pré-requisitos:

  • Configure o cluster do GKE para o checkpoint em vários níveis seguindo as instruções em Treinar modelos de machine learning em grande escala no GKE com o checkpoint em vários níveis:
    • Ative os complementos de checkpoint em vários níveis e do Cloud Storage FUSE (GcsFuseCsiDriver), além da federação de identidade da carga de trabalho para GKE, no cluster.
    • Crie um bucket do Cloud Storage com o namespace hierárquico ativado.
    • Conceda os papéis do Identity and Access Management (IAM) roles/storage.objectUser e roles/storage.insightsCollectorService à conta de serviço do Kubernetes de checkpoint de vários níveis.
    • Implante um recurso personalizado CheckpointConfiguration especificando o bucket do Cloud Storage, o seletor de nós de TPU e o tamanho do volume na memória (inMemoryVolumeSize).
  • Crie e envie uma imagem de contêiner secundário do Python colocalizada, por exemplo, usando o Dockerfile de amostra do contêiner secundário. Nas imagens de carga de trabalho e sidecar, use as versões jax e jaxlib correspondentes que se alinham à sua versão do Pathways e atendem aos seguintes requisitos mínimos de versão:
    • Versão do Orbax: >= 0.12.4
    • Versão do JAX (incluindo o Python colocalizado): >= 0.10.1

Suspender e retomar

O GKE normalmente envia um aviso de preempção para um pod acelerador antes de interrompê-lo. O Pathways ativa a tolerância à preempção por padrão em todas as implantações do Cloud Deploy, 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 possível, o Pathways tenta suspender de forma transparente a carga de trabalho de ML gravando o estado atual dela no 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 puder ser restaurada, o Pathways vai encerrar o job do acelerador e encaminhar a falha para a carga de trabalho 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 de ML típicas 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 e não podem ser restauradas.

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: se você definir minimum_slice_count como None ou como a contagem total de frações, o treinamento será pausado quando uma fração falhar e será retomado depois que todas as frações estiverem íntegras novamente. 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: se você definir minimum_slice_count como um valor menor que a contagem total de frações, o treinamento vai continuar nas frações ativas restantes quando uma fração falhar, reconstruindo a malha do dispositivo e recompilando para a topologia menor. Uma linha de execução de monitoramento em segundo plano verifica periodicamente as frações inativas e registra as frações que voltaram a ficar íntegras em elastic_manager.available_inactive_slices. O loop de treinamento pode reconstruir a malha do dispositivo com as frações 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/cloud-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. Por outro lado, ElasticIterator rastreia uma única posição global do conjunto de dados 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.
  • Aplicar @elastic_manager.elastic_retry a run_training_loop configura a função para aguardar fatias íntegras, definir jax.default_device como uma fatia ativa, manter a linha de execução do monitor em segundo plano ativa e tentar novamente em um jax.errors.JaxRuntimeError.
  • No início de uma tentativa de recuperação, seu código precisa recriar a malha do dispositivo com base em elastic_manager.active_devices e chamar snapshotter.load para restaurar o snapshot mais recente nas fragmentações de malha atualizadas. Deixe os buffers de dispositivo da tentativa com falha saírem do escopo para que o JAX possa liberar a memória deles.
  • 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 depois que o gerenciador atualiza active_slice_indices e antes que a função de loop de treinamento encapsulada seja executada.
  • on_elastic_event_callback: um callback opcional que é executado depois que elastic_retry captura um jax.errors.JaxRuntimeError de redução, antes do início da próxima tentativa. 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). Em uma função decorada com @elastic_manager.elastic_retry, 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 para verificar a integridade da fração até que pelo menos minimum_slice_count frações estejam ativas. Em seguida, ele 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.

Checkpoints em vários níveis

Os jobs de treinamento em grande escala sofrem interrupções frequentes, e a gravação ou restauração de checkpoints diretamente do Cloud Storage pode causar gargalos no progresso do treinamento. O checkpoint em várias camadas reduz a sobrecarga de checkpoint e o tempo de recuperação coordenando três camadas de armazenamento no GKE:

  • Disco RAM local: os workers da TPU gravam checkpoints de alta frequência em um volume CSI na memória (multitier-checkpoint.csi.storage.gke.io) montado em cada nó de trabalho. Esse volume persiste em reinicializações de pods no mesmo nó.
  • Replicação de peering no cluster: o serviço de replicação do GKE replica fragmentos de checkpoint locais entre workers da TPU de peering no cluster.
  • Backup do Cloud Storage: o serviço de replicação faz backup assíncrono de checkpoints locais em um bucket do Cloud Storage com o namespace hierárquico ativado, sem bloquear as etapas de treinamento da TPU.

Ao contrário do JAX de vários controladores, em que cada host de TPU executa o próprio script de treinamento, o Pathways executa um único controlador no pod pathways-head e envia o trabalho para os pods pathways-worker. Para ler e gravar fragmentos de checkpoint diretamente no disco de RAM local de cada host de TPU sem rotear dados de checkpoint pelo controlador único, o Pathways usa o Python colocalizado. Cada pod worker executa um contêiner colocated-python-sidecar ao lado do contêiner pathways-worker, e ambos os contêineres montam o volume CSI de checkpoint de várias camadas e um diretório de memória compartilhada (/tmp/sidecar). Durante a inicialização e o treinamento, o Orbax envia a configuração, as gravações e as restaurações de checkpoint do controlador do Pathways para os sidecars de worker em cada host de TPU. O checkpoint de vários níveis no Pathways também é compatível com o treinamento elástico ao reduzir escala vertical ou escalonar verticalmente o número de frações de TPU ativas.

O Pathways é compatível com o checkpoint em vários níveis com o MaxText. Para instruções de configuração, consulte o guia de checkpoint de vários níveis do MaxText.

Configurar a carga de trabalho do Pathways

É possível configurar uma carga de trabalho do Pathways para o checkpoint em vários níveis usando o Cluster Toolkit (gcluster job submit) ou configurando um manifesto personalizado do JobSet do Pathways.

Usar o Cluster Toolkit

Ao enviar um job com gcluster job submit --pathways, transmita as flags --gke-mtc-enabled, --gke-mtc-ramdisk-dir e --pathways-colocated-python-sidecar-image:

gcluster job submit \
    --pathways \
    --gke-mtc-enabled \
    --gke-mtc-ramdisk-dir="/tmp/mtc_checkpoints" \
    --pathways-colocated-python-sidecar-image="SIDECAR_IMAGE" \
    ...

Substitua SIDECAR_IMAGE pelo URI da imagem do contêiner secundário do Python colocada.

Quando você define essas flags, o gcluster configura automaticamente o seguinte:

  • Transmite --sidecar_name=external para o contêiner pathways-proxy.
  • Executa o contêiner colocated-python-sidecar em cada pod worker e monta um diretório /tmp/sidecar compartilhado (--cloud_pathways_sidecar_shm_directory=/tmp/sidecar) em pathways-worker e colocated-python-sidecar.
  • Ativa o volume CSI multitier-checkpoint.csi.storage.gke.io no caminho especificado por --gke-mtc-ramdisk-dir em pathways-worker e colocated-python-sidecar nos pods worker, mas não no pod pathways-head.

Usar um manifesto personalizado do JobSet do Pathways

Se você gerenciar seu próprio manifesto do JobSet do Pathways em vez de usar gcluster, transmita --sidecar_name=external para pathways-proxy. Na especificação do pod worker, monte o volume de memória do sidecar compartilhado e o volume CSI multitier-checkpoint.csi.storage.gke.io nos contêineres pathways-worker e colocated-python-sidecar:

containers:
- name: pathways-worker
  args:
  - --cloud_pathways_sidecar_shm_directory=/tmp/sidecar
  volumeMounts:
  - name: shared-tmp
    mountPath: /tmp/sidecar
  - name: checkpoint-mount
    mountPath: /tmp/mtc_checkpoints
- name: colocated-python-sidecar
  volumeMounts:
  - name: shared-tmp
    mountPath: /tmp/sidecar
  - name: checkpoint-mount
    mountPath: /tmp/mtc_checkpoints
volumes:
- name: shared-tmp
  emptyDir:
    medium: Memory
- name: checkpoint-mount
  csi:
    driver: multitier-checkpoint.csi.storage.gke.io

Usar o checkpoint de vários níveis do Orbax em uma carga de trabalho personalizada do JAX

Em um script de treinamento personalizado do JAX em execução no Pathways, use o pacote orbax.checkpoint.experimental.emergency.multi_tier_checkpointing e defina use_colocated_python=True em initialize_multi_tier_checkpointing e ReplicatorCheckpointManagerOptions:

from etils import epath
import jax
import orbax.checkpoint as ocp
from orbax.checkpoint.experimental.emergency.multi_tier_checkpointing import (
    initialization,
    replicator_checkpoint_manager,
)
import pathwaysutils

pathwaysutils.initialize()

local_checkpoint_dir = epath.Path("/tmp/mtc_checkpoints")

# Initialize multi-tier checkpointing across worker sidecars.
# When using elastic training, pass the active slice devices to `devices`
# on initial setup and whenever the active slice count changes.
initialization.initialize_multi_tier_checkpointing(
    local_checkpoint_directory=local_checkpoint_dir,
    backup_interval_minutes=10,
    num_slices=num_slices,
    run_name=run_name,
    data_parallelism=data_parallelism,
    use_colocated_python=True,
    devices=jax.devices(),
)

options = replicator_checkpoint_manager.ReplicatorCheckpointManagerOptions(
    save_interval_steps=20,
    use_colocated_python=True,
)
checkpoint_manager = (
    replicator_checkpoint_manager.ReplicatorCheckpointManager(
        local_checkpoint_dir,
        options=options,
        global_mesh=mesh,
    )
)

# Restore from the latest checkpoint if one exists.
start_step = 1
latest_step = checkpoint_manager.latest_step()
if latest_step is not None:
  restore_args = jax.tree.map(
      lambda x: ocp.type_handlers.ArrayRestoreArgs(sharding=x.sharding),
      state,
  )
  restored = checkpoint_manager.restore(
      latest_step,
      args=ocp.args.Composite(
          state=ocp.args.PyTreeRestore(
              item=state,
              restore_args=restore_args,
          )
      ),
  )
  state = restored.state
  start_step = latest_step + 1

# Step 0 is reserved by the replicator service to indicate "no checkpoint".
for step in range(start_step, final_step + 1):
  state = jitted_train_step(state)
  checkpoint_manager.save(
      step,
      args=ocp.args.Composite(
          state=ocp.args.PyTreeSave(item=state)
      ),
  )

checkpoint_manager.wait_until_finished()
checkpoint_manager.close()

Hot-swap

A troca rápida é um recurso da API JobSet do GKE que permite que um job de maior prioridade assuma rapidamente os recursos de um job de menor prioridade, minimizando o tempo de inatividade e acelerando a recuperação.

Quando você cria um JobSet, 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, o GKE vai marcar os pods afetados como com falha. Se o cluster do GKE tiver uma fração extra executando um job de prioridade mais baixa, o controlador 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 e salve-o como high-prior-job.yaml:

apiVersion: scheduling.k8s.io/v1
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 à carga de trabalho do Pathways. Se você enviar sua carga de trabalho usando o Cluster Toolkit, transmita --priority=high-prior-job para o comando gcluster job submit, que define priorityClassName nos pods pathways-head e worker. Se você escrever seu próprio manifesto do JobSet, defina priorityClassName no modelo de pod do job replicado worker, que executa o contêiner pathways-worker:

replicatedJobs:
- name: worker
  template:
    spec:
      template:
        spec:
          priorityClassName: high-prior-job

Se você criou o cluster usando um blueprint de TPU do Cluster Toolkit com o Pathways ativado, o cluster já inclui as classes de prioridade very-low, low, medium, high e very-high. Portanto, é possível transmitir uma delas para --priority em vez de criar uma própria.

A seguir