Entrene CodeFu-7B con veRL y Ray en Amazon SageMaker Trabajos de formación

El rápido avance de la inteligencia artificial (IA) ha creado una demanda sin precedentes de modelos especializados capaces de realizar tareas de razonamiento complejas, particularmente en programación competitiva donde los modelos deben generar código funcional mediante razonamiento algorítmico en lugar de memorización de patrones. El aprendizaje por refuerzo (RL) permite que los modelos aprendan mediante prueba y error al recibir recompensas basadas en la ejecución real del código, lo que lo hace particularmente adecuado para desarrollar capacidades genuinas de resolución de problemas en dominios algorítmicos.

Sin embargo, la implementación de capacitación de RL distribuida para la generación de código presenta importantes desafíos de infraestructura, como orquestar múltiples componentes heterogéneos, coordinar la compilación de código paralelo entre nodos y mantener la tolerancia a fallas para procesos de larga duración. Ray es uno de los marcos para cargas de trabajo distribuidas que aborda estos desafíos, debido a su sistema unificado que maneja todo el proceso de IA, su arquitectura basada en GPU y su perfecta integración con herramientas como Hugging Face Transformers y PyTorch.

Las cargas de trabajo se pueden ejecutar con el marco de Ray en trabajos de capacitación de SageMaker mediante la solución de trabajos de capacitación de Ray en Amazon SageMaker, que combina el marco de computación distribuida de Ray con la infraestructura totalmente administrada de SageMaker. Esta solución maneja automáticamente la inicialización del clúster Ray, la coordinación de múltiples nodos y la gestión de recursos distribuidos, lo que permite a los desarrolladores centrarse en el desarrollo de modelos mientras se benefician de las funciones de nivel empresarial de SageMaker.

En esta publicación, demostramos cómo entrenar CodeFu-7B, un modelo especializado de 7 mil millones de parámetros para programación competitiva, utilizando la Optimización de Políticas Relativas de Grupo (GRPO) con veRL, una biblioteca de capacitación flexible y eficiente para modelos de lenguaje grandes (LLM) que permite una extensión sencilla de diversos algoritmos RL y una integración perfecta con la infraestructura LLM existente, dentro de un clúster Ray distribuido administrado por trabajos de capacitación de SageMaker. Recorremos la implementación completa, que abarca la preparación de datos, la configuración de la capacitación distribuida y la observabilidad integral, mostrando cómo este enfoque unificado ofrece escala computacional y experiencia de desarrollador para cargas de trabajo de capacitación RL sofisticadas.

Acerca de CodeFu-7B

CodeFu-7B-v0.1 es un modelo de lenguaje de parámetros 7B entrenado específicamente para resolver problemas de programación competitiva (CP). Basado en el modelo base DeepSeek-R1-Distill-Qwen-7B, CodeFu demuestra cómo el aprendizaje por refuerzo puede desarrollar capacidades de razonamiento algorítmico y generación eficiente de código C++ más allá de los enfoques tradicionales de ajuste supervisado.

El modelo se entrena utilizando declaraciones de problemas del conjunto de datos de DeepMind CodeContest sin acceso a soluciones reales durante el entrenamiento, lo que lo obliga a aprender mediante prueba y error en función de los comentarios de ejecución del código. Este enfoque permite el desarrollo de capacidades genuinas de resolución de problemas en lugar de la memorización de patrones.

CodeFu está disponible públicamente en HuggingFace y se publica bajo la licencia del MIT, lo que lo hace accesible para investigadores y profesionales interesados ​​en la generación de código y el razonamiento algorítmico. La metodología de entrenamiento del modelo demuestra el potencial para aplicar técnicas de aprendizaje por refuerzo a tareas de razonamiento complejas más allá de la programación competitiva.

Ray en la solución de trabajos de formación de SageMaker

Ray en Amazon SageMaker Training jobs es una solución que permite el procesamiento de datos distribuidos y el entrenamiento de modelos utilizando Ray dentro del entorno de entrenamiento administrado de SageMaker. La solución proporciona capacidades clave que incluyen una arquitectura de iniciador universal para la configuración automática de clústeres de Ray, administración de clústeres de múltiples nodos con coordinación inteligente, soporte de clústeres heterogéneos para tipos de instancias mixtas y observabilidad integrada a través de la integración de Ray Dashboard, Prometheus, Grafana y Amazon CloudWatch.

La solución se integra perfectamente con SageMaker Python SDK utilizando la moderna API ModelTrainer. Esta solución disponible públicamente en GitHub permite a los desarrolladores utilizar las capacidades informáticas distribuidas de Ray mientras se benefician de la infraestructura administrada de SageMaker, lo que la hace ideal para cargas de trabajo complejas como capacitación de aprendizaje reforzado que requieren una coordinación distribuida y una gestión de recursos sofisticadas.

