Title: Leveraging Distributed Reinforcement Learning for Competitive Programming Code Generation with Ray on Amazon SageMaker
Introduction
The rapid advancement of artificial intelligence (AI) has created unprecedented demand for specialized models capable of complex reasoning tasks, particularly in competitive programming where models must generate functional code through algorithmic reasoning rather than pattern memorization. Reinforcement learning (RL) enables models to learn through trial and error by receiving rewards based on actual code execution, making it particularly well-suited for developing genuine problem-solving capabilities in algorithmic domains.
Challenges in Distributed RL Training
However, implementing distributed RL training for code generation presents significant infrastructure challenges such as orchestrating multiple heterogeneous components, coordinating parallel code compilation across nodes, and maintaining fault tolerance for long-running processes. Ray is one of the frameworks for distributed workloads that address these challenges, due to its unified system that handles the entire AI pipeline, GPU-first architecture, and seamless integration with tools like Hugging Face Transformers and PyTorch.
Integration with Amazon SageMaker
Workloads can be run with the Ray framework on SageMaker training jobs by using the Ray on Amazon SageMaker Training jobs solution, which combines Ray’s distributed computing framework with SageMaker’s fully managed infrastructure. This solution automatically handles Ray cluster initialization, multi-node coordination, and distributed resource management, enabling developers to focus on model development while benefiting from SageMaker’s enterprise-grade features.
Case Study: Training CodeFu-7B
In this post, we demonstrate how to train CodeFu-7B, a specialized 7-billion parameter model for competitive programming, using Group Relative Policy Optimization (GRPO) with veRL, a flexible and efficient training library for large language models (LLMs) that enables straightforward extension of diverse RL algorithms and seamless integration with existing LLM infrastructure, within a distributed Ray cluster managed by SageMaker training jobs.
About CodeFu-7B
CodeFu-7B-v0.1 is a 7B parameter language model specifically trained for solving Competitive Programming (CP) problems. Built upon the DeepSeek-R1-Distill-Qwen-7B base model, CodeFu demonstrates how reinforcement learning can develop capabilities in algorithmic reasoning and efficient C++ code generation beyond traditional supervised fine-tuning approaches.
Ray in SageMaker Training Jobs Solution
Ray on Amazon SageMaker Training jobs is a solution that enables distributed data processing and model training using Ray within SageMaker’s managed training environment. The solution provides capabilities such as universal launcher architecture for automatic Ray cluster setup and integrated observability, among others.
Solution Overview
The workflow for training CodeFu 7B with veRL and Ray on SageMaker training jobs consists of several important steps including: data preparation, training job submission, monitoring and observability, and automatic cleanup.
Prerequisites
Before running the notebook, certain prerequisites must be met, including increasing quotas for SageMaker AI instances and setting up necessary IAM roles.
Preparing the Dataset
The data preparation pipeline transforms the raw DeepMind CodeContest dataset into a suitable format for RL training, applying systematic filters and categorizing problems into various difficulty tiers.
GRPO Training Using veRL
The training process effectively utilizes Ray to orchestrate distributed execution and employs GRPO for stable training dynamics.
Ray Workload with SageMaker Training Jobs
To train CodeFu-7B, the ModelTrainer class from the SageMaker Python SDK is employed, enabling custom Docker containers and streamlined training workloads.
Experiment Tracking and Observability
The training pipeline integrates with Managed MLflow and third-party solutions for comprehensive experiment tracking and visualization of reinforcement learning metrics.
Clean-Up and Conclusion
Finally, it is essential to clean up resources after training to avoid ongoing charges. This post showcases the practical implementation of training specialized reasoning models using the Ray and Amazon SageMaker solutions, demonstrating the effective orchestration of distributed RL workloads.
About the Authors
The authors detail their backgrounds and experiences in the field of AI and machine learning, emphasizing their expertise in distributed systems and reinforcement learning.
Training Advanced AI for Competitive Programming: CodeFu-7B and Ray on SageMaker
Introduction
The rapid evolution of artificial intelligence (AI) has generated a pressing need for models that can tackle complex reasoning tasks, particularly within the sphere of competitive programming. In this domain, models must go beyond mere pattern recognition; they should generate functional code through algorithmic reasoning. The implementation of Reinforcement Learning (RL)—a technique where models learn via trial and error by receiving rewards based on code execution—has emerged as a powerful approach to cultivate genuine problem-solving abilities.
Despite its potential, distributed RL training for code generation introduces significant infrastructure challenges. Tasks such as orchestrating various heterogeneous components, coordinating code compilations across different nodes, and ensuring fault tolerance for long-running processes complicate the training landscape. Enter Ray, a framework designed for distributed workloads that effectively addresses these issues. With its unified system catering to the entire AI pipeline and seamless integration with various tools, Ray has become a go-to choice for developing sophisticated AI models.
In this post, we will explore how to train CodeFu-7B, a highly specialized 7-billion parameter model designed for competitive programming, using Group Relative Policy Optimization (GRPO) along with veRL—a flexible and efficient training library for large language models (LLMs). We’ll examine the entire process, from data preparation to distributed training setup, demonstrating how this unified approach can enhance both computational scale and the overall developer experience.
About CodeFu-7B
CodeFu-7B is a remarkable 7B parameter language model adept at solving Competitive Programming (CP) issues. Initially developed based on the DeepSeek-R1-Distill-Qwen-7B base model, CodeFu exemplifies how reinforcement learning can not only enhance algorithmic reasoning but also significantly improve C++ code generation—a feat that traditional supervised fine-tuning alone cannot achieve.
The training of CodeFu draws upon problem statements from the DeepMind CodeContest dataset without relying on ground-truth solutions. This unique approach forces the model to learn via trial and error, cultivating genuine problem-solving skills rather than rote memorization.
Moreover, CodeFu is publicly accessible on HuggingFace under the MIT license, which opens up opportunities for researchers and practitioners eager to venture into the realms of code generation and algorithmic reasoning.
Ray in SageMaker Training Jobs Solution
Ray on Amazon SageMaker Training jobs provides a robust solution for distributed data processing and model training using Ray in a fully managed training environment. This integrated system delivers key capabilities such as:
- Automatic Ray cluster setup
- Intelligent multi-node coordination
- Heterogeneous cluster support for mixed instance types
- Integrated observability with Ray Dashboard, Prometheus, and Amazon CloudWatch
This comprehensive integration allows developers to leverage Ray’s distributed computing capabilities while enjoying the managed infrastructure of SageMaker, making it an ideal choice for complex workloads like reinforcement learning training.
Solution Overview
The workflow for training CodeFu-7B with veRL and Ray on SageMaker involves a series of systematic steps:
- Data Preparation: Upload the preprocessed DeepMind CodeContest dataset and training configuration.
- Training Job Submission: Use the ModelTrainer class via the SageMaker Python SDK to submit a training job.
- Monitoring and Observability: Engage real-time monitoring through Ray Dashboard and supplementary tools like Prometheus and Grafana.
- Automatic Cleanup: SageMaker will manage model saving and resource decommissioning post-training.
This streamlined architecture focuses on delivering an entirely managed reinforcement learning training experience. It empowers developers to concentrate on model development while SageMaker and Ray take care of the complex infrastructure orchestration.
Prerequisites
Before executing the training notebook, several prerequisites must be completed:
- AWS Quota Increases: Request minimum quota increases for p4de instances suitable for training jobs.
- IAM Role Setup: Create an IAM role with the necessary permissions for SageMaker.
- (Optional) Studio Domain Creation: Set up an Amazon SageMaker Studio domain or use a local development environment for executing the training job.
Prepare the Dataset
The data preparation pipeline curates the raw DeepMind CodeContest dataset into a format suitable for reinforcement learning training. It applies systematic filters to identify compatible problems while removing those with insufficient ratings. Each problem is organized into components that shape the model’s learning environment, emphasizing learning from execution feedback instead of ground-truth solutions.
GRPO Training Using veRL
The training employs Ray to manage distributed execution and synchronization of the vLLM rollout, reward evaluation, and model optimization via GRPO. This nuanced approach enhances traditional proximal policy optimization by applying group-relative baselines, thereby stabilizing training and reducing variance in policy gradient estimates.
To facilitate efficient execution, the architecture is designed to support long-form reasoning and complex response generation through parallel code compilation and evaluation.
Ray Workload with SageMaker Training Jobs
The training of CodeFu-7B begins by leveraging the ModelTrainer class from the SageMaker Python SDK. This involves:
- Defining the instance type and container image for the training job.
- Creating a custom Docker container for the necessary dependencies.
- Setting up the training job using flexible input data channels from designated S3 paths.
Once the job is submitted, the process can be closely monitored through the SageMaker console, allowing for comprehensive insights into the job’s status and associated metrics.
Experiment Tracking
By integrating with Managed MLflow on Amazon SageMaker as well as other third-party solutions, the CodeFu training pipeline ensures efficient experiment tracking. Metrics such as reward progression, policy stability indicators, and validation performance allow for dynamic adjustments and refinements in the model training process.
Observability and Cleanup
To analyze training performance, Ray Dashboard and Grafana set up enable real-time insight into the workflow. Following completion, it’s crucial to clean up resources to avoid unnecessary costs, which involves deleting unused SageMaker resources and confirming that training jobs have concluded.
Conclusions
This post demonstrates the seamless integration of advanced AI training methodologies using Ray on Amazon SageMaker Training jobs. By simplifying the complexities inherent to orchestrating distributed RL workloads, organizations can leverage Ray’s advanced capabilities under the fully managed umbrella of SageMaker.
For those eager to embark on developing specialized reasoning models for competitive programming, the foundational solution framework and CodeFu-7B training implementation can be explored further via GitHub.
About the Authors
Bruno Pistone: Senior Worldwide Generative AI/ML Specialist Solutions Architect at AWS.
Giuseppe Angelo Porcelli: Principal Machine Learning Specialist Solutions Architect at AWS.
Yin Song: Senior Applied Scientist focused on tailored prototypes in AI.
Chen Wu: Principal Applied Scientist specializing in long-context language models and agentic systems.