AWS acelera carregamento de modelos LLM e amplia janelas de contexto com GPUDirect e Amazon FSx for Lustre
Para quem trabalha com grandes modelos de linguagem (LLMs) em instâncias GPU da AWS, o tempo de carregamento do modelo na memória de alta largura de banda (HBM) pode ser um gargalo significativo, especialmente à medida que os modelos ultrapassam centenas de bilhões de parâmetros. Pensando nisso, a AWS anunciou uma solução que combina o sistema de arquivos Amazon FSx for Lustre com a tecnologia NVIDIA GPUDirect Storage (GDS) e o recurso TurboQuant para acelerar esse processo e expandir as janelas de contexto durante a inferência.
O que foi lançado e como funciona
A novidade consiste na integração do Amazon FSx for Lustre com GPUDirect Storage, que permite a leitura direta dos dados armazenados para a memória HBM da GPU, eliminando a necessidade de passar pela memória do CPU. Tradicionalmente, o carregamento do modelo envolve a transferência sequencial dos pesos via CPU para cada GPU, um processo lento e ineficiente. Com GPUDirect, múltiplas GPUs podem ler em paralelo seus fragmentos do modelo diretamente do FSx for Lustre, reduzindo minutos de espera para segundos.
Além disso, o TurboQuant, uma tecnologia de compressão desenvolvida pelo Google Research, reduz significativamente o tamanho do cache de chave-valor (KV) usado durante a inferência, liberando memória da GPU para suportar janelas de contexto maiores ou cargas de trabalho mais intensas, sem perda de precisão.
Quem pode se beneficiar
Essa solução é ideal para desenvolvedores, pesquisadores e empresas que utilizam modelos LLM de larga escala, como o Llama 3.1 405B, especialmente em instâncias AWS EC2 das famílias P5en e P6 equipadas com GPUs NVIDIA Blackwell (H200 e B200). O método é particularmente útil para cenários de inferência distribuída, onde a redução do tempo de "cold start" impacta diretamente a latência, escalabilidade e custo operacional.
Disponibilidade e acesso
O Amazon FSx for Lustre com suporte a GPUDirect Storage está disponível para instâncias compatíveis na AWS. A configuração requer:
- Um sistema de arquivos FSx for Lustre configurado com Elastic Fabric Adapter (EFA) habilitado;
- Instâncias P5en ou P6 configuradas para GPUDirect Storage;
- Pré-processamento dos modelos para fragmentação tensor-paralela e quantização (por exemplo, para FP8) usando frameworks como vLLM;
- Configuração do cliente Lustre para otimização de throughput e alinhamento NUMA.
A AWS disponibiliza scripts e templates CloudFormation no repositório aws-samples para facilitar essa implantação.
Impacto prático para o usuário
Com essa integração, o tempo para carregar modelos gigantescos na memória GPU cai de dezenas de minutos para poucos segundos, otimizando o tempo até o primeiro token gerado (TTFT). Isso significa maior agilidade para iniciar serviços de inferência, melhor resposta em eventos de escalonamento automático e menor custo, já que os GPUs ficam menos tempo ociosos durante o carregamento.
Além disso, a compressão do KV cache via TurboQuant amplia a capacidade de contexto durante a inferência, permitindo processar sequências maiores ou aumentar o número de requisições simultâneas sem necessidade de hardware adicional.
Como começar
- Crie uma conta AWS se ainda não tiver: Cadastro AWS.
- Provisione um sistema FSx for Lustre com EFA habilitado e configure a capacidade para garantir o throughput desejado (exemplo: Persistent_2 SSD com 1000 MBps/TiB).
- Configure uma instância P5en ou P6 com suporte a GPUDirect Storage, seguindo as orientações do guia Configuring EFA clients.
- Pré-processe seu modelo para gerar shards tensor-paralelos e quantizados (FP8 recomendado) usando ferramentas como vLLM.
- Monte o sistema de arquivos FSx for Lustre e realize o carregamento paralelo direto para as GPUs utilizando a biblioteca fastsafetensors.