Descripción general de la solución

El flujo de trabajo para entrenar CodeFu 7B con veRL y Ray en trabajos de entrenamiento de SageMaker, como se ilustra en el diagrama adjunto, consta de los siguientes pasos:

Preparación de datos: cargue el conjunto de datos preprocesado de DeepMind CodeContest y la configuración de capacitación. Envío de trabajos de capacitación: envíe una solicitud API de trabajo de capacitación de SageMaker a través de la clase ModelTrainer del SDK de Python de SageMaker. Monitoreo y observabilidad: monitoree el progreso del entrenamiento en tiempo real a través de Ray Dashboard y, opcionalmente, con la recopilación de métricas de Prometheus, la visualización de Grafana y el seguimiento de experimentos. Limpieza automática: al finalizar la capacitación, SageMaker guarda automáticamente el modelo entrenado en S3, carga los registros de capacitación en CloudWatch y desmantela el clúster de computación.

Esta arquitectura optimizada ofrece una experiencia de capacitación de aprendizaje por refuerzo totalmente administrada, lo que permite a los desarrolladores centrarse en el desarrollo del modelo mientras SageMaker y Ray manejan la compleja orquestación de la infraestructura distribuida, dentro de un modelo de precios de pago por uso que factura solo por el tiempo de procesamiento real.

Requisitos previos

Se deben completar los siguientes requisitos previos antes de poder ejecutar el cuaderno:

Realice las siguientes solicitudes de aumento de cuota para SageMaker AI. Para este caso de uso, solicite un mínimo de 2 instancias p4de.24xlarge (con 8 GPU NVIDIA A100) y escale a más instancias p4de.24xlarge (según las compensaciones de tiempo y costo de capacitación para su caso de uso). También se admiten instancias P5 (con 8 GPU NVIDIA H100). En la consola de Cuotas de servicio, solicite las siguientes cuotas de SageMaker AI: instancias p4de (p4de.24xlarge) para el uso de trabajos de capacitación: 2 Cree un rol de AWS Identity and Access Management (IAM) con políticas administradas AmazonSageMakerFullAccess, AmazonS3FullAccess, AmazonSSMFullAccess para brindar el acceso necesario a SageMaker AI para ejecutar los ejemplos. Asigne la siguiente política como relación de confianza al rol de IAM creado:

{ "Versión":"2012-10-17", "Declaración":[ { "Sid":"", "Efecto":"Permitir", "Principal":{ "Servicio": "sagemaker.amazonaws.com" ] }, "Acción":"sts:AssumeRole" } ] }

(Opcional) Cree un dominio de Amazon SageMaker Studio (consulte Uso de la configuración rápida para Amazon SageMaker AI) para acceder a los cuadernos de Jupyter para ejecutar el código de capacitación. Alternativamente, JupyterLab se puede usar en una configuración local u otro entorno de desarrollo de Python para ejecutar el cuaderno y enviar el trabajo de capacitación de SageMaker.

Nota: Estos permisos otorgan un acceso amplio y no se recomienda su uso en entornos de producción. Consulte la Guía para desarrolladores de SageMaker para obtener orientación sobre cómo definir permisos más detallados.

El ejemplo de código se puede encontrar en este repositorio de GitHub.

Preparar el conjunto de datos

El proceso de preparación de datos transforma el conjunto de datos sin procesar de DeepMind CodeContest en un formato adecuado para la capacitación en aprendizaje por refuerzo. Aplicamos filtros sistemáticos para identificar problemas adecuados, eliminando aquellos con calificaciones de Codeforces inferiores a 800 e implementando controles de validación de calidad para casos de prueba faltantes, descripciones mal formadas y restricciones no válidas.

Clasificamos los problemas en tres niveles de dificultad: Fácil (800-1000 puntos), Difícil (1100-2200 puntos) y Experto (2300-3500 puntos). Esta publicación utiliza solo el conjunto de datos Easy para la capacitación. Cada problema está formateado con dos componentes: un mensaje de usuario que contiene la declaración del problema y una especificación de modelo de recompensa con casos de prueba, límites de tiempo y restricciones de memoria. Fundamentalmente, el campo ground_truth no contiene código de solución, solo casos de prueba, lo que obliga al modelo a aprender a través de señales de recompensa en lugar de memorizar soluciones.

