Armazenamento

O que “armazenamento” significa em ML em larga escala

No treinamento em larga escala, armazenamento não é apenas “onde os arquivos ficam”. É o sistema de ponta a ponta que alimenta aceleradores (GPUs/TPUs) com dados e preserva o estado de treinamento (checkpoints) de forma confiável e barata. As escolhas de armazenamento afetam diretamente:

  • Vazão (throughput) (você consegue manter as GPUs ocupadas ou está limitado por E/S?)
  • Tolerância a falhas (fault tolerance) (você consegue retomar após falhas de nó ou preempção?)
  • Reprodutibilidade (reproducibility) (você consegue recriar exatamente o mesmo conjunto de dados + código + checkpoint?)
  • Custo (SSD “quente” vs sistema de arquivos de rede vs armazenamento de objetos)
  • Complexidade operacional (fragmentação em shards, cache, consistência, políticas de retenção)

Este artigo se concentra em conjuntos de dados, checkpoints e padrões comuns de armazenamento para treinamento em larga escala.

Hierarquia de armazenamento e por que ela importa

O desempenho do treinamento é dominado pela movimentação de dados. Enquanto o cálculo acontece nos aceleradores, os dados frequentemente se originam bem longe. Um modelo mental útil é uma hierarquia de armazenamento, análoga aos caches de CPU:

  1. HBM / VRAM (na GPU/TPU) – extremamente rápida, muito pequena; mantém ativações/pesos do passo atual.
  2. RAM do host – usada para filas de prefetch, buffers de descompressão e preparação (staging) de lotes.
  3. SSD NVMe local – alta vazão/baixa latência; ideal para cache de shards do conjunto de dados e para escrever fragmentos frequentes de checkpoint.
  4. Sistema de arquivos em rede (NFS/Lustre/GPFS/Filestore/FSx) – semântica POSIX compartilhada; pode ser rápido, mas pode virar gargalo do cluster.
  5. Armazenamento de objetos (object storage) (S3/GCS/Azure Blob) – muito durável e escalável; maior latência e semântica diferente (consistência eventual e operações em “objetos” em vez de POSIX).

Um modo comum de falhar é projetar para correção (por exemplo, “todo mundo lê de um caminho NFS compartilhado”) e descobrir que isso não escala. Outro é projetar para vazão (caches locais em todo lugar) e perder reprodutibilidade ou desperdiçar dinheiro.

O desempenho de armazenamento interage fortemente com Memória e Largura de Banda e com a rede do cluster e tecidos RDMA como Interconexões (NVLink/InfiniBand).

Armazenamento de conjuntos de dados: de dados brutos a shards prontos para treinamento

Conjuntos de dados brutos vs processados

A maioria dos pipelines de treinamento tem pelo menos duas representações do conjunto de dados:

  • Bruto: arquivos originais (imagens, áudio, logs JSON, dumps de texto). Bom para proveniência, não para vazão.
  • Processado / pronto para treinamento: tokenizado, redimensionado, filtrado, aumentado e fragmentado em shards (sharded) em arquivos sequenciais maiores. Otimizado para leituras sequenciais e menos operações de metadados.

Ideia-chave: seu job de treinamento deve ler majoritariamente de forma sequencial a partir de um pequeno número de arquivos grandes, e não de forma aleatória a partir de milhões de arquivos minúsculos.

O problema de “arquivos pequenos”

Armazenar cada exemplo como seu próprio arquivo (por exemplo, 12345.jpg, 12346.jpg, …) causa:

  • Alto overhead de metadados (listagens de diretório, buscas de inode)
  • Muitas leituras aleatórias minúsculas (limitadas por latência)
  • Baixa escalabilidade em sistemas de arquivos compartilhados

Um remédio padrão é fragmentação em shards (sharding): empacotar muitas amostras em cada arquivo de shard.

Formatos comuns de conjunto de dados e por que são usados

