Treinar modelos de machine learning em grande escala no GKE com o Multi-Tier Checkpointing

Nesta página, mostramos como usar o checkpointing de várias camadas para armazenar e gerenciar de maneira confiável checkpoints durante o treinamento de modelo de machine learning no GKE. O armazenamento e o gerenciamento de pontos de verificação são cruciais para jobs de treinamento em grande escala, definidos como aqueles que utilizam mais de milhares de nós. As interrupções nesses jobs de grande escala são frequentes (potencialmente a cada hora), e a recuperação delas pode ser lenta.

Vantagens

O uso do checkpointing de várias camadas oferece os seguintes benefícios:

  • Gerenciamento de dados de ponto de verificação totalmente orquestrado, incluindo backups, replicação e restauração automática para as seguintes cargas de trabalho:
  • Recuperação rápida de jobs de treinamento de um checkpoint armazenado no nó local. Também é possível fazer a recuperação usando checkpoints armazenados em outro nó no cluster de treinamento.
  • Restauração rápida de jobs de treinamento de um checkpoint armazenado em um backup do Cloud Storage em cenários de pior caso, em que não há checkpoints no cluster.

Antes de começar

Antes de começar, verifique se você realizou as tarefas a seguir:

  • Ativar a API Google Kubernetes Engine.
    • Ativar a API Google Kubernetes Engine
  • Se você quiser usar a CLI do Google Cloud para essa tarefa, instale e inicialize a gcloud CLI. Se você instalou a gcloud CLI anteriormente, instale a versão mais recente executando gcloud components update.

Requisitos

O checkpointing de vários níveis requer o cluster do GKE versão 1.32.4-gke.1415000 ou posterior.

Limitações

  • Os clusters do Autopilot não são compatíveis.

Configurar nós do GKE para usar o checkpointing de vários níveis

Nesta seção, explicamos como configurar nós do GKE em clusters novos e existentes.

Configurar nós em um novo cluster

  1. Crie um cluster com o checkpointing de várias camadas, o driver CSI do Cloud Storage FUSE e a federação de identidade da carga de trabalho para o GKE ativados. Se você usar frações de TPU para sua carga de trabalho de machine learning, será necessário ajustar o comando de criação do cluster para incluir a configuração de um pool de nós de fração 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 fica a capacidade da TPU.
    • PROJECT_ID: o ID do projeto do Google Cloud .
    • CLUSTER_VERSION: a versão do cluster. A versão mínima compatível é 1.32.4-gke.1415000.
    • CLUSTER_LOCATION: a região em que você 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 checkpointing de várias camadas. Para treinamento em grande escala, recomendamos usar pelo menos e2-standard-4 máquinas. Você não vai usar esse tipo de máquina para treinamento de modelo. Em vez disso, vai criar pools de nós separados para essa finalidade, geralmente usando famílias de VMs otimizadas para aceleradores.
    • NUM_NODES: o número de nós a serem criados em cada uma das zonas do cluster.

Configurar nós em um cluster atual

Para usar o checkpointing de vários níveis com um cluster atual, ative-o junto com o driver CSI do Cloud Storage FUSE e a federação de identidade da carga de trabalho para o GKE com os seguintes comandos. A versão do cluster precisa ser posterior a 1.32.3-gke.1170000.

  1. Ative a federação de identidade da carga 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 ID do projeto do Google Cloud .
    • CLUSTER_LOCATION: a região do cluster.

  2. Ative o checkpointing de várias camadas e o driver CSI do Cloud Storage FUSE:

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

Configurar permissões para usar o checkpointing de vários níveis

Nesta seção, explicamos como configurar permissões para usar o checkpointing de vários níveis.

Conceder acesso a buckets do Cloud Storage

Os volumes efêmeros usados pelo driver CSI de checkpointing de vários níveis precisam usar buckets do Cloud Storage.

Para armazenar checkpoints em um bucket do Cloud Storage, o Multi-Tier Checkpointing precisa de acesso ao bucket. Conceda o papel de usuário de objetos do Storage (roles/storage.objectUser) no bucket à conta de serviço do Kubernetes para o checkpointing 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) Conceder acesso à conta de serviço padrão do Compute Engine

Se as instâncias do Compute Engine precisarem de acesso de leitura ao bucket do Cloud Storage, conceda o papel do IAM Leitor de objetos do Storage (roles/storage.objectViewer) à conta de serviço padrão 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

Implante o controlador do JobSet

O controlador JobSet é responsável por gerenciar os jobs em lote que executam o treinamento de modelo no GKE, e a alocação de recursos é ajustada para processar a carga de trabalho de maneira eficiente. Verifique se o iniciador de jobs de treinamento implanta e usa o JobSet.

Para aumentar a solicitação de memória para 1 Gi, o limite de memória para 2 Gi e a solicitação de CPU para 1 no contêiner do gerenciador na implantaçã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"}}}]'

Inicializar o driver CSI de checkpointing de vários níveis

Esta seção descreve como inicializar o driver CSI de checkpointing de vários níveis em nós onde suas cargas de trabalho serão executadas. O driver CSI é responsável por processar o armazenamento e o gerenciamento de checkpoints durante o processo de treinamento de modelo.