{ "data_source": "code_contests", "prompt": [ { "role": "user", "content": "Escribe una solución C++ para este problema: …" } ], "ability": "coding-cp", "reward_model": { "style": "rule", "ground_truth": { "name": "problem 1", "public_tests": { "input": ["test input 1", "test input 2"], "output": ["salida esperada 1", "salida esperada 2"] }, "private_tests": { "input": ["entrada privada 1", "entrada privada 2"], "salida": ["salida privada 1", "salida privada 2"] }, "time_limit": 2.0, "memory_limit_bytes": 268435456, "cf_rating": 1200 } } }

Para esta publicación, proporcionamos un subconjunto preprocesado del conjunto de datos de dificultad Fácil en el ejemplo de código para optimizar el ejemplo de capacitación, accesible desde el repositorio de GitHub.

Formación GRPO utilizando veRL

El proceso de capacitación utiliza Ray para orquestar la ejecución distribuida y la sincronización del lanzamiento de vLLM, la evaluación de recompensas (compilación y ejecución de código), el paralelismo del modelo FSDP y el paralelismo de secuencia Ulysses. Establecemos el grado de paralelismo de secuencia en 4 para el razonamiento de formato largo y las generaciones de código.

El marco veRL implementa una arquitectura sofisticada de múltiples componentes a través de su orquestador main_ppo.py, que coordina tres tipos principales de trabajadores distribuidos: ActorRolloutRefWorker para inferencias e implementaciones de políticas, CriticWorker para estimación de funciones de valor y RewardModelWorker para calificar soluciones generadas.

El algoritmo GRPO mejora la optimización de políticas próximas (PPO) tradicional al calcular las ventajas utilizando líneas de base relativas al grupo, lo que ayuda a estabilizar la capacitación al reducir la variación en las estimaciones del gradiente de políticas.

Ampliamos el repositorio de código de TinyZero utilizando Ray para administrar y distribuir el cálculo de la función de recompensa. Esto permite la compilación y evaluación de código C++ en paralelo en el mismo clúster para abordar la naturaleza de la ejecución del código que requiere un uso intensivo de computación y está ligada a la latencia. Todo el proceso se ejecuta como un trabajo de capacitación de SageMaker que se ejecuta en instancias ml.p4de.24xlarge. El proceso de capacitación consta de los siguientes pasos, como se muestra en la siguiente arquitectura:

Implementación: las indicaciones de los problemas de codificación se introducen en el motor de inferencia vLLM para implementar posibles soluciones. Generación de respuestas: vLLM genera múltiples respuestas (razonamiento + código) para cada mensaje. Ejecución de código: las soluciones de código se extraen de las respuestas y son compiladas y ejecutadas por trabajadores distribuidos (compiladores y tiempo de ejecución) administrados por Ray. Cálculo de recompensas: los resultados de la ejecución se utilizan para calcular las recompensas (es decir, índices de aprobación de casos de prueba) y las ventajas se calculan utilizando líneas de base relativas al grupo. Actualización de políticas: el actor utiliza ventajas y probabilidades de tokens para calcular la pérdida de PPO, que se utiliza para actualizar los parámetros de CodeFu mediante el descenso de gradiente. Iteración: el proceso se repite con lotes de ciclos de respuesta rápida y recompensa, con Ray administrando el muestreo distribuido, la ejecución y la sincronización del entrenamiento en todo el proceso.

La orquestación del proceso de capacitación implica varios componentes clave implementados en múltiples módulos. El ciclo de capacitación central de veRL se implementa en main_ppo.py, que inicializa a los trabajadores de Ray y administra el proceso de capacitación distribuido:

@ray.remote def main_task(config): # Inicializar tokenizer y descargar el modelo local_path = copy_local_path_from_hdfs(config.actor_rollout_ref.model.path) tokenizer = hf_tokenizer(local_path) # Definir roles de trabajadores distribuidos role_worker_mapping = { Role.ActorRollout: ray.remote(ActorRolloutRefWorker), Role.Critic: ray.remote(CriticWorker), Role.RefPolicy: ray.remote(ActorRolloutRefWorker), } # Inicializar el administrador de recompensas para la ejecución del código. Resource_pool_manager=resource_pool_manager, recompensa_fn=recompensa_fn,) entrenador.init_workers() entrenador.fit()

El sistema de evaluación de recompensas implementa la ejecución de código paralelo a través de funciones remotas de Ray, manejando la compilación de C++ y la ejecución de casos de prueba:

@ray.remote def Process_reward_item(idx, valid_response_length, secuencias_str, data_source, recompensa_model_data): # Extrae y compila código C++ de la respuesta del modelo ground_truth = json.loads(reward_model_data)["ground_truth"] # Selecciona la función de puntuación adecuada según la fuente de datos si data_source == "code_contests": Compute_score = code_contests.compute_score # Ejecute código contra casos de prueba y calcule la puntuación de la proporción de aprobación = Compute_score(solution_str=sequences_str, ground_truth=ground_truth) return idx, puntuación, valid_response_length, secuencias_str, data_source

El sistema de ejecución de casos de prueba en paralelo optimiza la eficiencia de la evaluación mediante el muestreo de casos de prueba y el uso de grupos de procesos:

def run_test_cases_parallel( bin_file: str, test_inputs: List[str], test_outputs: List[str], prob_name: str, junction_timeout: float, max_test_cases: int = 100, max_workers: int = 100) -> Tuple[int, int]: # Casos de prueba de muestra si hay demasiados disponibles if len(test_inputs) > max_test_cases: random_indices = np.random.choice(len(test_inputs), size=max_test_cases, replace=False) test_inputs = test_inputs[random_indices] test_outputs = test_outputs[random_indices] # Ejecutar casos de prueba en paralelo usando ProcessPoolExecutor con ProcessPoolExecutor(max_workers=min(max_workers, len(test_inputs))) como ejecutor: resultados = lista(executor.map(_process_test_case, args_list)) total_matches = suma(resultados) devuelve total_matches, len(test_inputs)

Esta implementación permite una capacitación distribuida eficiente al separar las preocupaciones: el orquestador main_ppo.py administra la coordinación de los trabajadores de Ray, mientras que el sistema de recompensa proporciona una evaluación de código escalable a través de la compilación y ejecución paralela en todo el clúster de SageMaker.

A continuación se muestra el pseudocódigo para el cálculo de recompensas utilizado en esta publicación para entrenar un modelo de codificación de programación competitivo. La función de recompensa es la parte más importante del aprendizaje por refuerzo, ya que define lo que se anima a lograr al modelo y lo que debe evitar. Esta implementación utiliza un sistema de penalización jerárquico que primero busca problemas fundamentales de ejecución de código, asignando penalizaciones severas para código no ejecutable (-1) y penalizaciones moderadas para fallas de compilación (-0,5). Las soluciones de código extraído se ejecutan con un límite de tiempo estricto: el código que excede el límite de tiempo especificado para el problema no recibe recompensa, lo que facilita condiciones de programación competitivas realistas. Para una solución C++ ejecutada con éxito, su recompensa se calcula como una función lineal basada en la fracción de casos de prueba privados aprobados, lo que anima al modelo a resolver tantos casos de prueba privados como sea posible y al mismo tiempo evita el sobreajuste a pruebas visibles públicamente. Este diseño prioriza la corrección del código y la validez de la ejecución, y el rendimiento de la prueba privada sirve como única señal para aprender soluciones de codificación óptimas.

def Compute_reward(code_output, ground_truth): # Maneja fallas de ejecución (igual para ambas etapas) si no es_executable(code_output): devuelve -1 si compilación_failed(code_output): devuelve -0.5 si excede_time_limit(code_output): devuelve 0 # Señal de recompensa principal: corrección en casos de prueba ocultos # Ejecuta código contra casos de prueba privados pass_private, total_private = run_private_tests(code_output, ground_truth, max_test_cases=1000) devuelve pass_private/total_private

Consulte scripts/verl/utils/reward_score/code_contests.py para obtener el código Python completo. La ejecución del código generado en entornos de producción requiere un espacio aislado adecuado. En esta configuración de demostración controlada, ejecutamos el código como un ejemplo rápido para evaluar su exactitud para asignar recompensas.

Carga de trabajo de Ray con trabajos de capacitación de SageMaker

Para entrenar CodeFu-7B usando veRL y Ray en trabajos de entrenamiento de SageMaker, usamos la clase ModelTrainer del SDK de Python de SageMaker. Comience configurando la carga de trabajo de capacitación distribuida con los siguientes pasos:

Seleccione el tipo de instancia y la imagen del contenedor para el trabajo de capacitación:

tipo_instancia = "ml.p4de.24xlarge" cuenta_instancia = 2 ID_cuenta = sts.get_caller_identity()["Cuenta"] región = sagemaker_session.boto_session.region_name nombre_repo = "codefu-pytorch" etiqueta = "último" image_uri = f"{account_id}.dkr.ecr.{region}.amazonaws.com/{repo_name}:{tag}"

La capacitación utiliza un contenedor Docker personalizado que incluye veRL, Ray y las dependencias necesarias para la capacitación de RL distribuida. Consulte el repositorio de GitHub para obtener la definición completa del contenedor y las instrucciones de compilación.

