Escribe tu propio bucle de entrenamiento — sin administrar GPUs — con las APIs de sesión de entrenamiento de Amazon SageMaker

Ajustar finamente un modelo de lenguaje de gran escala (LLM) hoy en día suele forzar una disyuntiva entre el control del bucle de entrenamiento y la simplicidad operativa. Con servicios completamente administrados, seleccionas un modelo base, proporcionas datos e hiperparámetros, y el servicio se encarga de todo lo demás, lo cual funciona bien cuando una receta estándar y preempaquetada se ajusta a tu problema. Los servicios menos administrados te dan más control, pero también debes administrar más de la infraestructura subyacente y los marcos de entrenamiento. A veces una receta estándar no es suficiente: necesitas una función de pérdida personalizada o un plan de datos para el ajuste fino supervisado (SFT), tu propia optimización de preferencias, una función de recompensa de aprendizaje por refuerzo (RL) o un algoritmo de alineación publicado la semana pasada. En cuanto necesitas dar forma al procedimiento de entrenamiento tú mismo, asumes ese trabajo de infraestructura para ganar libertad algorítmica.

Hoy, la personalización sin servidor de modelos en Amazon SageMaker AI se expande con las APIs de sesión de entrenamiento, que eliminan esa disyuntiva al trazar el límite de abstracción un nivel más abajo. En lugar de abstraer el algoritmo, abstraen la computación. Operaciones como forward_backward(), optim_step(), y sample() se convierten en llamadas a la API de SageMaker que compones en el procedimiento que tu técnica requiera: SFT, optimización de preferencias, RL o algo nuevo. Las llamadas se dirigen a una sesión activa y con estado, de modo que el modelo permanece cargado en las GPUs entre pasos, y puedes realizar tantas llamadas como necesites contra él. El servicio se encarga de la configuración del tiempo de ejecución, la ejecución distribuida, la sincronización de gradientes, la recuperación ante fallos, el aislamiento de inquilinos y las operaciones de ejecutar GPUs a escala. Tú escribes el bucle de entrenamiento, y SageMaker AI ejecuta la infraestructura.

Junto con el servicio, lanzamos el SageMaker training sessions cookbook, un conjunto de recetas SFT y RL listas para ejecutar que puedes usar como planos para tus propios bucles de entrenamiento.

En esta publicación, presentamos el modelo de programación y los tres clientes del SDK, y mostramos dónde encajan las APIs de sesión de entrenamiento entre las opciones de entrenamiento y personalización de modelos en SageMaker AI. Luego desarmamos un bucle RL mínimo en problemas de matemática de primaria (GSM8K) para mostrar cómo se ensamblan las operaciones, y terminamos con una carga de trabajo del mundo real del cookbook: aprendizaje por refuerzo en problemas de programación competitiva, donde el código generado se califica ejecutándolo contra pruebas unitarias.

¿Qué son las APIs de sesión de entrenamiento de SageMaker AI?

Con las APIs de sesión de entrenamiento de SageMaker AI, puedes controlar una sesión activa con GPUs desde tu propio proceso de Python en tu máquina local. Tú llamas a las operaciones, y SageMaker AI las ejecuta en aceleradores administrados cargados con tu modelo base elegido más tu adaptador Low-Rank Adaptation (LoRA). No hay un archivo de configuración que describa todo el trabajo por adelantado ni un paso de creación de clúster. Controlas lo que ocurre en cada paso, de forma interactiva, desde el código, para que puedas iterar y depurar rápidamente.

La división de responsabilidades se ve así:

La API se organiza en torno a tres clientes del SageMaker Python SDK, creados en secuencia. Un ServiceClient contiene la configuración a nivel de plataforma (rol de ejecución, grupo de paquetes de modelos, ubicación de salida de Amazon Simple Storage Service (Amazon S3) y, opcionalmente, una clave de AWS Key Management Service (AWS KMS) y una configuración de nube privada virtual (VPC), que también puedes sobrescribir por sesión) e inicia una sesión. El ServiceClient instancia un TrainingClient (operaciones de entrenamiento y checkpointing como forward, forward_backward, optim_step, y save_weights_for_sampler), que a su vez instancia un SamplingClient (sample y compute_logprobs, para generación y puntuación durante los rollouts y la evaluación). Cada llamada a la API se corresponde con una operación en tu bucle. Una sesión activa también sirve un endpoint de chat completions compatible con OpenAI, de modo que el código cliente de OpenAI existente puede muestrear a partir de él.

