A Amazon incluiu o particionamento de GPU no SageMaker Inference, segundo um post no AWS Machine Learning Blog. O recurso usa o NVIDIA Multi-Instance GPU (MIG) para cortar uma GPU física em partições. Cada partição recebe uma fatia fixa de computação e um bloco próprio de memória, e o isolamento é garantido por hardware.
O que mudou
Até agora, a menor unidade que dava para pedir era um acelerador inteiro, então uma instância com oito A100 hospedava no máximo oito modelos. Quem queria particionamento de verdade precisava operar a própria stack de Kubernetes. Agora o SageMaker cuida da configuração das GPUs, da alocação e do escalonamento assim que você escolhe um perfil de partição.
O recurso acrescenta dois campos opcionais à API, sem novas operações nem novas ações de IAM. Um layout na variante de produção, dentro da configuração do endpoint, declara como cada GPU é dividida. Cada inference component então indica o tipo de partição em que roda. A invocação não muda, e um endpoint que não usa nenhum dos dois campos se comporta como antes. A AWS mediu cerca de 17 segundos extras de provisionamento para a etapa de particionamento.
Onde está disponível
O particionamento vale para as famílias P4d, P5, G7 e G7e em sete Regiões: US East (N. Virginia e Ohio), US West (Oregon), Singapura, Tóquio, Frankfurt e Irlanda. Uma GPU rende no máximo sete partições. A G7 se divide em duas e a G7e em quatro. O endpoint precisa usar inference components. Endpoints de modelo único, endpoints serverless, pools heterogêneos de instâncias e invocação assíncrona ficam de fora.
Os limites
- O tipo de partição de um componente é fixo durante toda a vida dele. Para mudar o corte de um endpoint, é preciso criar outro e migrar o tráfego.
- Um componente que pede um perfil ausente do layout é rejeitado, sem recurso automático para uma GPU inteira.
- Modelos que exigem mais de uma GPU continuam precisando de dispositivos inteiros.
- A cobrança é pela instância completa. A economia vem do maior aproveitamento da GPU, já que o preço da hora não cai.
- O isolamento de GPU não é uma fronteira completa entre inquilinos, porque os modelos de uma mesma instância ainda dividem CPU do host, memória, armazenamento e rede.
- Na P4d, um dos perfis pequenos exige um driver NVIDIA mais novo que o da imagem padrão, e a AWS avisa que a mensagem de erro não cita o perfil.
O que a AWS mediu
A AWS comparou uma partição de 40 GB com três unidades de computação de uma H100 a uma A100-40GB inteira, ambas servindo o Llama-3.1-8B-Instruct, com uma réplica cada. No nível de carga em que o benefício de custo atingiu o pico, a partição gerou de 53% a 89% mais tokens de saída por segundo, com tempo até o primeiro token menor em todos os casos. Acima de uma concorrência de virada, saiu de 18% a 34% mais barata por milhão de tokens de saída. Essa virada foi em 64 para prompts de 100 tokens e em 4 para prompts de 8.000 tokens. Abaixo dela, a A100 inteira era mais barata. No topo da faixa de carga com prompts longos, a A100 falhou em um terço a metade das requisições, enquanto a partição não falhou em nenhuma. É um teste de um fornecedor com um único modelo, então vale tratá-lo como método para repetir com o seu próprio tráfego.
A AWS também lembra que a partição precisa comportar os pesos, o KV cache e de um a dois gigabytes de overhead do runtime. Um modelo de 8B em FP16 precisa de cerca de 16 GB só para os pesos.