Prepare modelos de aprendizagem automática de grande escala no GKE com a funcionalidade de pontos de verificação de vários níveis

Esta página mostra-lhe como usar a criação de pontos de verificação de vários níveis para armazenar e gerir de forma fiável os pontos de verificação durante a preparação do modelo de aprendizagem automática no GKE. O armazenamento e a gestão de pontos de verificação são cruciais para tarefas de preparação em grande escala, definidas como as que usam mais de milhares de nós. As interrupções destas tarefas de grande escala são frequentes (potencialmente de hora a hora) e a recuperação das mesmas pode ser lenta.

Vantagens

A utilização da funcionalidade de pontos de verificação de vários níveis oferece as seguintes vantagens:

  • Gestão de dados de pontos de verificação totalmente orquestrada, incluindo cópias de segurança, replicação e restauro automático para as seguintes cargas de trabalho:
  • Recuperação rápida de tarefas de preparação a partir de um ponto de verificação armazenado no nó local. Também pode fazer a recuperação através de pontos de verificação armazenados noutro nó no cluster de preparação.
  • Restauro rápido de tarefas de preparação a partir de um ponto de verificação armazenado numa cópia de segurança do Cloud Storage em cenários de pior caso, em que não existem pontos de verificação no cluster.

Antes de começar

Antes de começar, certifique-se de que realizou as seguintes tarefas:

  • Ative a API Google Kubernetes Engine.
    • Ative a API Google Kubernetes Engine
  • Se quiser usar a CLI gcloud para esta tarefa, instale-a e, em seguida, inicialize-a. Se instalou anteriormente a CLI gcloud, execute gcloud components update para obter a versão mais recente.

Requisitos

A funcionalidade de criação de pontos de verificação de vários níveis requer a versão 1.32.4-gke.1415000 ou posterior do cluster do GKE.

Limitações

  • Os clusters do Autopilot não são suportados.

Configure nós do GKE para usar a criação de pontos de verificação de vários níveis

Estas secções abordam a configuração de nós do GKE em clusters novos e existentes.

Configure nós num novo cluster

  1. Crie um cluster com a funcionalidade Multi-Tier Checkpointing, o controlador CSI do FUSE do Cloud Storage e a Federação de identidades de cargas de trabalho para o GKE ativados. Se usar fatias de TPU para a sua carga de trabalho de aprendizagem automática, tem de ajustar o comando de criação do cluster para incluir a configuração de um conjunto de nós de fatias de TPU.

    gcloud container clusters create CLUSTER_NAME \
    --addons=HighScaleCheckpointing,GcsFuseCsiDriver  \
    --node-locations=NODE_LOCATION \
    --workload-pool=PROJECT_ID.svc.id.goog \
    --cluster-version=CLUSTER_VERSION
    --location=CLUSTER_LOCATION \
    --machine-type=MACHINE_TYPE \
    --num-nodes=NUM_NODES

    Substitua os seguintes valores:

    • CLUSTER_NAME: o nome do cluster.
    • NODE_LOCATION: a zona dos nós do cluster. É aqui que se encontra a sua capacidade de TPU.
    • PROJECT_ID: o seu Google Cloud ID do projeto.
    • CLUSTER_VERSION: a versão do seu cluster. 1.32.4-gke.1415000 é a versão mínima suportada.
    • CLUSTER_LOCATION: a região na qual quer criar o cluster.
    • MACHINE_TYPE: o tipo de máquina usado para nós que executam componentes como o controlador JobSet e o controlador de pontos de verificação de vários níveis. Para a preparação em grande escala, recomendamos a utilização de, pelo menos, e2-standard-4 máquinas. Não vai usar este tipo de máquina para o preparação de modelos. Em alternativa, vai criar pools de nós separados para esse fim, usando frequentemente famílias de VMs otimizadas para aceleradores.
    • NUM_NODES: o número de nós a criar em cada uma das zonas do cluster.

Configure nós num cluster existente

Para usar a criação de pontos de verificação de vários níveis com um cluster existente, ative-a juntamente com o controlador CSI do FUSE do Cloud Storage e a federação de identidades de cargas de trabalho para o GKE com os seguintes comandos. A versão do cluster existente tem de ser posterior a 1.32.3-gke.1170000.

  1. Ative a federação de identidades de cargas de trabalho para o GKE:

    gcloud container clusters update CLUSTER_NAME \
    --workload-pool=PROJECT_ID.svc.id.goog \
    --location=CLUSTER_LOCATION

    Substitua os seguintes valores:

    • CLUSTER_NAME: o nome do cluster.
    • PROJECT_ID: o seu Google Cloud ID do projeto.
    • CLUSTER_LOCATION: a região do cluster.

  2. Ative a criação de pontos de verificação de vários níveis e o controlador CSI FUSE do Cloud Storage:

    gcloud container clusters update CLUSTER_NAME \
    --update-addons=HighScaleCheckpointing=ENABLED,GcsFuseCsiDriver=ENABLED \
    --location=CLUSTER_LOCATION

