Compreendendo o Desafio do Carregamento de Dados em ML
O Amazon Simple Storage Service (S3) é um serviço altamente elástico que se adapta automaticamente conforme a demanda das aplicações, oferecendo a performance de throughput necessária para cargas de trabalho modernas de aprendizagem de máquina. Conectores clientes de alto desempenho, como o Amazon S3 Connector para PyTorch e o Mountpoint para Amazon S3, possibilitam integração nativa ao S3 dentro de pipelines de treinamento, eliminando a necessidade de lidar diretamente com APIs REST do S3.
Este artigo examina técnicas práticas e recomendações para otimizar a throughput em cargas de trabalho de ML que leem dados diretamente de buckets S3 de propósito geral. Muitas das estratégias de otimização apresentadas são aplicáveis a diferentes arquiteturas de armazenamento. A AWS validou essas recomendações através de benchmarks em uma carga de trabalho representativa de visão computacional—uma tarefa de classificação de imagens com dezenas de milhares de arquivos JPEG pequenos.
Gargalos em Pipelines de Treinamento de ML
Embora GPUs desempenhem um papel vital na aceleração de computações de ML, o treinamento é um processo multifacetado com diversos estágios interdependentes—qualquer um deles pode se tornar um gargalo crítico.
Um pipeline típico de treinamento end-to-end passa por quatro etapas recorrentes de alto nível:
- Leitura de amostras de treinamento do armazenamento persistente para a memória
- Pré-processamento de amostras em memória (decodificação, transformação e aumento de dados)
- Atualização de parâmetros do modelo com base em gradientes computados e sincronizados entre GPUs
- Salvamento periódico de checkpoints para tolerância a falhas
A throughput efetiva de qualquer pipeline de ML é limitada pelo seu estágio mais lento. Enquanto a computação dos parâmetros do modelo (etapa 3) é o objetivo final, cargas de trabalho de ML em nuvem enfrentam desafios únicos. Em ambientes de nuvem, onde recursos de computação e armazenamento são tipicamente desacoplados por design, o pipeline de entrada de dados (etapas 1 e 2) frequentemente emerge como gargalo crítico.
Mesmo as GPUs mais modernas não conseguem acelerar o treinamento se ficarem ociosas esperando por dados. Quando ocorre escassez de dados, investimentos adicionais em hardware de computação mais poderoso geram retornos diminutos—uma ineficiência custosa em ambientes de produção. Maximizar a utilização de GPU requer otimização cuidadosa do pipeline de dados para garantir um fluxo contínuo de amostras de treinamento prontas para consumo pelos GPUs.
Entendendo Padrões de Acesso: Leitura Sequencial versus Aleatória
Um dos fatores mais importantes que influenciam a performance do carregamento de dados do S3 é o padrão de acesso aos dados durante o treinamento. A distinção entre leituras sequenciais e aleatórias desempenha um papel determinante na throughput e latência geral. Compreender como esses padrões de acesso interagem com as características subjacentes do S3 é fundamental para projetar pipelines de entrada eficientes.
A leitura de dados do S3 apresenta comportamento similar ao de unidades de disco rígido (HDDs) tradicionais com braços atuadores mecânicos. HDDs leem blocos de dados sequencialmente quando estão contíguos, permitindo que o braço minimize o movimento. Leituras aleatórias, por outro lado, exigem que o braço salte pela superfície do disco para acessar blocos espalhados, introduzindo atrasos pela reposição física do braço.
Ao acessar dados no S3, a situação é parcialmente similar. Cada requisição S3 incorre em uma sobrecarga de time-to-first-byte (TTFB) antes que a transferência de dados comece. Essa sobrecarga compreende vários componentes: estabelecimento de conexão, latência de round-trip de rede, operações internas do S3 (localização e acesso aos dados em disco) e tratamento de resposta no cliente. Enquanto o tempo de transferência escala com o tamanho dos dados, a sobrecarga TTFB é largamente fixa e independente do tamanho do objeto—um aspecto crucial para compreender o desempenho.
Em cargas de trabalho de ML, chamamos de padrão de leitura aleatória quando datasets consistem em numerosos arquivos pequenos armazenados no S3, cada arquivo contendo uma amostra de treinamento. Leitura aleatória também ocorre quando scripts de treinamento buscam amostras de diferentes partes dentro de um arquivo shard maior, usando requisições S3 GET com byte-range.
Padrões de leitura sequencial, por outro lado, ocorrem quando datasets são organizados em shards de arquivo grandes, cada shard contendo muitas amostras de treinamento, iteradas sequencialmente. Uma única requisição S3 GET pode recuperar múltiplas amostras, possibilitando throughput de dados muito mais alto que no cenário de leitura aleatória. Essa abordagem também simplifica o pré-carregamento de dados, permitindo antecipar, buscar e armazenar em buffer a próxima leva de amostras, deixando-as prontas para a GPU.
Caso de Estudo: Visão Computacional com Arquivos Pequenos
Para entender melhor como diferentes padrões de acesso afetam a performance, considere dois cenários em uma tarefa de visão computacional onde o dataset consiste em muitos arquivos de imagem relativamente pequenos (aproximadamente 100 KB cada).
Cenário 1 – Acesso Aleatório: O dataset é armazenado como está na classe de armazenamento S3 Standard, e o script de treinamento recupera cada imagem sob demanda. Isso cria um padrão de acesso aleatório, onde cada amostra de treinamento requer sua própria requisição S3 GET. Com latência TTFB na ordem de dezenas de milissegundos e tempo de download mínimo para arquivos pequenos, a performance do dataloader fica limitada por latência. Os threads do cliente gastam a maior parte do tempo ociosos aguardando a chegada dos dados.
Cenário 2 – Acesso Sequencial: O dataset é consolidado em shards de arquivo maiores (por exemplo, ~100 MB cada) antes de ser armazenado no S3. Agora o dataloader lê múltiplas amostras de treinamento sequencialmente com uma única requisição S3 GET. Isso muda a carga de trabalho para ser limitada por largura de banda, removendo o impacto TTFB por amostra e possibilitando streaming eficiente de amostras consecutivas durante a fase de download.
Técnicas Práticas de Otimização
Utilizar Clientes de Alto Desempenho Otimizados para S3
Escolher um cliente S3 performático pode ser desafiador dada a abundância de opções disponíveis. Para enfrentar isso, a AWS introduziu em 2023 dois clientes nativos de código aberto para S3: Mountpoint para Amazon S3 e Amazon S3 Connector para PyTorch. Ambos são construídos sobre o AWS Common Runtime (CRT), uma coleção de primitivas otimizadas em C que incluem um cliente S3 nativo implementando otimizações de performance baseadas em boas práticas, como paralelização de requisições, timeouts, retentativas e reutilização de conexões.
Mountpoint para Amazon S3 é um cliente de arquivo de código aberto que permite montar um bucket S3 em sua instância de computação e acessá-lo como um sistema de arquivos local sem necessidade de alterações no código existente. Isso o torna adequado para uma ampla gama de cargas de trabalho, incluindo treinamento de ML. Para ambientes Kubernetes, o Mountpoint para Amazon S3 Container Storage Interface (CSI) Driver estende essa capacidade apresentando um bucket S3 como um volume de armazenamento. Com o recente lançamento do Mountpoint para Amazon S3 CSI v2, o driver introduz cache compartilhado entre pods, permitindo que cargas de trabalho distribuídas de ML reutilizem dados localmente armazenados em cache, potencializando performance e eficiência de recursos.
Amazon S3 Connector para PyTorch oferece primitivas nativas do PyTorch que integram estritamente S3 com pipelines de treinamento. A integração possibilita acesso de alta throughput aos dados de treinamento e checkpointing eficiente diretamente ao Amazon S3, aplicando automaticamente otimizações de performance. O conector suporta datasets em estilo mapa para acesso aleatório e datasets em estilo iterável para acesso sequencial de streaming, adequando-se a diversos padrões de treinamento de ML. Inclui também interface de checkpointing integrada para salvar e carregar checkpoints do S3 sem depender de armazenamento local. A instalação é leve (usando pip, por exemplo), e o conector requer apenas mudanças mínimas no código de treinamento, com exemplos disponíveis no GitHub.
Fragmentar Datasets e Usar Padrões Sequenciais
Uma estratégia eficaz para otimizar o carregamento de dados do S3 é serializar datasets em fewer, shards de arquivo maiores, cada um contendo muitas amostras de treinamento, e ler essas amostras sequencialmente usando seu dataloader. Em micro-benchmarks S3, tamanhos de shard entre 100 MB a 1 GB tipicamente entregam excelente throughput. O tamanho ideal pode variar dependendo da carga de trabalho. Shards menores podem melhorar comportamento quasi-aleatório de buffers de pré-carregamento, enquanto shards maiores geralmente oferecem melhor throughput bruto.
Formatos comuns para fragmentação incluem tar (frequentemente usado em PyTorch através de bibliotecas como WebDataset) e TFRecord (usado com tf.data em TensorFlow). Fragmentar dados não garante leituras sequenciais. Se seu dataloader acessar aleatoriamente amostras dentro de um shard—comum com formatos como Parquet ou HDF5—os benefícios do acesso sequencial se perdem. Para completar os ganhos de performance, projete seu dataloader para que amostras sejam lidas em ordem dentro de cada shard.
Paralelização, Pré-carregamento e Cache
Otimizar os estágios de ingestão e pré-processamento de dados de um pipeline de ML é crítico para maximizar throughput de treinamento, especialmente quando padrões de acesso aleatório são inevitáveis.
Paralelização é uma das formas mais eficazes de melhorar throughput em pipelines de carregamento de dados, particularmente porque decodificação e pré-processamento de dados são frequentemente embarrassingly parallel—divisíveis em muitos processos independentes rodando simultaneamente sem necessidade de comunicação. Você pode usar frameworks como TensorFlow (tf.data) e PyTorch (DataLoader nativo) para ajustar o tamanho de seus pools de workers—threads ou processos CPU—para paralelizar ingestão de dados.
Para padrões de acesso sequencial, uma boa regra é corresponder o número de threads worker ao número de núcleos CPU disponíveis. Contudo, em instâncias com alto número de CPUs (por exemplo, mais de 20), usar um pool ligeiramente menor pode melhorar eficiência. Para padrões de acesso aleatório, particularmente ao ler diretamente do S3, tamanhos de pool maiores que o número de CPUs provaram ser benéficos. Por exemplo, em uma instância EC2 com 8 vCPUs, aumentar a configuração de num_workers do PyTorch para 64 ou mais melhorou significativamente throughput de dados.
Aumentar paralelismo não é uma solução universal. Over-paralelization pode sobrecarregar recursos de CPU e memória, deslocando o gargalo de I/O para pré-processamento. É importante fazer benchmark dentro do contexto de sua carga de trabalho específica para encontrar o equilíbrio correto.
Pré-carregamento complementa paralelização desacoplando carregamento de dados da computação de GPU. Usando um padrão produtor-consumidor, pré-carregamento permite que dados sejam preparados assincronamente e armazenados em buffer na memória, deixando o próximo lote pronto quando a GPU o necessite. Buffers de pré-carregamento bem dimensionados e tamanhos de pool worker adequadamente ajustados ajudam a amortizar latência de I/O e pré-processamento, melhorando throughput geral de treinamento.
Cache é particularmente eficaz para cargas de trabalho multi-epoch com padrões de acesso aleatório, onde as mesmas amostras de dados são lidas múltiplas vezes. Ferramentas como Mountpoint para Amazon S3 oferecem mecanismos de cache integrados que armazenam objetos do dataset localmente em armazenamento de instância (por exemplo, discos NVMe), volumes EBS ou memória. Removendo requisições S3 GET repetidas, cache melhora velocidade de treinamento e eficiência de custos.
Como o dataset de entrada típico permanece estático durante treinamento, recomenda-se configurar Mountpoint com indefinite metadata TTL (configurando –metadata-ttl indefinite, veja documentação Mountpoint para S3) para reduzir sobrecarga de requisição S3. Adicionalmente, nos benchmarks, cache de dados para NVMe foi habilitado, permitindo que Mountpoint armazene objetos localmente. O cache gerencia espaço automaticamente evictando os arquivos menos recentemente usados, mantendo pelo menos 5% de espaço disponível por padrão (configurável).
Validação Através de Benchmarks
A AWS conduziu uma série de benchmarks simulando uma carga de trabalho realista de visão computacional sob padrões de acesso aleatório e sequencial. Embora resultados exatos variem conforme seu caso específico, tendências e insights de performance são amplamente aplicáveis a pipelines de treinamento de ML.
Os benchmarks foram executados em uma instância Amazon EC2 g5.8xlarge equipada com GPU NVIDIA A10G e 32 vCPUs. A carga de trabalho usou o backbone google/vit-base-patch16-224-in21k ViT para classificação de imagens, treinando em um dataset de 10 GB contendo 100.000 imagens JPEG sintéticas (~115 KB cada), transmitidas diretamente do S3 Standard sob demanda.
Cada configuração de benchmark comparou diferentes clientes S3: dataloader baseado em fsspec, Mountpoint para Amazon S3 (sem cache de dados), Mountpoint para Amazon S3 (com cache de dados) e S3 Connector para PyTorch. Para benchmark com acesso sequencial, o dataset foi reorganizado em formato tar com tamanhos de shard variando de 4 MB a 256 MB.
Os resultados demonstraram que o S3 Connector para PyTorch alcançou a mais alta throughput de todos os clientes avaliados, atingindo aproximadamente 138 amostras/segundo com utilização próxima à saturação da GPU em acesso aleatório. Em cenários multi-epoch, cache de dados significativamente potencializou performance, com o dataset inteiro servido de disco a partir da segunda epoch, saturando completamente a GPU e maximizando throughput mesmo com pool de workers dataloader menor.
Para reproduzir benchmarks similares em seu próprio ambiente, a AWS fornece uma ferramenta de benchmark dedicada que suporta diversas configurações de carregamento de dados S3. Para resultados consistentes e significativos, use tipos idênticos de instância EC2 para cada cliente S3, coloque cada dataset de teste em buckets S3 separados e execute experimentos na mesma Região AWS que seus buckets.
Conclusão
Otimizar ingestão de dados é crucial para desbloquear completamente a performance de pipelines modernos de treinamento de ML em nuvem. Este artigo demonstrou como padrões de leitura aleatória e pequenos tamanhos de arquivo podem severamente limitar throughput devido a sobrecargas de latência, enquanto datasets consolidados com padrões de acesso sequencial podem maximizar largura de banda e manter GPUs plenamente utilizadas.
A AWS explorou como usar clientes Mountpoint para Amazon S3 e S3 Connector para PyTorch de alto desempenho pode fazer diferença significativa na performance de treinamento. Também demonstrou benefícios de fragmentar datasets em arquivos maiores, ajustar configurações de paralelização e aplicar cache para minimizar requisições S3 redundantes.
À medida que cargas de trabalho de treinamento crescem, revisite continuamente o design do seu pipeline de dados. Decisões cuidadosas sobre carregamento de dados podem entregar ganhos desproporciona is em eficiência de custos e tempo para resultados.
Fonte
Applying data loading best practices for ML training with Amazon S3 clients (https://aws.amazon.com/blogs/machine-learning/applying-data-loading-best-practices-for-ml-training-with-amazon-s3-clients/)