Não existe um único melhor formato; a escolha certa depende da modalidade e do ferramental.

  • WebDataset (shards tar): popular para visão + multimodal.
    • Prós: simples, streamable, compatível com armazenamento de objetos, boa vazão sequencial.
    • Contras: atualizações exigem reescrever shards; acesso aleatório é limitado.
  • TFRecord / RecordIO / LMDB: amplamente usados em pipelines de produção.
    • Prós: boas leituras sequenciais; ferramental maduro.
    • Contras: específico do ecossistema; o LMDB tem ressalvas de concorrência em sistemas de arquivos em rede.
  • Parquet / Arrow: comuns para dados tabulares + analytics, às vezes usados para texto.
    • Prós: esquema, poda de colunas, integração com data lakes.
    • Contras: o treinamento frequentemente precisa de streaming orientado a linhas; padrões de descompressão importam.
  • MDS / “conjuntos de dados de streaming” (streaming datasets) (por exemplo, estilo MosaicML): projetados para streaming escalável com indexação e fragmentação em shards.
    • Prós: feitos para streaming multi-worker e retomada.
    • Contras: específico do ecossistema.

Orientação geral:

  • Se você treina via streaming sequencial e quer simplicidade: tar fragmentado em shards (WebDataset) costuma ser uma boa opção.
  • Se sua organização já tem uma stack de data lake: Parquet/Arrow pode integrar melhor, mas avalie tamanhos de row-group e comportamento de descompressão.

Estratégia de fragmentação em shards (regras práticas)

A fragmentação em shards é tanto uma ferramenta de desempenho quanto de paralelismo (parallelism). Bons shards:

  • São grandes o suficiente para amortizar overhead (frequentemente 100MB–2GB por shard)
  • São numerosos o suficiente para distribuir entre workers (milhares de shards para corpora muito grandes)
  • Suportam aleatorização permutando a ordem dos shards e/ou a ordem das amostras dentro do shard

Fluxo de trabalho “regra de bolso”:

  1. Construa shards offline.
  2. Armazene os shards em armazenamento de objetos (durável) e opcionalmente espelhe em um sistema de arquivos compartilhado de alto desempenho.
  3. Use NVMe local para colocar em cache shards “quentes” em cada nó.

Compressão: trade-offs entre largura de banda e CPU

A compressão reduz o custo de armazenamento e a largura de banda de rede, mas custa tempo de CPU e pode virar gargalo.

  • Para texto: zstd costuma oferecer um bom equilíbrio entre velocidade e taxa.
  • Para imagens: JPEG/PNG já são comprimidos; compressão adicional pode não ajudar muito.
  • Para fluxos de tokens: codificação delta + zstd pode ser eficaz.

Na prática:

  • Faça benchmark da vazão de ponta a ponta (amostras/seg), não apenas MB/s.
  • Se a CPU ficar saturada, aumente os workers do dataloader, use codecs mais rápidos ou pré-descomprima para um cache.

Determinismo e versionamento do conjunto de dados

Reprodutibilidade exige mais do que “mesmos arquivos”:

  • Versão exata do conjunto de dados (snapshot imutável, endereçado por conteúdo se possível)
  • Versão do código de pré-processamento (versão do tokenizador, regras de filtragem)
  • Manifesto de shards (shard manifest) (lista de URLs de shards + checksums)
  • Política de seed de amostragem (seed global + derivação por worker)

Um padrão simples e robusto é armazenar um arquivo de manifesto junto aos shards:

{
  "dataset_name": "my-corpus-v3",
  "num_shards": 4096,
  "shards": [
    {"url": "s3://bucket/my-corpus-v3/000000.tar", "sha256": "..."},
    {"url": "s3://bucket/my-corpus-v3/000001.tar", "sha256": "..."}
  ],
  "tokenizer": {"name": "bpe-50k", "version": "1.2.0"},
  "created_at": "2026-01-01T12:00:00Z"
}

Ferramentas como DVC/lakeFS e registries internos de conjunto de dados formalizam isso, mas mesmo um manifesto mínimo melhora dramaticamente a confiabilidade.

Exemplo prático: streaming de shards WebDataset com cache local

Abaixo está um padrão simplificado (detalhes variam por ambiente) que:

  • faz streaming de shards a partir de armazenamento de objetos,
  • embaralha,
  • usa múltiplos workers,
  • e pode ser combinado com um diretório de cache por nó.
import webdataset as wds
from torch.utils.data import DataLoader

urls = "pipe:aws s3 cp s3://my-bucket/my-ds/{000000..004095}.tar -"

dataset = (
    wds.WebDataset(urls, shardshuffle=True)
      .shuffle(10_000)                 # in-sample shuffle buffer
      .decode("pil")
      .to_tuple("jpg", "cls")          # keys in the tar
)