Criar um CheckpointConfiguration

Um CheckpointConfiguration é um recurso personalizado do Kubernetes que especifica propriedades para implantação do driver CSI de checkpointing de várias camadas. Esse recurso tem escopo 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-cache:max-size-mb:-1
    - file-cache:cache-file-for-range-read:true
    - file-system:kernel-list-cache-ttl-secs:0
    - file-cache:enable-parallel-downloads:true
    - read_ahead_kb=1024
    - write:enable-streaming-writes:true

    Substitua:

    • MTC_CONFIG_NAME: o nome do CheckpointConfiguration. Esse nome é global para o cluster e não é específico do job.
    • GCS_BUCKET: o nome do bucket do Cloud Storage em que você vai armazenar os dados de checkpoint. Use o bucket que você configurou na etapa Configurar um bucket do Cloud Storage com permissões.

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

      Para mais informações sobre como executar cargas de trabalho distribuídas em GPUs com o GKE, consulte Como executar GPUs com várias instâncias. Para TPUs, consulte Criar o pool de nós de fração da TPU.

    • TOLERATION_KEY: esse campo permite que o driver CSI seja programado em nós com taints correspondentes. Para mais informações sobre como os taints funcionam em diferentes tipos de aceleradores, consulte estas páginas:

    • IN_MEMORY_VOLUME_SIZE: o tamanho do cache de checkpointing na memória. Especifique a quantidade e a unidade (por exemplo, 200 Gi).Esse valor precisa ser:

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

  2. Aplique o manifesto:

    kubectl apply -f checkpoint.yaml

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

    kubectl get pod -n gke-managed-checkpointing

    A resposta deverá ser semelhante a: Haverá 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

Desinstalar o driver CSI de checkpointing de várias camadas

Se você quiser cancelar a implantação do driver CSI de checkpointing de várias camadas, exclua o recurso CheckpointConfiguration. O controlador de checkpointing de várias camadas remove o driver CSI dos nós. Isso remove os discos RAM e libera memória para outras cargas de trabalho. Exemplo:

kubectl delete -f checkpoint.yaml

Gerenciar a retenção de dados e a coleta de lixo para backups do Cloud Storage

Você é responsável por implementar políticas de retenção para os backups de pontos de verificação do Cloud Storage. O checkpointing de vários níveis só grava backups de pontos de verificação no Cloud Storage e nunca os modifica ou exclui.

Muitas ferramentas de código aberto podem lidar com retenção e coleta de lixo, incluindo:

O exemplo a seguir usa backup-warden em que o diretório backup é ativado em um local de backup que usa o Cloud Storage FUSE:

# 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

Atualizar o manifesto do JobSet da carga de trabalho

Atualize o manifesto do JobSet para incluir o volume de checkpoints em grande escala. Os detalhes dependem da sua carga de trabalho.

Por exemplo, para estender o JobSet de amostra em Implantar Multislices de TPU no GKE, siga estas etapas:

  1. Adicione as seguintes linhas ao contêiner jax-tpu.

    volumeMounts:
- name: checkpoint
  mountPath: CHECKPOINT_DIR

    Substitua CHECKPOINT_DIR pelo caminho para o diretório de pontos de verificação. É o local em que o replicator.yaml é gerado e o Multi-Tier Checkpointing realiza a operação de salvamento do ponto de verificação. Para mais informações, consulte Integrar o checkpointing de várias camadas ao seu aplicativo.

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

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

Integrar o checkpointing de vários níveis ao seu aplicativo

Para compartilhar informações sobre locais de checkpoint e prontidão de replicação, modifique seu aplicativo para usar o seguinte protocolo de comunicação com o Multi-Tier Checkpointing.

Inicialização

Esta seção descreve as etapas iniciais que o aplicativo precisa seguir para interagir com o checkpointing de várias camadas.

O Replicator é um componente principal do Multi-Tier Checkpointing, executado em todos os nós como parte do driver CSI. O replicador gerencia a replicação de pontos de verificação em todos os níveis de armazenamento, do disco RAM local aos nós de mesmo nível e ao armazenamento externo, como o Cloud Storage.

O arquivo replicator.yaml atua como um plano de controle dinâmico entre seu trabalho de treinamento de ML (código do framework) e o componente Replicator. O aplicativo de ML gera esse arquivo de maneira programática no volume local (RAMDisk), que pode ser acessado pelo job de treinamento e pelo serviço Replicator. Esse manifesto permite que o framework de ML forneça instruções de configuração de tempo de execução e gerenciamento de ciclo de vida ao replicador, diferentes dos parâmetros de infraestrutura estática (por exemplo, frequência de upload do Cloud Storage) definidos durante a configuração do back-end.

Para um exemplo concreto dessa interação, consulte:

