Оптимизация обучения крупных ИИ‑моделей с помощью Amazon SageMaker AI и NVIDIA Blackwell позволяет решить ряд проблем: ограничение размеров батчей из‑за памяти GPU, сокращение длины последовательностей из‑за риска нехватки памяти, сложности шардинга модели, который увеличивает накладные расходы на связь при масштабировании.

Для тренировочных заданий подходят P6-B200 с 8 GPU Blackwell. Ёмкость можно забронировать — это обеспечит предсказуемый доступ, управление затратами и автоматизированное управление ресурсами. Amazon SageMaker AI автоматически выделяет и управляет вычислительной инфраструктурой, позволяя сосредоточиться на данных и алгоритмах.

NVIDIA Blackwell выделяется двухчиповой архитектурой и тензорными ядрами пятого поколения. NVLink 5 обеспечивает пропускную способность до 1,8 ТБ/с между GPU. Объём HBM и пропускная способность памяти B200 снижают нагрузку на память при больших батчах, длинных последовательностях и распределённых тренировочных нагрузках.

В примерах используется одноузловое обучение с 8 GPU и трансформер‑моделями от 1B до 64B параметров. Применяется техника распределённого обучения: параметры модели, градиенты и состояния оптимизатора разделяются между GPU.

Расширенная память Blackwell (180 ГБ на B200, 268 ГБ на B300) позволяет:
* увеличить размеры батчей (это сокращает количество шагов синхронизации градиентов между GPU и повышает общую пропускную способность);
* упростить шардинг модели (можно уменьшить степень параллелизма модели или полностью исключить его для некоторых моделей);
* увеличить длину последовательностей (модели могут обрабатывать больше контекста за один проход, что важно для задач с долгосрочными зависимостями).

Приоритеты оптимизации зависят от задачи:
* если цель — пропускная способность, начните с настройки размера батча;
* если узкое место — накладные расходы на связь, сначала упростите шардинг;
* если задача требует долгосрочного контекста, уделите внимание длине последовательности.

Размер пакета и длина последовательности увеличивают потребление памяти — важно найти баланс. Активация чекпоинтинга позволяет сбалансировать использование памяти и вычислений: она увеличивает время вычислений (обычно на 10–30 % в зависимости от архитектуры модели) за счёт пересчёта промежуточных активаций во время обратного прохода вместо их хранения, что снижает использование памяти GPU. Освобождённую память можно использовать для увеличения размера пакета или длины последовательностей. Чтобы оценить компромисс, нужно провести бенчмарки конкретной конфигурации.

Например, сравнивались три конфигурации обучения для LLM с 1B параметров с точностью MXFP8 и длиной последовательности 8K:
* без активации чекпоинтинга (BS=1) пропускная способность ~6K токенов/сек, пик памяти — 15,5 ГБ;
* при активации чекпоинтинга при том же размере пакета память снижается до 2,3 ГБ, но пропускная способность немного падает;
* при активации чекпоинтинга и увеличении размера пакета до 16 пропускная способность достигает ~51K токенов/сек (примерно в 8 раз выше базового уровня), пик памяти — 22,8 ГБ, что в пределах лимитов GPU.

Для решения о применении активации чекпоинтинга учитывайте размер модели и использование памяти:
* Малые модели (до ~14B параметров): активация чекпоинтинга обычно не нужна — с расширенной памятью Blackwell большинство малых моделей помещаются без неё. При работе на верхнем пределе диапазона и нехватке памяти чекпоинтинг добавляет вычислительные накладные расходы в обмен на существенную экономию памяти — её можно перенаправить на увеличение размеров пакетов.
* Большие модели (~14B+ параметров): потребление памяти — от 87 до 171 ГБ (зависит от размера пакета и длины последовательности). Без чекпоинтинга большинство конфигураций завершается ошибкой CUDA out-of-memory (OOM). С чекпоинтингом освобождённая память позволяет увеличить размер пакета — пропускная способность улучшается, несмотря на дополнительные вычислительные накладные расходы. Для больших моделей чекпоинтинг обязателен — это условие стабильного обучения.

Пятое поколение Tensor Cores в Blackwell обеспечивает аппаратное ускорение для форматов пониженной точности (FP8, MXFP8, NVFP4) — это в первую очередь оптимизация пропускной способности, а не экономия памяти.

При работе с вычислительно ограниченными нагрузками (обычно малые модели) скорость вычислений — ограничивающий фактор, накладные расходы на квантование частично компенсируют прирост пропускной способности от пониженной точности. При работе с нагрузками, ограниченными памятью (обычно большие модели), узким местом является перемещение данных, а уменьшенный объём памяти форматов пониженной точности напрямую устраняет ограничение, обеспечивая более значительный прирост.