loader = DataLoader(
    dataset,
    batch_size=256,
    num_workers=8,
    prefetch_factor=4,
    persistent_workers=True,
)

for images, labels in loader:
    # training step...
    pass

Para clusters grandes, é comum adicionar:

  • cache local por nó (node-local caching) de shards (evita re-download em retries),
  • limitação de taxa (rate limiting) e políticas de retry,
  • e atribuição de shards (shard assignment) por rank para reduzir leituras duplicadas.

Essas otimizações frequentemente importam tanto quanto melhorias no lado do modelo.

Checkpoints: o que salvar e como salvar com segurança

Um checkpoint é o estado mínimo necessário para retomar o treinamento como se não tivesse havido interrupção. Para aprendizado profundo moderno, isso geralmente inclui:

  • Pesos do modelo
  • Estado do otimizador (frequentemente 2–4× o tamanho do modelo para otimizadores da família Adam)
  • Estado do agendador de taxa de aprendizado (learning rate scheduler)
  • Estado do escalador de gradiente (gradient scaler) (para precisão mista)
  • Estados de RNG (Python/NumPy/framework + RNG por dispositivo)
  • Posição do data loader / época / índice de shard (para evitar duplicação ou omissão de dados)

Para treinamento distribuído em larga escala, o formato e o padrão de checkpointing importam tanto quanto o conteúdo.

Checkpoints completos vs incrementais

  • Checkpoint completo (full checkpoint): estado completo de treinamento em um passo.
    • Prós: mais simples, robusto.
    • Contras: caro em tempo e armazenamento.
  • Checkpoint incremental / delta (incremental / delta checkpoint): armazena mudanças desde o último completo.
    • Prós: menos largura de banda de escrita e menos armazenamento.
    • Contras: mais complexo; recuperação exige cadeia de deltas; risco de corrupção se acumula.

A maioria das equipes usa checkpoints completos periódicos (por exemplo, a cada N minutos ou passos) e gerencia custo via políticas de retenção.

Checkpoints fragmentados/distribuídos

Com paralelismo de dados e fragmentação do otimizador (optimizer sharding) (por exemplo, ZeRO, FSDP), cada rank pode possuir apenas uma fatia dos parâmetros/estado do otimizador. Uma abordagem escalável é:

  • Cada rank grava seu próprio shard (escritas paralelas).
  • Um pequeno arquivo de metadados registra a estrutura global e as localizações dos shards.

Benefícios:

  • Evita reunir o estado completo no rank 0 (o que pode causar OOM).
  • Acelera o checkpointing usando a largura de banda agregada.

Trade-offs:

  • Muitos arquivos por checkpoint (pressão de metadados).
  • Restauração requer world size consistente ou um mecanismo de re-sharding (alguns frameworks suportam isso).

Atomicidade: tornando checkpoints seguros sob falhas

Escritas de checkpoint podem ser interrompidas por preempção, falha de nó ou eventos de disco cheio. Um padrão robusto é:

  1. Escrever em um caminho temporário (por exemplo, ckpt_step_120000.tmp/)
  2. Validar (opcional, mas valioso): checksums, contagem de arquivos, verificações de sanidade de tamanho
  3. Publicar atomicamente:
    • FS POSIX: renomear diretório para ckpt_step_120000/
    • Armazenamento de objetos: escrever por último um pequeno objeto marcador de “commit” (por exemplo, _SUCCESS) e tratar checkpoints sem ele como incompletos

Isso evita retomar a partir de checkpoints parciais.

Checkpointing assíncrono

Checkpointing síncrono pausa o treinamento enquanto escreve. Em escala, essa pausa é dolorosa.

Checkpointing assíncrono encadeia (pipeline) a escrita:

  • Thread de treinamento enfileira o estado (ou referências) para uma thread/processo em segundo plano.
  • Writer em segundo plano faz flush para disco/object store.

Desafios:

  • Copiar tensores enormes para memória do host ainda pode causar travamento.
  • É preciso coordenar com o ciclo de vida de memória do framework.
  • É preciso lidar com “backpressure” (se o armazenamento for lento, a fila async cresce).

