wehub/ray-project--ray
1
0
Fork 0
Code Issues Pull Requests Actions Deployments Models Agent Notes Packages Projects Releases Wiki Activity

chore: import upstream snapshot with attribution

Browse Source
This commit is contained in:
wehub-resource-sync
2026-07-13 13:17:40 +08:00
commit f1825c8ceb
10096 changed files with 2364182 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
doc/source/train/examples/transformers/BUILD.bazel
+30
View File
@@ -0,0 +1,30 @@
load("//bazel:python.bzl", "py_test_run_all_notebooks")
filegroup(
name = "huggingface_transformers_examples",
srcs = glob(["*.ipynb"]),
visibility = ["//doc:__subpackages__"],
)
# GPU Tests
py_test_run_all_notebooks(
size = "large",
include = ["*.ipynb"],
data = ["//doc/source/train/examples/transformers:huggingface_transformers_examples"],
exclude = [],
tags = [
"exclusive",
"gpu",
"ray_air",
"team:ml",
],
)
filegroup(
name = "transformers_examples_ci_configs",
srcs = glob([
"**/ci/aws.yaml",
"**/ci/gce.yaml",
]),
visibility = ["//doc/source/train/examples:__pkg__"],
)
doc/source/train/examples/transformers/huggingface_text_classification.ipynb
+1338
View File
File diff suppressed because it is too large Load Diff
doc/source/train/examples/transformers/transformer_reinforcement_learning/README.ipynb
+417
View File
@@ -0,0 +1,417 @@
{
"cells": [
{
"cell_type": "markdown",
"id": "d31e727a103a6eb7",
"metadata": {},
"source": [
"# RL Post-Training using Hugging Face TRL with GRPO\n",
"\n",
"This notebook builds on the TRL GRPO Example, [GRPO Trainer](https://huggingface.co/docs/trl/en/grpo_trainer), to fine-tune a Qwen2.5 0.5B model to answer math questions using the DeepMath-103k dataset and Group Relative Policy Optimization (GRPO) algorithm.\n",
"\n",
"Ray Train scales this implementation efficiently across multiple GPUs without changing training logic.\n",
"\n",
"This notebook consists of the following steps:\n",
"\n",
"1. Quick Summary of GRPO and the DeepMath dataset\n",
"2. Package and Environment Setup\n",
"3. Running TRL with Ray Train\n",
"4. Scaling to more GPUs\n",
"\n",
"<div id=\"anyscale-note\" class=\"alert alert-block alert-warning\">\n",
"\n",
" <strong>Anyscale Specific Configuration</strong>\n",
"\n",
" <p><strong>Note:</strong> This tutorial is optimized for the Anyscale platform. When running on open source Ray, you need additional configuration. For example, you would need to manually:</p>\n",
"\n",
" <ul>\n",
" <li><strong>Configure your Ray Cluster</strong>: Set up your multi-node environment and manage resource allocation without Anyscale's automation.</li>\n",
" <li><strong>Manage Dependencies</strong>: Manually install and manage dependencies on each node.</li>\n",
" <li><strong>Set Up Storage</strong>: Configure your own distributed or shared storage system for model checkpointing.</li>\n",
" </ul>\n",
"\n",
" <p>All these configurations are handled automatically through the Anyscale platform.\n",
"</div>\n",
"\n",
"<style>\n",
" div#anyscale-note > p,\n",
" div#anyscale-note > ul,\n",
" div#anyscale-note > ul li {\n",
" color: black;\n",
" }\n",
"\n",
" div#anyscale-note {\n",
" background-color: rgb(255, 243, 205);\n",
" }\n",
"\n",
" div#anyscale-note {\n",
" border: 1px solid #ccc;\n",
" border-radius: 8px;\n",
" padding: 15px;\n",
" }\n",
"\n",
"</style>"
]
},
{
"cell_type": "markdown",
"id": "94faee5d4ef298c5",
"metadata": {},
"source": [
"## 1. Summary of the DeepMath dataset and GRPO algorithm\n",
"\n",
"### DeepMath dataset\n",
"\n",
"This example uses the DeepMath-103k dataset by [He et al., 2025](https://arxiv.org/abs/2504.11456), consisting of 103,000 challenging questions across a wide range of mathematical subjects including Algebra, Calculus, Number Theory, Geometry, Probability, and Discrete Mathematics. These questions are more heavily distributed towards challenging questions, in particular Levels 5 to 9, compared to alternative datasets, making them highly complex and difficult to solve without specialist training.\n",
"\n",
"Each example contains two fields:\n",
"- **`prompt`**: a list of chat-format message dicts; the question text is in `prompt[0][\"content\"]`.\n",
"- **`solution`**: a LaTeX-formatted ground-truth answer, used by the reward function to verify the model's response.\n",
"\n",
"Using a structured LaTeX format ensures the reward function can reliably parse and compare both the model's output and the ground truth. Here are a few examples from the dataset:\n",
"\n",
"| Subject | Prompt | Solution |\n",
"|---|---|---|\n",
"| Number Theory | Determine the number of 19th power residues modulo 229. | $12$ |\n",
"| Functional Analysis | Determine the norm $\\lVert T \\rVert$ of the linear operator $T: \\ell^2 \\rightarrow \\ell^2$ given by $(Tx)_1 = 0$, $(Tx)_n = -x_n + \\alpha x_{n+1}$ for $n \\geq 2$, where $\\alpha \\in \\mathbb{C}$. | $1 + \\lvert\\alpha\\rvert$ |\n",
"| Analysis | Determine whether there exists a Schwartz function $g \\in \\mathcal{S}(\\mathbb{R}^n)$ such that for a given continuous function $f: \\mathbb{R}^n \\to \\mathbb{R}$ with $f \\not\\equiv 0$, the integral $\\int_{\\mathbb{R}^n} g^2 f \\, dx \\neq 0$. | Yes |\n",
"\n",
"### Group Relative Policy Optimization Algorithm\n",
"\n",
"To RL post-train the model to answer math questions, this example uses TRL's GRPO (Group Relative Policy Optimization) implementation introduced by [Shao et al., 2024](https://arxiv.org/abs/2402.03300). GRPO is an online algorithm that iteratively improves by training on data the model generates itself. Each iteration proceeds through four steps:\n",
"\n",
"1. **Generate**: For each prompt in the batch, sample a group of G completions from the current policy.\n",
"2. **Score**: Apply the reward function to each completion, producing a scalar reward per completion.\n",
"3. **Advantage**: Normalize the rewards within each group by subtracting the group mean and dividing by the group standard deviation. This relative score becomes the advantage $\\hat{A}_{i,t}$ where positive means better than peers.\n",
"4. **Loss**: Update the policy to increase the probability of high-advantage completions and decrease it for low-advantage ones, while a KL penalty keeps the updated policy close to the reference.\n",
"\n",
"$$\n",
"\\mathcal{L}_{\\text{GRPO}}(\\theta) = -\\frac{1}{\\sum_{i=1}^{G} |o_i|} \\sum_{i=1}^{G} \\sum_{t=1}^{|o_i|} \\left[ \\frac{\\pi_\\theta(o_{i,t} \\mid q, o_{i,\\mathopen{<} t})}{\\left[\\pi_\\theta(o_{i,t} \\mid q, o_{i,\\mathopen{<} t})\\right]_{\\text{no grad}}} \\hat{A}_{i,t} - \\beta \\mathbb{D}_{\\text{KL}}\\left[\\pi_\\theta \\| \\pi_{\\text{ref}}\\right] \\right],\n",
"$$\n",
"\n",
"For a deeper, more technical, description of the algorithm, see the [TRL GRPO page](https://huggingface.co/docs/trl/en/grpo_trainer#looking-deeper-into-the-grpo-method) and the [original paper](https://arxiv.org/abs/2402.03300)."
]
},
{
"metadata": {},
"cell_type": "markdown",
"source": [
"## 2. Package and Environment Setup\n",
"\n",
"Install the required dependencies"
],
"id": "1c71d4c53875d94f"
},
{
"metadata": {},
"cell_type": "code",
"source": "# !pip install \"trl[vllm]\" \"math_verify\" \"transformers==4.57.6\"",
"id": "b78729cb7d9d2730",
"outputs": [],
"execution_count": null
},
{
"metadata": {},
"cell_type": "markdown",
"source": "Use `ray.init()` to initialize a local cluster. By default, this cluster contains only the machine you are running this notebook on. You can also run this notebook on an Anyscale cluster.",
"id": "a362d29fb3189847"
},
{
"metadata": {},
"cell_type": "code",
"outputs": [],
"execution_count": null,
"source": [
"import ray\n",
"\n",
"ray.init()"
],
"id": "3455573a2a2218f0"
},
{
"metadata": {},
"cell_type": "markdown",
"source": [
"Use `ray.cluster_resources()` to check which resources your cluster has access to.\n",
"If you're running this notebook on your local machine or Google Colab, you should see the number of CPU cores and GPUs available to you."
],
"id": "f7018eb1a753b41d"
},
{
"metadata": {},
"cell_type": "code",
"outputs": [],
"execution_count": null,
"source": [
"from pprint import pprint\n",
"\n",
"pprint(ray.cluster_resources())"
],
"id": "2690f780c2975613"
},
{
"metadata": {},
"cell_type": "markdown",
"source": "Change these two variables to control whether the training uses CPUs or GPUs, and how many workers to spawn. Each worker claims one CPU or GPU, so make sure not to request more resources than are available. By default, the training runs with four GPU worker.",
"id": "c7296dcf7a48bbaa"
},
{
"metadata": {},
"cell_type": "code",
"outputs": [],
"execution_count": null,
"source": [
"use_gpu = True # set this to `False` to run on CPUs\n",
"num_workers = 4 # set this to the number of GPUs or CPUs you want to use"
],
"id": "776d49988bfefc48"
},
{
"cell_type": "markdown",
"id": "43ad7b293617547a",
"metadata": {},
"source": [
"## 3. Using Hugging Face TRL (Transformer Reinforcement Learning) with Ray Train\n",
"\n",
"In comparison to supervised pre-training where models optimize to minimize the error for their next token, RL-based post-training aims to maximize their reward from a prompt. Therefore, it's crucial to define a reward function to measure the success and to train a model.\n",
"\n",
"This example uses the `trl.rewards.accuracy_reward` function to check whether the model's answer matches the answer in the dataset. As the answers use the LaTeX format, you must parse responses and solution before comparing them. The default `trl.rewards.accuracy_reward` implementation applies timeouts to the `parse` and `verify` functions, which are incompatible with Ray. The version defined below disables these timeouts by setting `parsing_timeout=0` and `timeout_seconds=0`."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "902e2cfaa92b68b1",
"metadata": {},
"outputs": [],
"source": [
"from latex2sympy2_extended import NormalizationConfig\n",
"from math_verify import LatexExtractionConfig, parse, verify\n",
"\n",
"def accuracy_reward(completions, solution, **kwargs):\n",
" \"\"\"Reward function that checks mathematical accuracy.\n",
"\n",
" This is a copy of `trl.rewards.accuracy_reward`.\n",
" The only difference is `parse(..., parsing_timeout=0)` and `verify(..., timeout_seconds=0)`\n",
" to avoid `signal.alarm()` issues with ray.\n",
" \"\"\"\n",
" contents = [completion[0][\"content\"] for completion in completions]\n",
" rewards = []\n",
" for content, sol in zip(contents, solution, strict=True):\n",
" gold_parsed = parse(sol, parsing_timeout=0)\n",
" if len(gold_parsed) != 0:\n",
" # We require the answer to be provided in correct latex (no malformed operators)\n",
" answer_parsed = parse(\n",
" content,\n",
" extraction_config=[\n",
" LatexExtractionConfig(\n",
" normalization_config=NormalizationConfig(units=True),\n",
" # Ensures that boxed is tried first\n",
" boxed_match_priority=0,\n",
" try_extract_without_anchor=False,\n",
" )\n",
" ],\n",
" extraction_mode=\"first_match\",\n",
" parsing_timeout=0,\n",
" )\n",
" reward = float(verify(gold_parsed, answer_parsed, timeout_seconds=0))\n",
" else:\n",
" # If the gold solution cannot be parsed, we assign `None` to skip this example\n",
" reward = None\n",
" rewards.append(reward)\n",
"\n",
" return rewards"
]
},
{
"cell_type": "markdown",
"id": "494cc96ea380247d",
"metadata": {},
"source": [
"Next, build the training function that's distributed across all the workers. The `GRPOTrainer` implements the GRPO algorithm for rolling out the model, computing the advantage and backpropagation the losses.\n",
"\n",
"With `vllm_mode=\"colocate\"`, each worker runs a vLLM instance that consumes about 30% of the GPU memory. The alternative `\"server\"` mode dedicates one worker to generation with vLLM while the others perform learning, but this introduces inter-GPU communication overhead that can reduce throughput. See [this blog post](https://huggingface.co/blog/vllm-colocate) for more details.\n",
"\n",
"Two Ray Train-specific additions integrate the HF Trainer into the Ray Train ecosystem:\n",
"\n",
"- **`RayTrainReportCallback`** — hooks into the HF Trainer's logging to forward metrics and checkpoints to Ray Train after each training step. This is what populates the `Result` object returned by `trainer.fit()` with metrics like `reward` and the best checkpoint.\n",
"- **`prepare_trainer`** — configures the HF Trainer for distributed execution across Ray workers. It disables HF's built-in distributed setup so that Ray Train controls the process group instead, and ensures correct device placement on each worker."
]
},
{
"cell_type": "code",
"id": "92c0cb05a65a9da6",
"metadata": {},
"source": [
"from ray.train.huggingface.transformers import RayTrainReportCallback, prepare_trainer\n",
"from trl import GRPOConfig, GRPOTrainer\n",
"from datasets import load_dataset\n",
"\n",
"def train_func(config):\n",
" # Load the DeepMath dataset with only 100 elements for this example.\n",
" # For larger datasets, use Ray Data.\n",
" dataset = load_dataset(\"trl-lib/DeepMath-103K\", split=\"train\").shuffle(seed=42).select(range(100))\n",
"\n",
" # Use vllm_mode=\"colocate\" to allow easy scaling as each GPU handles their own vLLM instance\n",
" training_args = GRPOConfig(\n",
" # Compute a training batch size of 4 for each prompt\n",
" per_device_train_batch_size=4,\n",
" # Use vLLM, colocated on each GPU which uses 30% of the vRAM.\n",
" use_vllm=True,\n",
" vllm_mode=\"colocate\",\n",
" vllm_gpu_memory_utilization=0.3,\n",
" # Run two training epochs over the whole dataset\n",
" num_train_epochs=2,\n",
" # Number of generations per prompt to sample, equivalent to G in the GRPO loss function.\n",
" num_generations=8,\n",
" # Save checkpoints and log metrics every epoch\n",
" save_strategy=\"epoch\",\n",
" logging_strategy=\"epoch\",\n",
" )\n",
"\n",
" # GRPO Trainer\n",
" trainer = GRPOTrainer(\n",
" model=\"Qwen/Qwen2.5-0.5B\",\n",
" args=training_args,\n",
" reward_funcs=accuracy_reward,\n",
" train_dataset=dataset,\n",
" )\n",
"\n",
" # Report metrics and checkpoints to Ray Train\n",
" trainer.add_callback(RayTrainReportCallback())\n",
"\n",
" # Prepare your trainer for TRL and Huggingface Transformer integration\n",
" trainer = prepare_trainer(trainer)\n",
"\n",
" # Start Training\n",
" trainer.train()"
],
"outputs": [],
"execution_count": null
},
{
"cell_type": "markdown",
"id": "5758d865e9d5b149",
"metadata": {},
"source": "With your `train_func` complete, you can now instantiate the `TorchTrainer`. Aside from calling the function, set the `scaling_config`, which controls the number of workers and resources used, and the `run_config` to configure checkpointing.\n\nThe `CheckpointConfig` controls how Ray Train saves and selects checkpoints during training:\n\n- **`num_to_keep=1`** — Ray Train retains only the single best checkpoint on disk, saving storage.\n- **`checkpoint_score_attribute=\"reward\"`** — Ray Train ranks checkpoints by the `\"reward\"` metric, which is the mean reward across the batch as reported by `RayTrainReportCallback`.\n- **`checkpoint_score_order=\"max\"`** — higher reward is better, so Ray Train keeps the checkpoint with the highest reward seen across all training steps."
},
{
"cell_type": "code",
"execution_count": null,
"id": "fff7fec396590de9",
"metadata": {},
"outputs": [],
"source": [
"from ray.train.torch import TorchTrainer\n",
"from ray.train import RunConfig, ScalingConfig, CheckpointConfig\n",
"\n",
"trainer = TorchTrainer(\n",
" train_func,\n",
" scaling_config=ScalingConfig(num_workers=num_workers, use_gpu=use_gpu),\n",
" run_config=RunConfig(\n",
" # On Anyscale this path is auto-mounted shared storage.\n",
" # For open-source Ray, use a shared NFS path or s3:// URI accessible from all nodes.\n",
" storage_path=\"/mnt/cluster_storage/\",\n",
" checkpoint_config=CheckpointConfig(\n",
" num_to_keep=1,\n",
" checkpoint_score_attribute=\"reward\",\n",
" checkpoint_score_order=\"max\",\n",
" ),\n",
" ),\n",
")"
]
},
{
"cell_type": "markdown",
"id": "b5b92f7546aef635",
"metadata": {},
"source": [
"Finally, call the `fit` method to start training with Ray Train. Save the `Result` object to a variable so you can access metrics and checkpoints."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "8ca549e43acc91b1",
"metadata": {},
"outputs": [],
"source": "results = trainer.fit()"
},
{
"cell_type": "markdown",
"id": "af3535209e06ad45",
"metadata": {},
"source": [
"You can use the returned `Result` object to access metrics and the Ray Train `Checkpoint` associated with the last iteration."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "fa7b49328202a7a5",
"metadata": {},
"outputs": [],
"source": "print(results.metrics_dataframe)"
},
{
"metadata": {},
"cell_type": "code",
"outputs": [],
"execution_count": null,
"source": "print(results.get_best_checkpoint(\"reward\", \"max\"))",
"id": "38431d92cf468f46"
},
{
"cell_type": "markdown",
"id": "ark4cvyvi2a",
"source": [
"## 4. Scaling to more GPUs or a larger model\n",
"\n",
"The preceding example trains with 4 workers on a small 100-sample subset of the dataset. To scale up, adjust `num_workers` in `ScalingConfig` and the `num_workers` variable at the top of the notebook, along with the dataset size and model.\n",
"\n",
"**Scaling workers:**\n",
"Each worker claims one GPU. With `vllm_mode=\"colocate\"`, each worker runs its own vLLM instance for generation and its own training process — generation requires no inter-GPU communication. Adding more workers increases the effective batch size and reduces time-to-convergence without any changes to training logic.\n",
"\n",
"**Scaling the model:**\n",
"Replace `\"Qwen/Qwen2.5-0.5B\"` with a larger checkpoint. Larger models require more GPU memory per worker. If a model doesn't fit on a single GPU, consider:\n",
"- Explore using DeepSpeed ZeRO to distribute the model weights across multiple GPUs.\n",
"- Reducing `vllm_gpu_memory_utilization` from `0.3` to leave more memory for model weights at the cost of a smaller vLLM KV cache.\n",
"\n",
"**Fault Tolerance:**\n",
"For longer runs where worker or node failures are possible, set `failure_config=FailureConfig(max_failures=1)` in `RunConfig`. Ray Train will automatically restart training up to twice on failure. Because this example already saves a checkpoint each epoch, a restarted run resumes from the last checkpoint rather than from scratch."
],
"metadata": {}
},
{
"cell_type": "markdown",
"id": "rnid76xr5u",
"metadata": {
"jp-MarkdownHeadingCollapsed": true
},
"source": [
"## Summary\n\nThis notebook demonstrated how to use Ray Train to scale RL post-training of a Qwen2.5 0.5B model across multiple GPUs using the TRL GRPO implementation and the DeepMath-103k dataset.\n\nThe key Ray Train integration points were:\n\n- **`RayTrainReportCallback`** — forwards HF Trainer metrics and checkpoints to Ray Train after each step.\n- **`prepare_trainer`** — configures the HF Trainer to run distributed training under Ray's process group.\n- **`TorchTrainer`** — launches `train_func` on all workers with the configured resources.\n- **`CheckpointConfig`** — automatically tracks and retains the best checkpoint based on the `\"reward\"` metric.\n\nFrom here, you can:\n- Scale to more GPUs or multiple nodes by increasing `num_workers` in `ScalingConfig`.\n- Swap in a different base model or dataset by updating the `model` argument and `load_dataset` call.\n- Experiment with different reward functions to train for other tasks."
]
}
],
"metadata": {
"kernelspec": {
"display_name": "Python 3 (ipykernel)",
"language": "python",
"name": "python3"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.10.11"
}
},
"nbformat": 4,
"nbformat_minor": 5
}
doc/source/train/examples/transformers/transformer_reinforcement_learning/README.md
+279
View File
@@ -0,0 +1,279 @@
# RL Post-Training using Hugging Face TRL with GRPO
This notebook builds on the TRL GRPO Example, [GRPO Trainer](https://huggingface.co/docs/trl/en/grpo_trainer), to fine-tune a Qwen2.5 0.5B model to answer math questions using the DeepMath-103k dataset and Group Relative Policy Optimization (GRPO) algorithm.
Ray Train scales this implementation efficiently across multiple GPUs without changing training logic.
This notebook consists of the following steps:
1. Quick Summary of GRPO and the DeepMath dataset
2. Package and Environment Setup
3. Running TRL with Ray Train
4. Scaling to more GPUs
<div id="anyscale-note" class="alert alert-block alert-warning">
<strong>Anyscale Specific Configuration</strong>
<p><strong>Note:</strong> This tutorial is optimized for the Anyscale platform. When running on open source Ray, you need additional configuration. For example, you would need to manually:</p>
<ul>
<li><strong>Configure your Ray Cluster</strong>: Set up your multi-node environment and manage resource allocation without Anyscale's automation.</li> <li><strong>Manage Dependencies</strong>: Manually install and manage dependencies on each node.</li> <li><strong>Set Up Storage</strong>: Configure your own distributed or shared storage system for model checkpointing.</li>
</ul>
<p>All these configurations are handled automatically through the Anyscale platform.
</div>
<style>
div#anyscale-note > p, div#anyscale-note > ul, div#anyscale-note > ul li { color: black; }
div#anyscale-note { background-color: rgb(255, 243, 205); }
div#anyscale-note { border: 1px solid #ccc; border-radius: 8px; padding: 15px; }
</style>
## 1. Summary of the DeepMath dataset and GRPO algorithm
### DeepMath dataset
This example uses the DeepMath-103k dataset by [He et al., 2025](https://arxiv.org/abs/2504.11456), consisting of 103,000 challenging questions across a wide range of mathematical subjects including Algebra, Calculus, Number Theory, Geometry, Probability, and Discrete Mathematics. These questions are more heavily distributed towards challenging questions, in particular Levels 5 to 9, compared to alternative datasets, making them highly complex and difficult to solve without specialist training.
Each example contains two fields:
- **`prompt`**: a list of chat-format message dicts; the question text is in `prompt[0]["content"]`.
- **`solution`**: a LaTeX-formatted ground-truth answer, used by the reward function to verify the model's response.
Using a structured LaTeX format ensures the reward function can reliably parse and compare both the model's output and the ground truth. Here are a few examples from the dataset:
| Subject | Prompt | Solution |
|---|---|---|
| Number Theory | Determine the number of 19th power residues modulo 229. | $12$ |
| Functional Analysis | Determine the norm $\lVert T \rVert$ of the linear operator $T: \ell^2 \rightarrow \ell^2$ given by $(Tx)_1 = 0$, $(Tx)_n = -x_n + \alpha x_{n+1}$ for $n \geq 2$, where $\alpha \in \mathbb{C}$. | $1 + \lvert\alpha\rvert$ |
| Analysis | Determine whether there exists a Schwartz function $g \in \mathcal{S}(\mathbb{R}^n)$ such that for a given continuous function $f: \mathbb{R}^n \to \mathbb{R}$ with $f \not\equiv 0$, the integral $\int_{\mathbb{R}^n} g^2 f \, dx \neq 0$. | Yes |
### Group Relative Policy Optimization Algorithm
To RL post-train the model to answer math questions, this example uses TRL's GRPO (Group Relative Policy Optimization) implementation introduced by [Shao et al., 2024](https://arxiv.org/abs/2402.03300). GRPO is an online algorithm that iteratively improves by training on data the model generates itself. Each iteration proceeds through four steps:
1. **Generate**: For each prompt in the batch, sample a group of G completions from the current policy.
2. **Score**: Apply the reward function to each completion, producing a scalar reward per completion.
3. **Advantage**: Normalize the rewards within each group by subtracting the group mean and dividing by the group standard deviation. This relative score becomes the advantage $\hat{A}_{i,t}$ where positive means better than peers.
4. **Loss**: Update the policy to increase the probability of high-advantage completions and decrease it for low-advantage ones, while a KL penalty keeps the updated policy close to the reference.
$$
\mathcal{L}_{\text{GRPO}}(\theta) = -\frac{1}{\sum_{i=1}^{G} |o_i|} \sum_{i=1}^{G} \sum_{t=1}^{|o_i|} \left[ \frac{\pi_\theta(o_{i,t} \mid q, o_{i,\mathopen{<} t})}{\left[\pi_\theta(o_{i,t} \mid q, o_{i,\mathopen{<} t})\right]_{\text{no grad}}} \hat{A}_{i,t} - \beta \mathbb{D}_{\text{KL}}\left[\pi_\theta \| \pi_{\text{ref}}\right] \right],
$$
For a deeper, more technical, description of the algorithm, see the [TRL GRPO page](https://huggingface.co/docs/trl/en/grpo_trainer#looking-deeper-into-the-grpo-method) and the [original paper](https://arxiv.org/abs/2402.03300).
## 2. Package and Environment Setup
Install the required dependencies
```python
# !pip install "trl[vllm]" "math_verify" "transformers==4.57.6"
```
Use `ray.init()` to initialize a local cluster. By default, this cluster contains only the machine you are running this notebook on. You can also run this notebook on an Anyscale cluster.
```python
import ray
ray.init()
```
Use `ray.cluster_resources()` to check which resources your cluster has access to. If you're running this notebook on your local machine or Google Colab, you should see the number of CPU cores and GPUs available to you.
```python
from pprint import pprint
pprint(ray.cluster_resources())
```
Change these two variables to control whether the training uses CPUs or GPUs, and how many workers to spawn. Each worker claims one CPU or GPU, so make sure not to request more resources than are available. By default, the training runs with four GPU worker.
```python
use_gpu = True # set this to `False` to run on CPUs
num_workers = 4 # set this to the number of GPUs or CPUs you want to use
```
## 3. Using Hugging Face TRL (Transformer Reinforcement Learning) with Ray Train
In comparison to supervised pre-training where models optimize to minimize the error for their next token, RL-based post-training aims to maximize their reward from a prompt. Therefore, it's crucial to define a reward function to measure the success and to train a model.
This example uses the `trl.rewards.accuracy_reward` function to check whether the model's answer matches the answer in the dataset. As the answers use the LaTeX format, you must parse responses and solution before comparing them. The default `trl.rewards.accuracy_reward` implementation applies timeouts to the `parse` and `verify` functions, which are incompatible with Ray. The version defined below disables these timeouts by setting `parsing_timeout=0` and `timeout_seconds=0`.
```python
from latex2sympy2_extended import NormalizationConfig
from math_verify import LatexExtractionConfig, parse, verify
def accuracy_reward(completions, solution, **kwargs):
"""Reward function that checks mathematical accuracy.
This is a copy of `trl.rewards.accuracy_reward`.
The only difference is `parse(..., parsing_timeout=0)` and `verify(..., timeout_seconds=0)`
to avoid `signal.alarm()` issues with ray.
"""
contents = [completion[0]["content"] for completion in completions]
rewards = []
for content, sol in zip(contents, solution, strict=True):
gold_parsed = parse(sol, parsing_timeout=0)
if len(gold_parsed) != 0:
# We require the answer to be provided in correct latex (no malformed operators)
answer_parsed = parse(
content,
extraction_config=[
LatexExtractionConfig(
normalization_config=NormalizationConfig(units=True),
# Ensures that boxed is tried first
boxed_match_priority=0,
try_extract_without_anchor=False,
)
],
extraction_mode="first_match",
parsing_timeout=0,
)
reward = float(verify(gold_parsed, answer_parsed, timeout_seconds=0))
else:
# If the gold solution cannot be parsed, we assign `None` to skip this example
reward = None
rewards.append(reward)
return rewards
```
Next, build the training function that's distributed across all the workers. The `GRPOTrainer` implements the GRPO algorithm for rolling out the model, computing the advantage and backpropagation the losses.
With `vllm_mode="colocate"`, each worker runs a vLLM instance that consumes about 30% of the GPU memory. The alternative `"server"` mode dedicates one worker to generation with vLLM while the others perform learning, but this introduces inter-GPU communication overhead that can reduce throughput. See [this blog post](https://huggingface.co/blog/vllm-colocate) for more details.
Two Ray Train-specific additions integrate the HF Trainer into the Ray Train ecosystem:
- **`RayTrainReportCallback`** — hooks into the HF Trainer's logging to forward metrics and checkpoints to Ray Train after each training step. This is what populates the `Result` object returned by `trainer.fit()` with metrics like `reward` and the best checkpoint.
- **`prepare_trainer`** — configures the HF Trainer for distributed execution across Ray workers. It disables HF's built-in distributed setup so that Ray Train controls the process group instead, and ensures correct device placement on each worker.
```python
from ray.train.huggingface.transformers import RayTrainReportCallback, prepare_trainer
from trl import GRPOConfig, GRPOTrainer
from datasets import load_dataset
def train_func(config):
# Load the DeepMath dataset with only 100 elements for this example.
# For larger datasets, use Ray Data.
dataset = load_dataset("trl-lib/DeepMath-103K", split="train").shuffle(seed=42).select(range(100))
# Use vllm_mode="colocate" to allow easy scaling as each GPU handles their own vLLM instance
training_args = GRPOConfig(
# Compute a training batch size of 4 for each prompt
per_device_train_batch_size=4,
# Use vLLM, colocated on each GPU which uses 30% of the vRAM.
use_vllm=True,
vllm_mode="colocate",
vllm_gpu_memory_utilization=0.3,
# Run two training epochs over the whole dataset
num_train_epochs=2,
# Number of generations per prompt to sample, equivalent to G in the GRPO loss function.
num_generations=8,
# Save checkpoints and log metrics every epoch
save_strategy="epoch",
logging_strategy="epoch",
)
# GRPO Trainer
trainer = GRPOTrainer(
model="Qwen/Qwen2.5-0.5B",
args=training_args,
reward_funcs=accuracy_reward,
train_dataset=dataset,
)
# Report metrics and checkpoints to Ray Train
trainer.add_callback(RayTrainReportCallback())
# Prepare your trainer for TRL and Huggingface Transformer integration
trainer = prepare_trainer(trainer)
# Start Training
trainer.train()
```
With your `train_func` complete, you can now instantiate the `TorchTrainer`. Aside from calling the function, set the `scaling_config`, which controls the number of workers and resources used, and the `run_config` to configure checkpointing.
The `CheckpointConfig` controls how Ray Train saves and selects checkpoints during training:
- **`num_to_keep=1`** — Ray Train retains only the single best checkpoint on disk, saving storage.
- **`checkpoint_score_attribute="reward"`** — Ray Train ranks checkpoints by the `"reward"` metric, which is the mean reward across the batch as reported by `RayTrainReportCallback`.
- **`checkpoint_score_order="max"`** — higher reward is better, so Ray Train keeps the checkpoint with the highest reward seen across all training steps.
```python
from ray.train.torch import TorchTrainer
from ray.train import RunConfig, ScalingConfig, CheckpointConfig
trainer = TorchTrainer(
train_func,
scaling_config=ScalingConfig(num_workers=num_workers, use_gpu=use_gpu),
run_config=RunConfig(
# On Anyscale this path is auto-mounted shared storage.
# For open-source Ray, use a shared NFS path or s3:// URI accessible from all nodes.
storage_path="/mnt/cluster_storage/",
checkpoint_config=CheckpointConfig(
num_to_keep=1,
checkpoint_score_attribute="reward",
checkpoint_score_order="max",
),
),
)
```
Finally, call the `fit` method to start training with Ray Train. Save the `Result` object to a variable so you can access metrics and checkpoints.
```python
results = trainer.fit()
```
You can use the returned `Result` object to access metrics and the Ray Train `Checkpoint` associated with the last iteration.
```python
print(results.metrics_dataframe)
```
```python
print(results.get_best_checkpoint("reward", "max"))
```
## 4. Scaling to more GPUs or a larger model
The preceding example trains with 4 workers on a small 100-sample subset of the dataset. To scale up, adjust `num_workers` in `ScalingConfig` and the `num_workers` variable at the top of the notebook, along with the dataset size and model.
**Scaling workers:** Each worker claims one GPU. With `vllm_mode="colocate"`, each worker runs its own vLLM instance for generation and its own training process — generation requires no inter-GPU communication. Adding more workers increases the effective batch size and reduces time-to-convergence without any changes to training logic.
**Scaling the model:** Replace `"Qwen/Qwen2.5-0.5B"` with a larger checkpoint. Larger models require more GPU memory per worker. If a model doesn't fit on a single GPU, consider:
- Explore using DeepSpeed ZeRO to distribute the model weights across multiple GPUs.
- Reducing `vllm_gpu_memory_utilization` from `0.3` to leave more memory for model weights at the cost of a smaller vLLM KV cache.
**Fault Tolerance:** For longer runs where worker or node failures are possible, set `failure_config=FailureConfig(max_failures=1)` in `RunConfig`. Ray Train will automatically restart training up to twice on failure. Because this example already saves a checkpoint each epoch, a restarted run resumes from the last checkpoint rather than from scratch.
## Summary
This notebook demonstrated how to use Ray Train to scale RL post-training of a Qwen2.5 0.5B model across multiple GPUs using the TRL GRPO implementation and the DeepMath-103k dataset.
The key Ray Train integration points were:
- **`RayTrainReportCallback`** — forwards HF Trainer metrics and checkpoints to Ray Train after each step.
- **`prepare_trainer`** — configures the HF Trainer to run distributed training under Ray's process group.
- **`TorchTrainer`** — launches `train_func` on all workers with the configured resources.
- **`CheckpointConfig`** — automatically tracks and retains the best checkpoint based on the `"reward"` metric.
From here, you can:
- Scale to more GPUs or multiple nodes by increasing `num_workers` in `ScalingConfig`.
- Swap in a different base model or dataset by updating the `model` argument and `load_dataset` call.
- Experiment with different reward functions to train for other tasks.
doc/source/train/examples/transformers/transformer_reinforcement_learning/ci/aws.yaml
+12
View File
@@ -0,0 +1,12 @@
cloud_id: {{env["ANYSCALE_CLOUD_ID"]}}
region: us-west-2
head_node_type:
name: head_node
instance_type: m5.2xlarge
worker_node_types:
- instance_type: g5.12xlarge # 4x A10G GPUs
name: '4xA10G:48CPU-192GB'
min_workers: 1
max_workers: 1
doc/source/train/examples/transformers/transformer_reinforcement_learning/ci/gce.yaml
+13
View File
@@ -0,0 +1,13 @@
cloud_id: {{env["ANYSCALE_CLOUD_ID"]}}
region: us-central1
head_node_type:
name: head
instance_type: n2-standard-8
worker_node_types:
- name: gpu_worker
instance_type: g2-standard-48-nvidia-l4-4 # 4x L4 GPUs
min_workers: 1
max_workers: 1
use_spot: false
doc/source/train/examples/transformers/transformer_reinforcement_learning/ci/nb2py.py
+72
View File
@@ -0,0 +1,72 @@
#!/usr/bin/env python3
import argparse
import nbformat
def convert_notebook(input_path: str, output_path: str) -> None:
"""
Read a Jupyter notebook and write a Python script, converting all %%bash
cells and IPython "!" commands into subprocess.run calls that raise on error.
Cells that load or autoreload extensions are ignored.
"""
nb = nbformat.read(input_path, as_version=4)
with open(output_path, "w") as out:
for cell in nb.cells:
# Only process code cells
if cell.cell_type != "code":
continue
lines = cell.source.splitlines()
# Skip cells that load or autoreload extensions
if any(
l.strip().startswith("%load_ext autoreload")
or l.strip().startswith("%autoreload all")
for l in lines
):
continue
# Detect a %%bash cell
if lines and lines[0].strip().startswith("%%bash"):
bash_script = "\n".join(lines[1:]).rstrip()
out.write("import subprocess\n")
out.write(
f"subprocess.run(r'''{bash_script}''',\n"
" shell=True,\n"
" check=True,\n"
" executable='/bin/bash')\n\n"
)
else:
# Detect any IPython '!' shell commands in code lines
has_bang = any(line.lstrip().startswith("!") for line in lines)
if has_bang:
out.write("import subprocess\n")
for line in lines:
stripped = line.lstrip()
if stripped.startswith("!"):
cmd = stripped[1:].lstrip()
out.write(
f"subprocess.run(r'''{cmd}''',\n"
" shell=True,\n"
" check=True,\n"
" executable='/bin/bash')\n"
)
else:
out.write(line.rstrip() + "\n")
out.write("\n")
else:
# Regular Python cell: dump as-is
out.write(cell.source.rstrip() + "\n\n")
def main() -> None:
parser = argparse.ArgumentParser(
description="Convert a Jupyter notebook to a Python script, preserving bash cells and '!' commands as subprocess calls."
)
parser.add_argument("input_nb", help="Path to the input .ipynb file")
parser.add_argument("output_py", help="Path for the output .py script")
args = parser.parse_args()
convert_notebook(args.input_nb, args.output_py)
if __name__ == "__main__":
main()
doc/source/train/examples/transformers/transformer_reinforcement_learning/ci/tests.sh
+5
View File
@@ -0,0 +1,5 @@
#!/bin/bash
set -euo pipefail
python ci/nb2py.py README.ipynb README.py # convert notebook to py script
python README.py # run the converted python script
rm README.py # remove the generated script
doc/source/train/examples/transformers/transformer_reinforcement_learning/configs/aws.yaml
+9
View File
@@ -0,0 +1,9 @@
head_node_type:
name: head_node
instance_type: m5.2xlarge
worker_node_types:
- instance_type: g5.12xlarge # 4x A10G GPUs
name: '4xA10G:48CPU-192GB'
min_workers: 1
max_workers: 1
doc/source/train/examples/transformers/transformer_reinforcement_learning/configs/gce.yaml
+10
View File
@@ -0,0 +1,10 @@
head_node_type:
name: head
instance_type: n2-standard-8
worker_node_types:
- name: gpu_worker
instance_type: g2-standard-48-nvidia-l4-4 # 4x L4 GPUs
min_workers: 1
max_workers: 1
use_spot: false
doc/source/train/examples/transformers/transformers_torch_trainer_basic.rst
+28
View File
@@ -0,0 +1,28 @@
:orphan:
.. _transformers_torch_trainer_basic_example:
Fine-tune a Text Classifier with Hugging Face Transformers
==========================================================
.. raw:: html
<a id="try-anyscale-quickstart-transformers_torch_trainer_basic" target="_blank" href="https://console.anyscale.com/register/ha?render_flow=ray&utm_source=ray_docs&utm_medium=docs&utm_campaign=transformers_torch_trainer_basic">
<img src="../../../_static/img/run-on-anyscale.svg" alt="Run on Anyscale" />
<br/><br/>
</a>
This basic example of distributed training with Ray Train and Hugging Face (HF) Transformers
fine-tunes a text classifier on the Yelp review dataset using HF Transformers and Ray Train.
Code example
------------
.. literalinclude:: /../../python/ray/train/examples/transformers/transformers_torch_trainer_basic.py
See also
--------
* :ref:`Get Started with Hugging Face Transformers <train-pytorch-transformers>` for a tutorial
* :doc:`Ray Train Examples <../../examples>` for more use cases