Configure autorizações para usar a verificação de pontos de controlo de vários níveis

Esta secção aborda a forma de configurar autorizações para usar a funcionalidade de validação de vários níveis.

Conceda acesso a contentores do Cloud Storage

Os volumes efémeros usados pelo controlador CSI de verificação de pontos de controlo de vários níveis têm de usar contentores do Cloud Storage existentes.

Para armazenar pontos de verificação num contentor do Cloud Storage, a funcionalidade de pontos de verificação de vários níveis precisa de acesso ao contentor. Conceda a função de IAM de utilizador de objetos de armazenamento (roles/storage.objectUser) no contentor à conta de serviço do Kubernetes para a criação de pontos de verificação de vários níveis.

gcloud storage buckets add-iam-policy-binding gs://GCS_BUCKET \
    --member "principal://iam.googleapis.com/projects/PROJECT_NUMBER/locations/global/workloadIdentityPools/PROJECT_ID.svc.id.goog/subject/ns/gke-managed-checkpointing/sa/gke-checkpointing-multitier-node" \
    --role "roles/storage.objectUser"

Substitua os seguintes valores:

(Opcional) Conceda acesso à conta de serviço predefinida do Compute Engine

Se as suas instâncias do Compute Engine precisarem de acesso de leitura ao contentor do Cloud Storage, conceda a função do IAM Storage Object Viewer (roles/storage.objectViewer) à conta de serviço predefinida do Compute Engine.

gcloud storage buckets add-iam-policy-binding gs://GCS_BUCKET \
    --member serviceAccount:PROJECT_NUMBER-compute@developer.gserviceaccount.com \
    --role roles/storage.objectViewer

Implemente o controlador JobSet

O controlador JobSet é responsável por gerir as tarefas em lote que executam a preparação do seu modelo no GKE, e a respetiva atribuição de recursos é ajustada para processar a carga de trabalho de forma eficiente. Certifique-se de que o iniciador de tarefas de preparação é implementado e usa o JobSet.

Para aumentar o pedido de memória para 1 Gi, o limite de memória para 2 Gi e o pedido de CPU para 1 para o contentor do gestor na sua implementação do JobSet, execute o seguinte comando de patch:

kubectl patch -n jobset-system deploy jobset-controller-manager --type json \
    --patch '[{"op": "add", "path": "/spec/template/spec/containers/0/resources", "value": {"limits": {"memory": "2Gi"}, "requests": {"cpu": "1", "memory": "1Gi"}}}]'

Inicialize o controlador CSI de pontos de verificação de vários níveis

Esta secção descreve como inicializar o controlador CSI de pontos de verificação de vários níveis em nós onde as suas cargas de trabalho vão ser executadas. O controlador CSI é responsável pelo processamento do armazenamento e da gestão dos pontos de verificação durante o processo de preparação do modelo.

Crie uma CheckpointConfiguration

Um CheckpointConfiguration é um recurso personalizado do Kubernetes que especifica propriedades para implementar o controlador CSI de pontos de verificação de vários níveis. Este recurso tem âmbito de cluster.

  1. Crie o seguinte manifesto checkpoint.yaml.

    apiVersion: checkpointing.gke.io/v1
kind: CheckpointConfiguration
metadata:
  name: MTC_CONFIG_NAME-configuration