Muitas stacks de produção combinam:

  • escritas frequentes e rápidas em NVMe local (async),
  • uploads mais lentos para object store com menor frequência (ou em segundo plano).

Exemplo prático: publicação segura de diretório de checkpoint

Um padrão POSIX simples:

STEP=120000
TMP=/checkpoints/runA/ckpt_${STEP}.tmp
FINAL=/checkpoints/runA/ckpt_${STEP}

mkdir -p "$TMP"
python save_checkpoint.py --out "$TMP"

# Optional validation
test -f "$TMP/metadata.json"

# Atomic publish on POSIX filesystems
mv "$TMP" "$FINAL"

# Update "latest" pointer (symlink if allowed)
ln -sfn "$FINAL" /checkpoints/runA/latest

Em armazenamento de objetos, substitua mv por “escrever o marcador por último” e use prefixes versionados.

Políticas de retenção

A retenção de checkpoints evita custo descontrolado:

  • Manter os N mais recentes (por exemplo, últimos 5)
  • Manter marcos (milestones) (por exemplo, a cada 10k passos ou fim de época)
  • Manter os melhores checkpoints por métrica de validação (para fine-tuning)
  • Fazer coleta de lixo de checkpoints incompletos (sem marcador de commit)

Para grandes modelos fundacionais, o custo de armazenamento pode ser dominado pelo estado do otimizador; se você só precisa de inferência, considere salvar artefatos apenas do modelo e descartar o estado do otimizador após o treinamento.

Padrões de armazenamento para treinamento em larga escala

Padrão 1: Object store como fonte da verdade + cache local por nó

O mais comum em ambientes de nuvem e híbridos.

  • Fonte da verdade: S3/GCS/Azure Blob
  • Cada nó coloca shards em cache no NVMe local
  • O treinamento lê do cache local quando disponível; caso contrário, faz download

Prós:

  • Alta durabilidade, distribuição fácil do conjunto de dados
  • Escala entre regiões/contas (com IAM adequado)
  • Evita gargalos de sistemas de arquivos compartilhados

Contras:

  • Complexidade de invalidação de cache e política de evicção
  • Primeira época pode ser mais lenta (“cache frio”)

Padrão 2: Sistema de arquivos paralelo compartilhado para conjuntos de dados quentes

Comum em HPC.

  • Conjuntos de dados ficam em Lustre/GPFS/FSx for Lustre
  • Muitos nós leem concorrentemente com alta largura de banda agregada

Prós:

  • Semântica POSIX simplifica o ferramental
  • Alta vazão se provisionado corretamente

Contras:

Padrão 3: Checkpointing em dois níveis (local e depois remoto)

Uma abordagem robusta de checkpoint:

  • Escrever checkpoints frequentes em NVMe local (rápido, baixa interrupção).
  • Sincronizar periodicamente para armazenamento remoto durável (object store).
  • Em caso de falha, restaurar do checkpoint durável mais recente; opcionalmente usar o local se o nó tiver sobrevivido.

Isso reduz stalls de treinamento mantendo segurança contra perda de nó.

Padrão 4: Preparar (staging) conjuntos de dados nos nós antes do treinamento começar

Às vezes chamado de prefetch ou data warmup:

  • Um pré-job copia um subconjunto ou todos os shards para cada nó.
  • O treinamento então lê localmente.

Funciona bem quando:

  • O conjunto de dados cabe no armazenamento local agregado
  • Jobs são longos (custo de warmup é amortizado)
  • Rede/object store tem limites de taxa de requisição

Isso pode ser orquestrado via init containers ou sidecars em Kubernetes para ML (Alto Nível).

Padrão 5: Streaming + pré-processamento online (com cuidado)

Para corpora massivos, você pode fazer streaming de dados quase brutos e tokenizar “on the fly”.

Prós:

  • Menos pré-processamento offline
  • Flexibilidade (mudar tokenizador/filtragem sem reescrever o conjunto de dados inteiro)

Contras:

  • Tokenização pode se tornar limitada por CPU
  • Reprodutibilidade é mais difícil a menos que você versione exatamente o container e o código de pré-processamento
  • Mais variância no tempo de passo (jitter)

Frequentemente, o “ponto doce” é semi-processado: armazenar texto em segmentos limpos e então tokenizar online com uma versão fixada do tokenizador, colocando em cache as saídas tokenizadas por shard.

