Пошаговый туториал: как обучить CodeFu-7B для решения задач по программированию с помощью veRL и Ray на Amazon SageMaker

Ты хочешь, чтобы твоя модель сама писала код, который проходит тесты на LeetCode. Не просто генерировала синтаксически верный мусор, а реально решала задачи. Fine-tuning на готовых решениях дает прирост, но модель все равно тупит на новых задачах. Ей не хватает логики, способности итеративно исправлять ошибки.

Reinforcement Learning (RL) выглядит логичным решением: модель пытается, получает награду за пройденные тесты, учится. Но классический RL для языковых моделей – это ад. Нестабильность, огромные вычислительные затраты, сложность масштабирования.

veRL (Vectorized Reinforcement Learning) решает часть проблем. Вместо последовательных взаимодействий со средой ты генерируешь целый батч разных решений одной задачи, оцениваешь их все сразу (векторно) и обучаешь модель на лучших. Это как отбор лучших клонов в параллельных вселенных. На практике это означает ускорение обучения в разы.

А теперь представь, что тебе нужно запустить это на кластере из 16 GPU. Оркестрация, синхронизация, обработка ошибок – голова болит уже на этапе планирования.

Здесь на сцену выходит дуэт Ray и Amazon SageMaker. Ray – фреймворк для распределенного Python, который превращает твой скрипт в распределенное приложение почти магией. SageMaker – управляемый сервис от AWS, который избавляет тебя от возни с инфраструктурой. Ты просто загружаешь код, указываешь конфигурацию и ждешь результат.

Этот гайд – пошаговая инструкция, как собрать этот пазл. От подготовки датасета до запуска распределенного обучения на кластере SageMaker. Мы не будем скрывать подводные камни – я укажу на каждую кочку, на которую сам наступал.

1Готовим полигон: данные и среду исполнения

Без качественного датасета и надежного способа исполнения кода ничего не выйдет. Тебе нужны задачи с тестами.

Берем APPS или CodeContests датасеты. Они содержат тысячи задач по программированию с открытыми тестами. Наш пайплайн:

• Скачиваем датасет.
• Фильтруем задачи: оставляем только те, где есть полный набор input/output тестов (несколько тестов на задачу).
• Подготавливаем в формате JSONL: один объект на задачу с полями problem_id, prompt (текст задачи), test_cases (список пар input/output).

Следующий критичный компонент – изолированная среда для исполнения кода. Запускать недоверенный код, сгенерированный ИИ, на своей машине – плохая идея. Используем Docker.

Пишем простой Python-сервис на FastAPI, который принимает код (скажем, на Python) и тестовые входные данные, запускает их в отдельном Docker-контейнере с ограничениями по времени и памяти, и возвращает результат (stdout, stderr, код возврата).

# Упрощенная логика исполнителя (executor.py)
import docker
import tempfile
import os

client = docker.from_env()

def run_code_in_container(code: str, input_data: str, timeout_seconds=5) -> dict:
    """Запускает код в изолированном контейнере и возвращает результат."""
    # Создаем временный Dockerfile с кодом
    with tempfile.TemporaryDirectory() as tmpdir:
        code_path = os.path.join(tmpdir, "solution.py")
        with open(code_path, 'w') as f:
            f.write(code)
        
        # Создаем и запускаем контейнер
        container = client.containers.run(
            image="python:3.9-slim",
            command=["python", "/app/solution.py"],
            stdin_open=True,
            mem_limit="128m",
            cpu_period=100000,
            cpu_quota=50000,
            network_mode="none",
            volumes={tmpdir: {'bind': '/app', 'mode': 'ro'}},
            working_dir="/app",
            detach=True
        )
        try:
            # Подаем входные данные и ждем результат
            result = container.wait(timeout=timeout_seconds)
            output = container.logs(stdout=True, stderr=False).decode()
            error = container.logs(stdout=False, stderr=True).decode()
            container.remove()
            return {
                "stdout": output,
                "stderr": error,
                "return_code": result['StatusCode'],
                "timed_out": False
            }
        except Exception as e:
            container.remove(force=True)
            return {"error": str(e), "timed_out": True}

Этот исполнитель станет ключевой частью нашей среды вознаграждения (reward environment). Модель генерирует код, мы запускаем его против батареи тестов, считаем процент пройденных – это и есть награда.

2Собираем тренировочный пайплайн с veRL и Ray

