SageMaker научился делить одну GPU на 7 частей: больше моделей на одном железе

Знакомая ситуация: у вас есть классификатор, реранкер, эмбеддер и пара дообученных моделей. Каждой нужно 10–20 ГБ, а SageMaker выдаёт по целой видеокарте. Остальная память простаивает, а счёт растёт. AWS наконец это исправил.

Что изменилось

SageMaker Inference теперь поддерживает GPU partitioning на базе NVIDIA Multi-Instance GPU (MIG). Физическая карта делится на независимые партиции, у каждой своя доля вычислительных блоков, кэша, памяти и пропускной способности памяти. Изоляция аппаратная, а не программная.

Раньше хотелось MIG, шли в самостоятельный Kubernetes: драйверы, device plugin, конфигурация партиций, планирование. Multi-model endpoints не выручали: модели делят устройство без изоляции, а ещё есть cold start при первой загрузке. Теперь вы выбираете профиль партиции, а остальное делает SageMaker.

Как это выглядит на практике

Добавляются два необязательных поля. AcceleratorPartitionLayouts указывается в production variant при создании конфигурации эндпоинта и описывает, как делится каждая карта. Например, на p4d.24xlarge раскладка mig-4g.20gb + mig-2g.10gb + mig-1g.5gb использует все семь слайсов A100. С восемью картами получается 24 партиции, по восемь каждого типа.

AcceleratorPartitionType указывается в inference component и выбирает, на какой партиции живёт модель. Новых API-операций и IAM-действий нет, InvokeEndpoint не меняется. Увеличив CopyCount, можно разместить дополнительные копии маленькой модели на свободных партициях. Создание нарезки добавило около 17 секунд к развёртыванию эндпоинта.

Доступные профили зависят от GPU. A100 и H100 делятся на седьмые по вычислениям и восьмые по памяти. G7 (RTX PRO 4500 Blackwell) делится пополам, G7e (RTX PRO 6000 Blackwell) на четверти. Профиль 3g хитрый: берёт 3 из 7 вычислительных блоков, но 4 из 8 блоков памяти, поэтому две такие партиции заполняют карту целиком.

Как выбирать профиль

Вычислительные блоки нужны для prefill, то есть обработки промпта: чем длиннее вход и больше параллельных запросов, тем важнее. Память нужна для decode и ёмкости: веса, KV-кэш и 1–2 ГБ накладных расходов на CUDA-контекст. Модель на 8B в FP16 весит около 16 ГБ, и на партиции 40 ГБ она комфортно держит контекст в 8000 токенов. Если считать только веса, модель загрузится, но под нагрузкой упадёт или тихо урежет контекст.

Где это выгодно

• Набор небольших и средних моделей: дообученные варианты, эмбеддеры, реранкеры.
• Мультитенантность: аппаратная граница между клиентами, но хост, сеть и маршрутизация остаются общими, так что это не полная изоляция.
• Пайплайны на одном инстансе: транскрибация, диаризация и LLM без сетевых прыжков.
• Guardrail-модель с предсказуемой задержкой.
• Канарейки: кандидат на маленькой партиции рядом с продакшеном.

Ограничения

Раскладку нужно выбрать при создании эндпоинта, менять её на лету нельзя. Работает только с inference components: одиночные, serverless и гетерогенные эндпоинты не поддерживаются. Модель, которой нужно больше одной карты, по-прежнему требует целых устройств. На P4d профиль mig-1g.10gb требует драйвер R525+, а образ по умолчанию идёт с 470, поэтому нужно указать InferenceAmiVersion al2023-ami-sagemaker-inference-gpu-4-1. Иначе нарезка упадёт без внятной ошибки.

Что делать

Если на SageMaker у вас несколько небольших моделей на отдельных GPU, посчитайте память по формуле «веса + KV-кэш + overhead» и попробуйте уплотнить их на одной карте. Заодно стоит прогнать нагрузочные тесты: производительность партиции не зависит от соседей, так что результаты можно переносить в продакшен.