Práticas recomendadas de armazenamento para cargas de trabalho de IA/ML em VMs de TPU

Para maximizar o desempenho e a eficiência de custo das suas cargas de trabalho de IA/ML em VMs de TPU, selecione e configure a solução de armazenamento certa para sua carga de trabalho. Ao remover gargalos de E/S, você reduz o tempo de inatividade dos aceleradores de TPU, o que diminui o tempo e os custos de treinamento.

Este documento oferece recomendações de armazenamento específicas para cargas de trabalho e práticas recomendadas de otimização para treinamento, checkpointing, disponibilização e armazenamento em cache em VMs de TPU. Antes de aplicar essas práticas, consulte as Opções de armazenamento disponíveis para dados de TPU. Neste documento, presumimos que você tenha familiaridade com as VMs da TPU e experiência básica no provisionamento de recursos do Cloud Storage.

Orientações específicas para a carga de trabalho

A tabela a seguir fornece recomendações de armazenamento, listadas por ordem de preferência, para diferentes cargas de trabalho:

Carga de trabalho Recomendação Otimização e ferramentas (se aplicável)
Conjuntos de dados de treinamento, incluindo preparação de dados
  1. Use buckets regionais do Cloud Storage com o Rapid Cache e ajuste do lado do cliente (como leituras antecipadas) para ter o menor custo.
  2. Use o Rapid Bucket do Cloud Storage para ter o melhor desempenho e escalonamento.
  3. Use o Managed Lustre se você estiver padronizando o Lustre (armazenamento do sistema de arquivos) devido à alta simultaneidade de metadados ou arquivos pequenos (menos de 1 MB).
  • Para buckets regionais do Cloud Storage, ative o namespace hierárquico.
  • Para buckets regionais do Cloud Storage ou Rapid Bucket, use o Cloud Storage FUSE e o gRPC. Para o GKE, use o perfil gcsfusecsi-training.
  • Para o Managed Lustre, use o nível dinâmico para reduzir custos e otimizar automaticamente a performance.
Criação de checkpoints e pesos de aprendizagem por reforço
  1. Use o Managed Lustre para alto desempenho sem ajuste extenso para requisitos de baixa latência (menos de 1 ms), como pontos de verificação síncronos e propagação de peso de alta velocidade para aprendizado por reforço.
  2. Use o Rapid Bucket do Cloud Storage para ter a maior capacidade de processamento para checkpoints assíncronos e de várias camadas e para paralelizar operações de checkpoint no armazenamento. O Rapid Bucket oferece a alta capacidade de processamento necessária para essas cargas de trabalho.
Para o Rapid Bucket, use o perfil gcsfusecsi-checkpointing.
Armazenamento e download de modelos
  1. Use o Hyperdisk ML para armazenamento em cache de modelos pequenos e até 2.500 instâncias.
  2. Use a mesma solução de armazenamento que você usa para conjuntos de dados de treinamento (Managed Lustre, Rapid Bucket ou buckets regionais) se quiser padronizar uma única solução de armazenamento.
Para fazer o download de modelos, use o GKE Run:ai Model Streamer ou o Cloud Storage FUSE com um ponto de montagem separado usando o perfil gcsfusecsi-serving.
Descarregamento de cache de chave-valor (KV)
  1. Sempre use a memória da CPU do host (RAM) como o nível principal para liberar VRAM sem adicionar latência significativa.
  2. Se a memória da CPU for insuficiente, use o Managed Lustre como uma segunda camada paralela e de baixa latência na hierarquia de cache para atender às necessidades de latência ultrabaixa (menos de 1 ms) de descarregamento de cache de chave-valor com alta capacidade de processamento (até 10 TB/s).

Otimizações do Cloud Storage

As seções a seguir descrevem as práticas recomendadas para otimizar o desempenho ao usar o Cloud Storage com VMs da TPU.

Ativar o namespace hierárquico para otimização de metadados

Para melhorar o desempenho dos metadados, ative o namespace hierárquico ao criar buckets regionais para cargas de trabalho de IA/ML. A performance de metadados se refere à rapidez com que o Cloud Storage pode processar operações que envolvem pesquisa, listagem ou modificação de caminhos e pastas de objetos, em vez de ler ou gravar o conteúdo dos arquivos.