Figura 1: La jerarquía de clientes del SageMaker Python SDK para las API de sesión de entrenamiento

Para aprendizaje por refuerzo, el bucle tiene una forma característica. La siguiente figura muestra un paso de entrenamiento y dónde se ejecuta cada parte: los pasos en verde llaman a API en GPU administradas por el servicio, y los pasos en morado, el control de tu bucle y la lógica de recompensa, se ejecutan en tu máquina.

Figura 2: Un paso de entrenamiento, con operaciones de muestreo, gradiente y checkpoint en GPU administradas y el control del bucle y la lógica de recompensa en tu máquina

Cada paso carga un lote de prompts y (1) muestrea a partir del modelo actual en la GPU. Tu código Python luego (2) puntúa esas salidas: un verificador, una llamada a un modelo de recompensa, un arnés de pruebas, cualquier cosa que puedas expresar en código. (3) forward_backward y optim_step conviertes las puntuaciones en una actualización de gradiente en la GPU, y (4) save_weights_and_get_sampling_client() actualiza el muestreador para que el siguiente paso extraiga del modelo actualizado.

Tú controlas los datos, el control del bucle y la recompensa; SageMaker AI ejecuta las operaciones de muestreo, gradiente y checkpoint vinculadas a la GPU. Las mismas primitivas cubren el ajuste fino supervisado (SFT) y otras técnicas de entrenamiento de LLM. Para SFT, elimina los pasos de muestreo y recompensa y llama a forward_backward con una cross_entropy pérdida sobre ejemplos etiquetados.

¿Para quién es esto y dónde encaja dentro de Amazon SageMaker AI?

Las APIs de sesión de entrenamiento son el punto de partida adecuado cuando necesitas ciclos de iteración rápidos y libertad algorítmica, como una pérdida personalizada, una nueva señal de recompensa de RL o un pipeline de varias etapas, pero no quieres aprovisionar ni operar la infraestructura de GPU tú mismo.

El uso se factura por token procesado, por lo que solo pagas por la computación que consume tu bucle. Si tu función de recompensa tarda 10 segundos en puntuar un lote en tu máquina, no se factura tiempo de GPU mientras se ejecuta. En comparación con la facturación horaria de GPU, esto tiene dos ventajas prácticas. Primera, no hay costo de GPU inactiva: el cálculo de recompensas, la carga de datos, el registro y la depuración ocurren fuera del reloj, y solo se te factura cuando llamas a una operación vinculada a la GPU. Segunda, el costo escala con tu conjunto de datos y tu volumen de generación, no con el tiempo transcurrido ni con el tamaño del clúster, lo que hace que el gasto sea más fácil de predecir y los experimentos más fáciles de comparar. El ahorro frente a reservar GPUs para toda la sesión puede ser significativo para cargas de trabajo con computación local intensiva, como modelos de recompensa complejos, despliegues agénticos multi-herramienta o evaluación humana en el bucle.

Las APIs de sesión de entrenamiento están diseñadas para equipos cuyo trabajo está definido por el algoritmo de entrenamiento, no solo por el conjunto de datos. Como compones el bucle tú mismo, las mismas primitivas admiten el ajuste fino supervisado (SFT) con una cross_entropy pérdida sobre datos etiquetados, el aprendizaje por refuerzo con ppo, importance_sampling, o cispo (GRPO, PPO y otros métodos relativos a grupos), y objetivos personalizados, como la optimización de preferencias o la destilación, construidos a partir de las mismas piezas. Encajan para:

• Investigadores de post-entrenamiento y alineación que implementan métodos de RL (GRPO, PPO, variantes de muestreo de importancia), estimadores de ventaja personalizados u objetivos nuevos.
• Equipos que realizan ajuste fino supervisado u optimización de preferencias y necesitan una pérdida no estándar, un currículo de datos o un pipeline de varias etapas que las recetas basadas en configuración no expresan.
• Desarrolladores de RL agéntico y multi-turno que necesitan lógica de despliegue personalizada, moldeado de recompensas y currículos.
• Ingenieros de aprendizaje automático (ML) que han superado el ajuste fino basado en configuración pero no quieren asumir la carga operativa de clústeres de GPU.
• Investigadores que trasladan código existente de bucles de entrenamiento, ya que el modelo de programación basado en operaciones se mapea directamente sobre bucles escritos con APIs de bajo nivel forward_backward, optim_step, y sample .

Los modelos base admitidos y las regiones de AWS donde cada uno está disponible se enumeran en la documentación de las APIs de sesión de entrenamiento.