Для малых моделей (до ~14B параметров):
* форматы пониженной точности (FP8, MXFP8, NVFP4) дают схожие, умеренные улучшения пропускной способности по сравнению с FP16 (накладные расходы на квантование съедают преимущество в скорости);
* настройка размера пакета обычно даёт более значимые улучшения, чем выбор формата точности;
* начните с FP8 для более высокой пропускной способности — у него ниже накладные расходы, чем у MXFP8 или NVFP4, и он часто подходит в качестве значения по умолчанию для большинства рабочих нагрузок с малыми моделями;
* при ограничении памяти и поддержке оптимизатором используйте quantizedmodelinit, чтобы хранить веса непосредственно в FP8 — это снизит потребление памяти ниже уровня FP16.

Для больших моделей (~14B+ параметров):
* пониженная точность даёт наибольший эффект;
* FP8 обычно обеспечивает хороший баланс пропускной способности и эффективности использования памяти;
* MXFP8 теоретически более эффективен по памяти, но накладные расходы на транспонирование частично компенсируют это преимущество на практике; если для рабочей нагрузки важны стабильность сходимости или числовая точность, MXFP8 может быть лучшим выбором — его более детальная схема квантования надёжнее сохраняет точность модели, чем FP8;
* для больших моделей, где узким местом является память, NVFP4 может обеспечить дополнительный прирост пропускной способности — преимущество в скорости матричного умножения растёт с увеличением размера модели.

Для достижения результатов требуются значительные инженерные инвестиции. Используйте рецепты на уровне фреймворка из Megatron Core с валидированными конфигурациями NVFP4. TransformerEngine от NVIDIA упрощает реализацию: обеспечивает автоматическое переключение смешанной точности, объединённые ядра и динамическое масштабирование потерь. Перед запуском в продакшн проверьте сходимость, отслеживая кривые потерь в разных форматах, чтобы убедиться, что выбранная точность соответствует требованиям к точности.

Не каждая рабочая нагрузка выигрывает от агрессивной оптимизации: если модель обучается в пределах лимитов памяти и удовлетворяет требованиям к пропускной способности при FP16, усложнение за счёт форматов с пониженной точностью может не стоить усилий. Начинайте с базовых измерений, оптимизируйте только измеримые узкие места.

Для начала работы с обучением на Blackwell в Amazon SageMaker AI нужно принять ключевые решения: объём доступной памяти, целесообразность активации чекпоинтинга для размера модели, подходящий формат точности для рабочей нагрузки.

Amazon SageMaker AI предоставляет полностью управляемую среду для распределённого обучения на инстансах Blackwell: управляет предоставлением инстансов, оркестрацией контейнеров и интеграцией с сервисами AWS.

Предварительные требования:
* аккаунт AWS с правами на создание заданий обучения Amazon SageMaker AI, доступ к Amazon ECR и создание ролей IAM для выполнения Amazon SageMaker AI;
* доступ к инстансам ml.p6-b200.48xlarge через Flexible Training Plan или Managed Spot Training (перед началом проверьте Service Quotas);
* Docker, установленный локально;
* Python 3.9 или новее;
* SageMaker Python SDK;
* знание PyTorch и FSDP (для новичков в FSDP см. Getting started with distributed training).