Теперь сердцевина. Мы используем библиотеку TRL (Transformers Reinforcement Learning) от Hugging Face, которая поддерживает veRL через класс RayTuneTrainer. Но TRL не заточена под специфику исполнения кода. Придется кастомизировать.

План:

1. Загружаем предобученную CodeFu-7B. Модель отлично подходит как база – она уже натренирована на кодекс.
2. Пишем функцию вознаграждения. Она принимает сгенерированные тексты (код) от модели, прогоняет через наш исполнитель и возвращает scores.
3. Конфигурируем Ray для распределения нагрузки. Функция вознаграждения будет запускаться на множестве воркеров параллельно.

Вот скелет нашего основного тренировочного скрипта (train_verl.py):

# Основные импорты
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
from trl import RayTuneTrainer, RLTrainer, RLTrainerCallback
from datasets import load_dataset
import ray
from ray import tune

# Инициализируем Ray
ray.init(address='auto')  # В среде SageMaker это подключится к кластеру

# Загружаем модель и токенизатор
model_name = "neulab/codefu-7b"  # Актуально на 2026 год
tokenizer = AutoTokenizer.from_pretrained(model_name)
tokenizer.pad_token = tokenizer.eos_token  # Важно для padding
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype=torch.bfloat16,  # Для экономии памяти и ускорения
    device_map="auto"
)

# Загружаем и подготавливаем датасет
def format_dataset(example):
    """Форматируем промпт для модели."""
    prompt = f"Реши следующую задачу на Python:\n{example['prompt']}\n### Решение:\n"
    return {"prompt": prompt}

dataset = load_dataset("json", data_files="my_coding_problems.jsonl", split="train")
dataset = dataset.map(format_dataset)

# Функция вознаграждения (запускается на воркерах Ray)
@ray.remote(num_gpus=0.25)  # Каждому воркеру по 0.25 GPU, можно масштабировать
def compute_reward(generated_codes, problem_ids):
    """Принимает батч сгенерированных кодов и ID задач, возвращает награды."""
    rewards = []
    for code, problem_id in zip(generated_codes, problem_ids):
        # Загружаем тесты для этой задачи (из кеша или базы)
        test_cases = load_tests_for_problem(problem_id)
        passed = 0
        for inp, expected_out in test_cases:
            result = code_executor.run(code, inp)  # Наш Docker-исполнитель
            if result["return_code"] == 0 and result["stdout"].strip() == expected_out.strip():
                passed += 1
        reward = passed / len(test_cases)  # Доля пройденных тестов
        rewards.append(reward)
    return rewards

# Конфигурация тренера
trainer_config = {
    "model": model,
    "args": {
        "num_train_epochs": 3,
        "per_device_train_batch_size": 4,
        "gradient_accumulation_steps": 8,
        "warmup_steps": 100,
        "learning_rate": 1e-5,
        "fp16": False,
        "bf16": True,  # Предпочтительнее для современных GPU (A100, H100)
        "logging_steps": 10,
        "output_dir": "/opt/ml/model",  # Стандартный путь для SageMaker
        "report_to": "none",
    },
    "train_dataset": dataset,
    "reward_func": compute_reward,  # Ray-remote функция
    "ray_config": {
        "num_workers": 16,  # Количество параллельных воркеров для оценки
        "local_dir": "/opt/ml/checkpoints",  # Для сохранения чекпоинтов
    }
}

# Создаем и запускаем тренер
trainer = RayTuneTrainer(**trainer_config)
trainer.train()

3Упаковываем все в контейнер для Amazon SageMaker

SageMaker не будет просто так запускать твой скрипт. Нужно создать Docker-образ со всеми зависимостями и правильной структурой.

Создаем Dockerfile:

# Используем официальный PyTorch образ от AWS
FROM 763104351884.dkr.ecr.us-east-1.amazonaws.com/pytorch-training:2.5.0-gpu-py310-cu124-ubuntu20.04

# Устанавливаем системные зависимости для Docker внутри Docker (dind)
RUN apt-get update && apt-get install -y docker.io

# Копируем требования и устанавливаем Python пакеты
COPY requirements.txt /opt/program/requirements.txt
RUN pip install --upgrade pip && \
    pip install -r /opt/program/requirements.txt

