Implementación de infraestructura RL de múltiples turnos para Amazon Nova en Amazon SageMaker HyperPod

Cuando crea agentes empresariales que ejecutan flujos de trabajo de varios pasos, enfrenta un desafío de capacitación fundamental. Estos agentes consultan bases de datos, llaman a API, hacen referencias cruzadas de resultados y se recuperan de fallas a mitad del proceso. La calidad de cualquier acción depende de lo que suceda varios pasos después.

El aprendizaje por refuerzo estándar a partir de la retroalimentación humana (RLHF) optimiza respuestas únicas de forma aislada. Este enfoque se queda corto para flujos de trabajo de varios pasos donde un agente que valida los datos antes de continuar evita una cascada de errores posteriores. El aprendizaje por refuerzo (RL) de múltiples turnos aborda esta brecha optimizando secuencias de interacción completas. Sus agentes aprenden a orquestar herramientas, recuperar errores y razonar en varios pasos mediante prueba y error. El ajuste fino supervisado (SFT), la generación aumentada de recuperación (RAG) y el entrenamiento previo continuo son técnicas complementarias, pero normalmente no enseñan estas capacidades de toma de decisiones secuenciales por sí solas.

Amazon SageMaker AI también ofrece RL de múltiples turnos como una capacidad sin servidor totalmente administrada, lo que lleva esta técnica a los trabajos de capacitación de SageMaker sin infraestructura que administrar. Cuando necesita control total sobre la pila de capacitación: su propio entorno de agente, orquestación personalizada o configuraciones de instancias específicas. Para estos casos, la infraestructura RL de múltiples turnos para Amazon Nova en Amazon SageMaker HyperPod le brinda las capas de computación, orquestación y enrutamiento de recompensas para capacitar a los agentes en estos complejos flujos de trabajo. Amazon Nova ofrece inteligencia de vanguardia y un rendimiento de precios líder en la industria, y Amazon Nova Forge lo amplía con capacidades de capacitación de RL de múltiples turnos.

En esta publicación, implementará una infraestructura de dos fases para RL de varios turnos utilizando Amazon Nova Forge en Amazon SageMaker HyperPod. Al final, tiene una canalización basada en eventos que comienza a entrenarse cuando carga datos en Amazon Simple Storage Service (Amazon S3). El trabajo de capacitación le enseña al modelo a jugar Wordle, un marcador de posición para su propia tarea de RL.

Descripción general de la solución

La solución es una canalización impulsada por eventos: usted carga un conjunto de datos en Amazon S3 y la infraestructura aprovisiona cálculos, enruta recompensas y ejecuta entrenamiento RL de varios turnos automáticamente. Tres capas hacen el trabajo. Un clúster de SageMaker HyperPod genera respuestas y aplica actualizaciones de peso GRPO (optimización de políticas relativas al grupo). ECS en AWS Fargate ejecuta su entorno de recompensas. El SDK de Nova Forge enruta mensajes entre el modelo y ese entorno mientras rastrea el estado de la conversación en los turnos. AWS Step Functions organiza la ejecución, activada por Amazon EventBridge cuando los datos llegan a S3.

La arquitectura se divide en dos fases: una implementación única del AWS Cloud Development Kit (AWS CDK) proporciona la base de larga duración (VPC, EKS/HyperPod, ECS, S3, IAM y la canalización), mientras que cada ejecución de capacitación genera sus propios recursos efímeros. Esto evita que el cálculo de la GPU permanezca inactivo entre ejecuciones y le permite iterar sin volver a implementarlo.

El siguiente diagrama muestra cómo interactúan estos componentes durante una ejecución de entrenamiento.

RL de múltiples turnos en la arquitectura de infraestructura HyperPod de Amazon SageMaker

La canalización organiza un bucle de conversación de varios turnos en tres superficies informáticas:

Amazon SageMaker HyperPod (EKS): las réplicas de generación de vLLM, pods de trabajo y primarias de entrenamiento se ejecutan en instancias P5. El modelo genera respuestas. Los módulos de entrenamiento realizan actualizaciones de peso GRPO mediante señales de recompensa. ECS en Fargate: los trabajadores de recompensa ejecutan su entorno (por ejemplo, Wordle o un entorno personalizado Bring Your Own Orchestrator (BYOO)). Reciben respuestas modelo a través de SQS, las califican según su rúbrica y devuelven señales de recompensa. Amazon Nova Forge: la capa de proxy del SDK enruta mensajes entre el modelo y el entorno de recompensa, rastreando el estado de la conversación en cada turno.