spec:
    cloudStorageBucketName: GCS_BUCKET
    nodeSelector:
        node.kubernetes.io/instance-type: MACHINE_TYPE
    tolerations:
    - key: TOLERATION_KEY
        operator: Exists
        effect: NoSchedule
    inMemoryVolumeSize: IN_MEMORY_VOLUME_SIZE
    gcsFuseMountOptions:
    - implicit-dirs
    - metadata-cache:negative-ttl-secs:0
    - metadata-cache:ttl-secs:-1
    - metadata-cache:stat-cache-max-size-mb:-1
    - metadata-cache:type-cache-max-size-mb:-1
    - file-system:kernel-list-cache-ttl-secs:0
    - read_ahead_kb=1024
    - write:enable-streaming-writes:true

    Substitua o seguinte:

    • MTC_CONFIG_NAME: o nome da sua CheckpointConfiguration. Este nome é global para o cluster e não é específico da tarefa.
    • GCS_BUCKET: o nome do contentor do Cloud Storage onde vai armazenar os dados de pontos de verificação. Use o contentor que configurou no passo Configure um contentor do Cloud Storage com autorizações.

    • MACHINE_TYPE: o tipo de máquina dos aceleradores correspondentes. O valor pode ser um dos seguintes:

      Para mais informações sobre a execução de cargas de trabalho distribuídas em GPUs com o GKE, consulte o artigo Executar GPUs de várias instâncias. Para TPUs, consulte o artigo Crie o node pool de fatia da TPU.

    • TOLERATION_KEY: este campo permite que o controlador CSI seja agendado em nós com manchas correspondentes. Para mais informações sobre como as falhas funcionam em diferentes tipos de aceleradores, consulte estas páginas:

    • IN_MEMORY_VOLUME_SIZE: o tamanho da cache de verificação de pontos na memória. Especifique a quantidade e a unidade (por exemplo, 200 Gi).Este valor deve:

      • O tamanho do ponto de restauro local para as TPUs multiplicado por 2,2
      • O tamanho do ponto de verificação local para GPUs com um único par multiplicado por 6,6.

  2. Aplique o manifesto:

    kubectl apply -f checkpoint.yaml

  3. Verifique se o controlador CSI está em execução:

    kubectl get pod -n gke-managed-checkpointing

    O resultado deve ser semelhante ao seguinte. Existem várias entradas, uma por nó acelerado.

    NAME                                                          READY   STATUS    RESTARTS   AGE
multitier-driver-e2b033a7-a4e7-496a-87a3-ffd7fcc2e57b-2d4fz   5/5     Running   0          114s

Desinstale o controlador CSI de pontos de verificação de vários níveis

Se quiser anular a implementação do controlador CSI de pontos de verificação de vários níveis, elimine o recurso CheckpointConfiguration. O controlador de pontos de verificação de vários níveis remove o controlador CSI dos nós. Esta ação remove os discos RAM e liberta memória para outras cargas de trabalho. Por exemplo:

kubectl delete -f checkpoint.yaml

Faça a gestão da retenção de dados e da recolha de lixo para as cópias de segurança do Cloud Storage

É responsável pela implementação de políticas de retenção para as cópias de segurança dos pontos de verificação do Cloud Storage. A criação de pontos de verificação de vários níveis apenas escreve cópias de segurança de pontos de verificação no Cloud Storage e nunca as modifica nem elimina.

Muitas ferramentas de código aberto podem processar a retenção e a recolha de lixo, incluindo as seguintes:

O exemplo seguinte usa backup-warden, em que o diretório backup está montado numa localização de cópia de segurança que usa o FUSE do Cloud Storage:

# Add --delete option to actually delete the backups, as is it only shows what would be deleted (dry-run)
backup-warden -p backup \
    --hourly 24 \
    --daily 7 \
    --weekly 5 \
    --monthly always \
    --yearly always \
    --prefer-recent

Atualize o manifesto do JobSet da carga de trabalho

Atualize o manifesto JobSet do seu trabalho para incluir o volume de pontos de verificação em grande escala. Os detalhes dependem da sua carga de trabalho.

Por exemplo, para estender o JobSet de exemplo de Implemente multislices de TPU no GKE, siga estes passos:

  1. Adicione as seguintes linhas ao contentor jax-tpu.

    volumeMounts:
- name: checkpoint
  mountPath: CHECKPOINT_DIR

    Substitua CHECKPOINT_DIR pelo caminho para o diretório do ponto de verificação. Esta é a localização onde o replicator.yaml é gerado e a funcionalidade de pontos de controlo de vários níveis executa a operação de guardar pontos de controlo. Para mais informações, consulte o artigo Integre a verificação de vários níveis na sua aplicação.

  2. Adicione as seguintes linhas ao campo spec.template.spec da especificação do trabalho.

    volumes:
- name: checkpoint
  csi:
    driver: multitier-checkpoint.csi.storage.gke.io

Integre a verificação de pontos de controlo de vários níveis na sua aplicação

Para partilhar informações sobre localizações de pontos de verificação e disponibilidade de replicação, modifique a sua aplicação para usar o seguinte protocolo para comunicar com a verificação de pontos de vários níveis.

Arranque

Esta secção descreve os passos iniciais que a aplicação tem de seguir para interagir com a validação de pontos de verificação de vários níveis.

O Replicator é um componente essencial da criação de pontos de verificação de vários níveis, executado em todos os nós como parte do controlador CSI. O replicador gere a replicação de pontos de verificação em todos os níveis de armazenamento, desde o disco RAM local aos nós de pares e ao armazenamento externo, como o Cloud Storage.