# Копируем наш тренировочный код и утилиты
COPY train_verl.py /opt/program/
COPY executor.py /opt/program/
COPY entrypoint.sh /opt/program/entrypoint.sh

# Делаем entrypoint исполняемым
RUN chmod +x /opt/program/entrypoint.sh

WORKDIR /opt/program
ENTRYPOINT ["./entrypoint.sh"]

entrypoint.sh – это скрипт, который запускается внутри контейнера SageMaker. Он должен запустить наш тренировочный скрипт с правильными аргументами.

#!/bin/bash
# entrypoint.sh

# Запускаем Docker daemon в фоне (нужен для нашего code executor)
sudo dockerd &
sleep 5  # Ждем инициализации

# Запускаем основной тренировочный скрипт
python train_verl.py

requirements.txt должен содержать:

torch>=2.5.0
transformers>=4.45.0
trl>=0.9.0
ray[tune]>=2.36.0
accelerate>=0.30.0
datasets>=2.20.0
fastapi
uvicorn
docker
sagemaker-training

4Запускаем распределенное обучение через SageMaker SDK

Все готово. Теперь создаем SageMaker Training Job из Python скрипта (например, launch_training.py), используя SageMaker Python SDK.

import sagemaker
from sagemaker.pytorch import PyTorch
from sagemaker import get_execution_role

# Инициализация сессии SageMaker
sagemaker_session = sagemaker.Session()
role = get_execution_role()

# Загружаем наш Docker образ в ECR (предполагается, что ты уже собрал и загрузил)
image_uri = ".dkr.ecr..amazonaws.com/codefu-verl-training:latest"

# Конфигурация распределенного обучения с Ray
instance_count = 4  # 4 инстанса
instance_type = "ml.g5.12xlarge"  # Каждый с 4x A10G GPU (всего 16 GPU)

# Определяем параметры обучения
train_args = {
    "entry_point": "train_verl.py",
    "source_dir": ".",  # Текущая директория со скриптами
    "image_uri": image_uri,
    "role": role,
    "instance_count": instance_count,
    "instance_type": instance_type,
    "volume_size": 200,  # ГБ для датасета и чекпоинтов
    "max_run": 172800,  # 48 часов максимум
    "environment": {
        "RAY_DEDUP_LOGS": "0",
        "NCCL_SOCKET_IFNAME": "eth0",
        "PYTHONUNBUFFERED": "TRUE",
    },
    "disable_profiler": True,
    "debugger_hook_config": False,
    "distribution": {
        "pytorchddp": {
            "enabled": False  # Мы используем Ray для распределения, а не PyTorch DDP
        },
        "mpi": {
            "enabled": False
        },
        "ray": {
            "enabled": True,
            "ray_processes_start_cmd": "ray start --head --port=6379 --object-manager-port=8076 --autoscaling-config=~/ray_bootstrap_config.yaml",
        }
    }
}

# Создаем и запускаем задание
trainer = PyTorch(**train_args)
trainer.fit(wait=False)  # Не ждем завершения, чтобы скрипт мог продолжить

print(f"Запущен Training Job: {trainer.latest_training_job.name}")
print(f"Следить за логами можно в CloudWatch или через: aws sagemaker describe-training-job --training-job-name {trainer.latest_training_job.name}")

После запуска SageMaker развернет кластер из 4 инстансов, загрузит твой Docker-образ, установит соединение между нодами через Ray и начнет обучение.

Где собака зарыта: типичные ошибки и нюансы

Теория гладкая, практика – это грабли. Вот на что спотыкаются 90% людей:

Самая большая ошибка – ожидать, что модель с нуля научится решать сложные задачи. Начни с простых задач (например, только ввод-вывод, базовые алгоритмы). Используй предварительный fine-tuning на готовых решениях (как в этом руководстве), а затем уже включай RL для полировки и адаптации к формату тестов.

Частые вопросы (FAQ)

Итог: veRL на SageMaker с Ray – это мощный, но сложный инструмент. Он не для хобби-проектов на Colab. Это для команд, которые серьезно нацелены на создание специализированных AI для программирования. Ты берешь на себя боль инфраструктуры, но получаешь контроль и масштабируемость. Главное – начни с малого, валидируй каждый компонент по отдельности, и только потом запускай монстра на полную мощность. Удачи, и пусть твоя модель наконец-то сдаст ту задачу с бинарным поиском, которая не давала тебе покоя все эти годы.

Подписаться на канал