O dilema entre controle algorítmico e simplicidade operacional
Fazer o ajuste fino de um Modelo de Linguagem de Grande Escala (LLM) costuma impor uma escolha difícil: ou você usa um serviço totalmente gerenciado — onde basta selecionar o modelo base, fornecer dados e hiperparâmetros — ou parte para uma abordagem menos gerenciada, ganhando controle sobre o algoritmo, mas herdando toda a complexidade de infraestrutura e frameworks de treinamento.
O problema aparece quando uma receita padrão não resolve: quando é preciso implementar uma função de perda customizada, um currículo de dados para Ajuste Fino Supervisionado (SFT), uma otimização de preferência própria, uma função de recompensa para Aprendizado por Reforço (RL) ou um algoritmo de alinhamento recém-publicado. Nesse momento, para ganhar liberdade algorítmica, a equipe acaba assumindo toda a gestão de infraestrutura GPU.
A AWS acaba de expandir o Serverless Model Customization do Amazon SageMaker AI com as Training Session APIs, que eliminam esse dilema ao mover a fronteira de abstração um nível abaixo. Em vez de abstrair o algoritmo, essas APIs abstraem o compute. Operações como forward_backward(), optim_step() e sample() tornam-se chamadas de API do SageMaker que podem ser compostas em qualquer procedimento que a técnica exigir: SFT, otimização de preferência, RL ou algo completamente novo.
As chamadas são direcionadas a uma sessão ativa e com estado (stateful), de modo que o modelo permanece carregado nas GPUs entre os passos. O serviço cuida de toda a configuração de runtime, execução distribuída, sincronização de gradientes, recuperação de falhas, isolamento de tenants e operação das GPUs em escala. Quem desenvolve escreve o loop de treinamento; o SageMaker AI roda a infraestrutura.
Junto com o serviço, a AWS também publicou o SageMaker training sessions cookbook, um conjunto de receitas prontas de SFT e RL que podem servir de ponto de partida para loops de treinamento personalizados.
O que são as Training Session APIs do SageMaker AI?
Com as Training Session APIs, é possível conduzir uma sessão ativa, respaldada por GPUs, a partir do seu próprio processo Python na máquina local. Você faz as chamadas de operação, e o SageMaker AI as executa em aceleradores gerenciados, já carregados com o modelo base escolhido e o adaptador de Adaptação de Baixa Classificação (LoRA). Não há arquivo de configuração descrevendo todo o job antecipadamente nem etapa de criação de cluster. Você controla o que acontece em cada passo, de forma interativa, direto do código — o que permite iterar e depurar rapidamente.
A divisão de responsabilidades fica assim:
- Você cuida de: seus dados, ambientes de RL, agentes e algoritmo. Você escreve um loop de treinamento em Python que chama operações como
forward_backward(),optim_step()esample()para os passos que precisam de GPUs. - O SageMaker AI cuida de: alocação de GPU, carregamento do modelo e LoRA, execução distribuída, recuperação de falhas, armazenamento de checkpoints e segurança.
Os três clientes do SDK
A API é organizada em três clientes no SageMaker Python SDK, criados em sequência:
- ServiceClient — armazena configurações de plataforma (role de execução, grupo de pacotes de modelo, localização de saída no Amazon S3 e, opcionalmente, chave AWS KMS e configuração de VPC) e inicia a sessão.
- TrainingClient — instanciado pelo ServiceClient, concentra operações de treinamento e checkpointing como
forward,forward_backward,optim_stepesave_weights_for_sampler. - SamplingClient — instanciado pelo TrainingClient, oferece
sampleecompute_logprobspara geração e scoring durante rollouts e avaliação.
Uma sessão ativa também expõe um endpoint de chat completions compatível com a API da OpenAI, permitindo que código existente que use o cliente OpenAI faça sampling diretamente.
O formato de um passo de treinamento com RL
Para Aprendizado por Reforço, o loop tem um formato característico com quatro estágios. As operações que rodam nas GPUs gerenciadas (sampling, gradiente e checkpoint) ficam no serviço, enquanto o controle do loop e a lógica de recompensa rodam na sua máquina:
- (1) Sample — gera respostas a partir do modelo atual nas GPUs.
- (2) Score — seu código Python pontua as saídas localmente (verificador, modelo de recompensa, harness de testes, etc.).
- (3) forward_backward + optim_step — transforma as pontuações em uma atualização de gradiente nas GPUs.
- (4) save_weights_and_get_sampling_client() — atualiza o sampler para que o próximo passo use o modelo atualizado.
As mesmas primitivas cobrem SFT e outras técnicas de treinamento de LLMs. Para SFT, basta remover os passos de sampling e recompensa e chamar forward_backward com uma perda cross_entropy sobre exemplos rotulados.
Para quem é e onde se encaixa no Amazon SageMaker AI?
As Training Session APIs são o ponto de partida certo quando a equipe precisa de ciclos rápidos de iteração e liberdade algorítmica — como uma perda customizada, um novo sinal de recompensa de RL ou um pipeline multi-estágio — mas não quer provisionar ou operar infraestrutura GPU.
A cobrança é feita por token processado, de modo que só se paga pelo compute que o loop efetivamente consome. Se a função de recompensa leva 10 segundos para pontuar um lote na máquina local, nenhum tempo de GPU é cobrado durante essa execução. Isso traz duas vantagens práticas em relação à cobrança por hora de GPU:
- Sem custo de GPU ociosa: computação de recompensa, carregamento de dados, logging e debugging acontecem fora do relógio de cobrança.
- Custo escala com o dataset e volume de geração, não com tempo de parede ou tamanho do cluster, o que facilita a previsão de gastos e a comparação de experimentos.
Os perfis de uso incluem:
- Pesquisadores de pós-treinamento e alinhamento implementando métodos de RL (GRPO, PPO, variantes de importance-sampling), estimadores de vantagem customizados ou novos objetivos.
- Equipes fazendo SFT ou otimização de preferência que precisam de uma perda não-padrão, currículo de dados ou pipeline multi-estágio.
- Construtores de RL agêntico e multi-turno que precisam de lógica de rollout customizada, reward shaping e currículos.
- Engenheiros de Aprendizado de Máquina (ML) que superaram o fine-tuning baseado em configuração, mas não querem o ônus operacional de clusters GPU.
- Pesquisadores portando código de loops de treinamento existentes, já que o modelo de programação baseado em operações mapeia diretamente para loops escritos com APIs de baixo nível.
Os modelos base suportados e as Regiões AWS disponíveis estão listados na documentação das Training Session APIs.
Posicionamento entre as opções de customização do SageMaker AI
As Training Session APIs ocupam um ponto distinto entre as opções de customização de modelos do SageMaker AI: oferecem substancialmente mais controle sobre o algoritmo de treinamento do que o Serverless Model Customization totalmente gerenciado, e muito menos responsabilidade operacional do que treinar com Amazon SageMaker Training Jobs ou Amazon SageMaker HyperPod.
Em resumo: o Serverless Model Customization com APIs gerenciadas oferece baixo controle (escolha modelo, configure hiperparâmetros, forneça dados) e infraestrutura serverless. As Training Session APIs oferecem alto controle (você escreve o loop e chama primitivas) e também são serverless. Já o SageMaker Training Jobs e o HyperPod oferecem controle total, mas exigem seleção de instâncias ou gerenciamento de clusters persistentes.
Anatomia de um loop de treinamento: GRPO em problemas de matemática
Para ilustrar como as operações se encaixam, o artigo original detalha um loop mínimo de RL sobre o GSM8K, um benchmark com aproximadamente 8.500 problemas de matemática de nível escolar, cada um com uma resposta numérica única. Exemplo:
Natalia sold clips to 48 of her friends in April, and then she sold half as many clips in May. How many clips did Natalia sell altogether in April and May? → 72
O treinamento usa um adaptador LoRA para o GPT-OSS 20B da OpenAI com Otimização de Política Relativa de Grupo (GRPO): para cada pergunta, o modelo gera várias respostas, cada uma é pontuada por correção, e respostas que superam a média do grupo são reforçadas. O exemplo completo e executável está no repositório de exemplos do Amazon SageMaker.
Abrindo a sessão
Antes de abrir uma sessão, é necessário completar os pré-requisitos na documentação. Em seguida, cria-se a sessão, obtém-se o tokenizer e um SamplingClient:
from sagemaker.train.training_session import ServiceClient, SamplingParams, AdamParams
svc = ServiceClient(role_arn=role, region=REGION)
training_client = svc.create_lora_training_client(
"openai/gpt-oss-20b", # a name from svc.list_base_models()
training_session_name=SESSION_NAME,
rank=32, alpha=64, # LoRA rank and scaling factor
model_package_group_arn=MODEL_PACKAGE_GROUP_ARN,
s3_output_path=OUTPUT_S3_PATH,
)
tokenizer = training_client.get_tokenizer()
sampling_client = training_client.create_sampling_client()Definindo recompensa e vantagem
A lógica específica da tarefa são duas funções curtas que rodam na máquina local:
def compute_reward(response_text, ground_truth):
predicted = extract_answer(response_text) # parse the number after "#### "
return 1.0 if predicted == ground_truth else 0.0
def compute_advantages(rewards):
rewards = np.array(rewards)
advantages = rewards - rewards.mean() # center within the group (the GRPO baseline)
return advantages / (rewards.std() + 1e-8) # normalizeO loop principal
Cada passo segue os quatro estágios descritos anteriormente:
for step in range(NUM_STEPS):
datums, pending = [], []
# (1) sample K responses per question in one request each (GPU); returns a future right away
for question, ground_truth in draw(PROMPTS_PER_STEP): # data loading is local
prompt = tokenizer.encode(build_prompt(question), add_special_tokens=False)
pending.append((prompt, ground_truth,
sampling_client.sample(
prompt, num_samples=K,
sampling_params=SamplingParams(max_tokens=512, temperature=0.7, top_p=0.9))))
for prompt, ground_truth, future in pending:
result = future.result() # wait once the batch is in flight
# (2) score each response, then center the rewards within the group (local)
rewards = [compute_reward(tokenizer.decode(s.tokens), ground_truth)
for s in result.sequences]
advantages = compute_advantages(rewards)
datums += [build_training_datum(prompt, s, a)
for s, a in zip(result.sequences, advantages)]
# (3) one gradient step over the batch (GPU); the optimizer step is sequenced behind it
fwd_bwd_op = training_client.forward_backward(datums, loss_fn="ppo")
optim_op = training_client.optim_step(AdamParams(learning_rate=1e-5))
fwd_bwd_op.result(), optim_op.result()
# (4) refresh the sampler so the next step's rollouts come from the updated policy (GPU)
sampling_client = training_client.save_weights_and_get_sampling_client()
training_client.stop()Detalhes importantes do loop
Três detalhes fazem esse loop funcionar:
- Operações retornam futures. Toda chamada GPU retorna um
APIFutureimediatamente: a requisição é enfileirada e enviada por um submitter em background, e.result()aguarda o valor. O loop submete todas as requisições de sampling de um passo antes de esperar qualquer uma, de modo que todo o lote gera concorrentemente. Operações de treinamento rodam na ordem em que são chamadas, entãooptim_steppode ser submetido antes deforward_backwardter terminado. - O datum carrega tudo que a perda precisa.
build_training_datumempacota os tokens do prompt e da resposta, os logprobs por token dosample(), a vantagem e uma máscara de pesos que mantém o prompt fora da perda. - Checkpoints são pacotes de modelo. Cada
saveregistra uma versão no grupo de pacotes de modelo. Após o treinamento, dois saves adicionais criam um modelo mesclado para deploy e um estado resumível com pesos e otimizador.
training_client.save_weights_for_sampler(merge_weights=True).result() # merged, deployable model
training_client.save_state().result() # resumable state (weights + optimizer)Para retomar, basta iniciar uma nova sessão a partir do último checkpoint com create_training_client_from_state_with_optimizer. Uma sessão ativa conta contra a cota de sessões concorrentes até que stop() seja chamado ou ela expire. O checkpoint mesclado pode ser implantado em um endpoint de tempo real. Para servir checkpoints como adaptadores LoRA, consulte o notebook de deploy do cookbook.
No teste do exemplo, 10 passos de treinamento sobre 40 problemas elevaram a acurácia em dados de validação de 56% para 75%.
Caso de uso real: RL em programação competitiva
O GSM8K mantém a mecânica visível, mas um workload de RL em produção adiciona mais complexidade: gerações mais longas, interação multi-turno, um avaliador externo e throughput suficiente para manter as GPUs ocupadas. O SageMaker training sessions cookbook empacota essa engenharia em loops de treinamento reutilizáveis construídos sobre as mesmas operações.
Para uma nova tarefa de RL, basta escrever um environment: o que o modelo vê, como a resposta é pontuada e se há outro turno. O cookbook cuida de batching, templates de chat, estimativa de vantagem, overlap entre sampling e treinamento, validação, checkpointing com resume e métricas por passo.
A receita code_rl usa essa estrutura para treinar um modelo em problemas de programação competitiva:
- A tarefa: os problemas vêm do dataset DeepCoder (PrimeIntellect, TACO e LiveCodeBench para treinamento; Codeforces e LiveCodeBench para teste). O modelo escreve uma solução em Python, que é executada contra os testes unitários do problema em um contêiner SandboxFusion na máquina local. A recompensa é 1 quando todos os testes passam e 0 caso contrário, com penalidade de 0.1 quando a resposta não contém bloco de código.
- Uso de ferramentas multi-turno: com dois turnos, o modelo pode chamar uma ferramenta
check_solution, ler os resultados dos testes e revisar a resposta antes da nota final. As chamadas de ferramenta e resultados são renderizados no formato de chat nativo do modelo.
O que a receita adiciona ao loop mínimo:
- Overlap entre sampling e treinamento: workers de rollout continuam gerando enquanto o trainer executa um passo. Um grupo de respostas amostrado de uma política mais antiga é usado se tiver no máximo k passos de defasagem (caso contrário é reenfileirado), e a perda
cispocorrige a diferença de política. - Computação de gradiente em stream:
forward_backwardé submetido em sub-lotes à medida que os grupos terminam, em vez de uma única vez por passo. - Validação em hold-out, resume de checkpoint e métricas por passo, opcionalmente publicadas no SageMaker managed MLflow.
No treinamento descrito, um adaptador LoRA (rank 32) para GPT-OSS 20B foi treinado por 50 passos, cada um com 128 problemas e 8 respostas por problema, usando dois turnos e limite de defasagem de k=3. A recompensa de treinamento subiu de 0.46 para 0.60, e a acurácia em 200 problemas de teste reservados subiu de 40% para 50.5%.
Para executar, basta iniciar o sandbox e lançar a receita:
bash tools/sandbox_up.sh
python -m training_sessions_cookbook.recipes.code_rl.train \
--config training_sessions_cookbook/recipes/code_rl/configs/coderl_oss20b_service.json \
--max_steps_off_policy 3 --stale_policy requeue --num_rollout_workers 80 --max_steps 50 \
--base_model_arn "$BASE_MODEL_ARN" \
--model_package_group_arn "$MODEL_PACKAGE_GROUP_ARN" \
--role_arn "$ROLE_ARN" \
--output_s3_path "$OUTPUT_S3_PATH"O cookbook traz mais receitas na mesma estrutura, incluindo SFT em datasets de chat (Tulu 3, UltraChat), RL no GSM8K em escala completa e alinhamento de preferência KTO a partir de feedback binário. O README de cada receita lista os comandos, hiperparâmetros e métricas esperadas.
Conclusão
As Training Session APIs do SageMaker AI movem a fronteira de abstração para o nível da operação de treinamento: você escreve o loop em Python na sua máquina, e o serviço o executa em GPUs gerenciadas, cuidando do carregamento de modelo e LoRA, execução distribuída, recuperação de falhas e checkpointing.
Um loop GRPO completo se resume a quatro operações — sample, forward_backward, optim_step e save_weights_and_get_sampling_client — com a lógica de recompensa e vantagem no seu próprio código. As mesmas operações escalam para RL multi-turno com uso de ferramentas e avaliação em sandbox, como demonstra a receita Code RL.
Recompensas customizadas, rollouts agênticos, debugging passo a passo e novas técnicas publicadas — tudo usa o mesmo modelo de programação: o algoritmo fica no seu loop, e o SageMaker AI roda as GPUs.
Para começar:
- Escolha uma receita do SageMaker training sessions cookbook, como a Code RL.
- Execute o exemplo GSM8K do repositório de exemplos do Amazon SageMaker e substitua pela sua própria função de recompensa.
- Leia a documentação das Training Session APIs para pré-requisitos, cotas, Regiões, modelos suportados e referência da API.
Fonte
Write your own training loop — without managing GPUs — with Amazon SageMaker training session APIs (https://aws.amazon.com/blogs/machine-learning/write-your-own-training-loop-without-managing-gpus-with-amazon-sagemaker-training-sessions/)
Leave a Reply