Engenharia de desempenho: diagnosticando gargalos de armazenamento

Sintomas de que você está limitado por E/S

  • Utilização de GPU é baixa com lacunas ociosas frequentes
  • Tempo de passo é instável (jitter)
  • Workers de CPU estão saturados em descompressão/tokenização
  • Picos de egress de rede do object store
  • Muitos arquivos abertos / altas operações de metadados em FS compartilhado

O que medir

No mínimo:

  • Amostras/seg ou tokens/seg de ponta a ponta
  • Tempo do dataloader vs tempo de computação (profiler do framework)
  • Largura de banda de leitura (MB/s) por nó
  • Taxa de acerto do cache (cache hit rate) (se houver cache)
  • Taxa de requisições ao object store e contagem de erros/retries

Um ponto conceitual-chave: desempenho de armazenamento é um pipeline. Melhorar um estágio (por exemplo, discos mais rápidos) não ajuda se outro estágio (por exemplo, CPU de descompressão) estiver limitando.

Confiabilidade, consistência e segurança

Semântica de consistência

  • Sistemas de arquivos POSIX oferecem atomicidade de rename e operações de diretório, o que simplifica padrões de “publicação”.
  • Object stores exigem padrões como:
    • escrever objetos do checkpoint,
    • depois escrever um objeto final de manifesto/marcador,
    • e tratar apenas checkpoints marcados como válidos.

Integridade de dados

Para grandes conjuntos de dados e checkpoints:

  • Armazene checksums para shards e partes de checkpoint.
  • Valide periodicamente (especialmente para arquivos de longa vida).
  • Prefira buckets/prefixes imutáveis e versionados.

Controle de acesso e privacidade

Conjuntos de dados podem conter dados sensíveis. Requisitos práticos frequentemente incluem:

  • Criptografia em repouso e em trânsito
  • IAM granular (princípio do menor privilégio)
  • Logs de auditoria para acesso
  • Políticas de ciclo de vida (apagar após janela de retenção)

Em muitos ambientes regulados, isso não é “bom de ter”.

Escolhas de armazenamento interagem com escolhas de modelo/runtime

  • Frameworks de treinamento distribuído em larga escala (FSDP/ZeRO/TP) moldam formatos de checkpoint e contagens de shards.
  • Precisão mista e escolha de otimizador mudam drasticamente o tamanho do checkpoint (estados do Adam são enormes).
  • Stacks de compilador/runtime (por exemplo, XLA/TensorRT) afetam quais artefatos você salva e como os empacota; veja Compiladores e Runtimes.
  • Formatos de pesos pós-treinamento (por exemplo, pesos quantizados) mudam artefatos de armazenamento e de implantação; veja Quantização.

Recomendações práticas (padrões testados em batalha)

Se você está construindo uma stack de treinamento e precisa de padrões sensatos:

  • Conjunto de dados
    • Converta dados brutos em formatos sequenciais fragmentados em shards (shards tar do WebDataset, TFRecord ou MDS).
    • Mire em centenas de MB até alguns GB por shard.
    • Armazene um manifesto com checksums e versione tudo.
  • Pipeline de entrada
    • Use leituras em streaming, prefetching, e workers persistentes (persistent workers).
    • Adicione cache local por nó se estiver lendo de object store.
    • Faça benchmark com contagens realistas de workers e com augmentation/tokenização.
  • Checkpoints
    • Use checkpoints fragmentados em shards para estados fragmentados de otimizador/modelo.
    • Implemente publicação atômica (atomic publish) (rename ou marcador de commit).
    • Adote uma estratégia de dois níveis: local frequente, remoto durável periódico.
    • Aplique retenção e limpe checkpoints incompletos.
  • Operações
    • Monitore jitter do tempo de passo, largura de banda de leitura, taxa de acerto do cache e contagens de erros/retries.
    • Planeje para falhas: preempção, escritas parciais e correção na retomada.

Armazenamento bem projetado é um multiplicador de força: mantém aceleradores caros ocupados, reduz tempo de treinamento desperdiçado e torna grandes experimentos reprodutíveis e recuperáveis. Em muitos sistemas reais, melhorar o layout do conjunto de dados e a estratégia de checkpoint traz ganhos maiores do que micro-otimizar código do modelo.