Skip to content
Modal ModalTraining Gym
Repo

Getting started with RL

View on GitHub

This tutorial introduces some core features of the Training Gym by walking through a simple example of Reinforcement Learning with Verifiable Rewards (RLVR), a foundational method of RL post-training. Here, we teach Qwen3.5-4B how to write correct haikus. Step by step, we’ll show the foundations of running training jobs on the Gym.

import nltk
from nltk.corpus import cmudict


import re


from modal_training_gym import (
    Endpoint,
    HuggingFaceDataset,
    Qwen3_5_4B,
    Qwen3_5_4B_Recipe,
    TrainConfig,
)

Deploy the base model

As with all training tasks, we need a baseline to decide how much training we need. To do that, we need a way to run inference on the base model so that we can try it out.

Luckily, Endpoints allows us to easily deploy a production-ready LLM inference endpoint on Modal’s managed infrastructure. It supports both open model weights in addition to custom fine tunes, sourced from either a Hugging Face repo or a Modal Volume. To use it, we provide a class to instantiate one programatically.

It will take a moment to download the model weights onto a Modal Volume and boot containers past the cold-start. Once you see the URL has been printed, you’re ready to move on!

model = Qwen3_5_4B()


base_deployment = Endpoint.launch(
    model, unauthenticated=True, recreate_if_existing=True
)
base_deployment.wait_until_ready(timeout=15 * 60)
print(f"base model deployed to {base_deployment.url}")

Define a scoring function

To evaluate the base model, we need a function that takes as input a haiku and outputs a score (a.k.a. reward when we’re training) to represent whether it follows the 5-7-5 syllable format. We can do that using NLTK’s CMU Pronouncing Dictionary,

How should we define our scoring function? We could give it a score of 0 if it doesn’t follow the format and 1 if it does, but that’s not very informative for both the models being trained, and, more importantly, the human training the models! Instead, we want the score to provide sufficient granularity such that it’s immediately obvious what the failure mode is (if any). Below, we implement the following function:

  • Return -10 if the model was so incompetent that it failed to return three lines.
  • Otherwise, return the negative sum of absolute differences between the predicted and target syllable count for each line.

What does this mean? That the model will receive increasingly negative scores the further off its haiku is, with a maximum score of 0. Let’s now see how it does.

_cmudict_cache = {}


def _get_cmudict() -> dict:
    if not _cmudict_cache:
        nltk.download("cmudict", quiet=True)
        _cmudict_cache.update(cmudict.dict())
    return _cmudict_cache


def _count_syllables(text: str) -> int:
    cmu = _get_cmudict()
    total = 0
    for word in re.findall(r"[a-zA-Z]+", text):
        phones = cmu.get(word.lower())
        if phones:
            total += sum(p[-1].isdigit() for p in phones[0])
        else:
            count = len(re.findall(r"[aeiouy]+", word.lower()))
            if word.lower().endswith("e") and count > 1:
                count -= 1
            total += max(count, 1)
    return total


def score_haiku(response: str) -> float:
    lines = [line.strip() for line in response.strip().split("\n") if line.strip()]
    if len(lines) != 3:
        return -10
    total_diff = sum(
        abs(_count_syllables(line) - target)
        for line, target in zip(lines, [5, 7, 5])
    )
    return -float(total_diff)

Get the dataset

Note that we’ve only qualitatively assessed its performance. Now, we should get concrete numbers. How do we do that? First, we’ll have to curate a dataset. Luckily, statworx/haiku from Huggingface already exists, so we don’t have to create one ourselves.

Note that for more complex tasks, it is almost certainly the case that you will be creating your own dataset. Why? Because the task you’re trying to get your model to do is either too expensive or simply too hard for a bigger model. In either case, this is because the task is sufficiently out-of-distribution, and no existing dataset will serve your needs.

See the multi-turn example for a basic example of creating your own dataset, or the DatasetConfig documentation for a deeper dive.

class HaikuDataset(HuggingFaceDataset):
    hf_repo = "statworx/haiku"
    input_column = "keywords"
    output_column = "text"
    output_format = "jsonl"
    apply_chat_template = True
    always_prepare = True
    prompt_template = "Write a haiku about {input}."


train_dataset = HaikuDataset(hf_split="train[:10]")


eval_dataset = HaikuDataset(hf_split="train[10:15]")

Evaluate the base model

All we need to do now is, for each sample in our eval dataset, call the Endpoint, score each response, and calculate the mean. By default, Endpoints can process multiple inputs concurrently, so we loop over samples in parallel to speed up eval.