Requisitos previos

Antes de implementar, asegúrese de tener lo siguiente:

Suscripción a Amazon Nova Forge: necesita esta suscripción para acceder al SDK de Nova Forge y a las API de entrenamiento de modelos. Cuota de instancias de SageMaker HyperPod: Mínimo de 10 × ml.p5.48xlarge para generate_replicas: 4. Solicite 12 a 14 para cargas de trabajo de producción. También se admite ml.p5en.48xlarge. Entorno CDK con arranque: ejecute cdk bootstrap antes de su primera implementación. Python 3.12+: la aplicación CDK y el tiempo de ejecución de AWS Lambda requieren esta versión. AWS CDK v2: instale con npm install -g aws-cdk. Esto sintetiza e implementa la pila de AWS CloudFormation. Interfaz de línea de comandos de AWS (AWS CLI) v2: configurada con credenciales que tienen permisos para crear VPC, clústeres de EKS, clústeres de SageMaker HyperPod, roles de IAM y Step Functions. Docker: utiliza Docker para crear las imágenes del contenedor Lambda que empaquetan el SDK de Nova Forge.

Importante: Esta infraestructura cuesta aproximadamente entre $786 y $1180 por hora cuando se ejecuta (10 a 12 instancias ml.p5.48xlarge). Revise la sección Desglose de costos y planee destruir la pila cuando no esté entrenando activamente.

Implementar la infraestructura

Clonar e instalar

Comience clonando el repositorio de muestra e instalando las dependencias de Python para la aplicación AWS CDK. Ejecute estos comandos desde su máquina local o entorno de desarrollo:

git clone https://github.com/aws-samples/sample-nova-multi-turn-rl-infra.git cd nova-multi-turn-rl-infra pip install -r requisitos.txt

Configurar

Configure todos los parámetros en cdk.json bajo la clave de contexto. Se requieren dos parámetros antes de su primera implementación:

Parámetro Descripción project_tag Prefijo único para todos los nombres de recursos (por ejemplo, my-nova-rl). Se convierte en parte del nombre del clúster EKS, el nombre del clúster HyperPod y las etiquetas del depósito S3. sdk_resource_prefix Prefijo para los recursos creados por el SDK (por ejemplo, nrl-myproject). El SDK de Nova Forge usa esto al nombrar sus pilas de CloudFormation, funciones Lambda y colas SQS.

Parámetros clave de infraestructura:

Parámetro Descripción predeterminada tipo_instancia ml.p5.48xlarge Tipo de instancia de HyperPod instancia_count 10 Número de instancias en el grupo de instancias restringidas nova_model NOVA_LITE_2 Modelo a entrenar: NOVA_MICRO, NOVA_LITE, NOVA_LITE_2 o NOVA_PRO vf_env_id wordle Entorno de recompensa integrado para validación use_custom_env / custom_env_id: establezca use_custom_env en "verdadero" y especifique el nombre del directorio en entornos personalizados/ para BYOO recompensa_cpu / recompensa_memoria CPU 2048, 4096 MiB Tamaño de tarea Fargate eks_kubernetes_version 1.32 Versión del clúster EKS

Parámetros de entrenamiento. Estos fluyen a través del evento de canalización hasta el trabajo de capacitación:

Parámetro Descripción predeterminada Training_method: RFT_MULTITURN_FULL o RFT_MULTITURN_LORA max_steps 10 Número de pasos de entrenamiento generate_replicas 4 réplicas de generación de vLLM global_batch_size 64 Muestras por paso de entrenamiento

Puede anular cualquier parámetro en el momento de la implementación:

implementación de cdk -c recuento_instancia=1 -c pasos_max=20

Desplegar

implementación de cdk –require-approval nunca

Modelo de implementación de dos fases

La infraestructura sigue un modelo de implementación de dos fases que separa los recursos fundamentales de larga duración de los recursos efímeros por ejecución. Esta separación reduce los costos al crear costosos recursos informáticos solo cuando los necesita, acelera la iteración al evitar la reimplementación completa para cada ejecución de capacitación y simplifica la administración al otorgar a cada capa un ciclo de vida independiente.