Cree ModelTrainer para encapsular la configuración de entrenamiento basada en Ray:

La clase ModelTrainer proporciona opciones de ejecución flexibles a través de su configuración SourceCode, lo que permite a los usuarios personalizar sus flujos de trabajo de capacitación con diferentes marcos y lanzadores. Especifique un Entry_Script para la ejecución directa del script Python o utilice el parámetro comando para comandos de ejecución personalizados, lo que permite la integración con marcos especializados como Ray, Hugging Face Accelerate o soluciones de capacitación distribuidas personalizadas.

… args = [ "–entrypoint", "train.py", "–config", "/opt/ml/input/data/config/args.yaml", ] # Definir el script que se ejecutará con Ray Launcher source_code = SourceCode( source_dir="./scripts", requisitos="requirements.txt", command=f"python launcher.py {' '.join(args)}", ) # Definir la configuración de cálculo Compute_configs = Compute( tipo_instancia=tipo_instancia, recuento_instancia=recuento_instancia, keep_alive_period_in_segundos=1800, ) nombre_trabajo = "train-codefu-verl-ray" ruta_salida = f"s3://{nombre_depósito}/{nombre_trabajo}" model_trainer = ModelTrainer( imagen_entrenamiento=imagen_uri, código_fuente=código_fuente, nombre_trabajo_base=nombre_trabajo, Compute=compute_configs, stop_condition=StoppingCondition(max_runtime_in_segundos=3600 * 24 * 5), output_data_config=OutputDataConfig(s3_output_path=output_path), checkpoint_config=CheckpointConfig( s3_uri=output_path + "/checkpoint", local_path="/opt/ml/checkpoints" ), entorno={ "RAY_PROMETHEUS_HOST": "", "RAY_GRAFANA_HOST": "", "RAY_PROMETHEUS_NAME": "prometheus", "BASE_MODEL": "deepseek-ai/DeepSeek-R1-Distill-Qwen-7B", "RUN_NAME": "sagemaker-training-run", … }, role=get_execution_role(), ).with_remote_debug_config(RemoteDebugConfig(enable_remote_debug=True))

El script launcher.py sirve como punto de entrada universal que detecta el entorno de SageMaker (clúster de un solo nodo o de múltiples nodos, homogéneo o heterogéneo), inicializa el clúster de Ray con la coordinación adecuada del nodo principal/trabajador y ejecuta su script de entrenamiento personalizado. Las funcionalidades clave de launcher.py son:

Configuración del clúster de Ray: detecta automáticamente el entorno del clúster e inicializa Ray con la selección adecuada del nodo principal. Coordinación de nodos: gestiona la comunicación entre los nodos principales y trabajadores en las instancias de SageMaker. Ejecución de script: ejecuta el script –entrypoint especificado (train.py) dentro del contexto del clúster Ray. Conectividad de Prometheus y Grafana: configura Ray para exportar métricas y establece una conexión a servidores externos de Prometheus y Grafana especificados por RAY_PROMETHEUS_HOST y RAY_GRAFANA_HOST para un monitoreo integral del clúster. Para obtener información adicional, consulte Ray sobre trabajos de capacitación de SageMaker: observabilidad con Prometheus y Grafana.

Para conocer la implementación completa de la configuración del clúster Ray con trabajos de capacitación de SageMaker, consulte launcher.py.

El script train.py sirve como orquestador de capacitación real que:

Carga la configuración de veRL desde el archivo YAML proporcionado Configura el entorno de entrenamiento distribuido con el tokenizador y la inicialización del modelo adecuados Construye y ejecuta el comando de entrenamiento de veRL con los parámetros necesarios Maneja la configuración de variables de entorno para los trabajadores de Ray y la comunicación de la Biblioteca de Comunicaciones Colectivas (NCCL) de NVIDIA Gestiona el ciclo de vida completo del entrenamiento desde la carga de datos hasta los puntos de control del modelo

Para conocer la implementación completa del script del punto de entrada, consulte train.py.

Configure los canales de entrada para ModelTrainer creando objetos InputData a partir de las rutas del depósito S3:

… train_input = InputData( nombre_canal="tren", data_source=S3DataSource( s3_data_type="S3Prefix", s3_uri=train_dataset_s3_path, s3_data_distribution_type="FullyReplicated", ), ) config_input = InputData( nombre_canal="config", data_source=S3DataSource( s3_data_type="S3Prefix", s3_uri=train_config_s3_path, s3_data_distribution_type="Completamente replicado", ), )

Envíe el trabajo de entrenamiento usando la llamada a la función de entrenamiento en el ModelTrainer creado:

model_trainer.train( input_data_config=[train_input, val_input, config_input], esperar=Falso)

El trabajo se puede monitorear directamente desde la salida del notebook o mediante la consola de SageMaker, que muestra el estado del trabajo y los registros de CloudWatch correspondientes.

Consola de trabajos de formación de SageMaker

Métricas del sistema de trabajos de capacitación de SageMaker

El script launcher.py organiza la inicialización del clúster Ray a través de los siguientes pasos automatizados, que se pueden monitorear en tiempo real a través de registros de CloudWatch:

Configurar trabajos de entrenamiento de SageMaker y variables de entorno de Ray: configura las variables de entorno necesarias tanto para la integración de SageMaker como para la comunicación del clúster de Ray:

__main__ – INFO – Argumento de punto de entrada proporcionado: train.py __main__ – INFO – Establecer source_dir=, Entry_script=train.py … __main__ – INFO – Entorno SageMaker encontrado con hosts: … __main__ – INFO – Host actual: algo-1 __main__ – INFO – Host Prometheus configurado: __main__ – INFO – Host Grafana configurado: __main__ – INFO – El entorno de ejecución de Ray contiene 137 variables de entorno totales __main__ – INFORMACIÓN – Entorno de ejecución de Ray: …

Identifique el tipo de clúster del trabajo de entrenamiento de SageMaker: detecta si la implementación es de un solo nodo o de varios nodos, y determina un clúster de uno o varios nodos, y si es una configuración de clúster homogénea o heterogénea:

__main__ – INFO – Configuración de clúster homogéneo: 2 hosts en total __main__ – INFO – Todos los hosts: ['algo-1', 'algo-2'] __main__ – INFO – Se encontraron varios hosts, inicializando Ray como un clúster de múltiples nodos

Configurar nodos principales y trabajadores: identifica qué instancia sirve como nodo principal de Ray y configura las instancias restantes como nodos trabajadores:

__main__ – INFO – Nodo principal: algo-1, Host actual: algo-1 __main__ – INFO – CPU para el nodo principal: 192 __main__ – INFO – GPU para el nodo principal: 8

Iniciar nodo Ray: inicializa el nodo principal de Ray y los nodos trabajadores con la asignación de recursos y la configuración del panel adecuadas, verificando que los nodos trabajadores se conectan correctamente al nodo principal antes de continuar:

#011INFO trabajador.py:1723 – Conexión al clúster de Ray existente en la dirección: … #011INFO trabajador.py:1908 – Conectado al clúster de Ray. Vea el panel en __main__ – INFORMACIÓN – ¡Todos los nodos conectados al clúster de Ray!

Ejecute el script de entrenamiento: inicia el script de punto de entrada especificado (train.py) dentro del contexto del clúster Ray completamente inicializado:

Ruta del script: /opt/ml/input/data/code/train.py… __main__ – INFORMACIÓN – Cargando y ejecutando el script Python usando importlib…

Una vez completado el trabajo, los pesos y puntos de control del modelo entrenado estarán disponibles en la ruta de salida de S3 especificada, listos para su implementación o evaluación adicional.

Seguimiento de experimentos

El canal de capacitación de CodeFu se integra perfectamente con Managed MLflow en Amazon SageMaker AI, así como con soluciones de terceros, para un seguimiento integral de experimentos y visualización de métricas de aprendizaje de refuerzo.

La siguiente imagen muestra las métricas que son particularmente útiles para monitorear durante el entrenamiento de CodeFu.

El gráfico de métricas muestra una progresión de aprendizaje prometedora de GRPO/PPO para el modelo de programación competitivo. Las señales de recompensa demuestran una mejora clara, con crítico/recompensa/media aumentando de -0,8 a 0,6 y crítico/recompensa/min recuperándose de las fallas iniciales -1,0 a un rendimiento moderado -0,5, mientras que crítico/recompensa/máximo mantiene puntuaciones perfectas de 1,0 durante todo el entrenamiento, lo que indica que el modelo puede lograr soluciones óptimas.

Las métricas de Actor revelan una dinámica de entrenamiento saludable: actor/ppo_kl permanece bajo ~0,0002 después de un pico inicial, lo que confirma actualizaciones de políticas estables, mientras que actor/pg_clipfrac se mantiene en un rango razonable ~0,002-0,004, lo que sugiere pasos de aprendizaje de tamaño apropiado.