Чтобы запустить задание обучения, выполните шаги:
1. Создайте скрипт: скачайте файл fsdp.py из FSDP example from the NVIDIA TransformerEngine repository (скрипт реализует обучение FSDP и принимает гиперпараметры в качестве аргументов командной строки).
2. Создайте скрипт точки входа: подготовьте файл train.sh для настройки torchrun и запуска скрипта обучения (```

!/bin/bash

SageMaker passes hyperparameters as environment variables (SMHP<NAME>)

PRECISION=${SMHPPRECISION:-"mxfp8"}
NUMLAYERS=${SMHPNUMLAYERS:-10}
BATCHSIZE=${SMHPBATCHSIZE:-8}
SEQLENGTH=${SMHPSEQLENGTH:-2048}

NUMGPUS=$(nvidia-smi --list-gpus | wc -l)
torchrun --standalone --nnodes=1 --nproc-per-node="$NUMGPUS" \
fsdp.py --no-defer-init --precision "$PRECISION" \
--num-layers "$NUMLAYERS" --checkpoint-layer "transformerlayer" \
--batch-size "$BATCHSIZE" --seq-length "$SEQ_LENGTH"
).
3. Соберите и отправьте контейнер: создайте пользовательский контейнер Docker, который расширяет AWS Deep Learning Containers (DLC), включает fsdp.py и train.sh, имеет установленный TransformerEngine 2.11. DLC предоставляет валидированный базовый образ с PyTorch и библиотеками CUDA, необходимыми для совместимости с Blackwell (Dockerfile:
FROM 763104351884.dkr.ecr.us-west-2.amazonaws.com/pytorch-training:2.9.0-gpu-py312-cu130-ubuntu22.04-sagemaker

Install Transformer Engine

RUN pip install --upgrade --no-build-isolation transformer_engine[pytorch]==2.11.0

Provide libcudart.so.12 for the pre-built flash-attn wheel

RUN pip install nvidia-cuda-runtime-cu12

Make the linker able to find it

ENV LDLIBRARYPATH=/usr/local/lib/python3.12/site-packages/nvidia/cudaruntime/lib:$LDLIBRARY_PATH

COPY fsdp.py /opt/ml/code/fsdp.py
COPY train.sh /opt/ml/code/train.sh

ENV SAGEMAKERSUBMITDIRECTORY /opt/ml/code
ENV SAGEMAKERPROGRAM train.sh
```).
4. Создайте приватный репозиторий Amazon ECR, если его ещё нет, и отправьте образ в Amazon ECR (запомнить repositoryUri из вывода для использования в командах docker tag и push).
5. Обеспечьте доступ к ресурсам — зарезервируйте ресурсы через Flexible Training Plan для предсказуемого доступа по стандартной ставке либо используйте Managed Spot Training для оптимизированных по стоимости рабочих нагрузок. Flexible Training Plans применять для производственных тренировочных запусков, требующих резервирования ресурсов и обеспечивающих непрерывную доступность; Spot — для экспериментов и отказоустойчивых рабочих нагрузок, где снижение затрат важнее риска прерывания. При использовании Spot‑инстансов (подверженных прерыванию) скрипт обучения должен сохранять контрольные точки в Amazon S3 через регулярные промежутки времени. Amazon SageMaker AI автоматически возобновит прерванное Spot‑задание, если в конфигурации оценщика указан checkpoints3uri.
6. Для предсказуемой доступности ресурсов в консоли SageMaker AI выберите «training-job» в качестве целевого ресурса. Если задание допускает перезапуски и приоритет — снижение затрат, можно выбрать Spot, предварительно увеличив квоту для использования тренировочного задания ml.p6-b200.48xlarge spot.
7. Отправьте тренировочное задание — замените значения‑заполнители на фактический ARN плана обучения и URI образа ECR, затем выполните следующий код из локальной среды разработки или экземпляра ноутбука SageMaker AI: ```
from sagemaker.estimator import Estimator
from sagemaker import getexecution_role
from sagemaker.debugger import ProfilerConfig

trainingplanarn = "<your-training-plan-arn>" # Replace with your training plan ARN
ecr_image = "<your-ecr-image-uri>" # Replace with your ECR image URI

Adjust these values to match your workload

precision = "mxfp8"
numlayers = 10
batchsize = 8
seq_length = 2048

estimator = Estimator(imageuri=ecrimage,
role=getexecutionrole,
basejobname='blackwell-training',
instancecount=1,
instancetype='ml.p6-b200.48xlarge',
hyperparameters={
"precision": precision,
"num-layers": numlayers,
"batch-size": batchsize,
"seq-length": seqlength,
},
profilerconfig=ProfilerConfig(disableprofiler=True),
trainingplan=trainingplanarn)

estimator.fit
```
8. Для Managed Spot Training замените trainingplan на usespotinstances=True, установите maxrun и maxwait, добавьте checkpoints3_uri для автоматического возобновления.

Мониторинг задания обучения в Amazon SageMaker AI:
* журналы автоматически передаются в Amazon CloudWatch;
* в SageMaker AI console перейдите в Training jobs, выберите задание, найдите группу журналов CloudWatch;
* откройте /aws/sagemaker/TrainingJobs, ищите строки [rank 0] со значениями потерь и пропускной способности;
* для подтверждения загрузки формата точности ищите сообщения «Using FP8 recipe» или «MXFP8 enabled».

При ошибке CUDA out-of-memory (OOM) в журнале отображается размер выделения. Уменьшите размер пакета или длину последовательности либо добавьте контрольную точку активации.

Для избежания дальнейших расходов остановите выполняющиеся задания обучения в консоли Amazon SageMaker AI. Выполните очистку: удалите образы контейнеров (Delete your Amazon ECR repository), артефакты обучения в Amazon S3 (данные обучения и контрольные точки), группы журналов CloudWatch, роль выполнения IAM, созданную для Amazon SageMaker AI.

Модели Transformer от 1B до 64B параметров показывают стабильное улучшение при сочетании этих оптимизаций. Важно определить, ограничена ли рабочая нагрузка вычислениями или памятью, и постепенно вносить изменения, чтобы измерить влияние каждого из них.

Для начала работы изучите Amazon SageMaker AI documentation (варианты экземпляров и детали конфигурации) или purchase a Flexible Training Plan (резервирование ёмкости Blackwell для следующего запуска обучения). При вопросах о конкретной рабочей нагрузке обратитесь к AWS Support или в команду аккаунта AWS.

Artificial Intelligence

Amazon SageMaker AI

NVIDIA Blackwell GPUs

available on Amazon SageMaker AI

Flexible Training Plan

PyTorch Fully Sharded Data Parallel (FSDP)

Amazon SageMaker AI training jobs

Amazon SageMaker AI

Amazon Simple Storage Service (Amazon S3)

Amazon CloudWatch

Amazon Elastic Container Registry (Amazon ECR)

AWS Identity and Access Management (AWS IAM)

Adapting your own Docker container to work with Amazon SageMaker AI

Create a Flexible Training Plan