Fase 1: implementación de AWS CDK (una sola vez): cuando ejecuta cdk implementación, aprovisiona la infraestructura fundamental:

La implementación tarda aproximadamente entre 30 y 40 minutos. La mayor parte del tiempo se dedica a la creación del clúster EKS (~15 minutos), seguido de la instalación del gráfico SageMaker HyperPod Helm a través de AWS CodeBuild (~5 minutos), el aprovisionamiento del clúster SageMaker HyperPod (~15-25 minutos, dependiendo de la capacidad de P5) y la creación de imágenes de contenedor Lambda (~5 minutos).

Fase 2: Tiempo de ejecución (por ejecución de entrenamiento): cuando carga un archivo .jsonl en el prefijo/datos de entrenamiento del depósito S3, EventBridge activa una canalización de Step Functions que crea los recursos de tiempo de ejecución para esa ejecución de entrenamiento. El SDK de Nova Forge implementa su propia pila de AWS CloudFormation que contiene funciones de AWS Lambda (proxy de conversación), colas FIFO de Amazon Simple Queue Service (Amazon SQS) (enrutamiento de mensajes entre el modelo y el entorno), una tabla de Amazon DynamoDB (seguimiento del estado de la conversación) y tareas de ECS Fargate (trabajadores de recompensa). El SDK gestiona el ciclo de vida de estos recursos.

Activar una carrera de entrenamiento

Después de implementar la infraestructura, comienza una ejecución de capacitación con una única carga de S3. EventBridge observa los datos de entrenamiento/prefijo del depósito e inicia la canalización de Step Functions automáticamente.

Una ejecución de capacitación exitosa muestra puntuaciones de recompensa crecientes en las métricas de Amazon CloudWatch en pasos sucesivos. Para el entorno Wordle, se espera que el modelo converja en 50-100 pasos, con una recompensa promedio mejorando de casi cero a 0,6-0,8 a medida que el modelo aprende a limitar las conjeturas en función de la retroalimentación.

Prepara tus datos de entrenamiento

El archivo .jsonl utiliza un formato basado en metadatos: cada línea contiene un mensaje y la respuesta real que el entorno de recompensa utiliza para puntuar:

{"id": "wordle_train_001", "metadata": {"prompt": "Adivina la palabra de 5 letras", "answer": "crane"}} {"id": "wordle_train_002", "metadata": {"prompt": "Adivina la palabra de 5 letras", "answer": "slate"}} {"id": "wordle_train_003", "metadata": {"prompt": "Adivina la palabra de 5 letras", "respuesta": "plomada"}}

El campo de identificación debe ser único en todos los registros. El modelo ve metadata.prompt al comienzo de cada conversación y el entorno de recompensa usa metadata.answer para calificar las respuestas del modelo en los turnos. Cada ejemplo de entrenamiento se convierte en una conversación de varios turnos: el modelo genera una suposición, recibe retroalimentación del entorno de recompensa e itera hasta que resuelve el rompecabezas o agota sus turnos.

Cargar y activar automáticamente

# Obtener el nombre del depósito de las salidas del CDK BUCKET=$(aws cloudformation describe-stacks –stack-name NovaMultiTurnRlStack –query "Stacks[0].Outputs[?OutputKey=='TrainingBucketName'].OutputValue" –texto de salida) # Carga: la canalización se inicia automáticamente aws s3 cp Training-data.jsonl s3://$BUCKET/training-data/training-data.jsonl

Cuando un archivo llega al prefijo datos de entrenamiento/, EventBridge detecta el evento PutObject e invoca la función Lambda S3TriggerFn. Esta función extrae la ruta del archivo, la fusiona con los parámetros de entrenamiento de cdk.json y pasa los datos del evento resultante a Step Functions, que ejecuta la canalización en cinco etapas. Puede observar el progreso en la consola de Step Functions, donde cada paso muestra su entrada, salida, duración y recuento de reintentos.

Activar manualmente para ejecuciones ad hoc

Para iterar sobre los parámetros sin volver a cargar datos, use el script de configuración para omitir EventBridge e invocar la canalización directamente:

# Anular los parámetros de entrenamiento para una ejecución de prueba rápida ./scripts/setup.sh –data-path s3://$BUCKET/training-data/training-data.jsonl –max-steps 5 –global-batch-size 32 –training-method RFT_MULTITURN_LORA

Esto ayuda durante el desarrollo: apunte a un conjunto de datos existente y experimente con diferentes recuentos de pasos, tamaños de lotes o el método de entrenamiento (ajuste completo versus LoRA) sin modificar cdk.json ni volver a implementar.

Monitorear y depurar

Puede monitorear y solucionar problemas de cada capa de este proceso de múltiples etapas en SageMaker HyperPod, ECS, Lambda y SQS. La infraestructura proporciona observabilidad en cada capa para que pueda aislar los problemas rápidamente.

Monitoreo de tuberías

La consola de Step Functions es el panel principal. Cada ejecución muestra el progreso paso a paso con tiempo, datos de entrada y salida e historial de reintentos. Cada paso escribe en su propio grupo de registros de Amazon CloudWatch.

Durante el funcionamiento normal, cada paso de Step Functions se completa dentro de su ventana de tiempo prevista: configuración de la infraestructura en 2 a 3 minutos, implementación de trabajadores de recompensa en 1 a 2 minutos, validación de datos en menos de 1 minuto y envío de capacitación en 3 a 5 minutos. Las colas SQS deben mostrar mensajes que fluyen de manera constante y sin retrasos sostenidos.

Alerta de falla

Un tema de Amazon Simple Notification Service (Amazon SNS), cifrado con AWS Key Management Service (AWS KMS), publica una notificación cuando falla la ejecución de Step Functions. Suscríbase para recibir alertas por correo electrónico, Slack o PagerDuty:

# Obtener el ARN del tema de las salidas del CDK TOPIC_ARN=$(aws cloudformation describe-stacks –stack-name NovaMultiTurnRlStack –query "Stacks[0].Outputs[?OutputKey=='AlertTopicArn'].OutputValue" –texto de salida) # Suscríbase y confirme por correo electrónico aws sns subscribe –topic-arn $TOPIC_ARN –protocol email –notification-endpoint your-team@example.com

Fallos de activación

Si una carga de S3 no inicia la canalización, verifique la cola de mensajes fallidos. Las invocaciones fallidas de EventBridge-to-Lambda se almacenan aquí con los datos del evento S3 original y los detalles del error:

DLQ_URL=$(aws cloudformation describe-stacks –stack-name NovaMultiTurnRlStack –query "Pilas[0].Outputs[?OutputKey=='TriggerDlqUrl'].OutputValue" –texto de salida) aws sqs mensaje de recepción –queue-url $DLQ_URL –número máximo de mensajes 5

Las causas comunes incluyen nombres de archivos con formato incorrecto (el desencadenante espera una extensión .jsonl), desviación de permisos de IAM y límites de simultaneidad de Lambda.

Depuración de puestos de entrenamiento

Si la canalización comienza pero el entrenamiento se detiene, la causa raíz suele estar en la capa de enrutamiento de mensajes entre el modelo y el entorno de recompensa. Verifique el estado de la cola SQS utilizando el diagnóstico integrado del SDK:

from rft_infra import check_all_queues # Devuelve recuentos de mensajes (en curso, disponibles, retrasados) para todas las colas FIFO check_all_queues()

Un retraso creciente en la cola de solicitudes con una cola de respuestas vacía indica que los trabajadores de recompensa no están procesando. Un retraso creciente en la cola de respuestas indica que el modelo no está consumiendo recompensas. Cualquiera de los patrones le indica qué componente investigar.

Salud del hiperpod

Verifique que todos los pods de capacitación y generación se estén ejecutando en el clúster SageMaker HyperPod:

kubectl obtiene vainas -n kubeflow

Todos los pods deben mostrar el estado En ejecución o Completado. Los pods bloqueados en Pendiente indican una capacidad de instancia insuficiente en el grupo de instancias restringido. Los pods en CrashLoopBackOff (un estado de Kubernetes que indica que el contenedor falla y se reinicia repetidamente) indican errores a nivel de contenedor. Consulte los registros del pod con kubectl logs -n kubeflow para obtener más detalles.

Limpiar y gestionar costos

