SageMaker agora divide uma GPU em 7 partes: mais modelos no mesmo hardware

Uma situação conhecida: você tem um classificador, um reranker, um embedder e alguns modelos com fine-tuning. Cada um precisa de 10–20 GB, mas o SageMaker entrega uma GPU inteira para cada modelo. O restante da memória fica ocioso, enquanto a conta continua crescendo. A AWS finalmente corrigiu isso.

O que mudou

O SageMaker Inference agora oferece particionamento de GPU baseado no NVIDIA Multi-Instance GPU (MIG). A placa física é dividida em partições independentes, cada uma com sua própria parcela de blocos computacionais, cache, memória e largura de banda de memória. O isolamento é feito por hardware, não por software.

Antes, quem quisesse usar MIG precisava operar seu próprio Kubernetes: drivers, device plugin, configuração de partições e agendamento. Os endpoints multimodelo não resolviam o problema: os modelos compartilhavam o dispositivo sem isolamento e ainda havia cold start no primeiro carregamento. Agora, você escolhe um perfil de partição e o SageMaker cuida do restante.

Como funciona na prática

Foram adicionados dois campos opcionais. AcceleratorPartitionLayouts é especificado na production variant durante a criação da configuração do endpoint e descreve como cada placa é dividida. Por exemplo, em uma p4d.24xlarge, o layout mig-4g.20gb + mig-2g.10gb + mig-1g.5gb usa todos os sete segmentos da A100. Com oito placas, são criadas 24 partições: oito de cada tipo.

AcceleratorPartitionType é especificado no inference component e define em qual partição o modelo será executado. Não há novas operações de API nem ações do IAM, e InvokeEndpoint permanece inalterado. Ao aumentar CopyCount, é possível colocar cópias adicionais de um modelo pequeno nas partições disponíveis. A criação do particionamento acrescentou cerca de 17 segundos à implantação do endpoint.

Os perfis disponíveis dependem da GPU. A100 e H100 são divididas em sétimos em termos de computação e em oitavos em termos de memória. G7 (RTX PRO 4500 Blackwell) é dividida pela metade, enquanto G7e (RTX PRO 6000 Blackwell) é dividida em quartos. O perfil 3g tem uma peculiaridade: usa 3 dos 7 blocos computacionais, mas 4 dos 8 blocos de memória. Por isso, duas partições desse tipo ocupam a placa inteira.

Como escolher um perfil

Os blocos computacionais são necessários para o prefill, ou seja, o processamento do prompt: quanto maior a entrada e maior o número de solicitações simultâneas, mais importantes eles são. A memória é necessária para o decode e para a capacidade: pesos, cache KV e 1–2 GB de overhead para o contexto CUDA. Um modelo 8B em FP16 ocupa cerca de 16 GB e, em uma partição de 40 GB, consegue manter confortavelmente um contexto de 8.000 tokens. Se apenas os pesos forem considerados, o modelo poderá ser carregado, mas falhará sob carga ou reduzirá silenciosamente a janela de contexto.

Onde vale a pena

• Um conjunto de modelos pequenos e médios: variantes com fine-tuning, embedders e rerankers.
• Multitenancy: existe uma fronteira de hardware entre os clientes, mas host, rede e roteamento continuam compartilhados, portanto não se trata de isolamento completo.
• Pipelines em uma única instância: transcrição, diarização e um LLM sem saltos de rede.
• Um modelo de guardrails com latência previsível.
• Implantações canary: um candidato em uma partição pequena ao lado do modelo de produção.

Limitações

O layout precisa ser escolhido durante a criação do endpoint e não pode ser alterado dinamicamente. O recurso funciona apenas com inference components: endpoints independentes, serverless e heterogêneos não são compatíveis. Um modelo que precise de mais de uma GPU ainda exige dispositivos inteiros. Na P4d, o perfil mig-1g.10gb exige o driver R525 ou posterior, enquanto a imagem padrão vem com a versão 470. Portanto, é necessário especificar InferenceAmiVersion al2023-ami-sagemaker-inference-gpu-4-1. Caso contrário, o particionamento falhará sem uma mensagem de erro clara.

O que fazer

Se você mantém vários modelos pequenos em GPUs separadas no SageMaker, calcule a memória usando a fórmula “pesos + cache KV + overhead” e tente consolidá-los em uma única placa. Também vale executar testes de carga: o desempenho de uma partição não depende das partições vizinhas, portanto os resultados podem ser levados para produção.