SageMaker ya puede dividir una GPU en 7 partes: más modelos sobre el mismo hardware

Una situación habitual: tiene un clasificador, un reranker, un embedder y un par de modelos ajustados. Cada uno necesita entre 10 y 20 GB, pero SageMaker asigna una GPU completa a cada uno. El resto de la memoria queda inactivo mientras la factura sigue creciendo. AWS por fin lo ha solucionado.

Qué ha cambiado

SageMaker Inference ahora admite la partición de GPU basada en NVIDIA Multi-Instance GPU (MIG). La tarjeta física se divide en particiones independientes, cada una con su propia cuota de bloques de cómputo, caché, memoria y ancho de banda de memoria. El aislamiento se implementa por hardware, no por software.

Antes, quien quisiera utilizar MIG tenía que desplegar su propio Kubernetes: controladores, device plugin, configuración de particiones y planificación. Los endpoints multimodelo tampoco resolvían el problema: los modelos comparten el dispositivo sin aislamiento y, además, la primera carga sufre un arranque en frío. Ahora basta con elegir un perfil de partición y SageMaker se ocupa del resto.

Cómo funciona en la práctica

Se han añadido dos campos opcionales. AcceleratorPartitionLayouts se especifica en la production variant al crear la configuración del endpoint y describe cómo se divide cada tarjeta. Por ejemplo, en una p4d.24xlarge, la distribución mig-4g.20gb + mig-2g.10gb + mig-1g.5gb utiliza los siete segmentos de la A100. Con ocho tarjetas se obtienen 24 particiones: ocho de cada tipo.

AcceleratorPartitionType se especifica en el inference component y determina en qué partición se ejecuta el modelo. No hay nuevas operaciones de API ni acciones de IAM, e InvokeEndpoint no cambia. Al aumentar CopyCount, se pueden colocar copias adicionales de un modelo pequeño en las particiones disponibles. La creación de la distribución de particiones añadió unos 17 segundos al despliegue del endpoint.

Los perfiles disponibles dependen de la GPU. Las A100 y H100 se dividen en séptimos en términos de cómputo y en octavos en términos de memoria. G7 (RTX PRO 4500 Blackwell) se divide en mitades y G7e (RTX PRO 6000 Blackwell), en cuartos. El perfil 3g tiene una particularidad: utiliza 3 de los 7 bloques de cómputo, pero 4 de los 8 bloques de memoria, por lo que dos particiones de este tipo ocupan la tarjeta completa.

Cómo elegir un perfil

Los bloques de cómputo son necesarios para el prefill, es decir, el procesamiento del prompt: cuanto más larga sea la entrada y más solicitudes simultáneas haya, más importantes resultan. La memoria es necesaria para el decode y la capacidad: los pesos, la caché KV y entre 1 y 2 GB de sobrecarga para el contexto de CUDA. Un modelo 8B en FP16 ocupa unos 16 GB y, en una partición de 40 GB, puede mantener cómodamente un contexto de 8.000 tokens. Si solo se tienen en cuenta los pesos, el modelo podrá cargarse, pero fallará bajo carga o reducirá silenciosamente la ventana de contexto.

Dónde resulta rentable

• Un conjunto de modelos pequeños y medianos: variantes ajustadas, embedders y rerankers.
• Multi-tenancy: existe una frontera de hardware entre clientes, pero el host, la red y el enrutamiento siguen siendo compartidos, por lo que no se trata de un aislamiento completo.
• Pipelines en una sola instancia: transcripción, diarización y un LLM sin saltos de red.
• Un modelo de guardrails con latencia predecible.
• Despliegues canary: un candidato en una partición pequeña junto al modelo de producción.

Limitaciones

La distribución debe elegirse al crear el endpoint y no puede modificarse sobre la marcha. Solo funciona con inference components: no se admiten endpoints independientes, serverless ni heterogéneos. Un modelo que necesite más de una GPU sigue requiriendo dispositivos completos. En P4d, el perfil mig-1g.10gb requiere el controlador R525 o posterior, mientras que la imagen predeterminada incluye la versión 470; por eso es necesario especificar InferenceAmiVersion al2023-ami-sagemaker-inference-gpu-4-1. De lo contrario, la partición fallará sin mostrar un mensaje de error claro.

Qué hacer

Si tiene varios modelos pequeños ejecutándose en GPU separadas dentro de SageMaker, calcule la memoria con la fórmula «pesos + caché KV + sobrecarga» e intente consolidarlos en una sola tarjeta. También conviene ejecutar pruebas de carga: el rendimiento de una partición no depende de sus vecinas, por lo que los resultados pueden trasladarse a producción.