La tendencia creciente de actor/kl_loss indica una divergencia creciente con respecto al modelo de referencia, como se esperaba durante el ajuste de RL. Lo más importante es que val/test_score/code_contests muestra una mejora constante de -0,6 a ~0,5, y la comparación de validación del tren revela una buena generalización con ambas curvas siguiendo de cerca, lo que indica que el modelo está aprendiendo a resolver problemas de codificación de manera efectiva sin sobreajuste.

La siguiente tabla explica las métricas clave de capacitación de GRPO y por qué es importante monitorear cada una de ellas para diagnosticar el estado y el desempeño de la capacitación:

Métrica Descripción Propósito crítico/recompensa/min Recompensa mínima lograda en el conjunto de entrenamiento Detectar fallas catastróficas: las recompensas extremadamente negativas indican que el modelo está produciendo resultados deficientes que necesitan atención crítico/recompensa/media Recompensa promedio en todo el conjunto de entrenamiento Indicador de progreso primario: muestra una mejora general del rendimiento del modelo; generalmente debe tener una tendencia ascendente durante un entrenamiento exitoso crítico/recompensa/máxima Recompensa máxima lograda en el conjunto de entrenamiento Seguimiento del desempeño en el mejor de los casos: muestra la capacidad máxima del modelo; ayuda a identificar si el modelo puede lograr resultados excelentes incluso si el promedio es bajo actor/ppo_kl Divergencia de KL entre la iteración de política actual y la anterior Monitoreo de la estabilidad de la capacitación: los valores altos indican cambios rápidos en las políticas que pueden desestabilizar la capacitación; debe permanecer moderado actor/pg_clipfrac Fracción de actualizaciones de políticas que alcanzan el límite de recorte Indicador de agresividad de actualización: los valores moderados indican un aprendizaje saludable; demasiado alto sugiere actualizaciones demasiado agresivas que pueden desestabilizar la capacitación, demasiado bajo (por ejemplo, cero) sugiere un aprendizaje ineficiente. Esto es válido sólo durante las actualizaciones de PPO fuera de la política. actor/kl_loss Divergencia de KL entre la política actual y el modelo de referencia fijo Prevención de deriva de referencia: ayuda a evitar que el modelo se desvíe demasiado del comportamiento original; importante para mantener las capacidades de codificación val/test_score/code_contests Recompensa/rendimiento en el conjunto de validación retenido Comprobación de generalización: métrica más importante para el rendimiento real; detecta sobreajuste y mide la verdadera mejora del modelo

(Opcional) Observabilidad con el panel de Ray y Grafana

Para acceder al Panel de Ray y habilitar la visualización de Grafana durante el entrenamiento, establezca el reenvío de puertos mediante AWS Systems Manager (SSM). Para obtener más información sobre la configuración de AWS SSM, consulte Configuración rápida de AWS Systems Manager.

Primero, identifique el nodo principal en su clúster de múltiples nodos examinando los registros de CloudWatch:

__main__ – INFORMACIÓN – Se encontraron varios hosts, inicializando Ray como un clúster de múltiples nodos __main__ – INFORMACIÓN – Nodo principal: algo-1, Host actual: algo-2

Acceda a Ray Dashboard reenviando el puerto 8265 desde el nodo principal:

aws ssm start-session —target sagemaker-training-job:train-codefu-verl-ray-20250821185206_algo-1 –region us-east-1 –document-name AWS-StartPortForwardingSession –parameters '{"portNumber":["8265"],"localPortNumber":["8265"]}'

Habilite Grafana para recopilar métricas de Ray reenviando el puerto 8080 (puerto de exportación de métricas de Ray):

aws ssm start-session —target sagemaker-training-job:train-codefu-verl-ray-20250821185206_algo-1 –region us-east-1 –document-name AWS-StartPortForwardingSession –parameters '{"portNumber":["8080"],"localPortNumber":[""]}'

Una vez que se establece el reenvío de puertos, se puede acceder al Ray Dashboard en localhost:8265 en su navegador, que proporciona información detallada sobre:

Utilización de los trabajadores en el clúster distribuido Estado de ejecución de tareas y métricas de rendimiento Consumo de recursos, incluido el uso de GPU y memoria Programación de actores y tareas entre los trabajadores de Ray

Los paneles integrados de Grafana brindan una visualización integral de las métricas de capacitación, el rendimiento del sistema y el estado del clúster en tiempo real:

Esta configuración de observabilidad es crucial para depurar problemas de capacitación de RL distribuida, optimizar la asignación de recursos y garantizar que el proceso de capacitación progrese de manera eficiente en todo el clúster de SageMaker de múltiples nodos.

Limpiar

Para limpiar sus recursos y evitar cargos continuos, siga estos pasos:

Eliminar los recursos de SageMaker Studio no utilizados (opcional) Eliminar el dominio de SageMaker Studio En la consola de SageMaker, elija Capacitación en el panel de navegación y verifique que su trabajo de capacitación ya no se esté ejecutando.

Conclusiones

Esta publicación demuestra cómo entrenar modelos de razonamiento especializados para programación competitiva utilizando la solución de trabajos de capacitación Ray en Amazon SageMaker combinada con el marco de aprendizaje por refuerzo de veRL.

La solución de trabajos de capacitación de Ray on SageMaker simplifica la complejidad de orquestar cargas de trabajo de RL distribuidas al manejar automáticamente la inicialización del clúster de Ray, la coordinación de múltiples nodos y la administración de recursos en entornos informáticos heterogéneos. Esta integración permite a las organizaciones utilizar las capacidades informáticas distribuidas avanzadas de Ray, incluido el soporte para arquitecturas complejas de múltiples componentes, asignación dinámica de recursos y ejecución tolerante a fallas, mientras se benefician de la infraestructura totalmente administrada, la seguridad de nivel empresarial y el modelo de precios de pago por uso de SageMaker.

El análisis detallado de métricas demostró cómo monitorear el estado de la capacitación a través de la progresión de recompensas, indicadores de estabilidad de políticas y desempeño de generalización, lo que permitió a los profesionales identificar configuraciones de capacitación óptimas y solucionar problemas de capacitación distribuida de manera efectiva.

Para comenzar a implementar la capacitación de RL distribuida con Ray en SageMaker, visite el repositorio de GitHub de trabajos de capacitación de Ray en Amazon SageMaker para conocer el marco de la solución fundamental. La implementación completa de la capacitación de CodeFu-7B, incluidos ejemplos de configuración e integración de veRL, está disponible en este repositorio de GitHub.

Sobre los autores

Bruno Pistone

Bruno Pistone es arquitecto senior especializado en soluciones de IA/ML generativa a nivel mundial en AWS con sede en Milán, Italia. Trabaja con equipos de productos de AWS y grandes clientes para ayudarlos a comprender completamente sus necesidades técnicas y diseñar soluciones de aprendizaje automático e inteligencia artificial que aprovechen al máximo la nube de AWS y la pila de aprendizaje automático de Amazon. Su experiencia incluye cargas de trabajo de inferencia y capacitación distribuida, personalización de modelos, IA generativa y aprendizaje automático de un extremo a otro. Le gusta pasar tiempo con amigos, explorar nuevos lugares y viajar a nuevos destinos.

Giuseppe Angelo Porcelli

Giuseppe Angelo Porcelli es arquitecto principal de soluciones especializado en aprendizaje automático para Amazon Web Services. Con varios años de ingeniería de software y experiencia en aprendizaje automático, trabaja con clientes de cualquier tamaño para comprender sus necesidades técnicas y comerciales y diseñar soluciones de inteligencia artificial y aprendizaje automático que aprovechen al máximo la nube de AWS y la pila de aprendizaje automático de Amazon. Ha trabajado en proyectos en diferentes dominios, incluidos MLOps, visión por computadora y PNL, que involucran un amplio conjunto de servicios de AWS. En su tiempo libre, a Giuseppe le gusta jugar al fútbol.

canción yin

Yin Song es científico aplicado senior en el equipo de creación de prototipos de AWS en Sydney, Australia, con más de cinco años de experiencia ayudando a los clientes a crear prototipos personalizados que demuestran casos de uso complejos de servicios de AWS. Su trabajo se centra en la investigación sobre el ajuste y el servicio de modelos de IA, lo que permite soluciones de IA impactantes de un extremo a otro. Yin, un apasionado defensor del código abierto, lidera iniciativas de IA generativa que han producido modelos ampliamente adoptados.

Chen Wu

Chen Wu es científico aplicado principal en el equipo de creación de prototipos de AWS, donde impulsa tanto la investigación aplicada como la interacción de alto impacto con los clientes. Se especializa en modelos de lenguaje de contexto largo, LLM de razonamiento, sistemas agentes y sistemas de inteligencia artificial de alto rendimiento. Chen lidera el desarrollo del Agent Training Kit, un marco de código abierto para agentes en continuo aprendizaje. Ha realizado compromisos estratégicos en modelos básicos genómicos, optimización de LLM, generación de imágenes multiescala y canales de IA volumétrica 3D/4D. Sus LLM abiertos sobre Hugging Face han logrado más de 1 millón de descargas y su investigación de contexto prolongado ha aparecido en NeurIPS 2024 y ACL 2025. Es finalista del Premio ACM Gordon Bell.