Dónde encajan las APIs de sesión de entrenamiento dentro de Amazon SageMaker AI

Las APIs de sesión de entrenamiento ocupan un punto distintivo entre las opciones de personalización de modelos en SageMaker AI: sustancialmente más control sobre el algoritmo de entrenamiento que la Personalización de modelos sin servidor, y mucha menos propiedad operativa que entrenar en Amazon SageMaker Training Jobs o Amazon SageMaker HyperPod.

Figura 3: Dónde encajan las APIs de sesión de entrenamiento entre las opciones de personalización de modelos en SageMaker AI

Anatomía de un bucle de entrenamiento: GRPO en problemas de matemáticas

Esta sección desmonta un bucle mínimo de aprendizaje por refuerzo para mostrar cómo encajan las operaciones. La tarea es GSM8K, un benchmark de aproximadamente 8,500 problemas de matemáticas de primaria con enunciado, cada uno con una única respuesta numérica, como en el siguiente ejemplo:

Natalia vendió clips a 48 de sus amigos en abril, y luego vendió la mitad de esa cantidad en mayo. ¿Cuántos clips vendió Natalia en total entre abril y mayo? → 72

Entrenamos un adaptador LoRA para el GPT-OSS 20B de OpenAI con Group Relative Policy Optimization (GRPO): para cada pregunta el modelo genera varias respuestas, cada respuesta se puntúa según su corrección, y se refuerzan las respuestas que superan el promedio del grupo. El ejemplo completo y ejecutable está en el repositorio de ejemplos de Amazon SageMaker.

Antes de abrir una sesión, complete los requisitos previos de la documentación.

Abrir una sesión

Cree una sesión y, a continuación, obtenga el tokenizador y un SamplingClient a partir de ella:

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()

Definir la recompensa y la ventaja

La lógica específica de la tarea son dos funciones breves que se ejecutan en su máquina:

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)     # normalize

El bucle

Cada paso sigue las cuatro etapas mostradas en la Figura 2. draw, build_prompt, y build_training_datum son ayudantes breves del ejemplo:

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()

Tres detalles hacen que este bucle funcione:

Las operaciones devuelven futuros. Cada llamada dependiente de la GPU devuelve un APIFuture inmediatamente: la solicitud se pone en cola y la envía un remitente en segundo plano, y .result() espera el valor. El bucle envía todas las solicitudes de muestreo de un paso antes de esperar cualquiera de ellas, de modo que todo el lote se genera de forma concurrente. Las operaciones de entrenamiento se ejecutan en el orden en que se las llama, por lo que optim_step puede enviarse antes de que forward_backward haya terminado.

El dato lleva todo lo que la pérdida necesita. build_training_datum empaqueta los tokens del prompt y de la respuesta, los logprobs por token de sample(), la ventaja y una weights máscara que mantiene el prompt fuera de la pérdida.

Los checkpoints son paquetes de modelos. Cada guardado registra una versión en su grupo de paquetes de modelos, incluida la actualización por pasos del muestreador en el bucle anterior. Después del entrenamiento, dos guardados más añaden un modelo fusionado para despliegue y un estado reanudable con pesos y optimizador.

training_client.save_weights_for_sampler(merge_weights=True).result()   # merged, deployable model
training_client.save_state().result()                                   # resumable state (weights + optimizer)

Para reanudar, inicie una nueva sesión desde el último checkpoint con create_training_client_from_state_with_optimizer.

Una sesión en ejecución consume su cuota de sesiones concurrentes hasta que llame a stop() o esta expire, así que deténgala cuando termine, también en las rutas de error, o abra el cliente en un bloque with . El checkpoint fusionado se despliega en un endpoint en tiempo real. Para servir los checkpoints como adaptadores LoRA en su lugar, consulte el cookbook de cuaderno de despliegue.

En nuestra ejecución de la muestra, 10 pasos de entrenamiento sobre 40 problemas elevaron la precisión en datos no vistos del 56 por ciento al 75 por ciento.

Una carga de trabajo del mundo real: RL en programación competitiva