O ficheiro replicator.yaml funciona como um plano de controlo dinâmico entre a tarefa de preparação de ML (código da framework) e o componente Replicator. A sua aplicação de ML gera este ficheiro programaticamente no volume local (RAMDisk), que é acessível tanto pela tarefa de preparação como pelo serviço Replicator. Este manifesto permite que a framework de ML forneça instruções de configuração de tempo de execução e gestão do ciclo de vida ao replicador, distintas dos parâmetros de infraestrutura estáticos (por exemplo, frequência de carregamento do Cloud Storage) definidos durante a configuração do back-end.

Para ver um exemplo concreto desta interação, consulte:

A sua aplicação deve seguir os seguintes passos durante o arranque:

  1. Aguarde até que o ficheiro replicator.yaml esteja ausente, o que indica que o replicador está pronto para ser configurado pela sua aplicação. O ficheiro replicator.yaml é gerado na localização CHECKPOINT_DIR que configurou na secção Atualize o manifesto do JobSet da carga de trabalho.

    Quando a tarefa de preparação de modelo é criada pela primeira vez, o ficheiro replicator.yaml não existe e a sua aplicação pode prosseguir imediatamente. No entanto, se o trabalho tiver sido reiniciado (por exemplo, devido a uma falha ou intervenção manual), o sistema ainda pode estar a processar a instância do trabalho anterior, e o replicator.yaml dessa instância ainda pode estar presente no volume local.

  2. A sua aplicação ou tarefa de ML cria o ficheiro replicator.yaml com a configuração semelhante à seguinte.

    Orbax

    job-name: orbax
framework: orbax
assume-data-parallelism: 3
node-rank: 0
nodes: 32
peer-ranks: [1, 16] or peers-per-node: 2
backup-interval-minutes: 30

    PyTorch

    job-name: nemo
framework: pytorch.distributed
node-rank: 0
nodes: 32
peer-ranks: [1, 16] or peers-per-node: 2
backup-interval-minutes: 30

    Esta configuração de exemplo tem os seguintes campos:

    • name: o nome da tarefa de preparação.
    • framework: a framework de AA usada pela tarefa de preparação.
    • node-rank: o identificador exclusivo do nó atual na tarefa de preparação distribuída. Representa a classificação do nó do nó que está a criar este ficheiro. Cada nó que participa na execução tem a sua própria classificação.
    • nodes: o número total de nós que participam na tarefa de preparação distribuída. Este valor é proveniente dos metadados do podcast. A tarefa de preparação de ML também pode ver este valor.
    • peer-ranks ou peers-per-node: duas formas alternativas de especificar a topologia de replicação. Apenas um destes dois parâmetros deve estar presente.
      • peer-ranks: classificações explícitas dos nós de pares para os quais os dados de pontos de verificação do nó atual devem ser replicados. Isto dá um controlo detalhado sobre os nós específicos que funcionam como parceiros de replicação.
      • peers-per-node: o número de nós pares por nó que o replicador deve selecionar automaticamente para replicação.
    • backup-interval-minutes: a frequência, em minutos, com que é feita uma cópia de segurança dos pontos de verificação no Cloud Storage. Recomendamos que defina este valor para 30 minutos ou mais.

  3. Aguarde até que o novo ficheiro replicator.yaml seja eliminado pelo sistema. Isto indica que o replicador foi reiniciado e efetuou a limpeza. Este passo permite-lhe evitar ficheiros desatualizados ou temporários no volume local quando a sua aplicação executa os passos na secção seguinte.

Restaurar a partir do último posto de controlo conhecido (LKG)

  1. Após a inicialização do replicador, a funcionalidade Multi-Tier Checkpointing cria um link simbólico por trabalhador de TPU ou GPU. Estes links simbólicos são criados no mesmo volume local montado que o ficheiro replicator.yaml, onde a tarefa guarda os pontos de verificação.

    Os links simbólicos têm o formato <job-name>-s{step}-n<node-rank>-w<worker-index>.restore.

  2. Restaure cada trabalhador a partir do ficheiro .restore correspondente. Para ver um exemplo, consulte o exemplo do gestor de pontos de verificação replicados do Orbax na secção seguinte.

Guarde o ponto de controlo

A sua aplicação executa estes passos várias vezes enquanto a tarefa de preparação progride. As operações de guardar ocorrem na localização CHECKPOINT_DIR que configurou em Atualize o manifesto JobSet da carga de trabalho.

Orbax

Crie pontos de verificação Orbax. O diretório tem o nome do número do passo. O replicador deteta o diretório de pontos de verificação recém-criado, faz a replicação ou a cópia de segurança conforme necessário e faz a limpeza automaticamente.