El script de limpieza detiene los recursos activos en este orden: ejecuciones de Step Functions, tareas de ECS, la pila de SDK CloudFormation y la pila de CDK.

Para evitar cargos, ejecute el siguiente comando para destruir toda la pila:

Para destruir la infraestructura y mantener tus datos de entrenamiento en S3 (aún se aplican cargos de almacenamiento):

./cleanup.sh –retener-datos

Importante: destruye la pila cuando no estés entrenando activamente. Los costos de inactividad de las instancias SageMaker HyperPod y NAT Gateway se acumulan rápidamente.

Desglose de costos

Recurso Costo por hora Notas SageMaker HyperPod 8 × ml.p5.48xlarge ~$786/hr Configuración mínima (10 instancias, 8 utilizadas para computación) SageMaker HyperPod 12 × ml.p5.48xlarge ~$1180/hr Configuración de producción (12 instancias con espacio libre) Plano de control EKS $0,10/hr NAT Gateway ~$0,045/hr ECS Fargate recompensa a los trabajadores ~$0,15-0,30/h + transferencia de datos S3, Lambda, CodeBuild, SQS, DynamoDB Insignificante

Para conocer los precios actuales de Amazon Nova Forge, consulte la página de precios de Amazon Nova.

Optimizar por costo

Puede reducir los costos inactivos con dos estrategias:

Escale a cero: proporcione un clúster de instancia única con políticas de escalado automático que se reducen a cero durante los períodos de inactividad, eliminando los costos de computación inactiva. Programación de tipos de instancias: utilice tipos de instancias de menor costo durante las fases de configuración y de menor actividad, luego realice la transición a instancias de mayor rendimiento cuando comiencen las cargas de trabajo de capacitación.

Conclusión

Ahora dispone de una infraestructura basada en eventos y lista para producción para el aprendizaje por refuerzo de múltiples turnos con Amazon Nova Forge en Amazon SageMaker HyperPod. Cargue un conjunto de datos en Amazon S3 y la canalización manejará el aprovisionamiento, la orquestación del entrenamiento y el enrutamiento de recompensas automáticamente. Para adaptar esto a su caso de uso, reemplace el entorno de Wordle con su propio agente de llamada API o flujo de trabajo empresarial.

Próximos pasos:

Clone el repositorio de muestra: aws-samples/nova-multi-turn-rl-infra. Lea la documentación de Amazon SageMaker HyperPod. Lea sobre las capacidades y los precios de Amazon Nova. Personalice su canalización con la guía para desarrolladores de AWS Step Functions.

Sobre los autores

María Masod

María Masod

María se especializa en IA agente, ajuste de refuerzos y capacitación de agentes de múltiples turnos. Tiene experiencia en aprendizaje automático, que abarca la personalización de grandes modelos de lenguaje, el modelado de recompensas y la creación de canales de capacitación de un extremo a otro para agentes de IA. María, una entusiasta de la sostenibilidad de corazón, disfruta de la jardinería y de preparar café con leche.

Nick Biso

Nick Biso

Nick es ingeniero sénior de aprendizaje automático en AWS Professional Services. Resuelve complejos desafíos organizativos y técnicos utilizando la ciencia y la ingeniería de datos. Además, crea e implementa modelos de IA/ML en la nube de AWS. Su pasión se extiende a su propensión a viajar y diversas experiencias culturales.

Christian Kamwangala

Christian Kamwangala

Christian es arquitecto de soluciones especializado en IA/ML e IA generativa en AWS, donde se asocia con clientes empresariales para diseñar, optimizar e implementar soluciones de IA de nivel de producción. Su experiencia radica en la optimización de la inferencia: equilibrar el rendimiento, el costo y la latencia para implementaciones a gran escala. Fuera del trabajo, le gusta explorar la naturaleza y pasar tiempo con familiares y amigos.

Manoj Gupta

Manoj Gupta

Manoj es arquitecto senior de soluciones en AWS, con sede en San Francisco. Con más de 4 años de experiencia en AWS, trabaja en estrecha colaboración con los clientes para crear soluciones optimizadas impulsadas por AI/ML e infraestructura en la nube. Sus principales áreas de enfoque son datos, IA/ML y seguridad, ayudando a las organizaciones a modernizar sus pilas de tecnología. Fuera del trabajo, le gustan las actividades al aire libre y viajar con la familia.