O que muda com o Blackwell no SageMaker AI
Quem treina modelos grandes hoje convive com um conjunto bem conhecido de limitações: batch sizes restritos pela memória da GPU, sequence lengths cortadas para evitar erros de falta de memória, e fragmentação de modelo (model sharding) que aumenta o overhead de comunicação conforme a escala cresce.
A Amazon SageMaker AI agora suporta instâncias P6-B200 com 8 GPUs NVIDIA Blackwell em jobs de treinamento. Essas instâncias estão disponíveis para uso geral no SageMaker AI, e a capacidade pode ser reservada via Flexible Training Plan — com acesso previsível, controle de custos e gerenciamento automatizado de recursos.
A arquitetura dual-chip do Blackwell, combinada com Tensor Cores de quinta geração, entrega ganhos mensuráveis para treinamento multi-GPU. O interconector NVLink 5 oferece até 1,8 TB/s de largura de banda bidirecional entre GPUs, enquanto a maior capacidade de Memória de Alta Largura de Banda (HBM) do B200 — 180 GB — e do B300 — 268 GB — reduz a pressão de memória em cargas de trabalho com grandes batches, sequências longas e treinamento distribuído.
Gerenciamento de memória: onde o Blackwell muda o jogo
A memória expandida do Blackwell abre espaço para otimização em três frentes: batches maiores, fragmentação de modelo simplificada e sequências mais longas.
- Batches maiores reduzem o número de etapas de sincronização de gradientes entre GPUs, aumentando o throughput geral.
- Fragmentação simplificada se torna viável porque mais memória por GPU pode eliminar ou reduzir o grau de paralelismo de modelo — menos fragmentos significam menos overhead de comunicação entre GPUs.
- Sequências mais longas permitem que o modelo processe mais contexto em uma única passagem, algo crítico para tarefas que dependem de dependências de longo alcance.
Se throughput é a prioridade, comece ajustando o batch size. Se o gargalo é comunicação, simplifique a fragmentação primeiro. Se a tarefa exige contexto longo, priorize a sequence length.
Activation Checkpointing: a troca certa entre memória e computação
Batch size e sequence length aumentam o consumo de memória juntos — encontrar um equilíbrio eficiente é fundamental. O activation checkpointing ajuda nessa balança: em vez de armazenar as ativações intermediárias durante o passo de avanço (forward pass), elas são recomputadas durante o passo de retropropagação (backward pass). Isso libera memória ao custo de mais computação — tipicamente 10 a 30% de overhead, dependendo da arquitetura.
Para ilustrar o impacto, considere três configurações de treinamento de um LLM de 1 bilhão de parâmetros com precisão MXFP8 e sequence length de 8K:
- Sem activation checkpointing (BS=1): throughput de aproximadamente 6 mil tokens/segundo, com memória de pico em 15,5 GB.
- Com activation checkpointing (BS=1): a memória cai drasticamente para 2,3 GB, mas o throughput também recua levemente pela recomputação.
- Com activation checkpointing e batch size 16: a memória liberada permite um batch muito maior, elevando o throughput para aproximadamente 51 mil tokens/segundo — cerca de 8 vezes o baseline — com memória de pico em 22,8 GB, ainda bem dentro dos limites da GPU.
A decisão de usar ou não activation checkpointing depende do tamanho do modelo:
- Modelos pequenos (até ~14B parâmetros): geralmente não precisam de activation checkpointing. A memória expandida do Blackwell acomoda a maioria desses modelos sem ele. Se houver pressão de memória na faixa superior, o checkpointing adiciona overhead de computação em troca de economia de memória que pode ser reinvestida em batches maiores.
- Modelos grandes (~14B+ parâmetros): o consumo de memória varia de 87 a 171 GB dependendo do batch size e da sequence length. Sem activation checkpointing, a maioria das configurações falha com erros de falta de memória CUDA (OOM). Com o checkpointing ativado, a memória liberada permite aumentar o batch size o suficiente para que o throughput melhore mesmo com o overhead adicional. Para modelos grandes, o checkpointing não é opcional — é um pré-requisito para treinamento estável.
Formatos de precisão: escolhendo o certo para cada workload
Os Tensor Cores de quinta geração do Blackwell oferecem aceleração em hardware para formatos de precisão reduzida: FP8, MXFP8 e NVFP4. Esses formatos são principalmente otimizações de throughput, não de economia de memória. Usar precisão menor reduz os requisitos de largura de banda de memória e aumenta o número de operações que a GPU executa por ciclo.
Porém, o treinamento com precisão reduzida é aproximadamente neutro em termos de memória por padrão — o TransformerEngine mantém os pesos primários em alta precisão para atualizações do otimizador e cópias quantizadas em paralelo. Isso significa que formatos de menor precisão não se traduzem diretamente em menor uso de memória.
O benefício líquido depende do tamanho do modelo e se o treinamento é limitado por computação ou por memória:
- Modelos pequenos (até ~14B parâmetros): FP8, MXFP8 e NVFP4 entregam melhorias modestas e similares em relação ao FP16, pois o overhead de quantização consome parte do ganho. O ajuste de batch size tende a trazer ganhos mais expressivos do que a escolha do formato de precisão. Comece com FP8 — ele tem overhead menor que MXFP8 ou NVFP4 e costuma ser um bom padrão para a maioria dos workloads com modelos pequenos.
- Modelos grandes (~14B+ parâmetros): é aqui que a precisão reduzida entrega seu maior impacto. FP8 oferece um bom equilíbrio entre throughput e eficiência de memória. MXFP8 é teoricamente mais eficiente em memória, mas o overhead de transposição parcialmente anula essa vantagem na prática. No entanto, se estabilidade de convergência ou precisão numérica for prioridade, MXFP8 pode ser a melhor escolha, pois seu esquema de quantização mais granular tende a preservar a acurácia do modelo de forma mais confiável. Para modelos grandes com memória como gargalo principal, NVFP4 pode entregar ganhos adicionais de throughput, mas exige investimento significativo de engenharia — recomenda-se usar as receitas validadas do Megatron Core.
O NVIDIA TransformerEngine lida com a complexidade de implementação: troca automática de precisão mista, kernels fundidos e escalonamento dinâmico de perda. Antes de ir para produção, valide a convergência acompanhando as curvas de perda entre os formatos para confirmar que a precisão escolhida atende aos requisitos de acurácia.
Configurando um job de treinamento no SageMaker AI com Blackwell
O Amazon SageMaker AI fornece um ambiente totalmente gerenciado para treinamento distribuído em instâncias Blackwell, cuidando do provisionamento de instâncias, orquestração de contêineres e integração com serviços AWS como Amazon Simple Storage Service (Amazon S3), Amazon CloudWatch, Amazon Elastic Container Registry (Amazon ECR) e AWS Identity and Access Management (AWS IAM).
Pré-requisitos
Antes de começar, confirme que você tem:
- Uma conta AWS com permissões para criar jobs de treinamento no SageMaker AI, acessar o ECR e criar roles de execução no IAM.
- Acesso a instâncias
ml.p6-b200.48xlargevia Flexible Training Plan ou Managed Spot Training. Verifique suas Service Quotas antes de começar. - Docker instalado localmente, Python 3.9 ou superior e o SageMaker Python SDK.
- Familiaridade com PyTorch e Paralelismo de Dados Totalmente Fragmentado do PyTorch (FSDP). Se você é novo no FSDP, consulte o guia de introdução ao treinamento distribuído.
Passo 1: Crie seu script de treinamento
Faça o download do arquivo fsdp.py do exemplo FSDP do repositório NVIDIA TransformerEngine. Esse script implementa o treinamento com FSDP e aceita hiperparâmetros como argumentos de linha de comando.
Passo 2: Crie o script de entrada
Prepare um arquivo train.sh para configurar o torchrun e iniciar o script de treinamento:
#!/bin/bash
# SageMaker passes hyperparameters as environment variables (SM_HP_)
PRECISION=${SM_HP_PRECISION:-"mxfp8"}
NUM_LAYERS=${SM_HP_NUM_LAYERS:-10}
BATCH_SIZE=${SM_HP_BATCH_SIZE:-8}
SEQ_LENGTH=${SM_HP_SEQ_LENGTH:-2048}
NUM_GPUS=$(nvidia-smi --list-gpus | wc -l)
torchrun --standalone --nnodes=1 --nproc-per-node="$NUM_GPUS" \
fsdp.py --no-defer-init --precision "$PRECISION" \
--num-layers "$NUM_LAYERS" --checkpoint-layer "transformerlayer" \
--batch-size "$BATCH_SIZE" --seq-length "$SEQ_LENGTH" Passo 3: Construa e envie seu contêiner
Construa um contêiner Docker personalizado que estenda os AWS Deep Learning Containers (DLC), inclua os arquivos fsdp.py e train.sh, e tenha o TransformerEngine 2.11 instalado. O DLC fornece uma imagem base validada com PyTorch e as bibliotecas CUDA necessárias para compatibilidade com o Blackwell. Abaixo está o 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 LD_LIBRARY_PATH=/usr/local/lib/python3.12/site-packages/nvidia/cuda_runtime/lib:$LD_LIBRARY_PATH
COPY fsdp.py /opt/ml/code/fsdp.py
COPY train.sh /opt/ml/code/train.sh
ENV SAGEMAKER_SUBMIT_DIRECTORY /opt/ml/code
ENV SAGEMAKER_PROGRAM train.shApós a construção, crie um repositório privado no ECR (se ainda não tiver um) e envie a imagem para lá. Para instruções detalhadas de build, consulte a documentação sobre como adaptar seu próprio contêiner Docker para o Amazon SageMaker AI.
Passo 4: Reserve capacidade
Reserve capacidade via Flexible Training Plan para acesso previsível à taxa padrão, ou use Managed Spot Training para workloads com tolerância a interrupções e foco em redução de custos. Use Flexible Training Plans para execuções de produção que exigem disponibilidade contínua; use Spot para experimentação e workloads tolerantes a falhas onde a redução de custo supera o risco de interrupção.
Instâncias Spot estão sujeitas a interrupções — certifique-se de que seu script de treinamento salva checkpoints no Amazon S3 em intervalos regulares. O SageMaker AI retoma automaticamente um job Spot interrompido se você fornecer um checkpoint_s3_uri na configuração do estimador.
Atenção: Flexible Training Plans reservam capacidade e geram cobranças pela duração do plano, independentemente de jobs estarem em execução. Revise os detalhes de preços antes de criar um plano.
Passo 5: Submeta o job de treinamento
Substitua os valores de exemplo pelo ARN do seu training plan e pela URI da imagem ECR, e execute o código abaixo no seu ambiente de desenvolvimento local ou em uma instância de notebook do SageMaker AI:
from sagemaker.estimator import Estimator
from sagemaker import get_execution_role
from sagemaker.debugger import ProfilerConfig
training_plan_arn = "" # Replace with your training plan ARN
ecr_image = "" # Replace with your ECR image URI
# Adjust these values to match your workload
precision = "mxfp8"
num_layers = 10
batch_size = 8
seq_length = 2048
estimator = Estimator(
image_uri=ecr_image,
role=get_execution_role(),
base_job_name='blackwell-training',
instance_count=1,
instance_type='ml.p6-b200.48xlarge',
hyperparameters={
"precision": precision,
"num-layers": num_layers,
"batch-size": batch_size,
"seq-length": seq_length,
},
profiler_config=ProfilerConfig(disable_profiler=True),
training_plan=training_plan_arn)
estimator.fit() Para Managed Spot Training, substitua training_plan por use_spot_instances=True, defina max_run e max_wait, e adicione um checkpoint_s3_uri para retomada automática.
Passo 6: Monitore o job de treinamento
O SageMaker AI transmite logs para o Amazon CloudWatch automaticamente. No console do SageMaker AI, navegue até Training jobs e selecione seu job para encontrar o grupo de logs do CloudWatch. Abra /aws/sagemaker/TrainingJobs e procure por linhas [rank 0] para valores de perda e throughput. Para confirmar que o formato de precisão foi carregado, procure por mensagens como "Using FP8 recipe" ou "MXFP8 enabled".
Se o job parar com um erro de falta de memória CUDA (OOM), o log mostrará o tamanho da alocação. Reduza o batch size ou a sequence length, ou adicione activation checkpointing se ainda não tiver feito isso.
Limpeza de recursos
Para evitar cobranças contínuas após seus testes, encerre quaisquer jobs de treinamento em execução no console do SageMaker AI (note que isso não cancela um Flexible Training Plan).
Atenção: Os passos de limpeza abaixo excluem recursos permanentemente e não podem ser desfeitos. Verifique se você fez backup de todos os dados necessários antes de prosseguir.
- Exclua seu repositório Amazon ECR (isso remove permanentemente as imagens de contêiner).
- Exclua quaisquer artefatos de treinamento armazenados no Amazon S3 (isso remove permanentemente dados de treinamento e checkpoints).
- Remova os grupos de logs do CloudWatch (isso remove permanentemente os logs de treinamento).
- Exclua a role de execução do IAM criada para o SageMaker AI.
Conclusão
A combinação de memória expandida, interconexão de alta largura de banda e suporte a formatos de precisão reduzida no NVIDIA Blackwell muda o que é viável para treinamento de modelos grandes. Modelos transformer de 1B a 64B parâmetros apresentam ganhos consistentes quando essas otimizações são combinadas nas instâncias P6-B200 do SageMaker AI.
A chave está em entender se o workload é limitado por computação ou por memória, e então aplicar as mudanças de forma incremental para medir o impacto de cada uma. Para quem quiser explorar as opções de instância e detalhes de configuração, a documentação do Amazon SageMaker AI é o ponto de partida, ou é possível adquirir um Flexible Training Plan para reservar capacidade Blackwell. Dúvidas sobre workloads específicos podem ser direcionadas ao AWS Support ou ao time de conta AWS.
Fonte
Optimize model training on Amazon SageMaker AI with NVIDIA Blackwell (https://aws.amazon.com/blogs/machine-learning/optimize-model-training-on-amazon-sagemaker-ai-with-nvidia-blackwell/)
Leave a Reply