Seu aplicativo precisa seguir estas etapas durante a inicialização:

  1. Aguarde até que o arquivo replicator.yaml não esteja mais presente, o que indica que o Replicator está pronto para ser configurado pelo aplicativo. O arquivo replicator.yaml é gerado no local CHECKPOINT_DIR que você configurou na seção Atualizar o manifesto do JobSet da carga de trabalho.

    Quando o job treinamento de modelo é criado, o arquivo replicator.yaml não existe, e o aplicativo pode continuar imediatamente. No entanto, se o job foi reiniciado (por exemplo, devido a uma falha ou intervenção manual), o sistema ainda pode estar processando a instância anterior do job, e o replicator.yaml dessa instância ainda pode estar presente no volume local.

  2. Seu aplicativo ou job de ML cria o arquivo 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 do job de treinamento.
    • framework: o framework de ML usado pelo job de treinamento.
    • node-rank: o identificador exclusivo do nó atual no job de treinamento distribuído. Isso representa a classificação do nó que criou este arquivo. Cada nó participante da execução terá uma classificação própria.
    • nodes: o número total de nós que participam do job de treinamento distribuído. Esse valor vem dos metadados do pod. O trabalho de treinamento de ML também pode conferir esse valor.
    • peer-ranks ou peers-per-node: duas maneiras alternativas de especificar a topologia de replicação. Apenas um desses dois parâmetros deve estar presente.
      • peer-ranks: classificações explícitas de nós semelhantes em que os dados de ponto de verificação do nó atual precisam ser replicados. Isso oferece controle refinado sobre quais nós específicos servem como parceiros de replicação.
      • peers-per-node: o número de nós iguais por nó que o Replicator seleciona automaticamente para replicação.
    • backup-interval-minutes: a frequência, em minutos, com que os checkpoints são armazenados em backup no Cloud Storage. Recomendamos que você defina esse valor como 30 minutos ou mais.

  3. Aguarde até que o novo arquivo replicator.yaml seja excluído pelo sistema. Isso indica que o replicador foi reiniciado e executou a limpeza. Essa etapa evita arquivos temporários ou desatualizados no volume local quando o aplicativo executa as etapas na próxima seção.

Restaurar do último ponto de verificação válido conhecido (LKG)

  1. Depois que o replicador é inicializado, o checkpointing de várias camadas cria um link simbólico por worker de TPU ou GPU. Esses links simbólicos são criados no mesmo volume local montado que o arquivo replicator.yaml, em que o job salva 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 worker do arquivo .restore correspondente. Por exemplo, confira o exemplo do gerenciador de pontos de verificação replicados do Orbax na próxima seção.

Salve o checkpoint.

O aplicativo executa essas etapas várias vezes enquanto o job de treinamento avança. As operações de salvamento acontecem no local CHECKPOINT_DIR que você configurou em Atualizar o manifesto do JobSet da carga de trabalho.

Orbax

Crie checkpoints do Orbax. O diretório é nomeado com o número da etapa. O replicador detecta o diretório de ponto de verificação recém-criado, realiza replicação ou backup conforme necessário e faz a limpeza automaticamente.

Para mais informações sobre como usar o gerenciador de checkpoints do replicador do Orbax, consulte o arquivo MaxtTest checkpointing. Para um exemplo de interação do serviço de replicação, consulte o arquivo max_utils do MaxText.

PyTorch

Use InClusterLocalCheckpointIO como um pytorch_lightning.CheckpointIO personalizado para ativar o checkpointing distribuído correto com armazenamento local. O comando de exemplo a seguir ativa o checkpointing de vários níveis usando uma implementação de referência criada no 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:

  • CHECKPOINT_DIR: o caminho para o diretório de pontos de verificação.
  • JOB_NAME: o nome da carga de trabalho do job de treinamento.

Upgrades de clusters

Para upgrades de cluster, é possível excluir e recriar o objeto CheckpointConfiguration antes ou depois do upgrade. Essa ação é necessária porque o daemonset do driver de checkpointing de nós, implantado dinamicamente por esse objeto, não será atualizado automaticamente.

Se você quiser manter a especificação do daemonset, não precisa fazer nada.

Resolver problemas

Nesta seção, fornecemos orientações para resolver problemas com o checkpointing de vários níveis. Para solução de problemas gerais de armazenamento, consulte Solução de problemas do Cloud Storage no GKE.

Checkpointing de vários níveis não ativado

O erro a seguir indica que o checkpointing de vários níveis não está ativado no cluster:

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

Esse erro pode aparecer depois de executar kubectl apply -f checkpoint.yaml na etapa Criar um CheckpointConfiguration.

Para resolver esse problema, verifique se você ativou o checkpointing de vários níveis no cluster com o seguinte comando:

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

Se o checkpointing de vários níveis estiver ativado, a saída será semelhante a esta:

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

Se o checkpointing de vários níveis não estiver ativado, atualize o cluster para ativar esse recurso.

O driver CSI de checkpointing de várias camadas não consegue montar volumes

Esse problema pode ocorrer se o driver CSI não conseguir montar o volume do Cloud Storage. Pode haver 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 esse problema, verifique os eventos do pod do driver CSI, conforme mostrado no exemplo a seguir:

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 ao erro PermissionDenied do bucket do Cloud Storage, como mostrado no exemplo, resolva o problema configurando as permissões corretamente.

A seguir