def run_eval(deployment, max_concurrency: int = 2) -> float:
    from concurrent.futures import ThreadPoolExecutor


    deployment.wait_until_ready(timeout=15 * 60)


    def _score_one(example):
        topic = str(example[eval_dataset.input_column])
        prompt = eval_dataset.prompt_template.format(input=topic)
        msg = deployment.chat(
            [{"role": "user", "content": prompt}],
            chat_template_kwargs={"enable_thinking": False},
        )
        return score_haiku(msg.get("content") or msg.get("reasoning_content") or "")


    with ThreadPoolExecutor(max_workers=max_concurrency) as executor:
        scores = list(executor.map(_score_one, eval_dataset.load()))
    return sum(scores) / len(scores) if scores else float("nan")


print("running base model evaluation...")
base_mean = run_eval(base_deployment)
print(f"average score: {base_mean:.1f}")

Creating a reward function

To make our scoring function a reward function, we just need to extract the text from the model’s response and pass it to our existing score_haiku. Simple enough.

async def haiku_rm(args, sample, **kwargs) -> float:
    response = model.parse_response(sample.response)
    return score_haiku(response.content)

Train the model

Finally, onto the training. The Gym supports both the Slime and Miles frameworks. Here, we use Slime for demonstration purposes.

Training is simple: pass in the model you intend to train, the dataset you wish to train on, and a recipe for how you want training to occur. The recipe wraps all framework-native flags in addition to providing Modal-specific ones.

Once we run the code below, training kicks off and we’ll immediately get a run ID, which we may use to watch the run’s progress in the dashboard.

config = TrainConfig(
    model=model,
    dataset=train_dataset,
    recipe=Qwen3_5_4B_Recipe(
        eval_interval=None,
        rollout_num_gpus=8,
        num_rollout=10,
        n_samples_per_prompt=8,
        save_interval=5,
        apply_chat_template_kwargs='{"enable_thinking": false}',
        custom_rm_function=haiku_rm,
        image_overlay=lambda image: image.run_commands(
            "uv pip install --system aiohttp 'nltk>=3.8.0'",
            "python -c \"import nltk; nltk.download('cmudict', quiet=True)\"",
        ),
    ),
)


run = config.launch()
print(f"run id: {run.training_run_id}")

Serve and evaluate the trained checkpoint

We’ll get the latest checkpoint and create a new Endpoint so we may evaluate it.

result = run.result()
checkpoint = result.checkpoints()[-1]
print(f"checkpoint: {checkpoint.path}")


trained_deployment = Endpoint.launch(
    model, checkpoint, unauthenticated=True, recreate_if_existing=True
)
trained_deployment.wait_until_ready(timeout=15 * 60)
print(f"checkpoint deployed to {trained_deployment.url}")

Now, let’s run the same eval as before.

print("running checkpoint evaluation...")
trained_mean = run_eval(trained_deployment)
print(f"average score: {trained_mean:.1f}")

Continue training off the checkpoint

Hmm, it looks like the trained model is still not doing very well. A likely cause is that it only trained for 10 iterations. Let’s continue training, starting from the last checkpoint.

new_config = TrainConfig(
    model=model,
    dataset=train_dataset,
    checkpoint=checkpoint,
    recipe=Qwen3_5_4B_Recipe(
        eval_interval=None,
        custom_rm_function=haiku_rm,
        rollout_num_gpus=8,
        num_rollout=20,
        n_samples_per_prompt=8,
        apply_chat_template_kwargs='{"enable_thinking": false}',
        image_overlay=lambda image: image.run_commands(
            "uv pip install --system aiohttp 'nltk>=3.8.0'",
            "python -c \"import nltk; nltk.download('cmudict', quiet=True)\"",
        ),
    ),
)


new_run = new_config.launch()
print(f"run id: {new_run.training_run_id}")

Evals Evals Evals

Once again, we’ll create a new Endpoint for the new checkpoint and run evals on it.

new_result = new_run.result()
new_checkpoint = new_result.checkpoints()[-1]
print(new_checkpoint.path)


new_deployment = Endpoint.launch(
    model, new_checkpoint, unauthenticated=True, recreate_if_existing=True
)
new_deployment.wait_until_ready(timeout=15 * 60)
print(f"new checkpoint deployed to {new_deployment.url}")


print("running new checkpoint evaluation...")
new_mean = run_eval(new_deployment)
print(f"average score: {new_mean:.1f}")