Em buckets sem namespace hierárquico ativado, as pastas não existem como recursos reais, mas são pastas simuladas representadas por prefixos de nome de objeto delimitados por barras (/). Isso torna operações como listar o conteúdo do diretório ou renomear diretórios lentas porque o sistema precisa verificar todos os objetos com esse prefixo. O namespace hierárquico oferece uma estrutura de sistema de arquivos real, o que é importante para cargas de trabalho de IA/ML por vários motivos:

  • Renomeações atômicas de diretório: os frameworks de ML usam renomeações de diretório para finalizar checkpoints. O namespace hierárquico oferece suporte a renomeações atômicas, garantindo que os checkpoints sejam finalizados rapidamente.
  • QPS inicial mais alta: o namespace hierárquico oferece suporte a até oito vezes mais consultas iniciais por segundo (QPS) para leituras e gravações em comparação com buckets sem namespace hierárquico ativado. Isso evita gargalos quando muitas TPUs acessam o armazenamento simultaneamente.
  • Operações eficientes no nível da pasta: encontrar e listar arquivos em diretórios específicos é muito mais rápido, reduzindo os tempos de resposta durante o treinamento e o carregamento de dados.

Os buckets zonais, oferecidos pelo Rapid Bucket, usam o namespace hierárquico por padrão. Para mais informações, consulte Visão geral do namespace hierárquico.

Usar o Cloud Storage FUSE com os perfis adequados

O Cloud Storage FUSE é um adaptador FUSE que permite ativar buckets como um sistema de arquivos local. Ao usar o Google Kubernetes Engine, recomendamos usar o driver CSI do Cloud Storage FUSE e os perfis do Cloud Storage FUSE para automatizar o ajuste de desempenho.

Para mais informações sobre as práticas recomendadas de uso do Cloud Storage FUSE, consulte Práticas recomendadas de ajuste de desempenho.

Personalizar o disco de inicialização da VM de TPU

É possível personalizar o ambiente do SO convidado em uma VM de TPU usando scripts de inicialização ou criando imagens personalizadas. Personalizar o disco de inicialização é útil para os seguintes cenários:

  • Pré-carregamento de software e bibliotecas: instale frameworks de ML, dependências ou software personalizado específicos para reduzir o tempo de inicialização da VM e garantir ambientes consistentes.
  • Usar distribuições de SO não padrão: use uma distribuição ou versão do SO não incluída na lista gerenciada pelo Google.
  • Aplicar configurações de segurança e monitoramento: aplique configurações de segurança personalizadas, instale agentes de monitoramento ou defina variáveis de ambiente.

No entanto, a recuperação do disco de inicialização para VMs de TPU é limitada. Não é possível separar ou criar um snapshot do disco de inicialização para reparo off-line. Por isso, tenha cuidado ao fazer mudanças que afetam o processo de inicialização. Seguindo essas práticas recomendadas, você pode reduzir o risco de falhas de inicialização ao personalizar os ambientes de VM da TPU.

Ao personalizar o disco de inicialização, lembre-se dos seguintes princípios:

  • Minimizar modificações no disco de inicialização: sempre que possível, instale aplicativos e armazene dados em volumes de Persistent Disk ou Hyperdisk em vez de modificar muito o disco de inicialização.

  • Use UUIDs para montagem: ao adicionar entradas ao arquivo /etc/fstab, sempre use UUIDs para identificar discos e partições (UUID=...) em vez de nomes de dispositivos, como /dev/sdb1. Não há garantia de que os nomes de dispositivos gerados automaticamente sejam estáveis em reinicializações.