GSM8K mantiene visibles los mecanismos, pero una carga de trabajo de RL en producción añade más: generaciones más largas, interacción multi-turno, un calificador externo y suficiente rendimiento para mantener las GPUs ocupadas. El cookbook de sesiones de entrenamiento de SageMaker empaqueta esa ingeniería en bucles de entrenamiento reutilizables construidos sobre las mismas operaciones. Para una nueva tarea de RL se escribe un entorno: qué ve el modelo, cómo se puntúa su respuesta y si obtiene otro turno. El cookbook se encarga del batching, las plantillas de chat, la estimación de ventaja, el muestreo y entrenamiento solapados, la validación, el checkpointing con reanudación y las métricas por paso, y registra la configuración de cada ejecución para que pueda repetirse. La code_rl receta lo usa para entrenar un modelo en problemas de programación competitiva.

La tarea. Los problemas provienen del conjunto de datos DeepCoder (PrimeIntellect, TACO y LiveCodeBench para entrenamiento; Codeforces y LiveCodeBench para pruebas). El modelo escribe una solución en Python, que se ejecuta contra las pruebas unitarias del problema en un contenedor de SandboxFusion en su máquina. La recompensa es 1 cuando todas las pruebas pasan y 0 en caso contrario, con una penalización de 0.1 cuando la respuesta no contiene ningún bloque de código.

Uso de herramientas multi-turno. Con dos turnos, el modelo puede llamar a una check_solution herramienta, leer los resultados de las pruebas y revisar su respuesta antes de la calificación final. Las llamadas a la herramienta y los resultados se representan en el formato de chat nativo del modelo, de modo que el modelo se entrena con el mismo formato de tokens que usa en la inferencia.

Lo que la receta añade al bucle mínimo:

• Muestreo y entrenamiento solapados: Los trabajadores de rollout siguen generando mientras el entrenador ejecuta un paso. Un grupo de respuestas muestreadas de una política anterior se usa si tiene como máximo k pasos de antigüedad y se vuelve a encolar en caso contrario, y la cispo pérdida corrige la diferencia de política.
• Cálculo de gradientes en streaming: forward_backward se envía en sublotes a medida que los grupos terminan, en lugar de una vez por paso.
• Validación en datos no vistos, reanudación de checkpoints y métricas por paso, opcionalmente publicadas en SageMaker managed MLflow.

Figura 4: La configuración de RL para código, con el bucle de entrenamiento local y el sandbox a la izquierda y las GPUs administradas a la derecha

Entrenamos un adaptador LoRA (rango 32) para GPT-OSS 20B durante 50 pasos, cada uno con 128 problemas y 8 respuestas por problema, usando dos turnos y un límite de obsolescencia de k=3. La recompensa de entrenamiento subió de 0.46 a 0.60, y la precisión en 200 problemas de prueba no vistos subió del 40 por ciento al 50.5 por ciento.

Para ejecutarlo, inicie el sandbox y lance la receta:

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"

El libro de recetas tiene más recetas sobre la misma estructura, incluyendo el ajuste fino supervisado en conjuntos de datos de chat (Tulu 3, UltraChat), RL en GSM8K a escala completa y alineación de preferencias KTO a partir de retroalimentación binaria. El README de cada receta enumera sus comandos, hiperparámetros y métricas esperadas.

Conclusión

Las APIs de sesiones de entrenamiento de SageMaker AI mueven el límite de abstracción hacia abajo, hasta la operación de entrenamiento: usted escribe el bucle de entrenamiento en Python en su propia máquina, y el servicio lo ejecuta en GPU administradas, gestionando la carga del modelo y de LoRA, la ejecución distribuida, la recuperación ante fallos y el checkpointing. Un bucle GRPO completo se reduce a cuatro operaciones, sample, forward_backward, optim_step, y save_weights_and_get_sampling_client, con la lógica de recompensa y ventaja en su propio código. Las mismas operaciones escalan a RL multiturno con uso de herramientas y evaluación en sandbox, como muestra la receta de Code RL.

Las recompensas personalizadas, los rollouts agénticos, la depuración paso a paso y las nuevas técnicas publicadas utilizan el mismo modelo de programación: el algoritmo permanece en su bucle, y SageMaker AI ejecuta las GPU.

Para comenzar:

• Elija una receta del cookbook de sesiones de entrenamiento de SageMaker, como Code RL.
• Ejecute el ejemplo de GSM8K del repositorio de ejemplos de Amazon SageMaker y sustituya su propia función de recompensa.
• Lea la documentación de las APIs de sesiones de entrenamiento para conocer los requisitos previos, las cuotas, las Regiones, los modelos compatibles y la referencia de la API.

Agradecimientos

Gracias a Arun Nagarajan, Xichao Wang y Shruti Sharma por sus contribuciones y orientación para esta publicación.

Sobre los autores