Para mais informações sobre como usar o gestor de pontos de verificação do replicador Orbax, consulte o ficheiro MaxtTest checkpointing. Para ver um exemplo de interação do serviço de replicador, consulte o ficheiro MaxText max_utils.

PyTorch

Use InClusterLocalCheckpointIO como um pytorch_lightning.CheckpointIO personalizado para ativar a criação de pontos de verificação distribuídos corretos com o armazenamento local. O seguinte comando de exemplo ativa a criação de pontos de verificação de vários níveis através de uma implementação de referência criada na framework NVIDIA NeMo:

torchrun train.py <other_train_flags> \
    --local-ckpt-dir=CHECKPOINT_DIR \
    --local-ckpt-interval=20 \
    --job-name=JOB_NAME \
    --enable-high-scale-ckpt

Substitua o seguinte:

  • CHECKPOINT_DIR: o caminho para o diretório do ponto de verificação.
  • JOB_NAME: o nome da carga de trabalho da tarefa de preparação.

Atualizações de clusters

Para atualizações de clusters, pode eliminar e recriar o objeto CheckpointConfiguration antes ou depois da atualização. Esta ação é necessária porque o conjunto de daemon do controlador de criação de pontos de verificação de nós, implementado dinamicamente por este objeto, não é atualizado automaticamente.

Se quiser especificamente manter a especificação do daemonset, não tem de fazer nada.

Resolver problemas

Esta secção fornece orientações de resolução de problemas para resolver problemas com a validação de pontos de controlo de vários níveis. Para a resolução de problemas gerais de armazenamento, consulte o artigo Resolução de problemas do Cloud Storage no GKE.

A verificação de pontos de vários níveis não está ativada

O seguinte erro indica que a verificação de pontos de controlo de vários níveis não está ativada no seu cluster:

error: unable to recognize "checkpoint.yaml": no matches for kind "CheckpointConfiguration" in version "checkpointing.gke.io/v1"

Pode encontrar este erro depois de executar kubectl apply -f checkpoint.yaml no passo Crie uma CheckpointConfiguration.

Para resolver este problema, verifique se ativou a funcionalidade Multi-Tier Checkpointing no seu cluster com o seguinte comando:

gcloud container clusters describe CLUSTER_NAME \
    --project PROJECT_ID
    --location CLUSTER_LOCATION

Se a funcionalidade de pontos de verificação de vários níveis estiver ativada, o resultado deve ser semelhante ao seguinte:

addonsConfig:
  gcePersistentDiskCsiDriverConfig:
    enabled: true
  gcsFuseCsiDriverConfig:
    enabled: true
  highScaleCheckpointingConfig:
    enabled: true
  kubernetesDashboard:
    disabled: true
  networkPolicyConfig:
    disabled: true

Se a funcionalidade Multi-Tier Checkpointing não estiver ativada, atualize o cluster para a ativar.

O controlador CSI de pontos de verificação de vários níveis não consegue montar volumes

Pode encontrar este problema se o controlador CSI não conseguir montar o volume do Cloud Storage. Podem existir várias linhas semelhantes a esta.

kubectl get pod -n gke-managed-checkpointing
NAME                                                          READY   STATUS     RESTARTS   AGE
multitier-driver-14694e4d-774f-4104-8bba-f0bd82fd7557-5vxr9   0/5     Init:0/1   0          6m32s

Para resolver este problema, verifique os eventos do pod do controlador CSI, conforme mostrado no seguinte exemplo:

kubectl describe pod multitier-driver-14694e4d-774f-4104-8bba-f0bd82fd7557-5vxr9 -n gke-managed-checkpointing

Events:
  Type     Reason       Age                 From               Message
  ----     ------       ----                ----               -------
  Normal   Scheduled    17m                 default-scheduler  Successfully assigned gke-managed-checkpointing/multitier-driver-14694e4d-774f-4104-8bba-f0bd82fd7557-5vxr9 to gke-my-cluster-default-pool-353c773f-6d8q
  Warning  FailedMount  82s (x16 over 17m)  kubelet            MountVolume.SetUp failed for volume "gcs" : rpc error: code = PermissionDenied desc = failed to get GCS bucket "checkpointing-test-bucket": googleapi: Error 403: Caller does not have storage.objects.list access to the Google Cloud Storage bucket. Permission 'storage.objects.list' denied on resource (or it may not exist)., forbidden

Se o problema ocorrer devido a um erro PermissionDenied do contentor do Cloud Storage, conforme mostrado no exemplo, pode resolver o problema configurando corretamente as autorizações.

O que se segue?