Siga estas recomendações para reduzir o risco de falhas de inicialização ao fazer mudanças no sistema:

  • Tratamento de erros: implemente uma verificação de erros robusta e modos de falha graceful nos seus scripts. Registre mensagens detalhadas no console serial e no Cloud Logging para ajudar na depuração.

  • Dependências críticas: tenha muito cuidado ao modificar arquivos essenciais para a inicialização, como o arquivo /etc/fstab, configurações de rede ou configurações do carregador de inicialização. Um erro de sintaxe ou uma entrada incorreta pode impedir a inicialização da VM.

  • Discos secundários: se o script depender de discos secundários, verifique se ele processa casos em que o disco pode não estar presente ou leva mais tempo para ser anexado do que o esperado. Evite tornar o processo de inicialização dependente de montagens de disco secundárias, a menos que seja absolutamente necessário.

    Confira abaixo exemplos de entradas /etc/fstab recomendadas e não recomendadas para montar um disco secundário:

    • Recomendado: UUID=a1b2c3d4-e5f6-7890-1234-567890abcdef /mnt/mydata ext4 defaults,nofail 0 2
    • Não recomendado: /dev/sdb1 /mnt/mydata ext4 defaults 0 2

    Usar a opção nofail pode impedir que o sistema pare se o disco não for encontrado, mas garanta que seu aplicativo possa lidar com o ponto de montagem indisponível.

  • Gerenciamento de pacotes: tenha cuidado ao adicionar repositórios de terceiros. Verifique se eles são confiáveis e compatíveis com a imagem do SO de base. Entenda as dependências de todos os pacotes instalados e o possível impacto deles nas bibliotecas do sistema.

  • Espaço em disco: monitora o uso do disco de inicialização. O registro em log extenso ou instalações de software grandes podem preencher o disco de inicialização, impedindo que a VM seja iniciada.

  • Geração de registros: configure seus aplicativos e scripts para gerar registros detalhados no console serial, já que essa é a principal ferramenta para diagnosticar problemas de inicialização em VMs de TPU.

Planejar a capacidade de armazenamento

É importante planejar a quantidade de capacidade de armazenamento que sua carga de trabalho vai precisar para usar totalmente os aceleradores. Isso inclui a capacidade de armazenamento e a largura de banda do checkpoint.

Estimar o armazenamento

Use as estimativas a seguir como ponto de partida para seus requisitos de armazenamento:

Tipo de carga de trabalho Armazenamento de conjuntos de dados Armazenamento de checkpoints
Pré-treinamento de LLM 2 TB por TPU 200 GB por TPU
Treinamento multimodal 12 TB por TPU 1 TB por TPU
Inferência 1 TB por TPU 1 GB por TPU

Estimar a largura de banda do checkpoint

Você pode estimar a largura de banda mínima de checkpointing necessária para cargas de trabalho de treinamento usando a seguinte fórmula. Para leituras de dados, várias execuções de treinamento ou treinamento e inferência, aumente proporcionalmente os requisitos estimados de largura de banda.

  1. Tamanho do checkpoint: número de parâmetros × bytes por parâmetro (aproximadamente 12 a 16 bytes por parâmetro para FP16 + estado do otimizador). Adicione um buffer (aproximadamente 3 vezes) para estados do otimizador e diferentes precisões.
  2. Intervalo de checkpoint: a frequência com que você salva um checkpoint (por exemplo, a cada 15 minutos).
  3. Largura de banda necessária: tamanho do ponto de verificação ÷ intervalo do ponto de verificação.

O exemplo a seguir mostra como estimar a largura de banda mínima de checkpointing para o Qwen3-72B:

  1. Tamanho do checkpoint: 72 bilhões de parâmetros × 12 bytes ≈ 864 GB por checkpoint. Com buffer, 3 × 864 GB ≈ 2,5 TB.
  2. Intervalo de checkpoint: 2 minutos = 120 segundos.
  3. Largura de banda necessária: 2,5 TB ÷ 120 segundos ≈ 20 GBps.

Receitas de referência

Para exemplos de configurações de armazenamento para hardware e cargas de trabalho específicos, consulte as seguintes receitas:

Cotas e limites de largura de banda

A largura de banda das ofertas do Cloud Storage e do Compute Engine é limitada por cotas padrão. Se você exceder uma cota, as solicitações de entrada e saída poderão ser limitadas.

Para informações sobre cotas do Cloud Storage e como solicitar aumentos, consulte Cotas e limites na documentação do Cloud Storage. Para informações sobre cotas do Compute Engine para hiperdisco e Persistent Disk, consulte Cotas de disco.

A seguir