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/BUILD.bazel
+13
View File
@@ -0,0 +1,13 @@
filegroup(
name = "train_examples_ci_configs",
srcs = glob([
"**/ci/aws.yaml",
"**/ci/gce.yaml",
]) + [
"//doc/source/train/examples/pytorch:pytorch_examples_ci_configs",
"//doc/source/train/examples/transformers:transformers_examples_ci_configs",
"//doc/source/train/examples/lightning:lightning_examples_ci_configs",
],
visibility = ["//doc:__pkg__"],
)
doc/source/train/examples/accelerate/accelerate_example.rst
+29
View File
@@ -0,0 +1,29 @@
:orphan:
Distributed Training with Hugging Face Accelerate
=================================================
.. raw:: html
<a id="try-anyscale-quickstart-accelerate_example" target="_blank" href="https://console.anyscale.com/register/ha?render_flow=ray&utm_source=ray_docs&utm_medium=docs&utm_campaign=accelerate_example">
<img src="../../../_static/img/run-on-anyscale.svg" alt="Run on Anyscale" />
<br/><br/>
</a>
This example does distributed data parallel training
with Hugging Face Accelerate, Ray Train, and Ray Data.
It fine-tunes a BERT model and is adapted from
https://github.com/huggingface/accelerate/blob/main/examples/nlp_example.py
Code example
------------
.. literalinclude:: /../../python/ray/train/examples/accelerate/accelerate_torch_trainer.py
See also
--------
* :ref:`Get Started with Hugging Face Accelerate <train-hf-accelerate>` for a tutorial on using Ray Train and HF Accelerate
* :doc:`Ray Train Examples <../../examples>` for more use cases
doc/source/train/examples/aws-trainium/llama3.rst
+109
View File
@@ -0,0 +1,109 @@
:orphan:
Distributed fine-tuning of Llama 3.1 8B on AWS Trainium with Ray and PyTorch Lightning
======================================================================================
.. raw:: html
<a id="try-anyscale-quickstart-aws-trainium-llama3" target="_blank" href="https://console.anyscale.com/register/ha?render_flow=ray&utm_source=ray_docs&utm_medium=docs&utm_campaign=aws-trainium-llama3">
<img src="../../../_static/img/run-on-anyscale.svg" alt="Run on Anyscale" />
<br/><br/>
</a>
This example demonstrates how to fine-tune the `Llama 3.1 8B <https://huggingface.co/NousResearch/Meta-Llama-3.1-8B/>`__ model on `AWS
Trainium <https://aws.amazon.com/ai/machine-learning/trainium/>`__ instances using Ray Train, PyTorch Lightning, and AWS Neuron SDK.
AWS Trainium is the machine learning (ML) chip that AWS built for deep
learning (DL) training of 100B+ parameter models. `AWS Neuron
SDK <https://aws.amazon.com/machine-learning/neuron/>`__ helps
developers train models on Trainium accelerators.
Prepare the environment
-----------------------
See `Setup EKS cluster and tools <https://github.com/aws-neuron/aws-neuron-eks-samples/tree/master/llama3.1_8B_finetune_ray_ptl_neuron#setupeksclusterandtools>`__ for setting up an Amazon EKS cluster leveraging AWS Trainium instances.
Create a Docker image
---------------------
When the EKS cluster is ready, create an Amazon ECR repository for building and uploading the Docker image containing artifacts for fine-tuning a Llama3.1 8B model:
1. Clone the repo.
::
git clone https://github.com/aws-neuron/aws-neuron-eks-samples.git
2. Go to the ``llama3.1_8B_finetune_ray_ptl_neuron`` directory.
::
cd aws-neuron-eks-samples/llama3.1_8B_finetune_ray_ptl_neuron
3. Trigger the script.
::
chmod +x 0-kuberay-trn1-llama3-finetune-build-image.sh
./0-kuberay-trn1-llama3-finetune-build-image.sh
4. Enter the zone your cluster is running in, for example: us-east-2.
5. Verify in the AWS console that the Amazon ECR service has the newly
created ``kuberay_trn1_llama3.1_pytorch2`` repository.
6. Update the ECR image ARN in the manifest file used for creating the Ray cluster.
Replace the <AWS_ACCOUNT_ID> and <REGION> placeholders with actual values in the ``1-llama3-finetune-trn1-create-raycluster.yaml`` file using commands below to reflect the ECR image ARN created above:
::
export AWS_ACCOUNT_ID=<enter_your_aws_account_id> # for ex: 111222333444
export REGION=<enter_your_aws_region> # for ex: us-east-2
sed -i "s/<AWS_ACCOUNT_ID>/$AWS_ACCOUNT_ID/g" 1-llama3-finetune-trn1-create-raycluster.yaml
sed -i "s/<REGION>/$REGION/g" 1-llama3-finetune-trn1-create-raycluster.yaml
Configuring Ray Cluster
-----------------------
The ``llama3.1_8B_finetune_ray_ptl_neuron`` directory in the AWS Neuron samples repository simplifies the
Ray configuration. KubeRay provides a manifest that you can apply
to the cluster to set up the head and worker pods.
Run the following command to set up the Ray cluster:
::
kubectl apply -f 1-llama3-finetune-trn1-create-raycluster.yaml
Accessing Ray Dashboard
-----------------------
Port forward from the cluster to see the state of the Ray dashboard and
then view it on `http://localhost:8265 <http://localhost:8265/>`__.
Run it in the background with the following command:
::
kubectl port-forward service/kuberay-trn1-head-svc 8265:8265 &
Launching Ray Jobs
------------------
The Ray cluster is now ready to handle workloads. Initiate the data preparation and fine-tuning Ray jobs:
1. Launch the Ray job for downloading the dolly-15k dataset and the Llama3.1 8B model artifacts:
::
kubectl apply -f 2-llama3-finetune-trn1-rayjob-create-data.yaml
2. When the job has executed successfully, run the following fine-tuning job:
::
kubectl apply -f 3-llama3-finetune-trn1-rayjob-submit-finetuning-job.yaml
3. Monitor the jobs via the Ray Dashboard
For detailed information on each of the steps above, see the `AWS documentation link <https://github.com/aws-neuron/aws-neuron-eks-samples/blob/master/llama3.1_8B_finetune_ray_ptl_neuron/README.md/>`__.
doc/source/train/examples/deepspeed/deepspeed_example.rst
+29
View File
@@ -0,0 +1,29 @@
:orphan:
Train with DeepSpeed ZeRO-3 and Ray Train
=========================================
.. raw:: html
<a id="try-anyscale-quickstart-deepspeed_example" target="_blank" href="https://console.anyscale.com/register/ha?render_flow=ray&utm_source=ray_docs&utm_medium=docs&utm_campaign=deepspeed_example">
<img src="../../../_static/img/run-on-anyscale.svg" alt="Run on Anyscale" />
<br/><br/>
</a>
This is an intermediate example that shows how to do distributed training with DeepSpeed ZeRO-3 and Ray Train.
It demonstrates how to use :ref:`Ray Data <data>` with DeepSpeed ZeRO-3 and Ray Train.
If you just want to quickly convert your existing TorchTrainer scripts into Ray Train, you can refer to the :ref:`Train with DeepSpeed <train-deepspeed>`.
Code example
------------
.. literalinclude:: /../../python/ray/train/examples/deepspeed/deepspeed_torch_trainer.py
See also
--------
* :doc:`Ray Train Examples <../../examples>` for more use cases.
* :ref:`Get Started with DeepSpeed <train-deepspeed>` for a tutorial.
doc/source/train/examples/deepspeed/gptj_deepspeed_fine_tuning.ipynb
+761
View File
@@ -0,0 +1,761 @@
{
"cells": [
{
"cell_type": "markdown",
"metadata": {},
"source": [
"# GPT-J-6B Fine-Tuning with Ray Train and DeepSpeed\n",
"\n",
"<a id=\"try-anyscale-quickstart-gptj_deepspeed_fine_tuning\" href=\"https://console.anyscale.com/register/ha?render_flow=ray&utm_source=ray_docs&utm_medium=docs&utm_campaign=gptj_deepspeed_fine_tuning\">\n",
" <img src=\"../../../_static/img/run-on-anyscale.svg\" alt=\"try-anyscale-quickstart\">\n",
"</a>\n",
"<br></br>\n",
"\n",
"This example showcases how to use Ray Train for **GPT-J fine-tuning**. GPT-J is a GPT-2-like causal language model trained on the Pile dataset. This particular model has 6 billion parameters. For more information, see [GPT-J](https://huggingface.co/docs/transformers/model_doc/gptj).\n",
"\n",
"This example uses the Ray Train \ud83e\udd17 Transformers integration and a pre-trained model from the Hugging Face Hub. Note that this example is adaptable to other similar models.\n",
"\n",
"This is an advanced example that focuses on the performance and distributed computing aspects of Ray Train. For a beginner-friendly introduction to the Ray Train \ud83e\udd17 Transformers integration, see {ref}`Basic Example for HuggingFace Transformers <transformers_torch_trainer_basic_example>`.\n",
"\n",
"Read [Ray Train Key Concepts](train-key-concepts) and [Ray Data Integration User Guides](data-ingest-torch) before starting this example.\n",
"\n",
"```{note}\n",
"To run this example, make sure your Ray cluster has access to at least one GPU with 16 or more GBs of memory. The required amount of memory depends on the model. This notebook is tested with 16 g4dn.4xlarge instances (including the head node).\n",
"```\n",
"\n",
"This notebook has the following steps:\n",
"1. [Set up Ray](#gptj-setup)\n",
"2. [Load the dataset](#gptj-load)\n",
"3. [Preprocess the dataset with Ray Data](#gptj-preprocess)\n",
"4. [Run the training with Ray Train](#gptj-train)\n",
"5. [Generate text from prompt](#gptj-predict)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Uncomment and run the following line in order to install all the necessary dependencies (this notebook was tested with `accelerate==0.33.0`, `transformers==4.36.2`, `deepspeed==0.17.2`):"
]
},
{
"cell_type": "code",
"execution_count": 1,
"metadata": {},
"outputs": [],
"source": [
"! pip install -q \"datasets\" \"evaluate\" \"accelerate==0.33.0\" \"transformers==4.36.2\" \"deepspeed==0.17.2\""
]
},
{
"cell_type": "code",
"execution_count": 1,
"metadata": {
"tags": []
},
"outputs": [],
"source": [
"import numpy as np\n",
"import pandas as pd\n",
"import os"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"(gptj-setup)=\n",
"## Set up Ray\n",
"\n",
"First, let's set some global variables. We will use 16 workers, each being assigned 1 GPU and 8 CPUs."
]
},
{
"cell_type": "code",
"execution_count": 3,
"metadata": {
"tags": []
},
"outputs": [],
"source": [
"model_name = \"EleutherAI/gpt-j-6B\"\n",
"use_gpu = True\n",
"num_workers = 16\n",
"cpus_per_worker = 8"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"We will use `ray.init()` to initialize a local cluster. By default, this cluster will be comprised of only the machine you are running this notebook on. You can also run this notebook on an Anyscale cluster.\n",
"\n",
"We define a {ref}`runtime environment <runtime-environments>` to ensure that the Ray workers have access to all the necessary packages. You can omit the `runtime_env` argument if you have all of the packages already installed on each node in your cluster."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"tags": []
},
"outputs": [],
"source": [
"import ray\n",
"\n",
"ray.init(\n",
" runtime_env={\n",
" \"pip\": [\n",
" \"datasets\",\n",
" \"evaluate\",\n",
" \"accelerate==0.33.0\",\n",
" \"transformers==4.36.2\",\n",
" \"deepspeed==0.17.2\",\n",
" ],\n",
" },\n",
")"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"tags": [
"hide-cell"
]
},
"outputs": [],
"source": [
"# THIS SHOULD BE HIDDEN IN DOCS AND ONLY RAN IN CI\n",
"# Download the model from our S3 mirror as it's faster\n",
"\n",
"import ray\n",
"import subprocess\n",
"import ray.util.scheduling_strategies\n",
"\n",
"\n",
"def force_on_node(node_id: str, remote_func_or_actor_class):\n",
" scheduling_strategy = ray.util.scheduling_strategies.NodeAffinitySchedulingStrategy(\n",
" node_id=node_id, soft=False\n",
" )\n",
" options = {\"scheduling_strategy\": scheduling_strategy}\n",
" return remote_func_or_actor_class.options(**options)\n",
"\n",
"\n",
"def run_on_every_node(remote_func_or_actor_class, **remote_kwargs):\n",
" refs = []\n",
" for node in ray.nodes():\n",
" if node[\"Alive\"] and node[\"Resources\"].get(\"GPU\", None):\n",
" refs.append(\n",
" force_on_node(node[\"NodeID\"], remote_func_or_actor_class).remote(\n",
" **remote_kwargs\n",
" )\n",
" )\n",
" return ray.get(refs)\n",
"\n",
"\n",
"@ray.remote(num_gpus=1)\n",
"def download_model():\n",
" from transformers.utils.hub import TRANSFORMERS_CACHE\n",
"\n",
" path = os.path.expanduser(\n",
" os.path.join(TRANSFORMERS_CACHE, \"models--EleutherAI--gpt-j-6B\")\n",
" )\n",
" subprocess.run([\"mkdir\", \"-p\", os.path.join(path, \"snapshots\", \"main\")])\n",
" subprocess.run([\"mkdir\", \"-p\", os.path.join(path, \"refs\")])\n",
" if os.path.exists(os.path.join(path, \"refs\", \"main\")):\n",
" return\n",
" subprocess.run(\n",
" [\n",
" \"aws\",\n",
" \"s3\",\n",
" \"sync\",\n",
" \"--no-sign-request\",\n",
" \"s3://large-dl-models-mirror/models--EleutherAI--gpt-j-6B/main/\",\n",
" os.path.join(path, \"snapshots\", \"main\"),\n",
" ]\n",
" )\n",
" with open(os.path.join(path, \"snapshots\", \"main\", \"hash\"), \"r\") as f:\n",
" f_hash = f.read().strip()\n",
" with open(os.path.join(path, \"refs\", \"main\"), \"w\") as f:\n",
" f.write(f_hash)\n",
" os.rename(\n",
" os.path.join(path, \"snapshots\", \"main\"), os.path.join(path, \"snapshots\", f_hash)\n",
" )\n",
"\n",
"\n",
"_ = run_on_every_node(download_model)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"(gptj-load)=\n",
"## Loading the dataset\n",
"\n",
"We will be fine-tuning the model on the [`tiny_shakespeare` dataset](https://huggingface.co/datasets/tiny_shakespeare), comprised of 40,000 lines of Shakespeare from a variety of Shakespeare's plays. The aim will be to make the GPT-J model better at generating text in the style of Shakespeare."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"tags": []
},
"outputs": [],
"source": [
"from datasets import load_dataset\n",
"\n",
"print(\"Loading tiny_shakespeare dataset\")\n",
"current_dataset = load_dataset(\"tiny_shakespeare\", trust_remote_code=True)\n",
"current_dataset"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"We will use [Ray Data](https://docs.ray.io/en/latest/data/data.html) for distributed preprocessing and data ingestion. We can easily convert the dataset obtained from Hugging Face Hub to Ray Data by using `ray.data.from_huggingface`."
]
},
{
"cell_type": "code",
"execution_count": 6,
"metadata": {
"tags": []
},
"outputs": [
{
"data": {
"text/plain": [
"{'train': MaterializedDataset(num_blocks=1, num_rows=1, schema={text: string}),\n",
" 'validation': MaterializedDataset(num_blocks=1, num_rows=1, schema={text: string})}"
]
},
"execution_count": 6,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"import ray.data\n",
"\n",
"ray_datasets = {\n",
" \"train\": ray.data.from_huggingface(current_dataset[\"train\"]),\n",
" \"validation\": ray.data.from_huggingface(current_dataset[\"validation\"]),\n",
"}\n",
"\n",
"ray_datasets"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"(gptj-preprocess)=\n",
"Note that the dataset is represented by a single line of large string, and needs some preprocessing. To do this, use the {meth}`~ray.data.Dataset.map_batches` API to apply transformation functions to batches of data.\n",
"\n",
"The `split_text` function takes the single string and splits it into separate lines, removing empty lines and character names ending with ':' (eg. 'ROMEO:'). The `tokenize` function takes the lines and tokenizes them using the \ud83e\udd17 Tokenizer associated with the model, ensuring each entry has the same length (`block_size`) by padding and truncating. This preprocessing is necessary for training.\n",
"\n",
"```{note}\n",
"This preprocessing can be done in other ways. A common pattern is to tokenize first, and then split the obtained tokens into equally-sized blocks.\n",
"```"
]
},
{
"cell_type": "code",
"execution_count": 7,
"metadata": {
"tags": []
},
"outputs": [],
"source": [
"block_size = 512"
]
},
{
"cell_type": "code",
"execution_count": 8,
"metadata": {
"tags": []
},
"outputs": [
{
"data": {
"text/plain": [
"{'train': MapBatches(tokenize)\n",
" +- MapBatches(split_text)\n",
" +- Dataset(num_blocks=1, num_rows=1, schema={text: string}),\n",
" 'validation': MapBatches(tokenize)\n",
" +- MapBatches(split_text)\n",
" +- Dataset(num_blocks=1, num_rows=1, schema={text: string})}"
]
},
"execution_count": 8,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"from transformers import AutoTokenizer\n",
"\n",
"\n",
"def split_text(batch: pd.DataFrame) -> pd.DataFrame:\n",
" text = list(batch[\"text\"])\n",
" flat_text = \"\".join(text)\n",
" split_text = [\n",
" x.strip()\n",
" for x in flat_text.split(\"\\n\")\n",
" if x.strip() and not x.strip()[-1] == \":\"\n",
" ]\n",
" return pd.DataFrame(split_text, columns=[\"text\"])\n",
"\n",
"\n",
"def tokenize(batch: pd.DataFrame) -> dict:\n",
" tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=False)\n",
" tokenizer.pad_token = tokenizer.eos_token\n",
" ret = tokenizer(\n",
" list(batch[\"text\"]),\n",
" truncation=True,\n",
" max_length=block_size,\n",
" padding=\"max_length\",\n",
" return_tensors=\"np\",\n",
" )\n",
" ret[\"labels\"] = ret[\"input_ids\"].copy()\n",
" return dict(ret)\n",
"\n",
"\n",
"processed_datasets = {\n",
" key: (\n",
" ds.map_batches(split_text, batch_format=\"pandas\")\n",
" .map_batches(tokenize, batch_format=\"pandas\")\n",
" )\n",
" for key, ds in ray_datasets.items()\n",
"}\n",
"processed_datasets"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"(gptj-train)=\n",
"### Fine-tuning the model with Ray Train\n",
"\n",
"Configure Ray Train's {class}`~ray.train.torch.TorchTrainer` to perform distributed fine-tuning of the model. Specify a `train_loop_per_worker` function, which defines the training logic to be distributed by Ray using Distributed Data Parallelism, which uses the PyTorch Distributed backend internally. Each worker has its own copy of the model, but operates on different data. At the end of each step, all the workers sync gradients.\n",
"\n",
"Because GPT-J is a relatively large model, it may not be possible to fit it on smaller GPU types (<=16 GB GRAM). To deal with that issue, this example uses [DeepSpeed](https://github.com/microsoft/DeepSpeed), a library to optimize the training process and to offload and partition optimizer and parameter states, reducing GRAM usage. Furthermore, DeepSpeed ZeRO Stage 3 can load large models without running out of memory.\n",
"\n",
"\ud83e\udd17 Transformers and Ray Train's {ref}`integrations <train-transformers-integration>` allow you to easily configure and use DDP and DeepSpeed. All you need to do is specify the DeepSpeed configuration in the [`TrainingArguments`](https://huggingface.co/docs/transformers/en/main_classes/trainer#transformers.TrainingArguments) object.\n",
"\n",
"```{tip}\n",
"There are many DeepSpeed settings that allow you to trade-off speed for memory usage. The settings used below are tailored to the cluster setup used (16 g4dn.4xlarge nodes) and per device batch size of 16. Some things to keep in mind:\n",
"- If your GPUs support bfloat16, use that instead of float16 mixed precision to get better performance and prevent overflows. Replace `fp16=True` with `bf16=True` in `TrainingArguments`.\n",
"- If you are running out of GRAM: try reducing batch size (defined in the cell below the next one), set `\"overlap_comm\": False` in DeepSpeed config.\n",
"- If you are running out of RAM, add more nodes to your cluster, use nodes with more RAM, set `\"pin_memory\": False` in the DeepSpeed config, reduce the batch size, and remove `\"offload_param\"` from the DeepSpeed config.\n",
"\n",
"For more information on DeepSpeed configuration, refer to [Hugging Face documentation](https://huggingface.co/docs/transformers/main_classes/deepspeed) and [DeepSpeed documentation](https://www.deepspeed.ai/docs/config-json/).\n",
"\n",
"Additionally, if you prefer a lower-level API, the logic below can be expressed as an [Accelerate training loop](https://github.com/huggingface/accelerate/blob/main/examples/by_feature/deepspeed_with_config_support.py) distributed by a Ray Train {class}`~ray.train.torch.torch_trainer.TorchTrainer`.\n",
"```\n",
"\n",
"#### Training speed\n",
"\n",
"As this example uses data parallelism, each worker operates on its own shard of the data. The batch size set in `train_ds.iter_torch_batches` is the **per device batch size** (per worker batch size). By changing the number of workers, you can change the **effective batch size** and thus the time needed for training to complete. Calculate the effective batch size as `per device batch size * number of workers * number of gradient accumulation steps`. As you add more workers, the effective batch size rises and thus less time is needed to complete a full epoch. While the speedup is not exactly linear due to extra communication overheads, in many cases it can be close to linear.\n",
"\n",
"The preprocessed dataset has 1348 examples. We have set per device batch size to 16.\n",
"\n",
"* With 16 g4dn.4xlarge nodes, the effective batch size was 256, which equals to 85 steps per epoch. One epoch took **~2440 seconds** (including initialization time).\n",
"\n",
"* With 32 g4dn.4xlarge nodes, the effective batch size was 512, which equals to 43 steps per epoch. One epoch took **~1280 seconds** (including initialization time)."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"tags": []
},
"outputs": [],
"source": [
"import evaluate\n",
"import torch\n",
"from transformers import (\n",
" Trainer,\n",
" TrainingArguments,\n",
" GPTJForCausalLM,\n",
" AutoTokenizer,\n",
" default_data_collator,\n",
")\n",
"from transformers.utils.logging import disable_progress_bar, enable_progress_bar\n",
"\n",
"from ray import train\n",
"from ray.train.huggingface.transformers import prepare_trainer, RayTrainReportCallback\n",
"\n",
"\n",
"def train_func(config):\n",
" # Use the actual number of CPUs assigned to this worker by Ray\n",
" runtime_ctx = ray.get_runtime_context()\n",
" assigned_cpus = runtime_ctx.get_assigned_resources().get(\"CPU\", 1)\n",
" os.environ[\"OMP_NUM_THREADS\"] = str(int(assigned_cpus))\n",
" # Enable tf32 for better performance\n",
" torch.backends.cuda.matmul.allow_tf32 = True\n",
"\n",
" batch_size = config.get(\"batch_size\", 4)\n",
" epochs = config.get(\"epochs\", 2)\n",
" warmup_steps = config.get(\"warmup_steps\", 0)\n",
" learning_rate = config.get(\"learning_rate\", 0.00002)\n",
" weight_decay = config.get(\"weight_decay\", 0.01)\n",
" steps_per_epoch = config.get(\"steps_per_epoch\")\n",
"\n",
" deepspeed = {\n",
" \"fp16\": {\n",
" \"enabled\": \"auto\",\n",
" \"initial_scale_power\": 8,\n",
" \"hysteresis\": 4,\n",
" \"consecutive_hysteresis\": True,\n",
" },\n",
" \"bf16\": {\"enabled\": \"auto\"},\n",
" \"optimizer\": {\n",
" \"type\": \"AdamW\",\n",
" \"params\": {\n",
" \"lr\": \"auto\",\n",
" \"betas\": \"auto\",\n",
" \"eps\": \"auto\",\n",
" },\n",
" },\n",
" \"zero_optimization\": {\n",
" \"stage\": 3,\n",
" \"offload_optimizer\": {\n",
" \"device\": \"cpu\",\n",
" \"pin_memory\": True,\n",
" },\n",
" \"overlap_comm\": True,\n",
" \"contiguous_gradients\": True,\n",
" \"gather_16bit_weights_on_model_save\": True,\n",
" \"round_robin_gradients\": True,\n",
" },\n",
" \"gradient_accumulation_steps\": \"auto\",\n",
" \"gradient_clipping\": \"auto\",\n",
" \"steps_per_print\": 10,\n",
" \"train_batch_size\": \"auto\",\n",
" \"train_micro_batch_size_per_gpu\": \"auto\",\n",
" \"wall_clock_breakdown\": False,\n",
" }\n",
"\n",
" print(\"Preparing training arguments\")\n",
" training_args = TrainingArguments(\n",
" \"output\",\n",
" logging_steps=1,\n",
" save_strategy=\"steps\",\n",
" save_steps=steps_per_epoch,\n",
" max_steps=steps_per_epoch * epochs,\n",
" per_device_train_batch_size=batch_size,\n",
" gradient_accumulation_steps=1,\n",
" learning_rate=learning_rate,\n",
" weight_decay=weight_decay,\n",
" warmup_steps=warmup_steps,\n",
" label_names=[\"input_ids\", \"attention_mask\"],\n",
" push_to_hub=False,\n",
" report_to=\"none\",\n",
" disable_tqdm=True, # declutter the output a little\n",
" fp16=True,\n",
" gradient_checkpointing=True,\n",
" deepspeed=deepspeed,\n",
" )\n",
" disable_progress_bar()\n",
"\n",
" tokenizer = AutoTokenizer.from_pretrained(model_name)\n",
" tokenizer.pad_token = tokenizer.eos_token\n",
"\n",
" print(\"Loading model\")\n",
"\n",
" model = GPTJForCausalLM.from_pretrained(model_name, use_cache=False)\n",
" model.resize_token_embeddings(len(tokenizer))\n",
"\n",
" print(\"Model loaded\")\n",
"\n",
" enable_progress_bar()\n",
"\n",
" metric = evaluate.load(\"accuracy\")\n",
"\n",
" train_ds = train.get_dataset_shard(\"train\")\n",
" eval_ds = train.get_dataset_shard(\"validation\")\n",
"\n",
" train_ds_iterable = train_ds.iter_torch_batches(\n",
" batch_size=batch_size,\n",
" local_shuffle_buffer_size=train.get_context().get_world_size() * batch_size,\n",
" )\n",
" eval_ds_iterable = eval_ds.iter_torch_batches(batch_size=batch_size)\n",
"\n",
" def compute_metrics(eval_pred):\n",
" logits, labels = eval_pred\n",
" predictions = np.argmax(logits, axis=-1)\n",
" return metric.compute(predictions=predictions, references=labels)\n",
"\n",
" trainer = Trainer(\n",
" model=model,\n",
" args=training_args,\n",
" train_dataset=train_ds_iterable,\n",
" eval_dataset=eval_ds_iterable,\n",
" compute_metrics=compute_metrics,\n",
" tokenizer=tokenizer,\n",
" data_collator=default_data_collator,\n",
" )\n",
"\n",
" # Add callback to report checkpoints to Ray Train\n",
" trainer.add_callback(RayTrainReportCallback())\n",
" trainer = prepare_trainer(trainer)\n",
" trainer.train()"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"After defining the training function, instantiate the {class}`~ray.train.torch.TorchTrainer`. Aside from the function, set the `scaling_config` to control the number of workers and amount of resources to use, and `datasets`(the preprocessed Ray Datasets) to use for training and evaluation.\n",
"\n",
"```{note}\n",
"Running with multiple nodes necessitates the persistence of checkpoints\n",
"and other outputs to some external storage for access after training has completed.\n",
"**You should set up cloud storage or NFS, then replace `storage_path` with your own cloud bucket URI or NFS path.**\n",
"\n",
"See {ref}`Configuration and Persistent Storage<persistent-storage-guide>` for more details.\n",
"```"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"storage_path = \"s3://your-bucket-here\" # TODO: Set up cloud storage\n",
"# storage_path=\"/mnt/path/to/nfs\" # TODO: Alternatively, set up NFS"
]
},
{
"cell_type": "code",
"execution_count": 10,
"metadata": {
"tags": [
"remove-cell"
]
},
"outputs": [],
"source": [
"import os, re\n",
"\n",
"artifact_storage = os.environ.get(\"ANYSCALE_ARTIFACT_STORAGE\", \"artifact_storage\")\n",
"user_name = re.sub(r\"\\s+\", \"__\", os.environ.get(\"ANYSCALE_USERNAME\", \"user\"))\n",
"storage_path = f\"{artifact_storage}/{user_name}/gptj-deepspeed-finetune\""
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"batch_size = 16\n",
"train_ds_size = processed_datasets[\"train\"].count()\n",
"steps_per_epoch = train_ds_size // (batch_size * num_workers)"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"tags": [
"remove-cell"
]
},
"outputs": [],
"source": [
"# SMOKE TEST SETTINGS FOR CI\n",
"steps_per_epoch = 10\n",
"num_workers = 8\n",
"batch_size = 1"
]
},
{
"cell_type": "code",
"execution_count": 12,
"metadata": {
"tags": []
},
"outputs": [],
"source": [
"from ray.train.torch import TorchTrainer\n",
"from ray.train import RunConfig, ScalingConfig\n",
"\n",
"trainer = TorchTrainer(\n",
" train_loop_per_worker=train_func,\n",
" train_loop_config={\n",
" \"epochs\": 1,\n",
" \"batch_size\": batch_size, # per device\n",
" \"steps_per_epoch\": steps_per_epoch,\n",
" },\n",
" scaling_config=ScalingConfig(\n",
" num_workers=num_workers,\n",
" use_gpu=use_gpu,\n",
" resources_per_worker={\"GPU\": 1, \"CPU\": cpus_per_worker},\n",
" ),\n",
" datasets=processed_datasets,\n",
" run_config=RunConfig(storage_path=storage_path),\n",
")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Finally, call the {meth}`~ray.train.torch.TorchTrainer.fit` method to start training with Ray Train. Save the {class}`~ray.train.Result` object to a variable to access metrics and checkpoints."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"tags": [
"hide-output"
]
},
"outputs": [],
"source": [
"results = trainer.fit()"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Use the returned {class}`~ray.train.Result` object to access metrics and the Ray Train {class}`~ray.train.Checkpoint` associated with the last iteration."
]
},
{
"cell_type": "code",
"execution_count": 15,
"metadata": {
"tags": []
},
"outputs": [
{
"data": {
"text/plain": [
"Checkpoint(filesystem=<pyarrow._s3fs.S3FileSystem object at 0x7f8c59d311b0>, path=anyscale-staging-data-cld-kvedzwag2qa8i5bjxuevf5i7/org_7c1Kalm9WcX2bNIjW53GUT/cld_kvedZWag2qA8i5BjxUevf5i7/artifact_storage/yunxuan__xiao/gptj-deepspeed-finetune/TorchTrainer_2023-08-18_18-09-11/TorchTrainer_01ea5_00000_0_2023-08-18_18-09-12/checkpoint_000000)"
]
},
"execution_count": 15,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"checkpoint = results.checkpoint\n",
"checkpoint"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"(gptj-predict)=\n",
"### Generate text from prompt\n",
"\n",
"First, download the persistent Ray Train checkpoint from a gpu node and load the fine-tuned model weights and tokenizer from the checkpoint. Then use \ud83e\udd17 Transformers [`pipeline`](https://huggingface.co/docs/transformers/en/main_classes/pipelines) to generate predictions from the fine-tuned model.\n",
"\n",
"```{tip}\n",
"For large scale batch inference, see {ref}`End-to-end: Offline Batch Inference <batch_inference_home>`.\n",
"```"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Set the `task` to `\"text-generation\"`, and also set `device_map=\"auto\"` for Ray Train to automatically place the model on the right device. "
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"tags": []
},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"[{'generated_text': 'Romeo and Juliet. This very night shall they come. A word with you, sir.'}]\n",
"[{'generated_text': 'Romeo! I know thee not. Lord Mercutio, is it you! Signior Montague.'}]\n",
"[{'generated_text': 'Juliet, look up in the vault, and there shalt find a grave; within the monument there is a table:'}]\n"
]
}
],
"source": [
"from transformers import pipeline, AutoTokenizer, GPTJForCausalLM\n",
"import os\n",
"\n",
"\n",
"@ray.remote(num_gpus=1)\n",
"def generate_text():\n",
" # Download the checkpoint\n",
" os.system(f\"aws s3 sync s3://{checkpoint.path} /mnt/local_storage/\")\n",
"\n",
" # Load the model and tokenizer \n",
" model = GPTJForCausalLM.from_pretrained(\"/mnt/local_storage/checkpoint\")\n",
" tokenizer = AutoTokenizer.from_pretrained(\"/mnt/local_storage/checkpoint\")\n",
"\n",
" pipe = pipeline(\n",
" model=model,\n",
" tokenizer=tokenizer,\n",
" task=\"text-generation\",\n",
" torch_dtype=torch.float16,\n",
" device_map=\"auto\",\n",
" )\n",
"\n",
" # Generate from prompts!\n",
" result = []\n",
" for sentence in pipe(\n",
" [\"Romeo and Juliet\", \"Romeo\", \"Juliet\"], do_sample=True, min_length=20\n",
" ):\n",
" result.append(sentence)\n",
" \n",
" return result\n",
"\n",
"ref = generate_text.remote()\n",
"print(ray.get(ref))"
]
}
],
"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.9.18"
},
"orphan": true,
"vscode": {
"interpreter": {
"hash": "3c0d54d489a08ae47a06eae2fd00ff032d6cddb527c382959b7b2575f6a8167f"
}
}
},
"nbformat": 4,
"nbformat_minor": 4
}
doc/source/train/examples/horovod/horovod_example.rst
+25
View File
@@ -0,0 +1,25 @@
:orphan:
Run Horovod Distributed Training with PyTorch and Ray Train
===========================================================
.. raw:: html
<a id="try-anyscale-quickstart-horovod_example" target="_blank" href="https://console.anyscale.com/register/ha?render_flow=ray&utm_source=ray_docs&utm_medium=docs&utm_campaign=horovod_example">
<img src="../../../_static/img/run-on-anyscale.svg" alt="Run on Anyscale" />
<br/><br/>
</a>
This basic example demonstrates how to run Horovod distributed training with PyTorch and Ray Train.
Code example
------------
.. literalinclude:: /../../python/ray/train/examples/horovod/horovod_example.py
See also
--------
* :ref:`Get Started with Horovod <train-horovod>` for a tutorial on using Horovod with Ray Train
* :doc:`Ray Train Examples <../../examples>` for more use cases
doc/source/train/examples/intel_gaudi/bert.ipynb
+416
View File
@@ -0,0 +1,416 @@
{
"cells": [
{
"cell_type": "markdown",
"metadata": {},
"source": [
"(hpu_bert_training)=\n",
"# BERT Model Training with Intel Gaudi\n",
"\n",
"<a id=\"try-anyscale-quickstart-intel_gaudi-bert\" href=\"https://console.anyscale.com/register/ha?render_flow=ray&utm_source=ray_docs&utm_medium=docs&utm_campaign=intel_gaudi-bert\">\n",
" <img src=\"../../../_static/img/run-on-anyscale.svg\" alt=\"try-anyscale-quickstart\">\n",
"</a>\n",
"<br></br>\n",
"\n",
"In this notebook, we will train a BERT model for sequence classification using the Yelp review full dataset. We will use the `transformers` and `datasets` libraries from Hugging Face, along with `ray.train` for distributed training.\n",
"\n",
"[Intel Gaudi AI Processors (HPUs)](https://habana.ai) are AI hardware accelerators designed by Intel Habana Labs. For more information, see [Gaudi Architecture](https://docs.habana.ai/en/latest/Gaudi_Overview/index.html) and [Gaudi Developer Docs](https://developer.habana.ai/).\n",
"\n",
"## Configuration\n",
"\n",
"A node with Gaudi/Gaudi2 installed is required to run this example. Both Gaudi and Gaudi2 have 8 HPUs. We will use 2 workers to train the model, each using 1 HPU.\n",
"\n",
"We recommend using a prebuilt container to run these examples. To run a container, you need Docker. See [Install Docker Engine](https://docs.docker.com/engine/install/) for installation instructions.\n",
"\n",
"Next, follow [Run Using Containers](https://docs.habana.ai/en/latest/Installation_Guide/Bare_Metal_Fresh_OS.html?highlight=installer#run-using-containers) to install the Gaudi drivers and container runtime.\n",
"\n",
"Next, start the Gaudi container:\n",
"```bash\n",
"docker pull vault.habana.ai/gaudi-docker/1.20.0/ubuntu22.04/habanalabs/pytorch-installer-2.6.0:latest\n",
"docker run -it --runtime=habana -e HABANA_VISIBLE_DEVICES=all -e OMPI_MCA_btl_vader_single_copy_mechanism=none --cap-add=sys_nice --net=host --ipc=host vault.habana.ai/gaudi-docker/1.20.0/ubuntu22.04/habanalabs/pytorch-installer-2.6.0:latest\n",
"```\n",
"\n",
"Inside the container, install the following dependencies to run this notebook.\n",
"```bash\n",
"pip install ray[train] notebook transformers datasets evaluate\n",
"```"
]
},
{
"cell_type": "code",
"execution_count": 1,
"metadata": {},
"outputs": [
{
"name": "stderr",
"output_type": "stream",
"text": [
"/usr/local/lib/python3.10/dist-packages/torch/distributed/distributed_c10d.py:252: UserWarning: Device capability of hccl unspecified, assuming `cpu` and `cuda`. Please specify it via the `devices` argument of `register_backend`.\n",
" warnings.warn(\n"
]
}
],
"source": [
"# Import necessary libraries\n",
"\n",
"import os\n",
"from typing import Dict\n",
"\n",
"import torch\n",
"from torch import nn\n",
"from torch.utils.data import DataLoader\n",
"from tqdm import tqdm\n",
"\n",
"import numpy as np\n",
"import evaluate\n",
"from datasets import load_dataset\n",
"import transformers\n",
"from transformers import (\n",
" Trainer,\n",
" TrainingArguments,\n",
" AutoTokenizer,\n",
" AutoModelForSequenceClassification,\n",
")\n",
"\n",
"import ray.train\n",
"from ray.train import ScalingConfig\n",
"from ray.train.torch import TorchTrainer\n",
"from ray.train.torch import TorchConfig\n",
"from ray.runtime_env import RuntimeEnv\n",
"\n",
"import habana_frameworks.torch.core as htcore"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Metrics Setup\n",
"\n",
"We will use accuracy as our evaluation metric. The `compute_metrics` function will calculate the accuracy of our model's predictions."
]
},
{
"cell_type": "code",
"execution_count": 2,
"metadata": {},
"outputs": [],
"source": [
"# Metrics\n",
"metric = evaluate.load(\"accuracy\")\n",
"\n",
"def compute_metrics(eval_pred):\n",
" logits, labels = eval_pred\n",
" predictions = np.argmax(logits, axis=-1)\n",
" return metric.compute(predictions=predictions, references=labels)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Training Function\n",
"\n",
"This function will be executed by each worker during training. It handles data loading, tokenization, model initialization, and the training loop. Compared to a training function for GPU, no changes are needed to port to HPU. Internally, Ray Train does these things:\n",
"\n",
"* Detect HPU and set the device.\n",
"\n",
"* Initializes the habana PyTorch backend.\n",
"\n",
"* Initializes the habana distributed backend."
]
},
{
"cell_type": "code",
"execution_count": 3,
"metadata": {},
"outputs": [],
"source": [
"def train_func_per_worker(config: Dict):\n",
" \n",
" # Datasets\n",
" dataset = load_dataset(\"yelp_review_full\")\n",
" tokenizer = AutoTokenizer.from_pretrained(\"bert-base-cased\")\n",
" \n",
" def tokenize_function(examples):\n",
" return tokenizer(examples[\"text\"], padding=\"max_length\", truncation=True)\n",
"\n",
" lr = config[\"lr\"]\n",
" epochs = config[\"epochs\"]\n",
" batch_size = config[\"batch_size_per_worker\"]\n",
"\n",
" train_dataset = dataset[\"train\"].select(range(1000)).map(tokenize_function, batched=True)\n",
" eval_dataset = dataset[\"test\"].select(range(1000)).map(tokenize_function, batched=True)\n",
"\n",
" # Prepare dataloader for each worker\n",
" dataloaders = {}\n",
" dataloaders[\"train\"] = torch.utils.data.DataLoader(\n",
" train_dataset, \n",
" shuffle=True, \n",
" collate_fn=transformers.default_data_collator, \n",
" batch_size=batch_size\n",
" )\n",
" dataloaders[\"test\"] = torch.utils.data.DataLoader(\n",
" eval_dataset, \n",
" shuffle=True, \n",
" collate_fn=transformers.default_data_collator, \n",
" batch_size=batch_size\n",
" )\n",
"\n",
" # Obtain HPU device automatically\n",
" device = ray.train.torch.get_device()\n",
"\n",
" # Prepare model and optimizer\n",
" model = AutoModelForSequenceClassification.from_pretrained(\n",
" \"bert-base-cased\", num_labels=5\n",
" )\n",
" model = model.to(device)\n",
" \n",
" optimizer = torch.optim.SGD(model.parameters(), lr=lr, momentum=0.9)\n",
"\n",
" # Start training loops\n",
" for epoch in range(epochs):\n",
" # Each epoch has a training and validation phase\n",
" for phase in [\"train\", \"test\"]:\n",
" if phase == \"train\":\n",
" model.train() # Set model to training mode\n",
" else:\n",
" model.eval() # Set model to evaluate mode\n",
"\n",
" # breakpoint()\n",
" for batch in dataloaders[phase]:\n",
" batch = {k: v.to(device) for k, v in batch.items()}\n",
"\n",
" # zero the parameter gradients\n",
" optimizer.zero_grad()\n",
"\n",
" # forward\n",
" with torch.set_grad_enabled(phase == \"train\"):\n",
" # Get model outputs and calculate loss\n",
" \n",
" outputs = model(**batch)\n",
" loss = outputs.loss\n",
"\n",
" # backward + optimize only if in training phase\n",
" if phase == \"train\":\n",
" loss.backward()\n",
" optimizer.step()\n",
" print(f\"train epoch:[{epoch}]\\tloss:{loss:.6f}\")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Main Training Function\n",
"\n",
"The `train_bert` function sets up the distributed training environment using Ray and starts the training process. To enable training using HPU, we only need to make the following changes:\n",
"* Require an HPU for each worker in ScalingConfig\n",
"* Set backend to \"hccl\" in TorchConfig"
]
},
{
"cell_type": "code",
"execution_count": 4,
"metadata": {},
"outputs": [],
"source": [
"def train_bert(num_workers=2):\n",
" global_batch_size = 8\n",
"\n",
" train_config = {\n",
" \"lr\": 1e-3,\n",
" \"epochs\": 10,\n",
" \"batch_size_per_worker\": global_batch_size // num_workers,\n",
" }\n",
"\n",
" # Configure computation resources\n",
" # In ScalingConfig, require an HPU for each worker\n",
" scaling_config = ScalingConfig(num_workers=num_workers, resources_per_worker={\"CPU\": 1, \"HPU\": 1})\n",
" # Set backend to hccl in TorchConfig\n",
" torch_config = TorchConfig(backend = \"hccl\")\n",
" \n",
" # start your ray cluster\n",
" ray.init()\n",
" \n",
" # Initialize a Ray TorchTrainer\n",
" trainer = TorchTrainer(\n",
" train_loop_per_worker=train_func_per_worker,\n",
" train_loop_config=train_config,\n",
" torch_config=torch_config,\n",
" scaling_config=scaling_config,\n",
" )\n",
"\n",
" result = trainer.fit()\n",
" print(f\"Training result: {result}\")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Start Training\n",
"\n",
"Finally, we call the `train_bert` function to start the training process. You can adjust the number of workers to use.\n",
"\n",
"Note: the following warning is fine, and is resolved in SynapseAI version 1.14.0+:\n",
"```text\n",
"/usr/local/lib/python3.10/dist-packages/torch/distributed/distributed_c10d.py:252: UserWarning: Device capability of hccl unspecified, assuming `cpu` and `cuda`. Please specify it via the `devices` argument of `register_backend`.\n",
"```"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"train_bert(num_workers=2)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Possible outputs\n",
"\n",
"``` text\n",
"Downloading builder script: 100%|██████████| 4.20k/4.20k [00:00<00:00, 27.0MB/s]\n",
"2025-03-03 03:37:08,776 INFO worker.py:1841 -- Started a local Ray instance.\n",
"/usr/local/lib/python3.10/dist-packages/ray/tune/impl/tuner_internal.py:125: RayDeprecationWarning: The `RunConfig` class should be imported from `ray.tune` when passing it to the Tuner. Please update your imports. See this issue for more context and migration options: https://github.com/ray-project/ray/issues/49454. Disable these warnings by setting the environment variable: RAY_TRAIN_ENABLE_V2_MIGRATION_WARNINGS=0\n",
" _log_deprecation_warning(\n",
"(RayTrainWorker pid=75123) Setting up process group for: env:// [rank=0, world_size=2]\n",
"(TorchTrainer pid=74734) Started distributed worker processes: \n",
"(TorchTrainer pid=74734) - (node_id=eef984cd0cd96cce50bad1b1dab12e19c809047f10be3c829524a3d1, ip=100.83.111.228, pid=75123) world_rank=0, local_rank=0, node_rank=0\n",
"(TorchTrainer pid=74734) - (node_id=eef984cd0cd96cce50bad1b1dab12e19c809047f10be3c829524a3d1, ip=100.83.111.228, pid=75122) world_rank=1, local_rank=1, node_rank=0\n",
"Generating train split: 0%| | 0/650000 [00:00<?, ? examples/s]\n",
"Generating train split: 7%|▋ | 45000/650000 [00:00<00:01, 435976.18 examples/s]\n",
"Generating train split: 15%|█▍ | 95000/650000 [00:00<00:01, 469481.51 examples/s]\n",
"Generating train split: 23%|██▎ | 150000/650000 [00:00<00:01, 477676.99 examples/s]\n",
"Generating train split: 31%|███ | 203000/650000 [00:00<00:00, 493746.70 examples/s]\n",
"Generating train split: 43%|████▎ | 279000/650000 [00:00<00:00, 499340.09 examples/s]\n",
"Generating train split: 55%|█████▍ | 355000/650000 [00:00<00:00, 498613.65 examples/s]\n",
"Generating train split: 66%|██████▋ | 431000/650000 [00:00<00:00, 497799.19 examples/s]\n",
"Generating train split: 78%|███████▊ | 506000/650000 [00:01<00:00, 495696.93 examples/s]\n",
"Generating train split: 86%|████████▌ | 556000/650000 [00:01<00:00, 494508.05 examples/s]\n",
"Generating train split: 94%|█████████▎| 609000/650000 [00:01<00:00, 490725.53 examples/s]\n",
"Generating train split: 100%|██████████| 650000/650000 [00:01<00:00, 494916.42 examples/s]\n",
"Generating test split: 0%| | 0/50000 [00:00<?, ? examples/s]\n",
"Generating test split: 100%|██████████| 50000/50000 [00:00<00:00, 509619.87 examples/s]\n",
"Map: 0%| | 0/1000 [00:00<?, ? examples/s]\n",
"Map: 100%|██████████| 1000/1000 [00:00<00:00, 3998.33 examples/s]\n",
"Map: 100%|██████████| 1000/1000 [00:00<00:00, 4051.80 examples/s]\n",
"Map: 100%|██████████| 1000/1000 [00:00<00:00, 3869.20 examples/s]\n",
"(RayTrainWorker pid=75123) Some weights of BertForSequenceClassification were not initialized from the model checkpoint at bert-base-cased and are newly initialized: ['classifier.bias', 'classifier.weight']\n",
"(RayTrainWorker pid=75123) You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.\n",
"Map: 0%| | 0/1000 [00:00<?, ? examples/s] [repeated 3x across cluster] (Ray deduplicates logs by default. Set RAY_DEDUP_LOGS=0 to disable log deduplication, or see https://docs.ray.io/en/master/ray-observability/user-guides/configure-logging.html#log-deduplication for more options.)\n",
"Map: 100%|██████████| 1000/1000 [00:00<00:00, 3782.66 examples/s] [repeated 2x across cluster]\n",
"(RayTrainWorker pid=75123) ============================= HABANA PT BRIDGE CONFIGURATION =========================== \n",
"(RayTrainWorker pid=75123) PT_HPU_LAZY_MODE = 1\n",
"(RayTrainWorker pid=75123) PT_HPU_RECIPE_CACHE_CONFIG = ,false,1024\n",
"(RayTrainWorker pid=75123) PT_HPU_MAX_COMPOUND_OP_SIZE = 9223372036854775807\n",
"(RayTrainWorker pid=75123) PT_HPU_LAZY_ACC_PAR_MODE = 1\n",
"(RayTrainWorker pid=75123) PT_HPU_ENABLE_REFINE_DYNAMIC_SHAPES = 0\n",
"(RayTrainWorker pid=75123) PT_HPU_EAGER_PIPELINE_ENABLE = 1\n",
"(RayTrainWorker pid=75123) PT_HPU_EAGER_COLLECTIVE_PIPELINE_ENABLE = 1\n",
"(RayTrainWorker pid=75123) PT_HPU_ENABLE_LAZY_COLLECTIVES = 0\n",
"(RayTrainWorker pid=75123) ---------------------------: System Configuration :---------------------------\n",
"(RayTrainWorker pid=75123) Num CPU Cores : 160\n",
"(RayTrainWorker pid=75123) CPU RAM : 1056374420 KB\n",
"(RayTrainWorker pid=75123) ------------------------------------------------------------------------------\n",
"2025-03-03 03:41:04,658 INFO tune.py:1009 -- Wrote the latest version of all result files and experiment state to '/root/ray_results/TorchTrainer_2025-03-03_03-37-11' in 0.0020s.\n",
"\n",
"View detailed results here: /root/ray_results/TorchTrainer_2025-03-03_03-37-11\n",
"To visualize your results with TensorBoard, run: `tensorboard --logdir /tmp/ray/session_2025-03-03_03-37-06_983992_65223/artifacts/2025-03-03_03-37-11/TorchTrainer_2025-03-03_03-37-11/driver_artifacts`\n",
"\n",
"Training started with configuration:\n",
"╭─────────────────────────────────────────────────╮\n",
"│ Training config │\n",
"├─────────────────────────────────────────────────┤\n",
"│ train_loop_config/batch_size_per_worker 4 │\n",
"│ train_loop_config/epochs 10 │\n",
"│ train_loop_config/lr 0.001 │\n",
"╰─────────────────────────────────────────────────╯\n",
"(RayTrainWorker pid=75123) train epoch:[0] loss:1.979938\n",
"(RayTrainWorker pid=75123) train epoch:[0] loss:1.756611 [repeated 36x across cluster]\n",
"(RayTrainWorker pid=75123) train epoch:[0] loss:1.643875 [repeated 180x across cluster]\n",
"(RayTrainWorker pid=75123) train epoch:[0] loss:1.416416 [repeated 177x across cluster]\n",
"(RayTrainWorker pid=75123) train epoch:[1] loss:1.272513 [repeated 107x across cluster]\n",
"(RayTrainWorker pid=75123) \n",
"(RayTrainWorker pid=75123) train epoch:[1] loss:2.086884 [repeated 155x across cluster]\n",
"(RayTrainWorker pid=75123) train epoch:[1] loss:1.426217 [repeated 178x across cluster]\n",
"(RayTrainWorker pid=75122) train epoch:[1] loss:0.991381 [repeated 160x across cluster]\n",
"(RayTrainWorker pid=75123) train epoch:[2] loss:1.294097 [repeated 28x across cluster]\n",
"(RayTrainWorker pid=75123) train epoch:[2] loss:1.386306 [repeated 169x across cluster]\n",
"(RayTrainWorker pid=75123) train epoch:[2] loss:1.190416 [repeated 181x across cluster]\n",
"(RayTrainWorker pid=75123) train epoch:[3] loss:1.171733 [repeated 130x across cluster]\n",
"(RayTrainWorker pid=75123) train epoch:[3] loss:1.287821 [repeated 152x across cluster]\n",
"(RayTrainWorker pid=75123) train epoch:[3] loss:1.055692 [repeated 179x across cluster]\n",
"(RayTrainWorker pid=75122) train epoch:[3] loss:1.677789 [repeated 162x across cluster]\n",
"(RayTrainWorker pid=75123) train epoch:[4] loss:0.942071 [repeated 19x across cluster]\n",
"(RayTrainWorker pid=75123) train epoch:[4] loss:1.592500 [repeated 167x across cluster]\n",
"(RayTrainWorker pid=75123) train epoch:[4] loss:0.936934 [repeated 180x across cluster]\n",
"(RayTrainWorker pid=75123) \n",
"(RayTrainWorker pid=75123) train epoch:[5] loss:2.465384 [repeated 141x across cluster]\n",
"(RayTrainWorker pid=75123) train epoch:[5] loss:1.659170 [repeated 156x across cluster]\n",
"(RayTrainWorker pid=75123) train epoch:[5] loss:1.850438 [repeated 180x across cluster]\n",
"(RayTrainWorker pid=75122) train epoch:[5] loss:1.101623 [repeated 160x across cluster]\n",
"(RayTrainWorker pid=75123) train epoch:[6] loss:2.125591 [repeated 18x across cluster]\n",
"(RayTrainWorker pid=75123) train epoch:[6] loss:1.612838 [repeated 170x across cluster]\n",
"(RayTrainWorker pid=75123) train epoch:[6] loss:1.759160 [repeated 177x across cluster]\n",
"(RayTrainWorker pid=75123) train epoch:[7] loss:1.338552 [repeated 139x across cluster]\n",
"(RayTrainWorker pid=75123) train epoch:[7] loss:1.467959 [repeated 157x across cluster]\n",
"(RayTrainWorker pid=75123) train epoch:[7] loss:1.682137 [repeated 181x across cluster]\n",
"(RayTrainWorker pid=75123) \n",
"(RayTrainWorker pid=75123) train epoch:[8] loss:1.395805 [repeated 162x across cluster]\n",
"(RayTrainWorker pid=75123) train epoch:[8] loss:1.527835 [repeated 153x across cluster]\n",
"(RayTrainWorker pid=75123) train epoch:[8] loss:1.672311 [repeated 177x across cluster]\n",
"(RayTrainWorker pid=75123) \n",
"(RayTrainWorker pid=75122) train epoch:[8] loss:1.093186 [repeated 166x across cluster]\n",
"(RayTrainWorker pid=75123) train epoch:[9] loss:1.457587 [repeated 13x across cluster]\n",
"(RayTrainWorker pid=75123) train epoch:[9] loss:1.727377 [repeated 171x across cluster]\n",
"(RayTrainWorker pid=75123) train epoch:[9] loss:1.694001 [repeated 182x across cluster]\n",
"\n",
"Training completed after 0 iterations at 2025-03-03 03:41:04. Total running time: 3min 53s\n",
"\n",
"Training result: Result(\n",
" metrics={},\n",
" path='/root/ray_results/TorchTrainer_2025-03-03_03-37-11/TorchTrainer_ca6cf_00000_0_2025-03-03_03-37-11',\n",
" filesystem='local',\n",
" checkpoint=None\n",
")\n",
"(RayTrainWorker pid=75122) Some weights of BertForSequenceClassification were not initialized from the model checkpoint at bert-base-cased and are newly initialized: ['classifier.bias', 'classifier.weight']\n",
"(RayTrainWorker pid=75122) You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.\n",
"(RayTrainWorker pid=75122) train epoch:[9] loss:0.417845 [repeated 136x across cluster]\n",
"```"
]
}
],
"metadata": {
"kernelspec": {
"display_name": "Python 3.10.12 64-bit",
"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.12"
},
"orphan": true,
"vscode": {
"interpreter": {
"hash": "31f2aee4e71d21fbe5cf8b01ff0e069b9275f58929596ceb00d14d90e3e16cd6"
}
}
},
"nbformat": 4,
"nbformat_minor": 4
}
doc/source/train/examples/intel_gaudi/llama.ipynb
+598
View File
@@ -0,0 +1,598 @@
{
"cells": [
{
"cell_type": "markdown",
"metadata": {},
"source": [
"# Fine-tuning Llama-2 Model with Intel Gaudi\n",
"\n",
"<a id=\"try-anyscale-quickstart-intel_gaudi-llama\" href=\"https://console.anyscale.com/register/ha?render_flow=ray&utm_source=ray_docs&utm_medium=docs&utm_campaign=intel_gaudi-llama\">\n",
" <img src=\"../../../_static/img/run-on-anyscale.svg\" alt=\"try-anyscale-quickstart\">\n",
"</a>\n",
"<br></br>\n",
"\n",
"In this Jupyter notebook, we will:\n",
"- fine-tuning a [Llama-2-7b](https://huggingface.co/meta-llama/Llama-2-7b-chat-hf) model by using Intel Gaudi accelerators with DDP method\n",
"- fine-tuning a [Llama-2-70b](https://huggingface.co/meta-llama/Llama-2-70b-chat-hf) model by using Intel Gaudi accelerators with DeepSpeed method\n",
"\n",
"We will use PyTorch for model training and Ray for distributed training. We will use dataset [tatsu-lab/alpaca](https://huggingface.co/datasets/tatsu-lab/alpaca).\n",
"\n",
"[Intel Gaudi AI Processors (HPUs)](https://habana.ai) are AI hardware accelerators designed by Habana Labs. For more information, see [Gaudi Architecture](https://docs.habana.ai/en/latest/Gaudi_Overview/index.html) and [Gaudi Developer Docs](https://developer.habana.ai/).\n",
"\n",
"Basic features for this fine-tuning example are:\n",
"- Running on HPUs, support three execution mode: [\"lazy\", \"eager\", \"eager.compile\"](https://docs.habana.ai/en/latest/PyTorch/Reference/PyTorch_Gaudi_Theory_of_Operations.html).\n",
"- LoRA training.\n",
"- DDP or DeepSpeed based method.\n",
"- [`GaudiTrainer`](https://github.com/huggingface/optimum-habana/blob/main/optimum/habana/transformers/trainer.py) based training.\n",
"- Llama-2-7b/Llama-2-70b model.\n",
"- Ray based resource scheduling and management."
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Prepare environment\n",
"This example run on single node with 4 HPUs.\n",
"\n",
"We recommend using a prebuilt container to run these examples. To run a container, you need Docker. See [Install Docker Engine](https://docs.docker.com/engine/install/) for installation instructions.\n",
"\n",
"Next, follow [Run Using Containers](https://docs.habana.ai/en/latest/Installation_Guide/Bare_Metal_Fresh_OS.html?highlight=installer#run-using-containers) to install the Habana drivers and container runtime.\n",
"\n",
"### Get docker image\n",
"``` bash\n",
"docker pull vault.habana.ai/gaudi-docker/1.20.0/ubuntu22.04/habanalabs/pytorch-installer-2.6.0:latest\n",
"```\n",
"### Run docker image\n",
"``` bash\n",
"docker run -it --runtime=habana -e HABANA_VISIBLE_DEVICES=all -e OMPI_MCA_btl_vader_single_copy_mechanism=none --cap-add=sys_nice --net=host --ipc=host vault.habana.ai/gaudi-docker/1.20.0/ubuntu22.04/habanalabs/pytorch-installer-2.6.0:latest\n",
"# maybe should mapping your workspace volumns\n",
"```\n",
"### Install dependency\n",
"``` bash\n",
"# \"optimum-habana>1.11.1\" if execution mode \"eager\" or \"eager.compile\" \n",
"# \"ray>=2.20.0\"\n",
"pip install ray[train] notebook transformers datasets evaluate peft accelerate scikit-learn optimum-habana\n",
"\n",
"# install deepspeed\n",
"pip install git+https://github.com/HabanaAI/DeepSpeed.git@1.20.0\n",
"\n",
"# this notebook verfied with packages' version:\n",
"# transformers==4.45.2\n",
"# datasets==3.3.2\n",
"# evaluate==0.4.3\n",
"# peft==0.14.0\n",
"# accelerate==0.33.0\n",
"# scikit-learn==1.6.1\n",
"# optimum-habana==1.15.0\n",
"\n",
"# deepspeed==0.16.1+hpu.synapse.v1.20.0\n",
"```"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Import necessary libraries"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"import os\n",
"import copy\n",
"from typing import Dict\n",
"\n",
"import torch\n",
"\n",
"import datasets\n",
"import transformers\n",
"from transformers import DataCollatorForLanguageModeling\n",
"\n",
"from tqdm import tqdm\n",
"\n",
"import peft\n",
"\n",
"from optimum.habana import GaudiTrainer, GaudiConfig, GaudiTrainingArguments\n",
"from optimum.habana.transformers.modeling_utils import adapt_transformers_to_gaudi"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Prepare Dataset Function\n",
"\n",
"Preprocessing the raw dataset's each line with specified format."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"\n",
"def preprocess_dataset(raw_datasets):\n",
"\n",
" PROMPT_DICT = {\n",
" \"prompt_with_input\": (\n",
" \"Below is an instruction that describes a task, paired with an input that provides further context. \"\n",
" \"Write a response that appropriately completes the request.\\n\\n\"\n",
" \"### Instruction:\\n{instruction}\\n\\n### Input:\\n{input}\\n\\n### Response:\"\n",
" ),\n",
" \"prompt_without_input\": (\n",
" \"Below is an instruction that describes a task. \"\n",
" \"Write a response that appropriately completes the request.\\n\\n\"\n",
" \"### Instruction:\\n{instruction}\\n\\n### Response:\"\n",
" ),\n",
" }\n",
"\n",
" def create_prompts(examples):\n",
" prompts = {}\n",
" prompts[\"source\"] = []\n",
" prompts[\"target\"] = []\n",
" for example in examples:\n",
" prompt_template = (\n",
" PROMPT_DICT[\"prompt_with_input\"] if example[\"input\"] != \"\" else PROMPT_DICT[\"prompt_without_input\"]\n",
" )\n",
" source = prompt_template.format_map(example)\n",
" prompts[\"source\"].append(source)\n",
" prompts[\"target\"].append(example[\"output\"])\n",
" return prompts\n",
"\n",
" # Preprocessing the datasets.\n",
" for key in raw_datasets:\n",
" prompts = create_prompts(raw_datasets[key])\n",
" columns_to_be_removed = list(raw_datasets[key].features.keys())\n",
" raw_datasets[key] = raw_datasets[key].add_column(\"prompt_sources\", prompts[\"source\"])\n",
" raw_datasets[key] = raw_datasets[key].add_column(\"prompt_targets\", prompts[\"target\"])\n",
" raw_datasets[key] = raw_datasets[key].remove_columns(columns_to_be_removed)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Dataset to Tokenizer Function\n",
"\n",
"Tokenize each line in dataset by model tokenizer.\n",
"\n",
"In example codes, we concatenate the dataset's line content to accelerate training speed.\n",
"\n",
"All datasets are processed as \"train\" datasets, no evaluation datasets are sampled from raw_datasets."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"\n",
"def preprocess_dataset_to_tokenizer(raw_datasets, tokenizer):\n",
" max_seq_length = 512\n",
" tokenizer.pad_token_id = 0\n",
" tokenizer.eos_token_id = 1\n",
" tokenizer.bos_token_id = 2\n",
"\n",
" def tokenize(prompt, add_eos_token=True):\n",
" results = tokenizer(\n",
" prompt,\n",
" truncation=True,\n",
" max_length=max_seq_length,\n",
" padding=False,\n",
" return_tensors=None,\n",
" )\n",
" for i in range(len(results[\"input_ids\"])):\n",
" if (\n",
" results[\"input_ids\"][i][-1] != tokenizer.eos_token_id\n",
" and len(results[\"input_ids\"][i]) < max_seq_length\n",
" and add_eos_token\n",
" ):\n",
" results[\"input_ids\"][i].append(tokenizer.eos_token_id)\n",
" results[\"attention_mask\"][i].append(1)\n",
"\n",
" results[\"labels\"] = copy.deepcopy(results[\"input_ids\"])\n",
" results[\"input_id_len\"] = [len(result) for result in results[\"input_ids\"]]\n",
" return results\n",
"\n",
" def preprocess_function(examples):\n",
" keys = list(examples.data.keys())\n",
" if len(keys) != 2:\n",
" raise ValueError(\"Unsupported dataset format\")\n",
"\n",
" st = [s + t for s, t in zip(examples[keys[0]], examples[keys[1]])]\n",
"\n",
" examples_tokenized = tokenize(st)\n",
" input_ids = examples_tokenized[\"input_ids\"]\n",
" labels = examples_tokenized[\"labels\"]\n",
" return {\n",
" \"input_ids\": input_ids,\n",
" \"labels\": labels,\n",
" \"attention_mask\": examples_tokenized[\"attention_mask\"],\n",
" }\n",
"\n",
" tokenized_datasets = raw_datasets.map(\n",
" preprocess_function,\n",
" batched=True,\n",
" load_from_cache_file=True,\n",
" )\n",
"\n",
" def concatenate_data(dataset, max_seq_length):\n",
" concatenated_dataset = {}\n",
" for column in dataset.features:\n",
" concatenated_data = [item for sample in dataset[column] for item in sample]\n",
" reshaped_data = [\n",
" concatenated_data[i * max_seq_length : (i + 1) * max_seq_length]\n",
" for i in range(len(concatenated_data) // max_seq_length)\n",
" ]\n",
" concatenated_dataset[column] = reshaped_data\n",
" return datasets.Dataset.from_dict(concatenated_dataset)\n",
"\n",
" tokenized_datasets_ = tokenized_datasets[\"train\"].remove_columns([\"prompt_sources\", \"prompt_targets\"])\n",
" tokenized_datasets[\"train\"] = concatenate_data(tokenized_datasets_, max_seq_length)\n",
"\n",
" return tokenized_datasets"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Prepare training arguments\n",
"\n",
"here some arguments are hard coded, you can pass arguments from `config`"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"\n",
"def prepare_training_args(config: Dict):\n",
" # prepare execution mode config\n",
" execution_mode = config[\"execution_mode\"]\n",
" use_lazy_mode = True if execution_mode == \"lazy\" else False\n",
" torch_compile_backend = \"hpu_backend\" if execution_mode == \"eager.compile\" else None\n",
"\n",
" deepspeed = config[\"deepspeed\"] if \"deepspeed\" in config else None\n",
"\n",
" return GaudiTrainingArguments(deepspeed=deepspeed,\n",
" output_dir=config[\"output\"],\n",
" do_train=True,\n",
" do_eval=False,\n",
" per_device_train_batch_size=config[\"batch_size_per_worker\"],\n",
" bf16=True,\n",
" learning_rate=config[\"lr\"],\n",
" save_strategy=\"no\",\n",
" torch_compile_backend=torch_compile_backend,\n",
" evaluation_strategy=\"no\",\n",
" lr_scheduler_type=\"cosine\",\n",
" num_train_epochs=config[\"epochs\"],\n",
" use_lazy_mode=use_lazy_mode,\n",
" use_habana=True,\n",
" pipelining_fwd_bwd=True,\n",
" save_only_model=True,\n",
" gradient_checkpointing=True,\n",
" warmup_ratio=0.03,\n",
" throughput_warmup_steps=3,\n",
" logging_steps=5)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Prepare model\n",
"\n",
"1. download model from huggingface or read model from local directory.\n",
"2. convert model to lora model.\n",
"3. move model to HPU device.\n",
"\n",
"If you doesn't want to fine-tune with LoRA, just remove LoRA conversion step."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"\n",
"def prepare_model(config: Dict, device):\n",
" # prepare from pretrained model\n",
" deepspeed = config[\"deepspeed\"] if \"deepspeed\" in config else None\n",
" if deepspeed is not None:\n",
" auto_config = transformers.AutoConfig.from_pretrained(config[\"model\"], use_cache=False, revision=\"main\", use_auth_token=None, trust_remote_code=None)\n",
" model = transformers.AutoModelForCausalLM.from_pretrained(config[\"model\"], config=auto_config, **config[\"model_config\"])\n",
" model.generation_config.attn_softmax_bf16 = True\n",
" model.generation_config.use_flash_attention = True\n",
" else:\n",
" model = transformers.AutoModelForCausalLM.from_pretrained(config[\"model\"], **config[\"model_config\"])\n",
" model.enable_input_require_grads()\n",
"\n",
" # convert to peft model for lora training\n",
" peft_config = peft.LoraConfig(**config[\"lora_config\"])\n",
" model = peft.get_peft_model(model, peft_config)\n",
"\n",
" model.to(dtype=config[\"model_config\"][\"torch_dtype\"], device=device)\n",
"\n",
" return model\n"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Training Function\n",
"\n",
"This function will be executed by each worker during training, with following steps:\n",
"\n",
"- preparing training args, an instance of `GaudiTrainingArguments`.\n",
"- loading datasets and preprocess datasets, just load the first 4096 item as training datasets.\n",
"- loading pretrained model as tokenizer, and process datasets to tokenizer.\n",
"- loading pretrained model.\n",
"- preparing data collator and gaidu_config.\n",
"- preparing instance of `GaudiTrainer`.\n",
"- calling `train()` to train model.\n",
"- saving model results.\n",
"\n",
"Compared to a training function for GPU, no changes are needed to port to HPU. Internally, Ray Train does these things:\n",
"\n",
"- Detect HPU and set the device.\n",
"- Initialize the habana PyTorch backend.\n",
"- Initialize the habana distributed backend."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"\n",
"def train_func_per_worker(config: Dict):\n",
" # adapt transformers to gaudi\n",
" adapt_transformers_to_gaudi()\n",
"\n",
" # prepare training arguments\n",
" training_args = prepare_training_args(config)\n",
"\n",
" # prepare datasets\n",
" # here we use dataset \"tatsu-lab/alpaca\" from huggingface\n",
" raw_datasets = datasets.DatasetDict({\"train\": datasets.load_dataset(\"tatsu-lab/alpaca\", split='train[0:4096]')})\n",
" preprocess_dataset(raw_datasets)\n",
"\n",
" # prepare tokenizer\n",
" tokenizer = transformers.AutoTokenizer.from_pretrained(config[\"model\"])\n",
" tokenized_datasets = preprocess_dataset_to_tokenizer(raw_datasets, tokenizer)\n",
"\n",
" # prepare model\n",
" model = prepare_model(config, training_args.device)\n",
"\n",
" # prepare data collator\n",
" data_collator = DataCollatorForLanguageModeling(tokenizer, pad_to_multiple_of=8, return_tensors=\"pt\", mlm=False)\n",
"\n",
" # prepare gaudi config\n",
" gaudi_config = GaudiConfig()\n",
" gaudi_config.use_fused_adam = True\n",
" gaudi_config.use_fused_clip_norm = True\n",
"\n",
" # instance GaudiTrainer\n",
" trainer = GaudiTrainer(\n",
" model=model,\n",
" gaudi_config=gaudi_config,\n",
" args=training_args,\n",
" train_dataset=tokenized_datasets[\"train\"],\n",
" eval_dataset=None,\n",
" tokenizer=tokenizer,\n",
" data_collator=data_collator,\n",
" compute_metrics=None,\n",
" preprocess_logits_for_metrics=None,\n",
" )\n",
"\n",
" train_result = trainer.train()\n",
" print(f\"train_result = {train_result}\")\n",
" trainer.save_model()\n",
"\n",
" return train_result"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Main Training Function\n",
"The `train_llama` function sets up the distributed training environment using Ray and starts the training process. To enable training using HPU, we only need to make the following changes:\n",
"- Set the exectuion mode for training, supported execution mode are:\n",
"\n",
" - \"lazy\": Deferred execution of graphs, comprising of ops delivered from script op by op similar to Eager mode. It gives the Eager mode experience with performance on Gaudi. Unlike Eager Mode with torch.compile, graph is analyzed in each iteration leading to a higher CPU usage.\n",
" - \"eager\": Op-by-op execution as defined in standard PyTorch Eager mode scripts.\n",
" - \"eager.compile\": Eager mode extended with `torch.compile` - Similar to Eager mode but extended with wrapping complete or part of model (such as a function) into a graph. Parts that are not wrapped are executed eagerly.\n",
"\n",
" More detail theory can be found [here](https://docs.habana.ai/en/latest/PyTorch/Reference/PyTorch_Gaudi_Theory_of_Operations.html), and detail performance results can be found [here](https://developer.habana.ai/get-started/habana-models-performance/)\n",
"- Set training method, supported method are:\n",
" - \"ddp\"\n",
" - \"deepspeed\"\n",
"- Require an HPU for each worker in ScalingConfig\n",
"- Set backend to `hccl` in TorchConfig"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"\n",
"def train_llama(num_workers, execution_mode, training_method):\n",
" import ray\n",
" from ray.train import ScalingConfig\n",
" from ray.train.torch import TorchTrainer, TorchConfig\n",
"\n",
" # deepspeed config, can also place it to config file\n",
" deepspeed_config = {\n",
" \"steps_per_print\": 64,\n",
" \"train_batch_size\": \"auto\",\n",
" \"train_micro_batch_size_per_gpu\": \"auto\",\n",
" \"gradient_accumulation_steps\": \"auto\",\n",
" \"bf16\": {\n",
" \"enabled\": True\n",
" },\n",
" \"gradient_clipping\": 1.0,\n",
" \"zero_optimization\": {\n",
" \"stage\": 3,\n",
" \"overlap_comm\": False,\n",
" \"contiguous_gradients\": False,\n",
" \"stage3_gather_16bit_weights_on_model_save\": True\n",
" }\n",
" }\n",
"\n",
" # Preparing train configurations\n",
" train_config = {\n",
" \"execution_mode\": execution_mode,\n",
" \"model\": \"/root/models/models--meta-llama--Llama-2-70b-chat-hf/snapshots/e9149a12809580e8602995856f8098ce973d1080/\",\n",
" \"model_config\": {\"torch_dtype\": torch.bfloat16, \"trust_remote_code\": False, \"use_auth_token\": None},\n",
" \"lora_config\": {\"task_type\": \"CAUSAL_LM\", \"r\": 8, \"lora_alpha\": 32, \"lora_dropout\": 0.1, \"target_modules\": [\"q_proj\", \"v_proj\"]},\n",
" \"lr\": 1e-4,\n",
" \"epochs\": 2,\n",
" \"batch_size_per_worker\": 8,\n",
" \"output\": \"/tmp/ray/\",\n",
" \"deepspeed\": deepspeed_config if training_method == \"deepspeed\" else None,\n",
" }\n",
"\n",
" # Configure computation resources\n",
" # In ScalingConfig, require an HPU for each worker\n",
" scaling_config = ScalingConfig(num_workers=num_workers, resources_per_worker={\"CPU\": 1, \"HPU\": 1})\n",
" # Set backend to hccl in TorchConfig\n",
" torch_config = TorchConfig(backend = \"hccl\")\n",
"\n",
" # start your ray cluster\n",
" ray.init()\n",
"\n",
" # Initialize a Ray TorchTrainer\n",
" trainer = TorchTrainer(\n",
" train_loop_per_worker=train_func_per_worker,\n",
" train_loop_config=train_config,\n",
" torch_config=torch_config,\n",
" scaling_config=scaling_config,\n",
" )\n",
"\n",
" result = trainer.fit()\n",
" print(f\"Training result: {result}\")\n"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Start Training\n",
"\n",
"Finally, we call the `train_llama` function to start the training process. You can adjust the number of workers to use, and the execution mode for HPU."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# set some environment variables\n",
"os.environ[\"RAY_EXPERIMENTAL_NOSET_HABANA_VISIBLE_MODULES\"] = \"0\"\n",
"# if using RAY_EXPERIMENTAL_NOSET_HABANA_VISIBLE_MODULES env var\n",
"# you must set HABANA_VISIBLE_DEVICES, such as\n",
"# os.environ[\"HABANA_VISIBLE_DEVICES\"] = \"0,1,2,3\"\n",
"\n",
"# execution_mode are [\"lazy\", \"eager\", \"eager.compile\"]\n",
"execution_mode = \"lazy\"\n",
"os.environ[\"PT_HPU_LAZY_MODE\"] = \"1\" if execution_mode == \"lazy\" else \"0\"\n",
"\n",
"# training_method are [\"ddp\", \"deepspeed\"]\n",
"training_method = \"deepspeed\"\n",
"if training_method == \"deepspeed\":\n",
" os.environ[\"PT_HPU_MAX_COMPOUND_OP_SIZE\"] = \"10\"\n",
" os.environ[\"DEEPSPEED_HPU_ZERO3_SYNC_MARK_STEP_REQUIRED\"] = \"1\"\n",
"\n",
"# here use 4 HPUs\n",
"train_llama(num_workers=4, execution_mode=execution_mode, training_method=training_method)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Final output\n",
"\n",
"### For DDP on HPUs\n",
"- Llama-2-70b-chat-hf\n",
"- 4 HPU\n",
"- LoRA\n",
"\n",
"``` bash\n",
"(RayTrainWorker pid=123181) {'loss': 1.8051, 'grad_norm': 0.6015625, 'learning_rate': 9.938441702975689e-05, 'epoch': 0.16, 'memory_allocated (GB)': 13.64, 'max_memory_allocated (GB)': 48.92, 'total_memory_available (GB)': 94.62}\n",
"(RayTrainWorker pid=123181) {'loss': 1.6754, 'grad_norm': 0.408203125, 'learning_rate': 9.567727288213005e-05, 'epoch': 0.32, 'memory_allocated (GB)': 13.64, 'max_memory_allocated (GB)': 48.92, 'total_memory_available (GB)': 94.62}\n",
"(RayTrainWorker pid=123181) {'loss': 1.568, 'grad_norm': 0.4453125, 'learning_rate': 8.885729807284856e-05, 'epoch': 0.48, 'memory_allocated (GB)': 13.64, 'max_memory_allocated (GB)': 48.92, 'total_memory_available (GB)': 94.62}\n",
"(RayTrainWorker pid=123181) {'loss': 1.4934, 'grad_norm': 0.4609375, 'learning_rate': 7.938926261462366e-05, 'epoch': 0.65, 'memory_allocated (GB)': 13.64, 'max_memory_allocated (GB)': 48.92, 'total_memory_available (GB)': 94.62}\n",
"(RayTrainWorker pid=123181) {'loss': 1.3965, 'grad_norm': 0.3515625, 'learning_rate': 6.7918397477265e-05, 'epoch': 0.81, 'memory_allocated (GB)': 13.64, 'max_memory_allocated (GB)': 48.92, 'total_memory_available (GB)': 94.62}\n",
"(RayTrainWorker pid=123181) {'loss': 1.3461, 'grad_norm': 0.34765625, 'learning_rate': 5.522642316338268e-05, 'epoch': 0.97, 'memory_allocated (GB)': 13.64, 'max_memory_allocated (GB)': 48.92, 'total_memory_available (GB)': 94.62}\n",
"(RayTrainWorker pid=123181) {'loss': 1.2924, 'grad_norm': 0.32421875, 'learning_rate': 4.2178276747988446e-05, 'epoch': 1.13, 'memory_allocated (GB)': 13.64, 'max_memory_allocated (GB)': 48.92, 'total_memory_available (GB)': 94.62}\n",
"(RayTrainWorker pid=123181) {'loss': 1.2643, 'grad_norm': 0.33203125, 'learning_rate': 2.9663167846209998e-05, 'epoch': 1.29, 'memory_allocated (GB)': 13.64, 'max_memory_allocated (GB)': 48.92, 'total_memory_available (GB)': 94.62}\n",
"(RayTrainWorker pid=123181) {'loss': 1.263, 'grad_norm': 0.318359375, 'learning_rate': 1.8533980447508137e-05, 'epoch': 1.45, 'memory_allocated (GB)': 13.64, 'max_memory_allocated (GB)': 48.92, 'total_memory_available (GB)': 94.62}\n",
"(RayTrainWorker pid=123181) {'loss': 1.2502, 'grad_norm': 0.275390625, 'learning_rate': 9.549150281252633e-06, 'epoch': 1.61, 'memory_allocated (GB)': 13.64, 'max_memory_allocated (GB)': 48.92, 'total_memory_available (GB)': 94.62}\n",
"(RayTrainWorker pid=123181) {'loss': 1.2161, 'grad_norm': 0.2734375, 'learning_rate': 3.3209786751399187e-06, 'epoch': 1.77, 'memory_allocated (GB)': 13.64, 'max_memory_allocated (GB)': 48.92, 'total_memory_available (GB)': 94.62}\n",
"(RayTrainWorker pid=123181) {'loss': 1.2517, 'grad_norm': 0.294921875, 'learning_rate': 2.7390523158633554e-07, 'epoch': 1.94, 'memory_allocated (GB)': 13.64, 'max_memory_allocated (GB)': 48.92, 'total_memory_available (GB)': 94.62}\n",
"```\n",
"\n",
"### For DeepSpeed on HPUs\n",
"- Llama-2-70b-chat-hf\n",
"- 4 HPU\n",
"- LoRA\n",
"\n",
"``` bash\n",
"(RayTrainWorker pid=50067) {'loss': 1.662, 'grad_norm': 0.36514782905578613, 'learning_rate': 9.938441702975689e-05, 'epoch': 0.16, 'memory_allocated (GB)': 32.86, 'max_memory_allocated (GB)': 94.46, 'total_memory_available (GB)': 94.62}\n",
"(RayTrainWorker pid=50067) {'loss': 1.6047, 'grad_norm': 0.396455317735672, 'learning_rate': 9.567727288213005e-05, 'epoch': 0.32, 'memory_allocated (GB)': 32.86, 'max_memory_allocated (GB)': 94.57, 'total_memory_available (GB)': 94.62}\n",
"(RayTrainWorker pid=50067) {'loss': 1.4974, 'grad_norm': 0.49250370264053345, 'learning_rate': 8.885729807284856e-05, 'epoch': 0.48, 'memory_allocated (GB)': 32.86, 'max_memory_allocated (GB)': 94.57, 'total_memory_available (GB)': 94.62}\n",
"(RayTrainWorker pid=50067) {'loss': 1.4078, 'grad_norm': 0.49840453267097473, 'learning_rate': 7.938926261462366e-05, 'epoch': 0.65, 'memory_allocated (GB)': 32.86, 'max_memory_allocated (GB)': 94.57, 'total_memory_available (GB)': 94.62}\n",
"(RayTrainWorker pid=50067) {'loss': 1.315, 'grad_norm': 0.3432576656341553, 'learning_rate': 6.7918397477265e-05, 'epoch': 0.81, 'memory_allocated (GB)': 32.86, 'max_memory_allocated (GB)': 94.59, 'total_memory_available (GB)': 94.62}\n",
"(RayTrainWorker pid=50067) {'loss': 1.2651, 'grad_norm': 0.32175061106681824, 'learning_rate': 5.522642316338268e-05, 'epoch': 0.97, 'memory_allocated (GB)': 32.86, 'max_memory_allocated (GB)': 94.59, 'total_memory_available (GB)': 94.62}\n",
"(RayTrainWorker pid=50067) {'loss': 1.1947, 'grad_norm': 0.3646097481250763, 'learning_rate': 4.2178276747988446e-05, 'epoch': 1.13, 'memory_allocated (GB)': 32.86, 'max_memory_allocated (GB)': 94.59, 'total_memory_available (GB)': 94.62}\n",
"(RayTrainWorker pid=50067) {'loss': 1.1534, 'grad_norm': 0.4598522186279297, 'learning_rate': 2.9663167846209998e-05, 'epoch': 1.29, 'memory_allocated (GB)': 32.86, 'max_memory_allocated (GB)': 94.59, 'total_memory_available (GB)': 94.62}\n",
"(RayTrainWorker pid=50067) {'loss': 1.1404, 'grad_norm': 0.2677183449268341, 'learning_rate': 1.8533980447508137e-05, 'epoch': 1.45, 'memory_allocated (GB)': 32.75, 'max_memory_allocated (GB)': 94.59, 'total_memory_available (GB)': 94.62}\n",
"(RayTrainWorker pid=50067) {'loss': 1.1283, 'grad_norm': 0.32087600231170654, 'learning_rate': 9.549150281252633e-06, 'epoch': 1.61, 'memory_allocated (GB)': 32.86, 'max_memory_allocated (GB)': 94.59, 'total_memory_available (GB)': 94.62}\n",
"(RayTrainWorker pid=50067) {'loss': 1.0877, 'grad_norm': 0.28305548429489136, 'learning_rate': 3.3209786751399187e-06, 'epoch': 1.77, 'memory_allocated (GB)': 32.86, 'max_memory_allocated (GB)': 94.59, 'total_memory_available (GB)': 94.62}\n",
"(RayTrainWorker pid=50067) {'loss': 1.1238, 'grad_norm': 0.25713953375816345, 'learning_rate': 2.7390523158633554e-07, 'epoch': 1.94, 'memory_allocated (GB)': 32.86, 'max_memory_allocated (GB)': 94.59, 'total_memory_available (GB)': 94.62}\n",
"\n",
"```"
]
}
],
"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.12"
},
"orphan": true
},
"nbformat": 4,
"nbformat_minor": 2
}
doc/source/train/examples/intel_gaudi/llama_pretrain.ipynb
+598
View File
@@ -0,0 +1,598 @@
{
"cells": [
{
"cell_type": "markdown",
"metadata": {},
"source": [
"# Llama model pre-training on Intel Gaudi\n",
"\n",
"<a id=\"try-anyscale-quickstart-intel_gaudi-llama_pretrain\" href=\"https://console.anyscale.com/register/ha?render_flow=ray&utm_source=ray_docs&utm_medium=docs&utm_campaign=intel_gaudi-llama_pretrain\">\n",
" <img src=\"../../../_static/img/run-on-anyscale.svg\" alt=\"try-anyscale-quickstart\">\n",
"</a>\n",
"<br></br>\n",
"\n",
"In this Jupyter notebook, we will pre-train a [huggyllama/llama-7b](https://huggingface.co/huggyllama/llama-7b) model by using Intel Gaudi accelerators.\n",
"\n",
"We will use PyTorch for model training and Ray for distributed training.\n",
"\n",
"[Intel Gaudi AI Processors (HPUs)](https://habana.ai) are AI hardware accelerators designed by Habana Labs. For more information, see [Gaudi Architecture](https://docs.habana.ai/en/latest/Gaudi_Overview/index.html) and [Gaudi Developer Docs](https://developer.habana.ai/).\n",
"\n",
"Basic features for this pre-training example are:\n",
"- Running on HPUs, support three execution mode: [\"lazy\", \"eager\", \"eager.compile\"](https://docs.habana.ai/en/latest/PyTorch/Reference/PyTorch_Gaudi_Theory_of_Operations.html).\n",
"- Pre-training llama model use configuration [huggyllama/llama-7b](https://huggingface.co/huggyllama/llama-7b)\n",
"- [`GaudiTrainer`](https://github.com/huggingface/optimum-habana/blob/main/optimum/habana/transformers/trainer.py) based training.\n",
"- DeepSpeed based pre-training.\n",
"- Ray based resource scheduling and management."
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Prepare environment\n",
"This example run on single node with 4 HPUs.\n",
"\n",
"We recommend using a prebuilt container to run these examples. To run a container, you need Docker. See [Install Docker Engine](https://docs.docker.com/engine/install/) for installation instructions.\n",
"\n",
"Next, follow [Run Using Containers](https://docs.habana.ai/en/latest/Installation_Guide/Bare_Metal_Fresh_OS.html?highlight=installer#run-using-containers) to install the Habana drivers and container runtime.\n",
"\n",
"### Get docker image\n",
"``` bash\n",
"# more available docker image can be found here: https://vault.habana.ai/ui/native/gaudi-docker\n",
"docker pull vault.habana.ai/gaudi-docker/1.20.0/ubuntu22.04/habanalabs/pytorch-installer-2.6.0:latest\n",
"```\n",
"### Run docker image\n",
"``` bash\n",
"docker run -it --runtime=habana -e HABANA_VISIBLE_DEVICES=all -e OMPI_MCA_btl_vader_single_copy_mechanism=none --cap-add=sys_nice --net=host --ipc=host vault.habana.ai/gaudi-docker/1.20.0/ubuntu22.04/habanalabs/pytorch-installer-2.6.0:latest\n",
"# maybe should mapping your workspace volumns\n",
"```\n",
"### Install dependency\n",
"``` bash\n",
"# \"optimum-habana>1.11.1\" if execution mode \"eager\" or \"eager.compile\" \n",
"# \"ray>=2.20.0\"\n",
"pip install ray[train] notebook transformers datasets evaluate peft accelerate scikit-learn optimum-habana\n",
"\n",
"# install deepspeed\n",
"pip install git+https://github.com/HabanaAI/DeepSpeed.git@1.20.0\n",
"\n",
"# this notebook verfied with packages' version:\n",
"# transformers==4.45.2\n",
"# datasets==3.3.2\n",
"# evaluate==0.4.3\n",
"# peft==0.14.0\n",
"# accelerate==0.33.0\n",
"# scikit-learn==1.6.1\n",
"# optimum-habana==1.15.0\n",
"\n",
"# deepspeed==0.16.1+hpu.synapse.v1.20.0\n",
"```"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Import necessary libraries"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"#!/usr/bin/env python\n",
"\n",
"import os\n",
"from typing import Any, Dict\n",
"from torch.utils.data import DataLoader\n",
"\n",
"import transformers\n",
"from itertools import chain\n",
"from datasets import load_dataset\n",
"from transformers import default_data_collator\n",
"from transformers.testing_utils import CaptureLogger\n",
"from optimum.habana import GaudiConfig, GaudiTrainer, GaudiTrainingArguments\n",
"from optimum.habana.utils import set_seed"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Build datasets\n",
"\n",
"Download and load dataset from huggingface.co"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"def load_datasets(config):\n",
" dataset_name = config[\"name\"] \n",
" dataset_config_name = config[\"config_name\"]\n",
"\n",
" # Downloading and loading a dataset from the hub.\n",
" raw_datasets = load_dataset(\n",
" dataset_name,\n",
" dataset_config_name,\n",
" cache_dir=None,\n",
" token=None,\n",
" streaming=False,\n",
" )\n",
" if \"validation\" not in raw_datasets.keys():\n",
" raw_datasets[\"validation\"] = load_dataset(\n",
" dataset_name,\n",
" dataset_config_name,\n",
" split=f\"train[:{data_args.validation_split_percentage}%]\",\n",
" cache_dir=None,\n",
" token=None,\n",
" streaming=False,\n",
" )\n",
" raw_datasets[\"train\"] = load_dataset(\n",
" dataset_name,\n",
" dataset_config_name,\n",
" split=f\"train[{data_args.validation_split_percentage}%:]\",\n",
" cache_dir=None,\n",
" token=None,\n",
" streaming=False,\n",
" )\n",
"\n",
" return raw_datasets"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Load tokenizer\n",
"\n",
"Download vocabulary from huggingface.co."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"def load_tokenizer(config):\n",
" name = config[\"name\"]\n",
" tokenizer_kwargs = {\n",
" \"cache_dir\": None,\n",
" \"use_fast\": True,\n",
" \"revision\": \"main\",\n",
" \"token\": None,\n",
" \"trust_remote_code\": False,\n",
" }\n",
" return transformers.AutoTokenizer.from_pretrained(name, **tokenizer_kwargs)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Tokenize dataset\n",
"\n",
"tokenize word to token ids."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"def tokenize_dataset(datasets, tokenizer):\n",
" column_names = list(datasets[\"train\"].features)\n",
" text_column_name = \"text\" if \"text\" in column_names else column_names[0]\n",
"\n",
" tok_logger = transformers.utils.logging.get_logger(\"transformers.tokenization_utils_base\")\n",
"\n",
" def tokenize_function(examples):\n",
" with CaptureLogger(tok_logger) as cl:\n",
" output = tokenizer(examples[text_column_name])\n",
" # clm input could be much much longer than block_size\n",
" if \"Token indices sequence length is longer than the\" in cl.out:\n",
" tok_logger.warning(\n",
" \"^^^^^^^^^^^^^^^^ Please ignore the warning above - this long input will be chunked into smaller bits\"\n",
" \" before being passed to the model.\"\n",
" )\n",
" return output\n",
"\n",
" tokenized_datasets = datasets.map(\n",
" tokenize_function,\n",
" batched=True,\n",
" num_proc=None,\n",
" remove_columns=column_names,\n",
" load_from_cache_file=True,\n",
" desc=\"Running tokenizer on dataset\",\n",
" )\n",
"\n",
" return tokenized_datasets"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Group dataset\n",
"\n",
"This preprocssing will concatenate all texts from our dataset and generate chunks of block_size, and will pre-train model much faster."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"def group_dataset(config, datasets, tokenizer):\n",
" config_name = config[\"name\"]\n",
" auto_config = transformers.AutoConfig.from_pretrained(config_name)\n",
" max_pos_embeddings = auto_config.max_position_embeddings\n",
" block_size = tokenizer.model_max_length\n",
" if block_size > max_pos_embeddings:\n",
" print(\n",
" f\"The tokenizer picked seems to have a very large `model_max_length` ({tokenizer.model_max_length}). \"\n",
" f\"Using block_size={min(1024, max_pos_embeddings)} instead. You can change that default value by passing --block_size xxx.\"\n",
" )\n",
" if max_pos_embeddings > 0:\n",
" block_size = min(1024, max_pos_embeddings)\n",
" else:\n",
" block_size = 1024\n",
"\n",
" # Main data processing function that will concatenate all texts from our dataset and generate chunks of block_size.\n",
" def group_texts(examples):\n",
" # Concatenate all texts.\n",
" concatenated_examples = {k: list(chain(*examples[k])) for k in examples.keys()}\n",
" total_length = len(concatenated_examples[list(examples.keys())[0]])\n",
" # We drop the small remainder, and if the total_length < block_size we exclude this batch and return an empty dict.\n",
" # We could add padding if the model supported it instead of this drop, you can customize this part to your needs.\n",
" total_length = (total_length // block_size) * block_size\n",
" # Split by chunks of max_len.\n",
" result = {\n",
" k: [t[i : i + block_size] for i in range(0, total_length, block_size)]\n",
" for k, t in concatenated_examples.items()\n",
" }\n",
" result[\"labels\"] = result[\"input_ids\"].copy()\n",
" return result\n",
"\n",
" lm_datasets = datasets.map(\n",
" group_texts,\n",
" batched=True,\n",
" num_proc=None,\n",
" load_from_cache_file=True,\n",
" desc=f\"Grouping texts in chunks of {block_size}\",\n",
" )\n",
" return lm_datasets"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Load model\n",
"\n",
"Download and load pre-configed model from huggingface.co, the detail model configurations in config.json"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"def load_model(config):\n",
" name = config[\"name\"]\n",
" model_config = config.get(\"config\", {})\n",
" auto_config = transformers.AutoConfig.from_pretrained(\n",
" pretrained_model_name_or_path=name, **model_config\n",
" )\n",
" model = transformers.AutoModelForCausalLM.from_config(auto_config, trust_remote_code=False)\n",
"\n",
" return model"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Prepare trainer\n",
"\n",
"Instance Trainer with `model`, `gaudi_config`, `training_args`, `tokenizer`\n",
"\n",
"No evaluation dataset passed, just training."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"def get_trainer(training_args, datasets, tokenizer, model):\n",
" gaudi_config = GaudiConfig.from_pretrained(\n",
" training_args.gaudi_config_name, revision=\"main\",\n",
" )\n",
"\n",
" trainer = GaudiTrainer(\n",
" model=model,\n",
" gaudi_config=gaudi_config,\n",
" args=training_args,\n",
" train_dataset=datasets[\"train\"],\n",
" eval_dataset=None,\n",
" tokenizer=tokenizer,\n",
" data_collator=default_data_collator,\n",
" )\n",
" return trainer"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Training Function\n",
"\n",
"This function will be executed by each worker during training, with following steps:\n",
"- prepare GaudiTrainingArguments object.\n",
"- load datasets from huggingface.co.\n",
"- load pre-configed tokenizer from huggingface.co.\n",
"- tokenize dataset with loaded model tokenizer.\n",
"- concatenate all texts from our dataset and generate chunks of block_size.\n",
"- instance object of `GaudiTrainer` with training_args, datasets, tokenizer, and model.\n",
"- call `train` of trainer.\n",
"- save model."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"def pretrain_llama(config: Dict[str, Any]):\n",
"\n",
" training_args = GaudiTrainingArguments(**config[\"training_args\"])\n",
" set_seed(training_args.seed)\n",
"\n",
" raw_datasets = load_datasets(config[\"datasets\"])\n",
"\n",
" tokenizer = load_tokenizer(config[\"tokenizer\"])\n",
"\n",
" tokenized_datasets = tokenize_dataset(raw_datasets, tokenizer)\n",
"\n",
" tokenized_datasets = group_dataset(config[\"model\"], tokenized_datasets, tokenizer)\n",
"\n",
" model = load_model(config[\"model\"])\n",
"\n",
" trainer = get_trainer(training_args, tokenized_datasets, tokenizer, model)\n",
"\n",
" result = trainer.train()\n",
" trainer.save_model()\n",
" print(result)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Main Training Function\n",
"\n",
"The `main` function sets up the distributed training environment using Ray and starts the training process. To enable training using HPU, we only need to make the following changes:\n",
"- Set the exectuion mode for training, supported execution mode are:\n",
"\n",
" - \"lazy\": Deferred execution of graphs, comprising of ops delivered from script op by op similar to Eager mode. It gives the Eager mode experience with performance on Gaudi. Unlike Eager Mode with torch.compile, graph is analyzed in each iteration leading to a higher CPU usage.\n",
" - \"eager\": Op-by-op execution as defined in standard PyTorch Eager mode scripts.\n",
" - \"eager.compile\": Eager mode extended with `torch.compile` - Similar to Eager mode but extended with wrapping complete or part of model (such as a function) into a graph. Parts that are not wrapped are executed eagerly.\n",
"\n",
" More detail theory can be found [here](https://docs.habana.ai/en/latest/PyTorch/Reference/PyTorch_Gaudi_Theory_of_Operations.html), and detail performance results can be found [here](https://developer.habana.ai/get-started/habana-models-performance/)\n",
"- Require an HPU for each worker in ScalingConfig\n",
"- Set backend to `hccl` in TorchConfig"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"def main(num_workers, execution_mode):\n",
" import ray\n",
" from ray.train import ScalingConfig\n",
" from ray.train.torch import TorchTrainer, TorchConfig\n",
"\n",
" pretrain_config = {\n",
" \"datasets\": {\n",
" \"name\": \"wikitext\",\n",
" \"config_name\": \"wikitext-2-raw-v1\",\n",
" },\n",
" \"tokenizer\": {\n",
" \"name\": \"huggyllama/llama-7b\",\n",
" \"config\": {}\n",
" },\n",
" \"model\": {\n",
" \"name\": \"huggyllama/llama-7b\",\n",
" \"config\": {\n",
" \"torch_dtype\": \"bfloat16\",\n",
" },\n",
" },\n",
" \"training_args\": {\n",
" \"per_device_train_batch_size\": 1,\n",
" \"do_train\": True,\n",
" \"save_strategy\": \"no\",\n",
" \"output_dir\": \"/tmp/ray/pretrain-llama-2\",\n",
" \"logging_steps\": 1,\n",
" \"gaudi_config_name\": \"Habana/llama\",\n",
" \"use_habana\": True,\n",
" \"throughput_warmup_steps\": 3,\n",
" \"use_lazy_mode\": True,\n",
" \"overwrite_output_dir\": True,\n",
" \"seed\": 42,\n",
" \"bf16\": True,\n",
" \"report_to\":'tensorboard',\n",
" \"deepspeed\": {\n",
" \"steps_per_print\": 64,\n",
" \"train_batch_size\": \"auto\",\n",
" \"train_micro_batch_size_per_gpu\": \"auto\",\n",
" \"gradient_accumulation_steps\": \"auto\",\n",
" \"bf16\": {\n",
" \"enabled\": True\n",
" },\n",
" \"gradient_clipping\": 1.0,\n",
" \"zero_optimization\": {\n",
" \"stage\": 3,\n",
" \"overlap_comm\": False,\n",
" \"reduce_scatter\": False,\n",
" \"contiguous_gradients\": False,\n",
" \"stage3_gather_16bit_weights_on_model_save\": True\n",
" }\n",
" },\n",
" },\n",
" }\n",
"\n",
" # if execution mode is eager with compile, must spcified with a compile backend\n",
" if execution_mode == \"eager.compile\":\n",
" pretrain_config[\"training_args\"].update({\"torch_compile_backend\": \"hpu_backend\"})\n",
"\n",
" scaling_config = ScalingConfig(num_workers=num_workers,\n",
" use_gpu=False,\n",
" resources_per_worker={\"CPU\": 1, \"HPU\": 1})\n",
"\n",
" # Set backend to hccl in TorchConfig\n",
" torch_config = TorchConfig(backend=\"hccl\")\n",
"\n",
" ray.init()\n",
"\n",
" # Initialize a Ray TorchTrainer\n",
" trainer = TorchTrainer(\n",
" train_loop_per_worker=pretrain_llama,\n",
" train_loop_config=pretrain_config,\n",
" torch_config=torch_config,\n",
" scaling_config=scaling_config\n",
" )\n",
"\n",
" result = trainer.fit()\n",
" print(result)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Start Training\n",
"\n",
"Finally, we call the `main` function to start the pre-training process.\n",
"\n",
"Before calling `main` function, you must set some environment variables.\n",
"\n",
"1. The visiable devices. Environment variable `HABANA_VISIBLE_DEVICES` and `HABANA_VISIBLE_MODULES` are used to control the HPU device visiable to application, you must set this two environment variable properly. For more detail usage of `HABANA_VISIBLE_DEVICES`, `HABANA_VISIBLE_MODULES`, please visit [here](https://docs.habana.ai/en/latest/PyTorch/Reference/PT_Multiple_Tenants_on_HPU/Multiple_Dockers_each_with_Single_Workload.html)\n",
"\n",
"2. The execution mode. Different execution mode has different runtime performance. The default execution mode is lazy mode."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# set some environment variables\n",
"os.environ[\"RAY_EXPERIMENTAL_NOSET_HABANA_VISIBLE_MODULES\"] = \"0\"\n",
"# if using RAY_EXPERIMENTAL_NOSET_HABANA_VISIBLE_MODULES env var\n",
"# you must set HABANA_VISIBLE_MODULES, such as\n",
"# os.environ[\"HABANA_VISIBLE_MODULES\"] = \"0,1,2,3\"\n",
"\n",
"# execution_mode are [\"lazy\", \"eager\", \"eager.compile\"]\n",
"execution_mode = \"lazy\"\n",
"os.environ[\"PT_HPU_LAZY_MODE\"] = \"1\" if execution_mode == \"lazy\" else \"0\"\n",
"\n",
"main(num_workers=4, execution_mode=execution_mode)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Possible outputs\n",
"\n",
"``` text\n",
"\n",
"...\n",
"\n",
"RayTrainWorker pid=36561) Setting up process group for: env:// [rank=0, world_size=4]\n",
"(TorchTrainer pid=36054) Started distributed worker processes: \n",
"(TorchTrainer pid=36054) - (node_id=409da2dba1dc3e5b8e58a2b766a4a19d90e7879c28c2fb13644148b8, ip=100.83.111.228, pid=36561) world_rank=0, local_rank=0, node_rank=0\n",
"(TorchTrainer pid=36054) - (node_id=409da2dba1dc3e5b8e58a2b766a4a19d90e7879c28c2fb13644148b8, ip=100.83.111.228, pid=36562) world_rank=1, local_rank=1, node_rank=0\n",
"(TorchTrainer pid=36054) - (node_id=409da2dba1dc3e5b8e58a2b766a4a19d90e7879c28c2fb13644148b8, ip=100.83.111.228, pid=36563) world_rank=2, local_rank=2, node_rank=0\n",
"(TorchTrainer pid=36054) - (node_id=409da2dba1dc3e5b8e58a2b766a4a19d90e7879c28c2fb13644148b8, ip=100.83.111.228, pid=36564) world_rank=3, local_rank=3, node_rank=0\n",
"\n",
"...\n",
"\n",
"(RayTrainWorker pid=36561) ============================= HABANA PT BRIDGE CONFIGURATION =========================== \n",
"(RayTrainWorker pid=36561) PT_HPU_LAZY_MODE = 1\n",
"(RayTrainWorker pid=36561) PT_HPU_RECIPE_CACHE_CONFIG = ,false,1024\n",
"(RayTrainWorker pid=36561) PT_HPU_MAX_COMPOUND_OP_SIZE = 9223372036854775807\n",
"(RayTrainWorker pid=36561) PT_HPU_LAZY_ACC_PAR_MODE = 0\n",
"(RayTrainWorker pid=36561) PT_HPU_ENABLE_REFINE_DYNAMIC_SHAPES = 0\n",
"(RayTrainWorker pid=36561) PT_HPU_EAGER_PIPELINE_ENABLE = 1\n",
"(RayTrainWorker pid=36561) PT_HPU_EAGER_COLLECTIVE_PIPELINE_ENABLE = 1\n",
"(RayTrainWorker pid=36561) PT_HPU_ENABLE_LAZY_COLLECTIVES = 0\n",
"(RayTrainWorker pid=36561) ---------------------------: System Configuration :---------------------------\n",
"(RayTrainWorker pid=36561) Num CPU Cores : 160\n",
"(RayTrainWorker pid=36561) CPU RAM : 1056374420 KB\n",
"(RayTrainWorker pid=36561) ------------------------------------------------------------------------------\n",
"\n",
"...\n",
"\n",
"(RayTrainWorker pid=36561) {'loss': 4.1052, 'grad_norm': 2.225008249282837, 'learning_rate': 8.26086956521739e-06, 'epoch': 2.5, 'memory_allocated (GB)': 28.87, 'max_memory_allocated (GB)': 94.26, 'total_memory_available (GB)': 94.62}\n",
"(RayTrainWorker pid=36561) {'loss': 4.0472, 'grad_norm': 2.0701019763946533, 'learning_rate': 8.212560386473431e-06, 'epoch': 2.51, 'memory_allocated (GB)': 28.87, 'max_memory_allocated (GB)': 94.26, 'total_memory_available (GB)': 94.62}\n",
"(RayTrainWorker pid=36561) {'loss': 4.097, 'grad_norm': 2.119075059890747, 'learning_rate': 8.164251207729469e-06, 'epoch': 2.51, 'memory_allocated (GB)': 28.87, 'max_memory_allocated (GB)': 94.26, 'total_memory_available (GB)': 94.62}\n",
"(RayTrainWorker pid=36561) {'loss': 3.7035, 'grad_norm': 2.1802899837493896, 'learning_rate': 8.115942028985508e-06, 'epoch': 2.51, 'memory_allocated (GB)': 28.87, 'max_memory_allocated (GB)': 94.26, 'total_memory_available (GB)': 94.62}\n",
"(RayTrainWorker pid=36561) {'loss': 4.242, 'grad_norm': 1.9516953229904175, 'learning_rate': 8.067632850241547e-06, 'epoch': 2.52, 'memory_allocated (GB)': 28.87, 'max_memory_allocated (GB)': 94.26, 'total_memory_available (GB)': 94.62}\n",
"(RayTrainWorker pid=36561) {'loss': 3.9594, 'grad_norm': 2.0580222606658936, 'learning_rate': 8.019323671497584e-06, 'epoch': 2.52, 'memory_allocated (GB)': 28.87, 'max_memory_allocated (GB)': 94.26, 'total_memory_available (GB)': 94.62}\n",
"(RayTrainWorker pid=36561) {'loss': 4.3415, 'grad_norm': 2.192605495452881, 'learning_rate': 7.971014492753623e-06, 'epoch': 2.52, 'memory_allocated (GB)': 28.87, 'max_memory_allocated (GB)': 94.26, 'total_memory_available (GB)': 94.62}\n",
"(RayTrainWorker pid=36561) {'loss': 3.9739, 'grad_norm': 2.0198025703430176, 'learning_rate': 7.922705314009662e-06, 'epoch': 2.52, 'memory_allocated (GB)': 28.87, 'max_memory_allocated (GB)': 94.26, 'total_memory_available (GB)': 94.62}\n",
"(RayTrainWorker pid=36561) {'loss': 4.1624, 'grad_norm': 2.0957565307617188, 'learning_rate': 7.874396135265701e-06, 'epoch': 2.53, 'memory_allocated (GB)': 28.87, 'max_memory_allocated (GB)': 94.26, 'total_memory_available (GB)': 94.62}\n",
"(RayTrainWorker pid=36561) {'loss': 3.9744, 'grad_norm': 2.1159448623657227, 'learning_rate': 7.82608695652174e-06, 'epoch': 2.53, 'memory_allocated (GB)': 28.87, 'max_memory_allocated (GB)': 94.26, 'total_memory_available (GB)': 94.62}\n",
"(RayTrainWorker pid=36561) {'loss': 4.1127, 'grad_norm': 2.159834623336792, 'learning_rate': 7.777777777777777e-06, 'epoch': 2.53, 'memory_allocated (GB)': 28.87, 'max_memory_allocated (GB)': 94.26, 'total_memory_available (GB)': 94.62}\n",
"(RayTrainWorker pid=36561) {'loss': 4.0588, 'grad_norm': 2.106534004211426, 'learning_rate': 7.729468599033817e-06, 'epoch': 2.54, 'memory_allocated (GB)': 28.87, 'max_memory_allocated (GB)': 94.26, 'total_memory_available (GB)': 94.62}\n",
"(RayTrainWorker pid=36561) {'loss': 3.8734, 'grad_norm': 2.445814371109009, 'learning_rate': 7.681159420289856e-06, 'epoch': 2.54, 'memory_allocated (GB)': 28.87, 'max_memory_allocated (GB)': 94.26, 'total_memory_available (GB)': 94.62}\n",
"(RayTrainWorker pid=36561) {'loss': 4.0278, 'grad_norm': 2.0376927852630615, 'learning_rate': 7.632850241545895e-06, 'epoch': 2.54, 'memory_allocated (GB)': 28.87, 'max_memory_allocated (GB)': 94.26, 'total_memory_available (GB)': 94.62}\n",
"(RayTrainWorker pid=36561) {'loss': 3.9643, 'grad_norm': 2.1097891330718994, 'learning_rate': 7.584541062801932e-06, 'epoch': 2.54, 'memory_allocated (GB)': 28.87, 'max_memory_allocated (GB)': 94.26, 'total_memory_available (GB)': 94.62}\n",
"(RayTrainWorker pid=36561) {'loss': 4.1384, 'grad_norm': 2.157325267791748, 'learning_rate': 7.536231884057972e-06, 'epoch': 2.55, 'memory_allocated (GB)': 28.87, 'max_memory_allocated (GB)': 94.26, 'total_memory_available (GB)': 94.62}\n",
"(RayTrainWorker pid=36561) {'loss': 3.9982, 'grad_norm': 2.230065107345581, 'learning_rate': 7.48792270531401e-06, 'epoch': 2.55, 'memory_allocated (GB)': 28.87, 'max_memory_allocated (GB)': 94.26, 'total_memory_available (GB)': 94.62}\n",
"(RayTrainWorker pid=36561) {'loss': 4.0988, 'grad_norm': 2.355875015258789, 'learning_rate': 7.439613526570048e-06, 'epoch': 2.55, 'memory_allocated (GB)': 28.87, 'max_memory_allocated (GB)': 94.26, 'total_memory_available (GB)': 94.62}\n",
"(RayTrainWorker pid=36561) {'loss': 4.0514, 'grad_norm': 2.1178295612335205, 'learning_rate': 7.391304347826088e-06, 'epoch': 2.56, 'memory_allocated (GB)': 28.87, 'max_memory_allocated (GB)': 94.26, 'total_memory_available (GB)': 94.62}\n",
"(RayTrainWorker pid=36561) {'loss': 3.9858, 'grad_norm': 2.089723825454712, 'learning_rate': 7.342995169082126e-06, 'epoch': 2.56, 'memory_allocated (GB)': 28.87, 'max_memory_allocated (GB)': 94.26, 'total_memory_available (GB)': 94.62}\n",
"(RayTrainWorker pid=36561) {'loss': 4.1548, 'grad_norm': 2.308490753173828, 'learning_rate': 7.294685990338164e-06, 'epoch': 2.56, 'memory_allocated (GB)': 28.87, 'max_memory_allocated (GB)': 94.26, 'total_memory_available (GB)': 94.62}\n",
"(RayTrainWorker pid=36561) {'loss': 4.0356, 'grad_norm': 1.9994627237319946, 'learning_rate': 7.246376811594203e-06, 'epoch': 2.57, 'memory_allocated (GB)': 28.87, 'max_memory_allocated (GB)': 94.26, 'total_memory_available (GB)': 94.62}\n",
"(RayTrainWorker pid=36561) {'loss': 3.7696, 'grad_norm': 1.9719663858413696, 'learning_rate': 7.1980676328502416e-06, 'epoch': 2.57, 'memory_allocated (GB)': 28.87, 'max_memory_allocated (GB)': 94.26, 'total_memory_available (GB)': 94.62}\n",
"(RayTrainWorker pid=36561) {'loss': 4.0157, 'grad_norm': 2.1598856449127197, 'learning_rate': 7.1497584541062814e-06, 'epoch': 2.57, 'memory_allocated (GB)': 28.87, 'max_memory_allocated (GB)': 94.26, 'total_memory_available (GB)': 94.62}\n",
"(RayTrainWorker pid=36561) {'loss': 4.0113, 'grad_norm': 1.997869849205017, 'learning_rate': 7.10144927536232e-06, 'epoch': 2.57, 'memory_allocated (GB)': 28.87, 'max_memory_allocated (GB)': 94.26, 'total_memory_available (GB)': 94.62}\n",
"(RayTrainWorker pid=36561) {'loss': 4.1048, 'grad_norm': 2.099222183227539, 'learning_rate': 7.053140096618358e-06, 'epoch': 2.58, 'memory_allocated (GB)': 28.87, 'max_memory_allocated (GB)': 94.26, 'total_memory_available (GB)': 94.62}\n",
"(RayTrainWorker pid=36561) {'loss': 4.0048, 'grad_norm': 2.100231885910034, 'learning_rate': 7.004830917874397e-06, 'epoch': 2.58, 'memory_allocated (GB)': 28.87, 'max_memory_allocated (GB)': 94.26, 'total_memory_available (GB)': 94.62}\n",
"(RayTrainWorker pid=36561) {'loss': 4.0302, 'grad_norm': 2.18204402923584, 'learning_rate': 6.956521739130435e-06, 'epoch': 2.58, 'memory_allocated (GB)': 28.87, 'max_memory_allocated (GB)': 94.26, 'total_memory_available (GB)': 94.62}\n",
"(RayTrainWorker pid=36561) {'loss': 3.7227, 'grad_norm': 2.190962553024292, 'learning_rate': 6.908212560386473e-06, 'epoch': 2.59, 'memory_allocated (GB)': 28.87, 'max_memory_allocated (GB)': 94.26, 'total_memory_available (GB)': 94.62}\n",
"(RayTrainWorker pid=36561) {'loss': 4.1111, 'grad_norm': 2.349518060684204, 'learning_rate': 6.859903381642513e-06, 'epoch': 2.59, 'memory_allocated (GB)': 28.87, 'max_memory_allocated (GB)': 94.26, 'total_memory_available (GB)': 94.62}\n",
"(RayTrainWorker pid=36561) {'loss': 4.024, 'grad_norm': 2.5497331619262695, 'learning_rate': 6.811594202898551e-06, 'epoch': 2.59, 'memory_allocated (GB)': 28.87, 'max_memory_allocated (GB)': 94.26, 'total_memory_available (GB)': 94.62}\n",
"(RayTrainWorker pid=36561) {'loss': 3.8844, 'grad_norm': 2.3125178813934326, 'learning_rate': 6.7632850241545894e-06, 'epoch': 2.59, 'memory_allocated (GB)': 28.87, 'max_memory_allocated (GB)': 94.26, 'total_memory_available (GB)': 94.62}\n",
"(RayTrainWorker pid=36561) {'loss': 4.2208, 'grad_norm': 2.1103923320770264, 'learning_rate': 6.7149758454106285e-06, 'epoch': 2.6, 'memory_allocated (GB)': 28.87, 'max_memory_allocated (GB)': 94.26, 'total_memory_available (GB)': 94.62}\n",
"\n",
"...\n",
"```"
]
}
],
"metadata": {
"language_info": {
"name": "python"
},
"orphan": true
},
"nbformat": 4,
"nbformat_minor": 2
}
doc/source/train/examples/intel_gaudi/resnet.ipynb
+610
View File
@@ -0,0 +1,610 @@
{
"cells": [
{
"cell_type": "markdown",
"metadata": {},
"source": [
"(hpu_resnet_training)=\n",
"# ResNet Model Training with Intel Gaudi\n",
"\n",
"<a id=\"try-anyscale-quickstart-intel_gaudi-resnet\" href=\"https://console.anyscale.com/register/ha?render_flow=ray&utm_source=ray_docs&utm_medium=docs&utm_campaign=intel_gaudi-resnet\">\n",
" <img src=\"../../../_static/img/run-on-anyscale.svg\" alt=\"try-anyscale-quickstart\">\n",
"</a>\n",
"<br></br>\n",
"\n",
"In this Jupyter notebook, we will train a ResNet-50 model to classify images of ants and bees using HPU. We will use PyTorch for model training and Ray for distributed training. The dataset will be downloaded and processed using torchvision's datasets and transforms.\n",
"\n",
"[Intel Gaudi AI Processors (HPUs)](https://habana.ai) are AI hardware accelerators designed by Intel Habana Labs. For more information, see [Gaudi Architecture](https://docs.habana.ai/en/latest/Gaudi_Overview/index.html) and [Gaudi Developer Docs](https://developer.habana.ai/).\n",
"\n",
"## Configuration\n",
"\n",
"A node with Gaudi/Gaudi2 installed is required to run this example. Both Gaudi and Gaudi2 have 8 HPUs. We will use 2 workers to train the model, each using 1 HPU.\n",
"\n",
"We recommend using a prebuilt container to run these examples. To run a container, you need Docker. See [Install Docker Engine](https://docs.docker.com/engine/install/) for installation instructions.\n",
"\n",
"Next, follow [Run Using Containers](https://docs.habana.ai/en/latest/Installation_Guide/Bare_Metal_Fresh_OS.html?highlight=installer#run-using-containers) to install the Gaudi drivers and container runtime.\n",
"\n",
"Next, start the Gaudi container:\n",
"```bash\n",
"docker pull vault.habana.ai/gaudi-docker/1.20.0/ubuntu22.04/habanalabs/pytorch-installer-2.6.0:latest\n",
"docker run -it --runtime=habana -e HABANA_VISIBLE_DEVICES=all -e OMPI_MCA_btl_vader_single_copy_mechanism=none --cap-add=sys_nice --net=host --ipc=host vault.habana.ai/gaudi-docker/1.20.0/ubuntu22.04/habanalabs/pytorch-installer-2.6.0:latest\n",
"```\n",
"\n",
"Inside the container, install Ray and Jupyter to run this notebook.\n",
"```bash\n",
"pip install ray[train] notebook\n",
"```"
]
},
{
"cell_type": "code",
"execution_count": 1,
"metadata": {},
"outputs": [],
"source": [
"import os\n",
"from typing import Dict\n",
"from tempfile import TemporaryDirectory\n",
"\n",
"import torch\n",
"from filelock import FileLock\n",
"from torch import nn\n",
"import torch.optim as optim\n",
"from torch.utils.data import DataLoader\n",
"from torchvision import datasets, transforms, models\n",
"from tqdm import tqdm\n",
"\n",
"import ray\n",
"import ray.train as train\n",
"from ray.train import ScalingConfig, Checkpoint\n",
"from ray.train.torch import TorchTrainer\n",
"from ray.train.torch import TorchConfig\n",
"from ray.runtime_env import RuntimeEnv\n",
"\n",
"import habana_frameworks.torch.core as htcore"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Define Data Transforms\n",
"\n",
"We will set up the data transforms for preprocessing images for training and validation. This includes random cropping, flipping, and normalization for the training set, and resizing and normalization for the validation set."
]
},
{
"cell_type": "code",
"execution_count": 2,
"metadata": {},
"outputs": [],
"source": [
"# Data augmentation and normalization for training\n",
"# Just normalization for validation\n",
"data_transforms = {\n",
" \"train\": transforms.Compose([\n",
" transforms.RandomResizedCrop(224),\n",
" transforms.RandomHorizontalFlip(),\n",
" transforms.ToTensor(),\n",
" transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),\n",
" ]),\n",
" \"val\": transforms.Compose([\n",
" transforms.Resize(256),\n",
" transforms.CenterCrop(224),\n",
" transforms.ToTensor(),\n",
" transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),\n",
" ]),\n",
"}"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Dataset Download Function\n",
"\n",
"We will define a function to download the Hymenoptera dataset. This dataset contains images of ants and bees for a binary classification problem."
]
},
{
"cell_type": "code",
"execution_count": 3,
"metadata": {},
"outputs": [],
"source": [
"def download_datasets():\n",
" os.system(\"wget https://download.pytorch.org/tutorial/hymenoptera_data.zip >/dev/null 2>&1\")\n",
" os.system(\"unzip hymenoptera_data.zip >/dev/null 2>&1\")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Dataset Preparation Function\n",
"\n",
"After downloading the dataset, we need to build PyTorch datasets for training and validation. The `build_datasets` function will apply the previously defined transforms and create the datasets."
]
},
{
"cell_type": "code",
"execution_count": 4,
"metadata": {},
"outputs": [],
"source": [
"def build_datasets():\n",
" torch_datasets = {}\n",
" for split in [\"train\", \"val\"]:\n",
" torch_datasets[split] = datasets.ImageFolder(\n",
" os.path.join(\"./hymenoptera_data\", split), data_transforms[split]\n",
" )\n",
" return torch_datasets"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Model Initialization Functions\n",
"\n",
"We will define two functions to initialize our model. The `initialize_model` function will load a pre-trained ResNet-50 model and replace the final classification layer for our binary classification task. The `initialize_model_from_checkpoint` function will load a model from a saved checkpoint if available."
]
},
{
"cell_type": "code",
"execution_count": 5,
"metadata": {},
"outputs": [],
"source": [
"def initialize_model():\n",
" # Load pretrained model params\n",
" model = models.resnet50(pretrained=True)\n",
"\n",
" # Replace the original classifier with a new Linear layer\n",
" num_features = model.fc.in_features\n",
" model.fc = nn.Linear(num_features, 2)\n",
"\n",
" # Ensure all params get updated during finetuning\n",
" for param in model.parameters():\n",
" param.requires_grad = True\n",
" return model"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Evaluation Function\n",
"\n",
"To assess the performance of our model during training, we define an `evaluate` function. This function computes the number of correct predictions by comparing the predicted labels with the true labels."
]
},
{
"cell_type": "code",
"execution_count": 6,
"metadata": {},
"outputs": [],
"source": [
"def evaluate(logits, labels):\n",
" _, preds = torch.max(logits, 1)\n",
" corrects = torch.sum(preds == labels).item()\n",
" return corrects"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Training Loop Function\n",
"\n",
"This function defines the training loop that will be executed by each worker. It includes downloading the dataset, preparing data loaders, initializing the model, and running the training and validation phases. Compared to a training function for GPU, no changes are needed to port to HPU. Internally, Ray Train does these things:\n",
"\n",
"* Detect HPU and set the device.\n",
"\n",
"* Initializes the habana PyTorch backend.\n",
"\n",
"* Initializes the habana distributed backend."
]
},
{
"cell_type": "code",
"execution_count": 7,
"metadata": {},
"outputs": [],
"source": [
"def train_loop_per_worker(configs):\n",
" import warnings\n",
"\n",
" warnings.filterwarnings(\"ignore\")\n",
"\n",
" # Calculate the batch size for a single worker\n",
" worker_batch_size = configs[\"batch_size\"] // train.get_context().get_world_size()\n",
"\n",
" # Download dataset once on local rank 0 worker\n",
" if train.get_context().get_local_rank() == 0:\n",
" download_datasets()\n",
" torch.distributed.barrier()\n",
"\n",
" # Build datasets on each worker\n",
" torch_datasets = build_datasets()\n",
"\n",
" # Prepare dataloader for each worker\n",
" dataloaders = dict()\n",
" dataloaders[\"train\"] = DataLoader(\n",
" torch_datasets[\"train\"], batch_size=worker_batch_size, shuffle=True\n",
" )\n",
" dataloaders[\"val\"] = DataLoader(\n",
" torch_datasets[\"val\"], batch_size=worker_batch_size, shuffle=False\n",
" )\n",
"\n",
" # Distribute\n",
" dataloaders[\"train\"] = train.torch.prepare_data_loader(dataloaders[\"train\"])\n",
" dataloaders[\"val\"] = train.torch.prepare_data_loader(dataloaders[\"val\"])\n",
"\n",
" # Obtain HPU device automatically\n",
" device = train.torch.get_device()\n",
"\n",
" # Prepare DDP Model, optimizer, and loss function\n",
" model = initialize_model()\n",
" model = model.to(device)\n",
"\n",
" optimizer = optim.SGD(\n",
" model.parameters(), lr=configs[\"lr\"], momentum=configs[\"momentum\"]\n",
" )\n",
" criterion = nn.CrossEntropyLoss()\n",
"\n",
" # Start training loops\n",
" for epoch in range(configs[\"num_epochs\"]):\n",
" # Each epoch has a training and validation phase\n",
" for phase in [\"train\", \"val\"]:\n",
" if phase == \"train\":\n",
" model.train() # Set model to training mode\n",
" else:\n",
" model.eval() # Set model to evaluate mode\n",
"\n",
" running_loss = 0.0\n",
" running_corrects = 0\n",
"\n",
" for inputs, labels in dataloaders[phase]:\n",
" inputs = inputs.to(device)\n",
" labels = labels.to(device)\n",
"\n",
" # zero the parameter gradients\n",
" optimizer.zero_grad()\n",
"\n",
" # forward\n",
" with torch.set_grad_enabled(phase == \"train\"):\n",
" # Get model outputs and calculate loss\n",
" outputs = model(inputs)\n",
" loss = criterion(outputs, labels)\n",
"\n",
" # backward + optimize only if in training phase\n",
" if phase == \"train\":\n",
" loss.backward()\n",
" optimizer.step()\n",
"\n",
" # calculate statistics\n",
" running_loss += loss.item() * inputs.size(0)\n",
" running_corrects += evaluate(outputs, labels)\n",
"\n",
" size = len(torch_datasets[phase]) // train.get_context().get_world_size()\n",
" epoch_loss = running_loss / size\n",
" epoch_acc = running_corrects / size\n",
"\n",
" if train.get_context().get_world_rank() == 0:\n",
" print(\n",
" \"Epoch {}-{} Loss: {:.4f} Acc: {:.4f}\".format(\n",
" epoch, phase, epoch_loss, epoch_acc\n",
" )\n",
" )\n",
"\n",
" # Report metrics and checkpoint every epoch\n",
" if phase == \"val\":\n",
" train.report(\n",
" metrics={\"loss\": epoch_loss, \"acc\": epoch_acc},\n",
" )"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Main Training Function\n",
"\n",
"The `train_resnet` function sets up the distributed training environment using Ray and starts the training process. It specifies the batch size, number of epochs, learning rate, and momentum for the SGD optimizer. To enable training using HPU, we only need to make the following changes:\n",
"* Require an HPU for each worker in ScalingConfig\n",
"* Set backend to \"hccl\" in TorchConfig"
]
},
{
"cell_type": "code",
"execution_count": 8,
"metadata": {},
"outputs": [],
"source": [
"def train_resnet(num_workers=2):\n",
" global_batch_size = 16\n",
"\n",
" train_loop_config = {\n",
" \"input_size\": 224, # Input image size (224 x 224)\n",
" \"batch_size\": 32, # Batch size for training\n",
" \"num_epochs\": 10, # Number of epochs to train for\n",
" \"lr\": 0.001, # Learning Rate\n",
" \"momentum\": 0.9, # SGD optimizer momentum\n",
" }\n",
" # Configure computation resources\n",
" # In ScalingConfig, require an HPU for each worker\n",
" scaling_config = ScalingConfig(num_workers=num_workers, resources_per_worker={\"CPU\": 1, \"HPU\": 1})\n",
" # Set backend to hccl in TorchConfig\n",
" torch_config = TorchConfig(backend = \"hccl\")\n",
" \n",
" ray.init()\n",
" \n",
" # Initialize a Ray TorchTrainer\n",
" trainer = TorchTrainer(\n",
" train_loop_per_worker=train_loop_per_worker,\n",
" train_loop_config=train_loop_config,\n",
" torch_config=torch_config,\n",
" scaling_config=scaling_config,\n",
" )\n",
"\n",
" result = trainer.fit()\n",
" print(f\"Training result: {result}\")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Start Training\n",
"\n",
"Finally, we call the `train_resnet` function to start the training process. You can adjust the number of workers to use. Before running this cell, ensure that Ray is properly set up in your environment to handle distributed training.\n",
"\n",
"Note: the following warning is fine, and is resolved in SynapseAI version 1.14.0+:\n",
"```text\n",
"/usr/local/lib/python3.10/dist-packages/torch/distributed/distributed_c10d.py:252: UserWarning: Device capability of hccl unspecified, assuming `cpu` and `cuda`. Please specify it via the `devices` argument of `register_backend`.\n",
"```"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"scrolled": true
},
"outputs": [],
"source": [
"train_resnet(num_workers=2) "
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Possible outputs\n",
"\n",
"``` text\n",
"2025-03-03 03:32:12,620 INFO worker.py:1841 -- Started a local Ray instance.\n",
"/usr/local/lib/python3.10/dist-packages/ray/tune/impl/tuner_internal.py:125: RayDeprecationWarning: The `RunConfig` class should be imported from `ray.tune` when passing it to the Tuner. Please update your imports. See this issue for more context and migration options: https://github.com/ray-project/ray/issues/49454. Disable these warnings by setting the environment variable: RAY_TRAIN_ENABLE_V2_MIGRATION_WARNINGS=0\n",
" _log_deprecation_warning(\n",
"(RayTrainWorker pid=63669) Setting up process group for: env:// [rank=0, world_size=2]\n",
"(TorchTrainer pid=63280) Started distributed worker processes: \n",
"(TorchTrainer pid=63280) - (node_id=9f2c34ea47fe405f3227e9168aa857f81655a83e95fd6be359fd76db, ip=100.83.111.228, pid=63669) world_rank=0, local_rank=0, node_rank=0\n",
"(TorchTrainer pid=63280) - (node_id=9f2c34ea47fe405f3227e9168aa857f81655a83e95fd6be359fd76db, ip=100.83.111.228, pid=63668) world_rank=1, local_rank=1, node_rank=0\n",
"(RayTrainWorker pid=63669) ============================= HABANA PT BRIDGE CONFIGURATION =========================== \n",
"(RayTrainWorker pid=63669) PT_HPU_LAZY_MODE = 1\n",
"(RayTrainWorker pid=63669) PT_HPU_RECIPE_CACHE_CONFIG = ,false,1024\n",
"(RayTrainWorker pid=63669) PT_HPU_MAX_COMPOUND_OP_SIZE = 9223372036854775807\n",
"(RayTrainWorker pid=63669) PT_HPU_LAZY_ACC_PAR_MODE = 1\n",
"(RayTrainWorker pid=63669) PT_HPU_ENABLE_REFINE_DYNAMIC_SHAPES = 0\n",
"(RayTrainWorker pid=63669) PT_HPU_EAGER_PIPELINE_ENABLE = 1\n",
"(RayTrainWorker pid=63669) PT_HPU_EAGER_COLLECTIVE_PIPELINE_ENABLE = 1\n",
"(RayTrainWorker pid=63669) PT_HPU_ENABLE_LAZY_COLLECTIVES = 0\n",
"(RayTrainWorker pid=63669) ---------------------------: System Configuration :---------------------------\n",
"(RayTrainWorker pid=63669) Num CPU Cores : 160\n",
"(RayTrainWorker pid=63669) CPU RAM : 1056374420 KB\n",
"(RayTrainWorker pid=63669) ------------------------------------------------------------------------------\n",
"(RayTrainWorker pid=63668) Downloading: \"https://download.pytorch.org/models/resnet50-0676ba61.pth\" to /root/.cache/torch/hub/checkpoints/resnet50-0676ba61.pth\n",
" 0%| | 0.00/97.8M [00:00<?, ?B/s]\n",
" 9%|▊ | 8.38M/97.8M [00:00<00:01, 87.7MB/s]\n",
"100%|██████████| 97.8M/97.8M [00:00<00:00, 193MB/s]\n",
"100%|██████████| 97.8M/97.8M [00:00<00:00, 203MB/s]\n",
"\n",
"View detailed results here: /root/ray_results/TorchTrainer_2025-03-03_03-32-15\n",
"To visualize your results with TensorBoard, run: `tensorboard --logdir /tmp/ray/session_2025-03-03_03-32-10_695011_53838/artifacts/2025-03-03_03-32-15/TorchTrainer_2025-03-03_03-32-15/driver_artifacts`\n",
"\n",
"Training started with configuration:\n",
"╭──────────────────────────────────────╮\n",
"│ Training config │\n",
"├──────────────────────────────────────┤\n",
"│ train_loop_config/batch_size 32 │\n",
"│ train_loop_config/input_size 224 │\n",
"│ train_loop_config/lr 0.001 │\n",
"│ train_loop_config/momentum 0.9 │\n",
"│ train_loop_config/num_epochs 10 │\n",
"╰──────────────────────────────────────╯\n",
"(RayTrainWorker pid=63669) Epoch 0-train Loss: 0.6574 Acc: 0.6066\n",
"\n",
"Training finished iteration 1 at 2025-03-03 03:32:45. Total running time: 29s\n",
"╭───────────────────────────────╮\n",
"│ Training result │\n",
"├───────────────────────────────┤\n",
"│ checkpoint_dir_name │\n",
"│ time_this_iter_s 24.684 │\n",
"│ time_total_s 24.684 │\n",
"│ training_iteration 1 │\n",
"│ acc 0.71053 │\n",
"│ loss 0.51455 │\n",
"╰───────────────────────────────╯\n",
"(RayTrainWorker pid=63669) Epoch 0-val Loss: 0.5146 Acc: 0.7105\n",
"(RayTrainWorker pid=63669) Epoch 1-train Loss: 0.5016 Acc: 0.7541\n",
"\n",
"Training finished iteration 2 at 2025-03-03 03:32:46. Total running time: 31s\n",
"╭───────────────────────────────╮\n",
"│ Training result │\n",
"├───────────────────────────────┤\n",
"│ checkpoint_dir_name │\n",
"│ time_this_iter_s 1.39649 │\n",
"│ time_total_s 26.0805 │\n",
"│ training_iteration 2 │\n",
"│ acc 0.93421 │\n",
"│ loss 0.30218 │\n",
"╰───────────────────────────────╯\n",
"(RayTrainWorker pid=63669) Epoch 1-val Loss: 0.3022 Acc: 0.9342\n",
"(RayTrainWorker pid=63669) Epoch 2-train Loss: 0.3130 Acc: 0.9180\n",
"\n",
"Training finished iteration 3 at 2025-03-03 03:32:47. Total running time: 32s\n",
"╭───────────────────────────────╮\n",
"│ Training result │\n",
"├───────────────────────────────┤\n",
"│ checkpoint_dir_name │\n",
"│ time_this_iter_s 1.37042 │\n",
"│ time_total_s 27.4509 │\n",
"│ training_iteration 3 │\n",
"│ acc 0.93421 │\n",
"│ loss 0.22201 │\n",
"╰───────────────────────────────╯\n",
"(RayTrainWorker pid=63669) Epoch 2-val Loss: 0.2220 Acc: 0.9342\n",
"(RayTrainWorker pid=63669) Epoch 3-train Loss: 0.2416 Acc: 0.9262\n",
"\n",
"Training finished iteration 4 at 2025-03-03 03:32:49. Total running time: 34s\n",
"╭───────────────────────────────╮\n",
"│ Training result │\n",
"├───────────────────────────────┤\n",
"│ checkpoint_dir_name │\n",
"│ time_this_iter_s 1.38353 │\n",
"│ time_total_s 28.8345 │\n",
"│ training_iteration 4 │\n",
"│ acc 0.96053 │\n",
"│ loss 0.17815 │\n",
"╰───────────────────────────────╯\n",
"(RayTrainWorker pid=63669) Epoch 3-val Loss: 0.1782 Acc: 0.9605\n",
"(RayTrainWorker pid=63669) Epoch 4-train Loss: 0.1900 Acc: 0.9508\n",
"\n",
"Training finished iteration 5 at 2025-03-03 03:32:50. Total running time: 35s\n",
"╭───────────────────────────────╮\n",
"│ Training result │\n",
"├───────────────────────────────┤\n",
"│ checkpoint_dir_name │\n",
"│ time_this_iter_s 1.37318 │\n",
"│ time_total_s 30.2077 │\n",
"│ training_iteration 5 │\n",
"│ acc 0.93421 │\n",
"│ loss 0.17063 │\n",
"╰───────────────────────────────╯\n",
"(RayTrainWorker pid=63669) Epoch 4-val Loss: 0.1706 Acc: 0.9342\n",
"(RayTrainWorker pid=63669) Epoch 5-train Loss: 0.1346 Acc: 0.9672\n",
"\n",
"Training finished iteration 6 at 2025-03-03 03:32:52. Total running time: 36s\n",
"╭───────────────────────────────╮\n",
"│ Training result │\n",
"├───────────────────────────────┤\n",
"│ checkpoint_dir_name │\n",
"│ time_this_iter_s 1.37999 │\n",
"│ time_total_s 31.5876 │\n",
"│ training_iteration 6 │\n",
"│ acc 0.96053 │\n",
"│ loss 0.1552 │\n",
"╰───────────────────────────────╯\n",
"(RayTrainWorker pid=63669) Epoch 5-val Loss: 0.1552 Acc: 0.9605\n",
"(RayTrainWorker pid=63669) Epoch 6-train Loss: 0.1184 Acc: 0.9672\n",
"\n",
"Training finished iteration 7 at 2025-03-03 03:32:53. Total running time: 38s\n",
"╭───────────────────────────────╮\n",
"│ Training result │\n",
"├───────────────────────────────┤\n",
"│ checkpoint_dir_name │\n",
"│ time_this_iter_s 1.39198 │\n",
"│ time_total_s 32.9796 │\n",
"│ training_iteration 7 │\n",
"│ acc 0.94737 │\n",
"│ loss 0.14702 │\n",
"╰───────────────────────────────╯\n",
"(RayTrainWorker pid=63669) Epoch 6-val Loss: 0.1470 Acc: 0.9474\n",
"(RayTrainWorker pid=63669) Epoch 7-train Loss: 0.0864 Acc: 0.9836\n",
"\n",
"Training finished iteration 8 at 2025-03-03 03:32:54. Total running time: 39s\n",
"╭───────────────────────────────╮\n",
"│ Training result │\n",
"├───────────────────────────────┤\n",
"│ checkpoint_dir_name │\n",
"│ time_this_iter_s 1.3736 │\n",
"│ time_total_s 34.3532 │\n",
"│ training_iteration 8 │\n",
"│ acc 0.94737 │\n",
"│ loss 0.14443 │\n",
"╰───────────────────────────────╯\n",
"(RayTrainWorker pid=63669) Epoch 7-val Loss: 0.1444 Acc: 0.9474\n",
"(RayTrainWorker pid=63669) Epoch 8-train Loss: 0.1085 Acc: 0.9590\n",
"\n",
"Training finished iteration 9 at 2025-03-03 03:32:56. Total running time: 40s\n",
"╭───────────────────────────────╮\n",
"│ Training result │\n",
"├───────────────────────────────┤\n",
"│ checkpoint_dir_name │\n",
"│ time_this_iter_s 1.37868 │\n",
"│ time_total_s 35.7319 │\n",
"│ training_iteration 9 │\n",
"│ acc 0.94737 │\n",
"│ loss 0.14194 │\n",
"╰───────────────────────────────╯\n",
"(RayTrainWorker pid=63669) Epoch 8-val Loss: 0.1419 Acc: 0.9474\n",
"(RayTrainWorker pid=63669) Epoch 9-train Loss: 0.0829 Acc: 0.9754\n",
"\n",
"2025-03-03 03:32:58,628 INFO tune.py:1009 -- Wrote the latest version of all result files and experiment state to '/root/ray_results/TorchTrainer_2025-03-03_03-32-15' in 0.0028s.\n",
"Training finished iteration 10 at 2025-03-03 03:32:57. Total running time: 42s\n",
"╭───────────────────────────────╮\n",
"│ Training result │\n",
"├───────────────────────────────┤\n",
"│ checkpoint_dir_name │\n",
"│ time_this_iter_s 1.36497 │\n",
"│ time_total_s 37.0969 │\n",
"│ training_iteration 10 │\n",
"│ acc 0.96053 │\n",
"│ loss 0.14297 │\n",
"╰───────────────────────────────╯\n",
"(RayTrainWorker pid=63669) Epoch 9-val Loss: 0.1430 Acc: 0.9605\n",
"\n",
"Training completed after 10 iterations at 2025-03-03 03:32:58. Total running time: 43s\n",
"\n",
"Training result: Result(\n",
" metrics={'loss': 0.1429688463869848, 'acc': 0.9605263157894737},\n",
" path='/root/ray_results/TorchTrainer_2025-03-03_03-32-15/TorchTrainer_19fd8_00000_0_2025-03-03_03-32-15',\n",
" filesystem='local',\n",
" checkpoint=None\n",
")\n",
"(RayTrainWorker pid=63669) Downloading: \"https://download.pytorch.org/models/resnet50-0676ba61.pth\" to /root/.cache/torch/hub/checkpoints/resnet50-0676ba61.pth\n",
" 0%| | 0.00/97.8M [00:00<?, ?B/s]\n",
" 68%|██████▊ | 66.1M/97.8M [00:00<00:00, 160MB/s] [repeated 6x across cluster] (Ray deduplicates logs by default. Set RAY_DEDUP_LOGS=0 to disable log deduplication, or see https://docs.ray.io/en/master/ray-observability/user-guides/configure-logging.html#log-deduplication for more options.)\n",
"```"
]
}
],
"metadata": {
"kernelspec": {
"display_name": "Python 3.10.12 64-bit",
"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.12"
},
"orphan": true,
"vscode": {
"interpreter": {
"hash": "31f2aee4e71d21fbe5cf8b01ff0e069b9275f58929596ceb00d14d90e3e16cd6"
}
}
},
"nbformat": 4,
"nbformat_minor": 4
}
doc/source/train/examples/intel_gaudi/sd.ipynb
+494
View File
File diff suppressed because one or more lines are too long
doc/source/train/examples/jax/intro_to_jax_trainer/README.ipynb
+878
View File
@@ -0,0 +1,878 @@
{
"cells": [
{
"cell_type": "markdown",
"id": "0ebcd4f7",
"metadata": {},
"source": [
"# Train a GPT-2 model with Ray Train `JaxTrainer`\n",
"\n",
"**Time to complete**: 15 min\n",
"\n",
"This template shows you how to distribute a JAX/Flax training loop with [Ray Train](https://docs.ray.io/en/latest/train/train.html)'s `JaxTrainer`. You’ll train a small GPT-2-style Transformer from scratch on the [OpenWebText](https://openwebtext2.readthedocs.io/en/latest/) dataset.\n",
"\n",
"Ray Train lets you keep your training code in a normal Python function, then runs that function on a set of Ray workers. `JaxTrainer` handles the orchestration (starting workers, setting up the distributed context, and collecting metrics/checkpoints) so you can focus on the model and the input pipeline.\n",
"\n",
"This tutorial is inspired by Andrej Karpathy’s [nanoGPT](https://github.com/karpathy/nanoGPT/tree/master) and Google’s [Train a GPT-2 model with JAX on TPU for free](https://developers.googleblog.com/en/train-gpt2-model-with-jax-on-tpu/).\n",
"\n",
"In this tutorial, you will:\n",
"\n",
"1. Prepare and the `OpenWebText` dataset and wrap with Ray Data.\n",
"2. Define a basic GPT2 model and train step in Jax/Flax.\n",
"3. Wrap the training loop in a `train_loop_per_worker` function and scale it out using Ray Train `JaxTrainer` with GPUs or TPUs!\n",
"\n"
]
},
{
"cell_type": "markdown",
"id": "a9fe11b6",
"metadata": {},
"source": [
"<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, additional configuration is required. For example, you 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",
"</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": "bf58e5fd",
"metadata": {},
"source": [
"## Step 1: Install dependencies and prepare the dataset\n"
]
},
{
"cell_type": "markdown",
"id": "1704e5cd",
"metadata": {},
"source": [
"First, install the required Python packages.\n",
"\n",
"If you’re running on GPUs, install `jax[cuda]`. If you’re running on Google TPUs, install `jax[tpu]`. For platform-specific requirements, see the [JAX installation guide](https://docs.jax.dev/en/latest/installation.html).\n",
"\n",
"This notebook uses `jax[cuda]` in the examples for simplicity.\n"
]
},
{
"cell_type": "markdown",
"id": "39b71121",
"metadata": {},
"source": [
"Run below if you plan to use GPUs."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "f8a6d4fc",
"metadata": {},
"outputs": [],
"source": [
"%%bash\n",
"pip install pandas numpy jax[cuda] flax tiktoken datasets transformers orbax optax\n"
]
},
{
"cell_type": "markdown",
"id": "d862bb66",
"metadata": {},
"source": [
"Run below if you plan to use TPUs."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "f5bf2907",
"metadata": {},
"outputs": [],
"source": [
"# %%bash\n",
"# pip install pandas numpy jax[tpu] flax tiktoken datasets transformers orbax optax"
]
},
{
"cell_type": "markdown",
"id": "193fafe8",
"metadata": {},
"source": [
"Next, prepare the data that you’ll feed into the training loop.\n",
"\n",
"This notebook uses [OpenWebText](https://openwebtext2.readthedocs.io/en/latest/), an open reproduction of OpenAI’s (private) WebText. The goal of this step is to tokenize the dataset once and write it to disk as two files:\n",
"\n",
"- `train.bin`: training tokens.\n",
"- `val.bin`: validation tokens.\n",
"\n",
"You can use Karpathy’s nanoGPT prep script ([`prepare.py`](https://github.com/karpathy/nanoGPT/blob/master/data/openwebtext/prepare.py)) or download preprocessed data from Kaggle ([OpenWebText GPT-2](https://www.kaggle.com/datasets/windmaple/openwebtext-gpt2)).\n",
"\n",
"If running on the Anyscale workspace, the following code adapts the nanoGPT approach and writes the output to the shared storage path used in an Anyscale workspace (`/mnt/cluster_storage`). If you already have `train.bin` and `val.bin`, you can skip this step.\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "90dcb1ba",
"metadata": {},
"outputs": [],
"source": [
"import os\n",
"from tqdm import tqdm\n",
"import numpy as np\n",
"import tiktoken\n",
"from datasets import load_dataset # huggingface datasets\n",
"\n",
"# number of workers in .map() call\n",
"# good number to use is ~order number of cpu cores // 2\n",
"num_proc = 8\n",
"storage_path = \"/mnt/cluster_storage/openwebtext\"\n",
"\n",
"# number of workers in load_dataset() call\n",
"# best number might be different from num_proc above as it also depends on NW speed.\n",
"# it is better than 1 usually though\n",
"num_proc_load_dataset = num_proc\n",
"\n",
"enc = tiktoken.get_encoding(\"gpt2\")\n",
"\n",
"dataset = load_dataset(\"openwebtext\", num_proc=num_proc_load_dataset)\n",
"\n",
"# owt by default only contains the 'train' split, so create a test split\n",
"split_dataset = dataset[\"train\"].train_test_split(test_size=0.0005, seed=2357, shuffle=True)\n",
"split_dataset['val'] = split_dataset.pop('test') # rename the test split to val\n",
"# we now want to tokenize the dataset. first define the encoding function (gpt2 bpe)\n",
"def process(example):\n",
" ids = enc.encode_ordinary(example['text']) # encode_ordinary ignores any special tokens\n",
" ids.append(enc.eot_token) # add the end of text token, e.g. 50256 for gpt2 bpe\n",
" # note: I think eot should be prepended not appended... hmm. it's called \"eot\" though...\n",
" out = {'ids': ids, 'len': len(ids)}\n",
" return out\n",
"\n",
"# tokenize the dataset\n",
"tokenized = split_dataset.map(\n",
" process,\n",
" remove_columns=['text'],\n",
" desc=\"tokenizing the splits\",\n",
" num_proc=num_proc,\n",
")\n",
"\n",
"# concatenate all the ids in each dataset into one large file we can use for training\n",
"for split, dset in tokenized.items():\n",
" arr_len = np.sum(dset['len'], dtype=np.uint64)\n",
" filename = os.path.join(storage_path, f'{split}.bin')\n",
" dtype = np.uint16 # (can do since enc.max_token_value == 50256 is < 2**16)\n",
" arr = np.memmap(filename, dtype=dtype, mode='w+', shape=(arr_len,))\n",
" total_batches = 1024\n",
"\n",
" idx = 0\n",
" for batch_idx in tqdm(range(total_batches), desc=f'writing {filename}'):\n",
" # Batch together samples for faster write\n",
" batch = dset.shard(num_shards=total_batches, index=batch_idx, contiguous=True).with_format('numpy')\n",
" arr_batch = np.concatenate(batch['ids'])\n",
" # Write into mmap\n",
" arr[idx : idx + len(arr_batch)] = arr_batch\n",
" idx += len(arr_batch)\n",
" arr.flush()\n",
"\n",
"# train.bin is ~18GB, val.bin ~8.8MB\n",
"# train has ~9B tokens \n",
"# val has ~4M tokens"
]
},
{
"cell_type": "markdown",
"id": "d89cd3f4",
"metadata": {},
"source": [
"After running the script, you should have two files in shared storage:\n",
"\n",
"1. Training dataset: `/mnt/cluster_storage/openwebtext/train.bin`\n",
"2. Validation dataset: `/mnt/cluster_storage/openwebtext/val.bin`\n",
"\n"
]
},
{
"cell_type": "markdown",
"id": "72f1caa8",
"metadata": {},
"source": [
"## Step 2: Load the dataset with Ray Data."
]
},
{
"cell_type": "markdown",
"id": "75c698e5",
"metadata": {},
"source": [
"Next, let's load these files with Ray Data. Ray Data can read and preprocess data in parallel, then shard it across workers so each training process streams its own batches. This keeps the input pipeline scalable as you add more GPUs or TPU hosts.\n",
"\n",
"For more details about Ray Data, check out the [Ray Data documentation](https://docs.ray.io/en/latest/data/data.html).\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "e4f234bb",
"metadata": {},
"outputs": [],
"source": [
"import ray\n",
"import ray.data\n",
"\n",
"def make_bin_xy_dataset(\n",
" bin_path: str,\n",
" seqlen: int,\n",
" *,\n",
" # How many sequences to generate per epoch-like pass.\n",
" # You can make this very large and then re-iterate over batches in the\n",
" # training loop.\n",
" num_sequences: int,\n",
" seed: int = 0,\n",
" dtype=np.uint16,\n",
" concurrency: int = 64,\n",
"):\n",
" \"\"\"\n",
" Build a Ray Dataset of (x,y) sequences sampled randomly from a .bin token file.\n",
"\n",
" Produces rows:\n",
" - x: int32[seqlen]\n",
" - y: int32[seqlen]\n",
" \"\"\"\n",
" if not os.path.exists(bin_path):\n",
" raise FileNotFoundError(bin_path)\n",
"\n",
" # Open memmap on driver just to get length.\n",
" data = np.memmap(bin_path, dtype=dtype, mode=\"r\")\n",
" n = int(len(data))\n",
" if n <= seqlen + 1:\n",
" raise ValueError(f\"{bin_path} too small: len={n}, seqlen={seqlen}\")\n",
"\n",
" rng = np.random.default_rng(seed)\n",
" # Each start index uses [i : i+seqlen+1]\n",
" starts = rng.integers(0, n - (seqlen + 1), size=num_sequences, dtype=np.int64)\n",
"\n",
" # Create a dataset of start indices\n",
" ds = ray.data.from_items([{\"i\": int(i)} for i in starts])\n",
"\n",
" def read_xy(batch):\n",
" # Open memmap inside worker process\n",
" mm = np.memmap(bin_path, dtype=dtype, mode=\"r\")\n",
" idx = batch[\"i\"].astype(np.int64, copy=False)\n",
"\n",
" # Allocate fixed arrays\n",
" bs = idx.shape[0]\n",
" x = np.empty((bs, seqlen), dtype=np.int32)\n",
" y = np.empty((bs, seqlen), dtype=np.int32)\n",
"\n",
" # Slice per row (still Python loop, but runs in parallel across Ray workers)\n",
" for j, start in enumerate(idx):\n",
" window = mm[start : start + seqlen + 1].astype(np.int32, copy=False)\n",
" x[j] = window[:-1]\n",
" y[j] = window[1:]\n",
"\n",
" return {\"x\": x, \"y\": y}\n",
"\n",
" # Batch reading for efficiency\n",
" ds = ds.map_batches(\n",
" read_xy,\n",
" batch_format=\"numpy\",\n",
" batch_size=32,\n",
" compute=ray.data.TaskPoolStrategy(size=concurrency),\n",
" zero_copy_batch=True,\n",
" )\n",
"\n",
" return ds\n",
"\n",
"train_ds = make_bin_xy_dataset(\n",
" \"/mnt/cluster_storage/openwebtext/train.bin\",\n",
" seqlen=1024,\n",
" num_sequences=5_000_000,\n",
" seed=2357,\n",
" )\n",
"val_ds = make_bin_xy_dataset(\n",
" \"/mnt/cluster_storage/openwebtext/val.bin\",\n",
" seqlen=1024,\n",
" num_sequences=5_000_000,\n",
" seed=2357,\n",
")"
]
},
{
"cell_type": "markdown",
"id": "87878c88",
"metadata": {},
"source": [
"## Step 3: Define a JAX/Flax GPT-2-style model\n"
]
},
{
"cell_type": "markdown",
"id": "9d3a679f",
"metadata": {},
"source": [
"Now define the model and the core training step with JAX and Flax.\n",
"\n",
"The JAX ecosystem is modular: JAX provides the array programming and compilation primitives, and libraries such as [Flax](https://github.com/google/flax) (neural network building blocks), [Optax](https://github.com/google-deepmind/optax) (optimizers and losses), and [Orbax](https://github.com/google/orbax) (checkpointing) provide higher-level components. You’ll use all three in this tutorial.\n",
"\n",
"In this section, nothing is Ray-specific yet—you’re building a normal single-process JAX/Flax training step that you’ll scale out with `JaxTrainer` in the next section.\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "25f1fde5",
"metadata": {},
"outputs": [],
"source": [
"import os\n",
"import time\n",
"import numpy as np\n",
"from dataclasses import dataclass\n",
"\n",
"import jax\n",
"import jax.numpy as jnp\n",
"from jax.experimental import mesh_utils\n",
"from jax.sharding import Mesh\n",
"\n",
"import flax.nnx as nnx\n",
"import optax\n",
"import orbax.checkpoint as orbax\n",
"\n",
"import tiktoken\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "267cb3bf",
"metadata": {},
"outputs": [],
"source": [
"@dataclass(frozen=True)\n",
"class TrainingConfig:\n",
" # Model config (GPT-2 base model configuration).\n",
" tokenizer = tiktoken.get_encoding(\"gpt2\") # We use gpt2 tokenizer for this tutorial.\n",
" vocab_size = tokenizer.n_vocab\n",
"\n",
" num_transformer_blocks: int = 12\n",
" seqlen: int = 1024\n",
" embed_dim: int = 768\n",
"\n",
" num_heads: int = 12\n",
" dropout_rate: float = 0.1\n",
"\n",
" dtype = jnp.bfloat16 # change to jnp.float32 for older GPUs\n",
" param_dtype = jnp.float32\n",
"\n",
" @property\n",
" def feed_forward_dim(self) -> int:\n",
" return 4 * self.embed_dim\n",
"\n",
" # Optimizer config.\n",
" init_learning_rate: float = 5e-4\n",
" weight_decay: float = 1e-1\n",
"\n",
" # Training loop config.\n",
" global_batch_size: int = 32\n",
" max_steps: int = 10_000\n",
" log_every_n_steps: int = 10\n",
" val_every_n_steps: int = 100\n",
" checkpoint_every_n_steps: int = 100\n",
"\n",
" # Data/config paths.\n",
" openwebtext_root: str = \"/mnt/cluster_storage/openwebtext\""
]
},
{
"cell_type": "markdown",
"id": "df1bdea2",
"metadata": {},
"source": [
"This tutorial provides a GPT-2-style Transformer model implemented with [Flax NNX](https://flax.readthedocs.io/en/v0.8.3/experimental/nnx/index.html). The model code is standard JAX/Flax—Ray Train doesn’t require any special model wrappers.\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "1d9103cb",
"metadata": {},
"outputs": [],
"source": [
"# --- Model Definitions (Safe to keep global) ---\n",
"# Keep dtype settings in one place so the model code doesn't rely on undefined globals.\n",
"dtype = TrainingConfig.dtype\n",
"param_dtype = TrainingConfig.param_dtype\n",
"\n",
"\n",
"def causal_attention_mask(seq_len):\n",
" return jnp.tril(jnp.ones((seq_len, seq_len)))\n",
"\n",
"class TransformerBlock(nnx.Module):\n",
" def __init__(self, embed_dim: int, num_heads: int, ff_dim: int, dropout_rate: float, rngs: nnx.Rngs):\n",
" self.layer_norm1 = nnx.LayerNorm(epsilon=1e-6,\n",
" num_features=embed_dim,\n",
" scale_init=nnx.with_partitioning(nnx.initializers.ones_init(), ('model',)),\n",
" bias_init=nnx.with_partitioning(nnx.initializers.zeros_init(), ('model',)),\n",
" dtype=dtype,\n",
" param_dtype=param_dtype,\n",
" rngs=rngs)\n",
" self.mha = nnx.MultiHeadAttention(num_heads=num_heads,\n",
" in_features=embed_dim,\n",
" kernel_init=nnx.with_partitioning(nnx.initializers.xavier_uniform(), (None, 'model')),\n",
" bias_init=nnx.with_partitioning(nnx.initializers.zeros_init(), ('model',)),\n",
" dtype=dtype,\n",
" param_dtype=param_dtype,\n",
" rngs=rngs)\n",
" self.dropout1 = nnx.Dropout(rate=dropout_rate, rngs=rngs)\n",
" self.layer_norm2 = nnx.LayerNorm(epsilon=1e-6,\n",
" num_features=embed_dim,\n",
" scale_init=nnx.with_partitioning(nnx.initializers.ones_init(), ('model',)),\n",
" bias_init=nnx.with_partitioning(nnx.initializers.zeros_init(), ('model',)),\n",
" dtype=dtype,\n",
" param_dtype=param_dtype,\n",
" rngs=rngs)\n",
" self.linear1 = nnx.Linear(in_features=embed_dim,\n",
" out_features=ff_dim,\n",
" kernel_init=nnx.with_partitioning(nnx.initializers.xavier_uniform(), (None, 'model')),\n",
" bias_init=nnx.with_partitioning(nnx.initializers.zeros_init(), ('model',)),\n",
" dtype=dtype,\n",
" param_dtype=param_dtype,\n",
" rngs=rngs)\n",
" self.linear2 = nnx.Linear(in_features=ff_dim,\n",
" out_features=embed_dim,\n",
" kernel_init=nnx.with_partitioning(nnx.initializers.xavier_uniform(), (None, 'model')),\n",
" bias_init=nnx.with_partitioning(nnx.initializers.zeros_init(), ('model',)),\n",
" dtype=dtype,\n",
" param_dtype=param_dtype,\n",
" rngs=rngs)\n",
" self.dropout2 = nnx.Dropout(rate=dropout_rate, rngs=rngs)\n",
"\n",
" def __call__(self, inputs, training: bool = False):\n",
" input_shape = inputs.shape\n",
" bs, seq_len, emb_sz = input_shape\n",
" attention_output = self.mha(\n",
" inputs_q=self.layer_norm1(inputs),\n",
" mask=causal_attention_mask(seq_len),\n",
" decode=False,\n",
" )\n",
" x = inputs + self.dropout1(attention_output, deterministic=not training)\n",
" mlp_output = self.linear1(self.layer_norm2(x))\n",
" mlp_output = nnx.gelu(mlp_output)\n",
" mlp_output = self.linear2(mlp_output)\n",
" mlp_output = self.dropout2(mlp_output, deterministic=not training)\n",
" return x + mlp_output\n",
"\n",
"class TokenAndPositionEmbedding(nnx.Module):\n",
" def __init__(self, seqlen: int, vocab_size: int, embed_dim: int, rngs: nnx.Rngs):\n",
" self.token_emb = nnx.Embed(num_embeddings=vocab_size, features=embed_dim, dtype=dtype, param_dtype=param_dtype, rngs=rngs)\n",
" self.pos_emb = nnx.Embed(num_embeddings=seqlen, features=embed_dim, dtype=dtype, param_dtype=param_dtype, rngs=rngs)\n",
"\n",
" def __call__(self, x):\n",
" positions = jnp.arange(0, x.shape[1])[None, :]\n",
" position_embedding = self.pos_emb(positions)\n",
" token_embedding = self.token_emb(x)\n",
" return self.token_emb, token_embedding+position_embedding\n",
"\n",
"class GPT2(nnx.Module):\n",
" def __init__(\n",
" self,\n",
" seqlen: int,\n",
" vocab_size: int,\n",
" embed_dim: int,\n",
" num_heads: int,\n",
" dropout_rate: float,\n",
" feed_forward_dim: int,\n",
" num_transformer_blocks: int,\n",
" rngs: nnx.Rngs,\n",
" ):\n",
" self.embedding_layer = TokenAndPositionEmbedding(seqlen, vocab_size, embed_dim, rngs=rngs)\n",
" self.dropout = nnx.Dropout(rate=dropout_rate, rngs=rngs)\n",
" self.transformer_blocks = nnx.List([\n",
" TransformerBlock(embed_dim, num_heads, feed_forward_dim, dropout_rate, rngs=rngs)\n",
" for _ in range(num_transformer_blocks)\n",
" ])\n",
" self.layer_norm = nnx.LayerNorm(\n",
" epsilon=1e-6,\n",
" num_features=embed_dim,\n",
" scale_init=nnx.with_partitioning(nnx.initializers.ones_init(), (\"model\",)),\n",
" bias_init=nnx.with_partitioning(nnx.initializers.zeros_init(), (\"model\",)),\n",
" dtype=dtype,\n",
" param_dtype=param_dtype,\n",
" rngs=rngs,\n",
" )\n",
"\n",
" def __call__(self, inputs, training: bool = False):\n",
" token_embedding, x = self.embedding_layer(inputs)\n",
" x = self.dropout(x, deterministic=not training)\n",
" for transformer_block in self.transformer_blocks:\n",
" x = transformer_block(x, training=training)\n",
" x = self.layer_norm(x)\n",
" outputs = token_embedding.attend(x)\n",
" return outputs\n",
"\n",
"\n",
"def create_model(*, rngs: nnx.Rngs, config: TrainingConfig):\n",
" return GPT2(\n",
" seqlen=config.seqlen,\n",
" vocab_size=config.vocab_size,\n",
" embed_dim=config.embed_dim,\n",
" num_heads=config.num_heads,\n",
" dropout_rate=config.dropout_rate,\n",
" feed_forward_dim=config.feed_forward_dim,\n",
" num_transformer_blocks=config.num_transformer_blocks,\n",
" rngs=rngs,\n",
" )"
]
},
{
"cell_type": "markdown",
"id": "3d453706",
"metadata": {},
"source": [
"Next, define the `loss_fn_train`, `loss_fn_eval`, and `train_step` functions.\n",
"\n",
"`train_step()` computes the loss, takes gradients, and updates model parameters through the optimizer. The training loop will call this function repeatedly.\n",
"\n",
"For performance, this notebook JIT-compiles these functions with `@nnx.jit`.\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "a819d3ba",
"metadata": {},
"outputs": [],
"source": [
"\n",
"@nnx.jit\n",
"def loss_fn_train(model, batch):\n",
" logits = model(batch[0], training=True)\n",
" loss = optax.softmax_cross_entropy_with_integer_labels(\n",
" logits=logits, labels=batch[1]\n",
" ).mean()\n",
" return loss, logits\n",
"\n",
"\n",
"@nnx.jit\n",
"def loss_fn_eval(model, batch):\n",
" logits = model(batch[0], training=False)\n",
" loss = optax.softmax_cross_entropy_with_integer_labels(\n",
" logits=logits, labels=batch[1]\n",
" ).mean()\n",
" return loss, logits\n",
"\n",
"\n",
"@nnx.jit\n",
"def train_step(model, optimizer, metrics, batch):\n",
" grad_fn = nnx.value_and_grad(loss_fn_train, has_aux=True)\n",
" (loss, logits), grads = grad_fn(model, batch)\n",
" metrics.update(loss=loss, logits=logits, labels=batch[1])\n",
" optimizer.update(model, grads)\n",
" return loss\n"
]
},
{
"cell_type": "markdown",
"id": "d9cccbed",
"metadata": {},
"source": [
"Next, let's wrap the JAX training logic in a `train_loop_per_worker` function.\n",
"\n",
"Each Ray Train worker runs the same Python function with a different world rank, and Ray sets device visibility per worker (for example, one GPU per worker). Inside this function, you can:\n",
"\n",
"- Read the distributed context (`world_rank`, `world_size`).\n",
"- Get the per-worker dataset shard (`train.get_dataset_shard(...)`) to stream batches.\n",
"- Stream sharded JAX Arrays directly to training devices using [`Dataset.iter_jax_batches`](https://docs.ray.io/en/latest/data/api/doc/ray.data.Dataset.iter_jax_batches.html), which handles device sharding natively.\n",
"- Report metrics and checkpoints back to the trainer with `ray.train.report(...)`."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "f4543c82",
"metadata": {},
"outputs": [],
"source": [
"\n",
"import ray\n",
"from ray import train\n",
"from ray.train import Checkpoint"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "57ea38fa",
"metadata": {},
"outputs": [],
"source": [
"def train_loop_per_worker(config_dict: dict) -> None:\n",
"\n",
" config = TrainingConfig(**config_dict)\n",
"\n",
" world_rank = ray.train.get_context().get_world_rank()\n",
" world_size = ray.train.get_context().get_world_size()\n",
" print(f\"Worker rank {world_rank}/{world_size} sees devices: {jax.devices()}\")\n",
"\n",
"\n",
" # Create a mesh per worker process. \n",
" device_mesh = mesh_utils.create_device_mesh((jax.process_count(), 1))\n",
" mesh = Mesh(device_mesh, axis_names=(\"data\", \"model\"))\n",
" jax.set_mesh(mesh)\n",
"\n",
"\n",
" # Initialize the model locally.\n",
" # Include a dropout RNG stream so nnx.Dropout can run when training=True.\n",
" model = create_model(rngs=nnx.Rngs(params=0, dropout=1), config=config)\n",
"\n",
"\n",
" # We use Ray data to load the training and validation datasets.\n",
" train_it = ray.train.get_dataset_shard(\"train\")\n",
" val_it = ray.train.get_dataset_shard(\"val\")\n",
" if train_it is None or val_it is None:\n",
" raise RuntimeError(\"No Ray Train datasets provided. Pass datasets={...} to JaxTrainer.\")\n",
" \n",
" local_batch_size = config.global_batch_size // jax.process_count()\n",
" \n",
" # A generator helper that yields JAX Array batches indefinitely.\n",
" # This removes the need for try/except StopIteration blocks in the training loop.\n",
" def batch_generator(dataset_shard):\n",
" while True:\n",
" # iter_jax_batches yields globally sharded JAX Array batches directly,\n",
" # mapped onto all process JAX-addressable training devices.\n",
" yield from dataset_shard.iter_jax_batches(\n",
" batch_size=local_batch_size,\n",
" prefetch_batches=2,\n",
" drop_last=True,\n",
" )\n",
"\n",
" train_batches = batch_generator(train_it)\n",
" val_batches = batch_generator(val_it)\n",
"\n",
" # Initialize the optimizer.\n",
" schedule = optax.cosine_decay_schedule(\n",
" init_value=config.init_learning_rate,\n",
" decay_steps=config.max_steps,\n",
" )\n",
" optax_chain = optax.chain(optax.adamw(learning_rate=schedule, weight_decay=config.weight_decay))\n",
" optimizer = nnx.Optimizer(model, optax_chain, wrt=nnx.Param)\n",
"\n",
" checkpointer = orbax.PyTreeCheckpointer()\n",
" start_time = time.time()\n",
"\n",
" train_metrics = nnx.metrics.Average('loss')\n",
" val_metrics = nnx.metrics.Average('val_loss')\n",
"\n",
" for step in range(config.max_steps):\n",
" batch = next(train_batches)\n",
" global_x, global_y = batch[\"x\"], batch[\"y\"]\n",
"\n",
" train_loss = train_step(model, optimizer, train_metrics, (global_x, global_y))\n",
"\n",
" if (step + 1) % config.log_every_n_steps == 0:\n",
" elapsed = time.time() - start_time\n",
" # Report metrics through Ray Train.\n",
" ray.train.report({\"step\": step + 1, \"train_loss\": float(train_loss), \"elapsed_s\": elapsed})\n",
" start_time = time.time()\n",
" \n",
" if (step + 1) % config.val_every_n_steps == 0:\n",
" validation_batch = next(val_batches)\n",
" global_val_input, global_val_target = validation_batch[\"x\"], validation_batch[\"y\"]\n",
" \n",
" loss, logits = loss_fn_eval(model, (global_val_input, global_val_target))\n",
" val_metrics.update(val_loss=loss, logits=logits)\n",
" val_loss = float(val_metrics.compute())\n",
" metrics = {\"step\": step + 1, \"train_loss\": float(train_loss),\"val_loss\": float(val_loss)}\n",
" \n",
" checkpoint = None\n",
" if (step + 1) % config.checkpoint_every_n_steps == 0:\n",
" \n",
" # Orbax checkpointing is a barrier.\n",
" train_state = nnx.to_pure_dict(nnx.state(model))\n",
" checkpoint_path = os.path.join(\"/mnt/cluster_storage/checkpoint/jax_gpt2_ray_data\", str(step + 1))\n",
" checkpointer.save(checkpoint_path, train_state)\n",
" \n",
" # Save a checkpoint and report validation metrics through Ray Train.\n",
" # The controller persists the checkpoint to the RunConfig storage path.\n",
" checkpoint = Checkpoint.from_directory(checkpoint_path) \n",
" if world_rank == 0:\n",
" train.report(metrics, checkpoint=checkpoint)\n",
" else:\n",
" train.report(metrics, checkpoint=None)"
]
},
{
"cell_type": "markdown",
"id": "c9b59b73",
"metadata": {},
"source": [
"> [!NOTE]\n",
"> [`iter_jax_batches`](https://docs.ray.io/en/latest/data/api/doc/ray.data.Dataset.iter_jax_batches.html) uses an internal 1D mesh for sharding. If your training loop uses a complex multi-dimensional mesh, JAX might perform implicit resharding. Ensure your mesh aligns with the device ordering to minimize overhead.\n"
]
},
{
"cell_type": "markdown",
"id": "36f206e9",
"metadata": {},
"source": [
"## Step 5: Define the `ScalingConfig`\n",
"\n",
"Let's define the `ScalingConfig` that we want to scale the training process. \n",
"\n",
"`JaxTrainer` now supports both GPU training and TPU training. \n",
"\n",
"For a walkthrough on configuring `ScalingConfig`, see [Get Started with Distributed Training using JAX](https://docs.ray.io/en/latest/train/getting-started-jax.html). "
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "260fbd24",
"metadata": {},
"outputs": [],
"source": [
"from ray.train import ScalingConfig"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "f1a1d22d",
"metadata": {},
"outputs": [],
"source": [
"# In this example, we use 2 GPUs.\n",
"scaling_config = ScalingConfig(\n",
" use_gpu=True,\n",
" num_workers=2, # Change this to match on your GPU cluster setting, by default, each worker uses one GPU.\n",
")\n",
"# If you plan to use TPUs, see an example below.\n",
"# This ScalingConfig requires a KubeRay cluster configured for a TPU v6e 4x4 slice with 4 TPU VMs.\n",
"# For more information about TPU clusters with Ray on Kubernetes, see the\n",
"# KubeRay TPU guide: https://docs.ray.io/en/master/cluster/kubernetes/user-guides/tpu.html#kuberay-tpu\n",
"# scaling_config = ScalingConfig(\n",
"# use_tpu=True,\n",
"# num_workers=4,\n",
"# topology=\"4x4\",\n",
"# accelerator_type=\"TPU-V6E\",\n",
"# resources_per_worker={\"TPU\": 4},\n",
"# )"
]
},
{
"cell_type": "markdown",
"id": "dadca9ae",
"metadata": {},
"source": [
"## Step 6: Launch with `JaxTrainer`\n",
"\n",
"To run `train_loop_per_worker` on a Ray cluster, you'll construct a `JaxTrainer` with:\n",
"\n",
"- `train_loop_per_worker`: the training function you've defined earlier. Each Ray Train worker runs this function.\n",
"- `train_loop_config`: a hyperparameter dictionary passed into the function.\n",
"- `scaling_config`: the `ScalingConfig` you've defined earlier.\n",
"- `datasets`: the Ray Datasets to ingest for training. Datasets are keyed by name (`{name: dataset}`). Each dataset can be accessed from within the `train_loop_per_worker` by calling `ray.train.get_dataset_shard(name)`. Sharding and additional configuration can be done by passing in a `dataset_config`.\n",
"- `run_config`: runtime configuration including where to write outputs such as checkpoints.\n",
"\n",
"If your workers hit CUDA or XLA library load errors, clear `LD_LIBRARY_PATH` in the runtime env to avoid picking up incompatible system libraries.\n",
"\n",
"`trainer.fit` spawns a controller process to orchestrate the training run and worker processes to execute the JAX training code."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "42b00efd",
"metadata": {},
"outputs": [],
"source": [
"from ray.train import RunConfig\n",
"from ray.train.v2.jax import JaxTrainer"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "fc12fdfb",
"metadata": {},
"outputs": [],
"source": [
"storage_path = \"/mnt/cluster_storage\"\n",
"\n",
"trainer = JaxTrainer(\n",
" train_loop_per_worker=train_loop_per_worker,\n",
" train_loop_config={\n",
" \"global_batch_size\": 32,\n",
" },\n",
" scaling_config=scaling_config,\n",
" run_config=RunConfig(\n",
" name=\"jax_gpt2\",\n",
" storage_path=storage_path,\n",
" # Make sure to unset ``LD_LIBRARY_PATH`` if you're using CUDA devices, \n",
" # since ``LD_LIBRARY_PATH`` can override the CUDA libraries.\n",
" worker_runtime_env={\"env_vars\": {\"LD_LIBRARY_PATH\": \"\"}},\n",
" ),\n",
" datasets={\"train\": train_ds, \"val\": val_ds},\n",
")\n",
"\n",
"result = trainer.fit()\n",
"print(result)\n",
"\n"
]
},
{
"cell_type": "markdown",
"id": "4a0d85e1",
"metadata": {},
"source": [
"## Summary\n",
"\n",
"In this notebook, you:\n",
"\n",
"1. Prepared and the `OpenWebText` dataset and wrapped with Ray Data.\n",
"2. Defined a basic GPT2 model and train step in Jax/Flax.\n",
"3. Wrapped the training loop in a `train_loop_per_worker` function and scaled it out using Ray Train `JaxTrainer` with GPUs or TPUs!"
]
},
{
"cell_type": "markdown",
"id": "c37928bd",
"metadata": {},
"source": []
}
],
"metadata": {
"language_info": {
"name": "python"
},
"orphan": true
},
"nbformat": 4,
"nbformat_minor": 5
}
doc/source/train/examples/jax/intro_to_jax_trainer/README.md
+684
View File
@@ -0,0 +1,684 @@
# Train a GPT-2 model with Ray Train `JaxTrainer`
**Time to complete**: 15 min
This template shows you how to distribute a JAX/Flax training loop with [Ray Train](https://docs.ray.io/en/latest/train/train.html)'s `JaxTrainer`. You’ll train a small GPT-2-style Transformer from scratch on the [OpenWebText](https://openwebtext2.readthedocs.io/en/latest/) dataset.
Ray Train lets you keep your training code in a normal Python function, then runs that function on a set of Ray workers. `JaxTrainer` handles the orchestration (starting workers, setting up the distributed context, and collecting metrics/checkpoints) so you can focus on the model and the input pipeline.
This tutorial is inspired by Andrej Karpathy’s [nanoGPT](https://github.com/karpathy/nanoGPT/tree/master) and Google’s [Train a GPT-2 model with JAX on TPU for free](https://developers.googleblog.com/en/train-gpt2-model-with-jax-on-tpu/).
In this tutorial, you will:
1. Prepare and the `OpenWebText` dataset and wrap with Ray Data.
2. Define a basic GPT2 model and train step in Jax/Flax.
3. Wrap the training loop in a `train_loop_per_worker` function and scale it out using Ray Train `JaxTrainer` with GPUs or TPUs!
<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, additional configuration is required. For example, you 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>
</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>
## Step 1: Install dependencies and prepare the dataset
First, install the required Python packages.
If you’re running on GPUs, install `jax[cuda]`. If you’re running on Google TPUs, install `jax[tpu]`. For platform-specific requirements, see the [JAX installation guide](https://docs.jax.dev/en/latest/installation.html).
This notebook uses `jax[cuda]` in the examples for simplicity.
Run below if you plan to use GPUs.
```bash
%%bash
pip install pandas numpy jax[cuda] flax tiktoken datasets transformers orbax optax
```
Run below if you plan to use TPUs.
```python
# %%bash
# pip install pandas numpy jax[tpu] flax tiktoken datasets transformers orbax optax
```
Next, prepare the data that you’ll feed into the training loop.
This notebook uses [OpenWebText](https://openwebtext2.readthedocs.io/en/latest/), an open reproduction of OpenAI’s (private) WebText. The goal of this step is to tokenize the dataset once and write it to disk as two files:
- `train.bin`: training tokens.
- `val.bin`: validation tokens.
You can use Karpathy’s nanoGPT prep script ([`prepare.py`](https://github.com/karpathy/nanoGPT/blob/master/data/openwebtext/prepare.py)) or download preprocessed data from Kaggle ([OpenWebText GPT-2](https://www.kaggle.com/datasets/windmaple/openwebtext-gpt2)).
If running on the Anyscale workspace, the following code adapts the nanoGPT approach and writes the output to the shared storage path used in an Anyscale workspace (`/mnt/cluster_storage`). If you already have `train.bin` and `val.bin`, you can skip this step.
```python
import os
from tqdm import tqdm
import numpy as np
import tiktoken
from datasets import load_dataset # huggingface datasets
# number of workers in .map() call
# good number to use is ~order number of cpu cores // 2
num_proc = 8
storage_path = "/mnt/cluster_storage/openwebtext"
# number of workers in load_dataset() call
# best number might be different from num_proc above as it also depends on NW speed.
# it is better than 1 usually though
num_proc_load_dataset = num_proc
enc = tiktoken.get_encoding("gpt2")
dataset = load_dataset("openwebtext", num_proc=num_proc_load_dataset)
# owt by default only contains the 'train' split, so create a test split
split_dataset = dataset["train"].train_test_split(test_size=0.0005, seed=2357, shuffle=True)
split_dataset['val'] = split_dataset.pop('test') # rename the test split to val
# we now want to tokenize the dataset. first define the encoding function (gpt2 bpe)
def process(example):
ids = enc.encode_ordinary(example['text']) # encode_ordinary ignores any special tokens
ids.append(enc.eot_token) # add the end of text token, e.g. 50256 for gpt2 bpe
# note: I think eot should be prepended not appended... hmm. it's called "eot" though...
out = {'ids': ids, 'len': len(ids)}
return out
# tokenize the dataset
tokenized = split_dataset.map(
process,
remove_columns=['text'],
desc="tokenizing the splits",
num_proc=num_proc,
)
# concatenate all the ids in each dataset into one large file we can use for training
for split, dset in tokenized.items():
arr_len = np.sum(dset['len'], dtype=np.uint64)
filename = os.path.join(storage_path, f'{split}.bin')
dtype = np.uint16 # (can do since enc.max_token_value == 50256 is < 2**16)
arr = np.memmap(filename, dtype=dtype, mode='w+', shape=(arr_len,))
total_batches = 1024
idx = 0
for batch_idx in tqdm(range(total_batches), desc=f'writing {filename}'):
# Batch together samples for faster write
batch = dset.shard(num_shards=total_batches, index=batch_idx, contiguous=True).with_format('numpy')
arr_batch = np.concatenate(batch['ids'])
# Write into mmap
arr[idx : idx + len(arr_batch)] = arr_batch
idx += len(arr_batch)
arr.flush()
# train.bin is ~18GB, val.bin ~8.8MB
# train has ~9B tokens
# val has ~4M tokens
```
After running the script, you should have two files in shared storage:
1. Training dataset: `/mnt/cluster_storage/openwebtext/train.bin`
2. Validation dataset: `/mnt/cluster_storage/openwebtext/val.bin`
## Step 2: Load the dataset with Ray Data.
Next, let's load these files with Ray Data. Ray Data can read and preprocess data in parallel, then shard it across workers so each training process streams its own batches. This keeps the input pipeline scalable as you add more GPUs or TPU hosts.
For more details about Ray Data, check out the [Ray Data documentation](https://docs.ray.io/en/latest/data/data.html).
```python
import ray
import ray.data
def make_bin_xy_dataset(
bin_path: str,
seqlen: int,
*,
# How many sequences to generate per epoch-like pass.
# You can make this very large and then re-iterate over batches in the
# training loop.
num_sequences: int,
seed: int = 0,
dtype=np.uint16,
concurrency: int = 64,
):
"""
Build a Ray Dataset of (x,y) sequences sampled randomly from a .bin token file.
Produces rows:
- x: int32[seqlen]
- y: int32[seqlen]
"""
if not os.path.exists(bin_path):
raise FileNotFoundError(bin_path)
# Open memmap on driver just to get length.
data = np.memmap(bin_path, dtype=dtype, mode="r")
n = int(len(data))
if n <= seqlen + 1:
raise ValueError(f"{bin_path} too small: len={n}, seqlen={seqlen}")
rng = np.random.default_rng(seed)
# Each start index uses [i : i+seqlen+1]
starts = rng.integers(0, n - (seqlen + 1), size=num_sequences, dtype=np.int64)
# Create a dataset of start indices
ds = ray.data.from_items([{"i": int(i)} for i in starts])
def read_xy(batch):
# Open memmap inside worker process
mm = np.memmap(bin_path, dtype=dtype, mode="r")
idx = batch["i"].astype(np.int64, copy=False)
# Allocate fixed arrays
bs = idx.shape[0]
x = np.empty((bs, seqlen), dtype=np.int32)
y = np.empty((bs, seqlen), dtype=np.int32)
# Slice per row (still Python loop, but runs in parallel across Ray workers)
for j, start in enumerate(idx):
window = mm[start : start + seqlen + 1].astype(np.int32, copy=False)
x[j] = window[:-1]
y[j] = window[1:]
return {"x": x, "y": y}
# Batch reading for efficiency
ds = ds.map_batches(
read_xy,
batch_format="numpy",
batch_size=32,
compute=ray.data.TaskPoolStrategy(size=concurrency),
zero_copy_batch=True,
)
return ds
train_ds = make_bin_xy_dataset(
"/mnt/cluster_storage/openwebtext/train.bin",
seqlen=1024,
num_sequences=5_000_000,
seed=2357,
)
val_ds = make_bin_xy_dataset(
"/mnt/cluster_storage/openwebtext/val.bin",
seqlen=1024,
num_sequences=5_000_000,
seed=2357,
)
```
## Step 3: Define a JAX/Flax GPT-2-style model
Now define the model and the core training step with JAX and Flax.
The JAX ecosystem is modular: JAX provides the array programming and compilation primitives, and libraries such as [Flax](https://github.com/google/flax) (neural network building blocks), [Optax](https://github.com/google-deepmind/optax) (optimizers and losses), and [Orbax](https://github.com/google/orbax) (checkpointing) provide higher-level components. You’ll use all three in this tutorial.
In this section, nothing is Ray-specific yet—you’re building a normal single-process JAX/Flax training step that you’ll scale out with `JaxTrainer` in the next section.
```python
import os
import time
import numpy as np
from dataclasses import dataclass
import jax
import jax.numpy as jnp
from jax.experimental import mesh_utils
from jax.sharding import Mesh
import flax.nnx as nnx
import optax
import orbax.checkpoint as orbax
import tiktoken
```
```python
@dataclass(frozen=True)
class TrainingConfig:
# Model config (GPT-2 base model configuration).
tokenizer = tiktoken.get_encoding("gpt2") # We use gpt2 tokenizer for this tutorial.
vocab_size = tokenizer.n_vocab
num_transformer_blocks: int = 12
seqlen: int = 1024
embed_dim: int = 768
num_heads: int = 12
dropout_rate: float = 0.1
dtype = jnp.bfloat16 # change to jnp.float32 for older GPUs
param_dtype = jnp.float32
@property
def feed_forward_dim(self) -> int:
return 4 * self.embed_dim
# Optimizer config.
init_learning_rate: float = 5e-4
weight_decay: float = 1e-1
# Training loop config.
global_batch_size: int = 32
max_steps: int = 10_000
log_every_n_steps: int = 10
val_every_n_steps: int = 100
checkpoint_every_n_steps: int = 100
# Data/config paths.
openwebtext_root: str = "/mnt/cluster_storage/openwebtext"
```
This tutorial provides a GPT-2-style Transformer model implemented with [Flax NNX](https://flax.readthedocs.io/en/v0.8.3/experimental/nnx/index.html). The model code is standard JAX/Flax—Ray Train doesn’t require any special model wrappers.
```python
# --- Model Definitions (Safe to keep global) ---
# Keep dtype settings in one place so the model code doesn't rely on undefined globals.
dtype = TrainingConfig.dtype
param_dtype = TrainingConfig.param_dtype
def causal_attention_mask(seq_len):
return jnp.tril(jnp.ones((seq_len, seq_len)))
class TransformerBlock(nnx.Module):
def __init__(self, embed_dim: int, num_heads: int, ff_dim: int, dropout_rate: float, rngs: nnx.Rngs):
self.layer_norm1 = nnx.LayerNorm(epsilon=1e-6,
num_features=embed_dim,
scale_init=nnx.with_partitioning(nnx.initializers.ones_init(), ('model',)),
bias_init=nnx.with_partitioning(nnx.initializers.zeros_init(), ('model',)),
dtype=dtype,
param_dtype=param_dtype,
rngs=rngs)
self.mha = nnx.MultiHeadAttention(num_heads=num_heads,
in_features=embed_dim,
kernel_init=nnx.with_partitioning(nnx.initializers.xavier_uniform(), (None, 'model')),
bias_init=nnx.with_partitioning(nnx.initializers.zeros_init(), ('model',)),
dtype=dtype,
param_dtype=param_dtype,
rngs=rngs)
self.dropout1 = nnx.Dropout(rate=dropout_rate, rngs=rngs)
self.layer_norm2 = nnx.LayerNorm(epsilon=1e-6,
num_features=embed_dim,
scale_init=nnx.with_partitioning(nnx.initializers.ones_init(), ('model',)),
bias_init=nnx.with_partitioning(nnx.initializers.zeros_init(), ('model',)),
dtype=dtype,
param_dtype=param_dtype,
rngs=rngs)
self.linear1 = nnx.Linear(in_features=embed_dim,
out_features=ff_dim,
kernel_init=nnx.with_partitioning(nnx.initializers.xavier_uniform(), (None, 'model')),
bias_init=nnx.with_partitioning(nnx.initializers.zeros_init(), ('model',)),
dtype=dtype,
param_dtype=param_dtype,
rngs=rngs)
self.linear2 = nnx.Linear(in_features=ff_dim,
out_features=embed_dim,
kernel_init=nnx.with_partitioning(nnx.initializers.xavier_uniform(), (None, 'model')),
bias_init=nnx.with_partitioning(nnx.initializers.zeros_init(), ('model',)),
dtype=dtype,
param_dtype=param_dtype,
rngs=rngs)
self.dropout2 = nnx.Dropout(rate=dropout_rate, rngs=rngs)
def __call__(self, inputs, training: bool = False):
input_shape = inputs.shape
bs, seq_len, emb_sz = input_shape
attention_output = self.mha(
inputs_q=self.layer_norm1(inputs),
mask=causal_attention_mask(seq_len),
decode=False,
)
x = inputs + self.dropout1(attention_output, deterministic=not training)
mlp_output = self.linear1(self.layer_norm2(x))
mlp_output = nnx.gelu(mlp_output)
mlp_output = self.linear2(mlp_output)
mlp_output = self.dropout2(mlp_output, deterministic=not training)
return x + mlp_output
class TokenAndPositionEmbedding(nnx.Module):
def __init__(self, seqlen: int, vocab_size: int, embed_dim: int, rngs: nnx.Rngs):
self.token_emb = nnx.Embed(num_embeddings=vocab_size, features=embed_dim, dtype=dtype, param_dtype=param_dtype, rngs=rngs)
self.pos_emb = nnx.Embed(num_embeddings=seqlen, features=embed_dim, dtype=dtype, param_dtype=param_dtype, rngs=rngs)
def __call__(self, x):
positions = jnp.arange(0, x.shape[1])[None, :]
position_embedding = self.pos_emb(positions)
token_embedding = self.token_emb(x)
return self.token_emb, token_embedding+position_embedding
class GPT2(nnx.Module):
def __init__(
self,
seqlen: int,
vocab_size: int,
embed_dim: int,
num_heads: int,
dropout_rate: float,
feed_forward_dim: int,
num_transformer_blocks: int,
rngs: nnx.Rngs,
):
self.embedding_layer = TokenAndPositionEmbedding(seqlen, vocab_size, embed_dim, rngs=rngs)
self.dropout = nnx.Dropout(rate=dropout_rate, rngs=rngs)
self.transformer_blocks = nnx.List([
TransformerBlock(embed_dim, num_heads, feed_forward_dim, dropout_rate, rngs=rngs)
for _ in range(num_transformer_blocks)
])
self.layer_norm = nnx.LayerNorm(
epsilon=1e-6,
num_features=embed_dim,
scale_init=nnx.with_partitioning(nnx.initializers.ones_init(), ("model",)),
bias_init=nnx.with_partitioning(nnx.initializers.zeros_init(), ("model",)),
dtype=dtype,
param_dtype=param_dtype,
rngs=rngs,
)
def __call__(self, inputs, training: bool = False):
token_embedding, x = self.embedding_layer(inputs)
x = self.dropout(x, deterministic=not training)
for transformer_block in self.transformer_blocks:
x = transformer_block(x, training=training)
x = self.layer_norm(x)
outputs = token_embedding.attend(x)
return outputs
def create_model(*, rngs: nnx.Rngs, config: TrainingConfig):
return GPT2(
seqlen=config.seqlen,
vocab_size=config.vocab_size,
embed_dim=config.embed_dim,
num_heads=config.num_heads,
dropout_rate=config.dropout_rate,
feed_forward_dim=config.feed_forward_dim,
num_transformer_blocks=config.num_transformer_blocks,
rngs=rngs,
)
```
Next, define the `loss_fn_train`, `loss_fn_eval`, and `train_step` functions.
`train_step()` computes the loss, takes gradients, and updates model parameters through the optimizer. The training loop will call this function repeatedly.
For performance, this notebook JIT-compiles these functions with `@nnx.jit`.
```python
@nnx.jit
def loss_fn_train(model, batch):
logits = model(batch[0], training=True)
loss = optax.softmax_cross_entropy_with_integer_labels(
logits=logits, labels=batch[1]
).mean()
return loss, logits
@nnx.jit
def loss_fn_eval(model, batch):
logits = model(batch[0], training=False)
loss = optax.softmax_cross_entropy_with_integer_labels(
logits=logits, labels=batch[1]
).mean()
return loss, logits
@nnx.jit
def train_step(model, optimizer, metrics, batch):
grad_fn = nnx.value_and_grad(loss_fn_train, has_aux=True)
(loss, logits), grads = grad_fn(model, batch)
metrics.update(loss=loss, logits=logits, labels=batch[1])
optimizer.update(model, grads)
return loss
```
## Step 4: Wrap training logic in `train_loop_per_worker`
Next, let's wrap the JAX training logic in a `train_loop_per_worker` function.
Each Ray Train worker runs the same Python function with a different world rank, and Ray sets device visibility per worker (for example, one GPU per worker). Inside this function, you can:
- Read the distributed context (`world_rank`, `world_size`).
- Get the per-worker dataset shard (`train.get_dataset_shard(...)`) to stream batches.
- Stream sharded JAX Arrays directly to training devices using [`Dataset.iter_jax_batches`](https://docs.ray.io/en/latest/data/api/doc/ray.data.Dataset.iter_jax_batches.html), which handles device sharding natively.
- Report metrics and checkpoints back to the trainer with `ray.train.report(...)`.
```python
import ray
from ray import train
from ray.train import Checkpoint
```
```python
def train_loop_per_worker(config_dict: dict) -> None:
config = TrainingConfig(**config_dict)
world_rank = ray.train.get_context().get_world_rank()
world_size = ray.train.get_context().get_world_size()
print(f"Worker rank {world_rank}/{world_size} sees devices: {jax.devices()}")
# Create a mesh per worker process.
device_mesh = mesh_utils.create_device_mesh((jax.process_count(), 1))
mesh = Mesh(device_mesh, axis_names=("data", "model"))
jax.set_mesh(mesh)
# Initialize the model locally.
# Include a dropout RNG stream so nnx.Dropout can run when training=True.
model = create_model(rngs=nnx.Rngs(params=0, dropout=1), config=config)
# We use Ray data to load the training and validation datasets.
train_it = ray.train.get_dataset_shard("train")
val_it = ray.train.get_dataset_shard("val")
if train_it is None or val_it is None:
raise RuntimeError("No Ray Train datasets provided. Pass datasets={...} to JaxTrainer.")
local_batch_size = config.global_batch_size // jax.process_count()
# A generator helper that yields JAX Array batches indefinitely.
# This removes the need for try/except StopIteration blocks in the training loop.
def batch_generator(dataset_shard):
while True:
# iter_jax_batches yields globally sharded JAX Array batches directly,
# mapped onto all process JAX-addressable training devices.
yield from dataset_shard.iter_jax_batches(
batch_size=local_batch_size,
prefetch_batches=2,
drop_last=True,
)
train_batches = batch_generator(train_it)
val_batches = batch_generator(val_it)
# Initialize the optimizer.
schedule = optax.cosine_decay_schedule(
init_value=config.init_learning_rate,
decay_steps=config.max_steps,
)
optax_chain = optax.chain(optax.adamw(learning_rate=schedule, weight_decay=config.weight_decay))
optimizer = nnx.Optimizer(model, optax_chain, wrt=nnx.Param)
checkpointer = orbax.PyTreeCheckpointer()
start_time = time.time()
train_metrics = nnx.metrics.Average('loss')
val_metrics = nnx.metrics.Average('val_loss')
for step in range(config.max_steps):
batch = next(train_batches)
global_x, global_y = batch["x"], batch["y"]
train_loss = train_step(model, optimizer, train_metrics, (global_x, global_y))
if (step + 1) % config.log_every_n_steps == 0:
elapsed = time.time() - start_time
# Report metrics through Ray Train.
ray.train.report({"step": step + 1, "train_loss": float(train_loss), "elapsed_s": elapsed})
start_time = time.time()
if (step + 1) % config.val_every_n_steps == 0:
validation_batch = next(val_batches)
global_val_input, global_val_target = validation_batch["x"], validation_batch["y"]
loss, logits = loss_fn_eval(model, (global_val_input, global_val_target))
val_metrics.update(val_loss=loss, logits=logits)
val_loss = float(val_metrics.compute())
metrics = {"step": step + 1, "train_loss": float(train_loss),"val_loss": float(val_loss)}
checkpoint = None
if (step + 1) % config.checkpoint_every_n_steps == 0:
# Orbax checkpointing is a barrier.
train_state = nnx.to_pure_dict(nnx.state(model))
checkpoint_path = os.path.join("/mnt/cluster_storage/checkpoint/jax_gpt2_ray_data", str(step + 1))
checkpointer.save(checkpoint_path, train_state)
# Save a checkpoint and report validation metrics through Ray Train.
# The controller persists the checkpoint to the RunConfig storage path.
checkpoint = Checkpoint.from_directory(checkpoint_path)
if world_rank == 0:
train.report(metrics, checkpoint=checkpoint)
else:
train.report(metrics, checkpoint=None)
```
> [!NOTE]
> [`iter_jax_batches`](https://docs.ray.io/en/latest/data/api/doc/ray.data.Dataset.iter_jax_batches.html) uses an internal 1D mesh for sharding. If your training loop uses a complex multi-dimensional mesh, JAX might perform implicit resharding. Ensure your mesh aligns with the device ordering to minimize overhead.
## Step 5: Define the `ScalingConfig`
Let's define the `ScalingConfig` that we want to scale the training process.
`JaxTrainer` now supports both GPU training and TPU training.
For a walkthrough on configuring `ScalingConfig`, see [Get Started with Distributed Training using JAX](https://docs.ray.io/en/latest/train/getting-started-jax.html).
```python
from ray.train import ScalingConfig
```
```python
# In this example, we use 2 GPUs.
scaling_config = ScalingConfig(
use_gpu=True,
num_workers=2, # Change this to match on your GPU cluster setting, by default, each worker uses one GPU.
)
# If you plan to use TPUs, see an example below.
# This ScalingConfig requires a KubeRay cluster configured for a TPU v6e 4x4 slice with 4 TPU VMs.
# For more information about TPU clusters with Ray on Kubernetes, see the
# KubeRay TPU guide: https://docs.ray.io/en/master/cluster/kubernetes/user-guides/tpu.html#kuberay-tpu
# scaling_config = ScalingConfig(
# use_tpu=True,
# num_workers=4,
# topology="4x4",
# accelerator_type="TPU-V6E",
# resources_per_worker={"TPU": 4},
# )
```
## Step 6: Launch with `JaxTrainer`
To run `train_loop_per_worker` on a Ray cluster, you'll construct a `JaxTrainer` with:
- `train_loop_per_worker`: the training function you've defined earlier. Each Ray Train worker runs this function.
- `train_loop_config`: a hyperparameter dictionary passed into the function.
- `scaling_config`: the `ScalingConfig` you've defined earlier.
- `datasets`: the Ray Datasets to ingest for training. Datasets are keyed by name (`{name: dataset}`). Each dataset can be accessed from within the `train_loop_per_worker` by calling `ray.train.get_dataset_shard(name)`. Sharding and additional configuration can be done by passing in a `dataset_config`.
- `run_config`: runtime configuration including where to write outputs such as checkpoints.
If your workers hit CUDA or XLA library load errors, clear `LD_LIBRARY_PATH` in the runtime env to avoid picking up incompatible system libraries.
`trainer.fit` spawns a controller process to orchestrate the training run and worker processes to execute the JAX training code.
```python
from ray.train import RunConfig
from ray.train.v2.jax import JaxTrainer
```
```python
storage_path = "/mnt/cluster_storage"
trainer = JaxTrainer(
train_loop_per_worker=train_loop_per_worker,
train_loop_config={
"global_batch_size": 32,
},
scaling_config=scaling_config,
run_config=RunConfig(
name="jax_gpt2",
storage_path=storage_path,
# Make sure to unset ``LD_LIBRARY_PATH`` if you're using CUDA devices,
# since ``LD_LIBRARY_PATH`` can override the CUDA libraries.
worker_runtime_env={"env_vars": {"LD_LIBRARY_PATH": ""}},
),
datasets={"train": train_ds, "val": val_ds},
)
result = trainer.fit()
print(result)
```
## Summary
In this notebook, you:
1. Prepared and the `OpenWebText` dataset and wrapped with Ray Data.
2. Defined a basic GPT2 model and train step in Jax/Flax.
3. Wrapped the training loop in a `train_loop_per_worker` function and scaled it out using Ray Train `JaxTrainer` with GPUs or TPUs!
doc/source/train/examples/lightning/BUILD.bazel
+33
View File
@@ -0,0 +1,33 @@
load("//bazel:python.bzl", "py_test_run_all_notebooks")
filegroup(
name = "lightning_examples",
srcs = glob(["*.ipynb"]),
visibility = ["//doc:__subpackages__"],
)
# GPU tests
py_test_run_all_notebooks(
size = "large",
include = ["*.ipynb"],
data = ["//doc/source/train/examples/lightning:lightning_examples"],
exclude = [
"dolly_lightning_fsdp_finetuning.ipynb", # Release Test
"vicuna_13b_lightning_deepspeed_finetune.ipynb", # Release Test
],
tags = [
"exclusive",
"gpu",
"ray_air",
"team:ml",
],
)
filegroup(
name = "lightning_examples_ci_configs",
srcs = glob([
"**/ci/aws.yaml",
"**/ci/gce.yaml",
]),
visibility = ["//doc/source/train/examples:__pkg__"],
)
doc/source/train/examples/lightning/dolly_lightning_fsdp_finetuning.ipynb
+842
View File
@@ -0,0 +1,842 @@
{
"cells": [
{
"cell_type": "markdown",
"metadata": {},
"source": [
"# Fine-tune `dolly-v2-7b` with Ray Train, PyTorch Lightning and FSDP\n",
"\n",
"<a id=\"try-anyscale-quickstart-dolly_lightning_fsdp_finetuning\" href=\"https://console.anyscale.com/register/ha?render_flow=ray&utm_source=ray_docs&utm_medium=docs&utm_campaign=dolly_lightning_fsdp_finetuning\">\n",
" <img src=\"../../../_static/img/run-on-anyscale.svg\" alt=\"try-anyscale-quickstart\">\n",
"</a>\n",
"<br></br>"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"In this example, we demonstrate how to use Ray Train to fine-tune a [`dolly-v2-7b`](https://huggingface.co/databricks/dolly-v2-7b) model. `dolly-v2-7b` is a 7 billion parameter causal language model created by Databricks, derived from EleutherAI’s [Pythia-6.9b](https://huggingface.co/EleutherAI/pythia-6.9b), and fine-tuned on a [~15K record instruction corpus](https://github.com/databrickslabs/dolly/tree/master/data).\n",
"\n",
"We load the pre-trained model from the HuggingFace model hub into a LightningModule and launch an FSDP fine-tuning job across 16 T4 GPUs with the help of {class}`Ray TorchTrainer <ray.train.torch.TorchTrainer>`. It is also straightforward to fine-tune other similar large language models in a similar manner as shown in this example.\n",
"\n",
"Before starting this example, we highly recommend reading [Ray Train Key Concepts](train-key-concepts) and [Ray Data Quickstart](data_quickstart)."
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"In this , we demonstrate how to use Ray Train to fine-tune a [`dolly-v2-7b`](https://huggingface.co/databricks/dolly-v2-7b) model. `dolly-v2-7b` is a 7 billion parameter causal language model created by Databricks, derived from EleutherAI’s [Pythia-6.9b](https://huggingface.co/EleutherAI/pythia-6.9b), and fine-tuned on a [~15K record instruction corpus](https://github.com/databrickslabs/dolly/tree/master/data).\n",
"\n",
"We load the pre-trained model from the HuggingFace model hub into a LightningModule and launch an FSDP fine-tuning job across 16 T4 GPUs with the help of {class}`Ray TorchTrainer <ray.train.torch.TorchTrainer>`. It is also straightforward to fine-tune other similar large language models in a similar manner as shown in this example.\n",
"\n",
"Before starting this example, we highly recommend reading [Ray Train Key Concepts](train-key-concepts) and [Ray Data Quickstart](data_quickstart)."
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Set up ray cluster \n",
"In this example, we are using a Ray cluster with a `m5.2xlarge` head node and 4 `g4dn.12xlarge` worker nodes. Each `g4dn.12xlarge has four Tesla T4 GPUs. "
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"tags": []
},
"outputs": [],
"source": [
"import ray\n",
"ray.init()"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"We then install the necessary dependencies on each node:"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"%%bash\n",
"pip install datasets\n",
"pip install evaluate\n",
"pip install \"transformers>=4.26.0\"\n",
"pip install \"torch>=1.12.0\"\n",
"pip install \"lightning>=2.0\"\n",
"pip install \"pydantic>=2,<3\""
]
},
{
"cell_type": "code",
"execution_count": 4,
"metadata": {
"tags": []
},
"outputs": [],
"source": [
"MODEL_NAME = \"databricks/dolly-v2-7b\""
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Prepare your data \n",
"We are using tiny_shakespeare for fine-tuning, which contains 40,000 lines of Shakespeare from a variety of Shakespeare's plays. Featured in Andrej Karpathy's blog post ['The Unreasonable Effectiveness of Recurrent Neural Networks'](http://karpathy.github.io/2015/05/21/rnn-effectiveness/). \n",
"\n",
"Dataset samples:\n",
"```\n",
"BAPTISTA:\n",
"I know him well: you are welcome for his sake.\n",
"\n",
"GREMIO:\n",
"Saving your tale, Petruchio, I pray,\n",
"Let us, that are poor petitioners, speak too:\n",
"Baccare! you are marvellous forward.\n",
"\n",
"PETRUCHIO:\n",
"O, pardon me, Signior Gremio; I would fain be doing.\n",
"```\n",
"\n",
"Here, we have adopted similar pre-processing logic from another demo: {doc}`GPT-J-6B Fine-Tuning with Ray Train and DeepSpeed <../deepspeed/gptj_deepspeed_fine_tuning>`."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"tags": []
},
"outputs": [],
"source": [
"import ray\n",
"import pandas as pd\n",
"from datasets import load_dataset\n",
"from transformers import AutoTokenizer, AutoModelForCausalLM\n",
"\n",
"def split_text(batch: pd.DataFrame) -> pd.DataFrame:\n",
" text = list(batch[\"text\"])\n",
" flat_text = \"\".join(text)\n",
" split_text = [\n",
" x.strip()\n",
" for x in flat_text.split(\"\\n\")\n",
" if x.strip() and not x.strip()[-1] == \":\"\n",
" ]\n",
" return pd.DataFrame(split_text, columns=[\"text\"])\n",
"\n",
"\n",
"def tokenize(batch: pd.DataFrame) -> dict:\n",
" tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, padding_side=\"left\")\n",
" tokenizer.pad_token = tokenizer.eos_token\n",
" ret = tokenizer(\n",
" list(batch[\"text\"]),\n",
" truncation=True,\n",
" max_length=256,\n",
" padding=\"max_length\",\n",
" return_tensors=\"np\",\n",
" )\n",
" ret[\"labels\"] = ret[\"input_ids\"].copy()\n",
" return dict(ret)\n",
"\n",
"hf_dataset = load_dataset(\"tiny_shakespeare\", trust_remote_code=True)\n",
"train_ds = ray.data.from_huggingface(hf_dataset[\"train\"])"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"We first split the original paragraphs into multiple sentences, then tokenize them. Here are some samples:"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"tags": []
},
"outputs": [],
"source": [
"# First split the dataset into multiple sentences.\n",
"train_ds = train_ds.map_batches(split_text, batch_format=\"pandas\", batch_size=\"auto\")\n",
"train_ds.take(10)"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# Then tokenize the dataset.\n",
"train_ds = train_ds.map_batches(tokenize, batch_format=\"pandas\", batch_size=\"auto\")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Define your lightning model\n",
"\n",
"In this example, we use the [dolly-v2-7b](https://huggingface.co/databricks/dolly-v2-7b) model for finetuning. It is an instruction-following large language model trained on the Databricks machine learning platform that is licensed for commercial use. We load the model weights from Huggingface Model Hub and encapsulate it into a `pl.LightningModule`.\n",
"\n",
":::{note}\n",
"Make sure you pass the FSDP wrapped model parameters `self.trainer.model.parameters()` into the optimizer, instead of `self.model.parameters()`. \n",
":::\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"tags": []
},
"outputs": [],
"source": [
"import torch\n",
"import lightning.pytorch as pl\n",
"\n",
"class DollyV2Model(pl.LightningModule):\n",
" def __init__(self, lr=2e-5, eps=1e-8):\n",
" super().__init__()\n",
" self.save_hyperparameters()\n",
" self.lr = lr\n",
" self.eps = eps\n",
" self.model = AutoModelForCausalLM.from_pretrained(MODEL_NAME)\n",
"\n",
" def forward(self, batch):\n",
" outputs = self.model(\n",
" batch[\"input_ids\"], \n",
" attention_mask=batch[\"attention_mask\"], \n",
" labels=batch[\"labels\"]\n",
" )\n",
" return outputs.loss\n",
"\n",
" def training_step(self, batch, batch_idx):\n",
" loss = self.forward(batch)\n",
" self.log(\"train_loss\", loss, prog_bar=True, on_step=True)\n",
" return loss\n",
"\n",
" def configure_optimizers(self):\n",
" if self.global_rank == 0:\n",
" print(self.trainer.model)\n",
" return torch.optim.AdamW(self.trainer.model.parameters(), lr=self.lr, eps=self.eps)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Configure your FSDP strategy\n",
"As `dolly-v2-7b` is a relatively large model, it cannot be properly fit into a single commercial GPU. In this example, we use the FSDP strategy to shard model parameters across multiple workers. This allows us to avoid GPU out-of-memory issues and support a larger global batch size.\n",
"\n",
"![](https://user-images.githubusercontent.com/26745457/236892936-d4b91751-4689-421e-ac5f-edfd2eeeb635.png)\n",
"Image source: [Fully Sharded Data Parallel: faster AI training with fewer GPUs](https://engineering.fb.com/2021/07/15/open-source/fsdp/)\n",
"\n",
":::{note}\n",
"FSDP is a type of data parallelism that shards model parameters, optimizer states and gradients across DDP ranks. This was inspired by Xu et al. as well as the ZeRO Stage 3 from DeepSpeed. You may refer to these blogs for more information:\n",
"\n",
"- [Fully Sharded Data Parallel: faster AI training with fewer GPUs](https://engineering.fb.com/2021/07/15/open-source/fsdp/)\n",
"- [Getting Started with Fully Sharded Data Parallel(FSDP)](https://pytorch.org/tutorials/intermediate/FSDP_tutorial.html#:~:text=FSDP%20is%20a%20type%20of,sizes%20for%20our%20training%20job.)\n",
"- [PyTorch FSDP Tutorial](https://www.youtube.com/watch?v=8_k76AHu__s&list=PL_lsbAsL_o2BT6aerEKgIoufVD_fodnuT)\n",
":::\n",
"\n",
"To start training with Lightning's [FSDPStrategy](https://lightning.ai/docs/pytorch/stable/api/lightning.pytorch.strategies.FSDPStrategy.html#lightning.pytorch.strategies.FSDPStrategy), you only need to create a {class}`~ray.train.lightning.RayFSDPStrategy` with the same initialization arguments.\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"tags": []
},
"outputs": [],
"source": [
"import functools\n",
"import lightning.pytorch as pl \n",
"\n",
"from torch.distributed.fsdp.wrap import transformer_auto_wrap_policy\n",
"from torch.distributed.fsdp import ShardingStrategy, BackwardPrefetch\n",
"from transformers.models.gpt_neox.modeling_gpt_neox import GPTNeoXLayer\n",
"\n",
"from ray.train.lightning import RayFSDPStrategy\n",
"\n",
"\n",
"# Define the model sharding policy:\n",
"# Wrap every GPTNeoXLayer as its own FSDP instance\n",
"auto_wrap_policy = functools.partial(\n",
" transformer_auto_wrap_policy,\n",
" transformer_layer_cls = {GPTNeoXLayer}\n",
")\n",
"\n",
"fsdp_strategy = RayFSDPStrategy(\n",
" sharding_strategy=ShardingStrategy.FULL_SHARD,\n",
" backward_prefetch=BackwardPrefetch.BACKWARD_PRE,\n",
" forward_prefetch=True,\n",
" auto_wrap_policy=auto_wrap_policy,\n",
" limit_all_gathers=True,\n",
" activation_checkpointing=[GPTNeoXLayer],\n",
" cpu_offload=True\n",
")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
":::{tip}\n",
"\n",
"Some tips for FSDP configuration:\n",
"- `sharding_strategy`:\n",
" - `ShardingStrategy.NO_SHARD`: Parameters, gradients, and optimizer states are not sharded. Similar to DDP.\n",
" - `ShardingStrategy.SHARD_GRAD_OP`: Gradients and optimizer states are sharded during computation, and additionally, parameters are sharded outside computation. Similar to ZeRO stage-2.\n",
" - `ShardingStrategy.FULL_SHARD`: Parameters, gradients, and optimizer states are sharded. It has minimal GRAM usage among the 3 options. Similar to ZeRO stage-3.\n",
"- `auto_wrap_policy`:\n",
" - Model layers are often wrapped with FSDP in a layered fashion. This means that only the layers in a single FSDP instance are required to aggregate all parameters to a single device during forwarding or backward calculations.\n",
" - Use `transformer_auto_wrap_policy` to automatically wrap each Transformer Block into a single FSDP instance. \n",
"- `backward_prefetch` and `forward_prefetch`:\n",
" - Overlap the upcoming all-gather while executing the current forward/backward pass. It can improve throughput but may slightly increase peak memory usage.\n",
":::"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Fine-tune with Ray TorchTrainer\n",
"\n",
"Ray TorchTrainer allows you to scale your PyTorch Lightning training workload over multiple nodes. See {ref}`train_scaling_config` for more details."
]
},
{
"cell_type": "code",
"execution_count": 10,
"metadata": {},
"outputs": [],
"source": [
"num_workers = 16\n",
"batch_size_per_worker = 5"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"tags": [
"remove-cell"
]
},
"outputs": [],
"source": [
"# To accelerate release tests\n",
"train_ds = train_ds.limit(num_workers * batch_size_per_worker * 5) # each worker has 5 batches"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Additionally, remember to define a Lightning callback that saves and reports checkpoints. Ray Train offers a simple implementation, {meth}`~ray.train.lightning.RayTrainReportCallback`, which persists your checkpoint and metrics in remote storage at the end of each training epoch. \n",
"\n",
"Note you can also implement your own report callback with customized logics, such as saving customized checkpoint files or reporting at a different frequency."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"tags": [
"remove-cell"
]
},
"outputs": [],
"source": [
"from lightning.pytorch.callbacks import TQDMProgressBar\n",
"\n",
"# Create a customized progress bar for Ray Data Iterable Dataset\n",
"class DollyV2ProgressBar(TQDMProgressBar):\n",
" def __init__(self, num_iters_per_epoch, *args, **kwargs):\n",
" super().__init__(*args, **kwargs)\n",
" self.num_iters_per_epoch = num_iters_per_epoch\n",
" \n",
" def on_train_epoch_start(self, trainer, *_):\n",
" super().on_train_epoch_start(trainer, *_)\n",
" self.train_progress_bar.reset(self.num_iters_per_epoch)\n",
"\n",
"num_iters_per_epoch = train_ds.count() // (num_workers * batch_size_per_worker)\n",
"prog_bar = DollyV2ProgressBar(num_iters_per_epoch)"
]
},
{
"cell_type": "code",
"execution_count": 13,
"metadata": {},
"outputs": [],
"source": [
"from ray.train import Checkpoint\n",
"from ray.train.lightning import RayLightningEnvironment, RayTrainReportCallback, prepare_trainer\n",
"\n",
"# Training function for each worker\n",
"def train_func(config):\n",
" lr = config[\"lr\"]\n",
" eps = config[\"eps\"]\n",
" strategy = config[\"strategy\"]\n",
" batch_size_per_worker = config[\"batch_size_per_worker\"]\n",
"\n",
" # Model\n",
" model = DollyV2Model(lr=lr, eps=eps)\n",
"\n",
" # Ray Data Ingestion\n",
" train_ds = ray.train.get_dataset_shard(\"train\")\n",
" train_dataloader = train_ds.iter_torch_batches(batch_size=batch_size_per_worker)\n",
"\n",
" # Lightning Trainer\n",
" trainer = pl.Trainer(\n",
" max_epochs=1, \n",
" devices=\"auto\",\n",
" accelerator=\"auto\", \n",
" precision=\"16-mixed\",\n",
" strategy=strategy,\n",
" plugins=[RayLightningEnvironment()],\n",
" callbacks=[RayTrainReportCallback()],\n",
" enable_checkpointing=False,\n",
" )\n",
"\n",
" trainer = prepare_trainer(trainer)\n",
"\n",
" trainer.fit(model, train_dataloaders=train_dataloader)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"```{note}\n",
"Since this example runs with multiple nodes, we need to persist checkpoints\n",
"and other outputs to some external storage for access after training has completed.\n",
"**You should set up cloud storage or NFS, then replace `storage_path` with your own cloud bucket URI or NFS path.**\n",
"\n",
"See the [storage guide](tune-storage-options) for more details.\n",
"```"
]
},
{
"cell_type": "code",
"execution_count": 14,
"metadata": {},
"outputs": [],
"source": [
"storage_path=\"s3://your-bucket-here\" # TODO: Set up cloud storage\n",
"# storage_path=\"/mnt/path/to/nfs\" # TODO: Alternatively, set up NFS"
]
},
{
"cell_type": "code",
"execution_count": 15,
"metadata": {
"tags": [
"remove-cell"
]
},
"outputs": [],
"source": [
"storage_path = \"/mnt/cluster_storage\""
]
},
{
"cell_type": "code",
"execution_count": 16,
"metadata": {
"tags": []
},
"outputs": [
{
"name": "stderr",
"output_type": "stream",
"text": [
"\u001b[36m(TrainController pid=6878)\u001b[0m /home/ray/anaconda3/lib/python3.10/site-packages/transformers/utils/generic.py:441: UserWarning: torch.utils._pytree._register_pytree_node is deprecated. Please use torch.utils._pytree.register_pytree_node instead.\n",
"\u001b[36m(TrainController pid=6878)\u001b[0m _torch_pytree._register_pytree_node(\n",
"\u001b[36m(TrainController pid=6878)\u001b[0m /home/ray/anaconda3/lib/python3.10/site-packages/transformers/utils/generic.py:309: UserWarning: torch.utils._pytree._register_pytree_node is deprecated. Please use torch.utils._pytree.register_pytree_node instead.\n",
"\u001b[36m(TrainController pid=6878)\u001b[0m _torch_pytree._register_pytree_node(\n",
"\u001b[36m(TrainController pid=6878)\u001b[0m [State Transition] INITIALIZING -> SCHEDULING.\n",
"\u001b[36m(TrainController pid=6878)\u001b[0m Attempting to start training worker group of size 16 with the following resources: [{'GPU': 1}] * 16\n",
"\u001b[36m(RayTrainWorker pid=4349, ip=10.0.157.249)\u001b[0m /home/ray/anaconda3/lib/python3.10/site-packages/transformers/utils/generic.py:441: UserWarning: torch.utils._pytree._register_pytree_node is deprecated. Please use torch.utils._pytree.register_pytree_node instead.\n",
"\u001b[36m(RayTrainWorker pid=4349, ip=10.0.157.249)\u001b[0m _torch_pytree._register_pytree_node(\n",
"\u001b[36m(RayTrainWorker pid=4350, ip=10.0.157.249)\u001b[0m /home/ray/anaconda3/lib/python3.10/site-packages/transformers/utils/generic.py:441: UserWarning: torch.utils._pytree._register_pytree_node is deprecated. Please use torch.utils._pytree.register_pytree_node instead.\n",
"\u001b[36m(RayTrainWorker pid=4350, ip=10.0.157.249)\u001b[0m _torch_pytree._register_pytree_node(\n",
"\u001b[36m(RayTrainWorker pid=4088, ip=10.0.163.141)\u001b[0m Setting up process group for: env:// [rank=0, world_size=16]\n",
"\u001b[36m(RayTrainWorker pid=4349, ip=10.0.157.249)\u001b[0m /home/ray/anaconda3/lib/python3.10/site-packages/huggingface_hub/file_download.py:795: FutureWarning: `resume_download` is deprecated and will be removed in version 1.0.0. Downloads always resume when possible. If you want to force a new download, use `force_download=True`.\n",
"\u001b[36m(RayTrainWorker pid=4349, ip=10.0.157.249)\u001b[0m warnings.warn(\n",
"\u001b[36m(TrainController pid=6878)\u001b[0m Started training worker group of size 16: \n",
"\u001b[36m(TrainController pid=6878)\u001b[0m - (ip=10.0.163.141, pid=4088) world_rank=0, local_rank=0, node_rank=0\n",
"\u001b[36m(TrainController pid=6878)\u001b[0m - (ip=10.0.163.141, pid=4089) world_rank=1, local_rank=1, node_rank=0\n",
"\u001b[36m(TrainController pid=6878)\u001b[0m - (ip=10.0.163.141, pid=4090) world_rank=2, local_rank=2, node_rank=0\n",
"\u001b[36m(TrainController pid=6878)\u001b[0m - (ip=10.0.163.141, pid=4091) world_rank=3, local_rank=3, node_rank=0\n",
"\u001b[36m(TrainController pid=6878)\u001b[0m - (ip=10.0.166.248, pid=4338) world_rank=4, local_rank=0, node_rank=1\n",
"\u001b[36m(TrainController pid=6878)\u001b[0m - (ip=10.0.166.248, pid=4337) world_rank=5, local_rank=1, node_rank=1\n",
"\u001b[36m(TrainController pid=6878)\u001b[0m - (ip=10.0.166.248, pid=4340) world_rank=6, local_rank=2, node_rank=1\n",
"\u001b[36m(TrainController pid=6878)\u001b[0m - (ip=10.0.166.248, pid=4339) world_rank=7, local_rank=3, node_rank=1\n",
"\u001b[36m(TrainController pid=6878)\u001b[0m - (ip=10.0.191.43, pid=4090) world_rank=8, local_rank=0, node_rank=2\n",
"\u001b[36m(TrainController pid=6878)\u001b[0m - (ip=10.0.191.43, pid=4248) world_rank=9, local_rank=1, node_rank=2\n",
"\u001b[36m(TrainController pid=6878)\u001b[0m - (ip=10.0.191.43, pid=4246) world_rank=10, local_rank=2, node_rank=2\n",
"\u001b[36m(TrainController pid=6878)\u001b[0m - (ip=10.0.191.43, pid=4247) world_rank=11, local_rank=3, node_rank=2\n",
"\u001b[36m(TrainController pid=6878)\u001b[0m - (ip=10.0.157.249, pid=4349) world_rank=12, local_rank=0, node_rank=3\n",
"\u001b[36m(TrainController pid=6878)\u001b[0m - (ip=10.0.157.249, pid=4350) world_rank=13, local_rank=1, node_rank=3\n",
"\u001b[36m(TrainController pid=6878)\u001b[0m - (ip=10.0.157.249, pid=4347) world_rank=14, local_rank=2, node_rank=3\n",
"\u001b[36m(TrainController pid=6878)\u001b[0m - (ip=10.0.157.249, pid=4348) world_rank=15, local_rank=3, node_rank=3\n",
"\u001b[36m(TrainController pid=6878)\u001b[0m [State Transition] SCHEDULING -> RUNNING.\n",
"\u001b[36m(RayTrainWorker pid=4246, ip=10.0.191.43)\u001b[0m /home/ray/anaconda3/lib/python3.10/site-packages/transformers/utils/generic.py:309: UserWarning: torch.utils._pytree._register_pytree_node is deprecated. Please use torch.utils._pytree.register_pytree_node instead.\u001b[32m [repeated 31x across cluster] (Ray deduplicates logs by default. Set RAY_DEDUP_LOGS=0 to disable log deduplication, or see https://docs.ray.io/en/master/ray-observability/user-guides/configure-logging.html#log-deduplication for more options.)\u001b[0m\n",
"\u001b[36m(RayTrainWorker pid=4246, ip=10.0.191.43)\u001b[0m _torch_pytree._register_pytree_node(\u001b[32m [repeated 31x across cluster]\u001b[0m\n"
]
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"\u001b[33m(raylet)\u001b[0m WARNING: 4 PYTHON worker processes have been started on node: 3ffdd02cd1be6b69a64a97e08f75fc0a80eddcf0caa627f3f4266c95 with address: 10.0.150.21. This could be a result of using a large number of actors, or due to tasks blocked in ray.get() calls (see https://github.com/ray-project/ray/issues/3644 for some discussion of workarounds).\n"
]
},
{
"name": "stderr",
"output_type": "stream",
"text": [
"\u001b[36m(RayTrainWorker pid=4088, ip=10.0.163.141)\u001b[0m GPU available: True (cuda), used: True\n",
"\u001b[36m(RayTrainWorker pid=4088, ip=10.0.163.141)\u001b[0m TPU available: False, using: 0 TPU cores\n",
"\u001b[36m(RayTrainWorker pid=4088, ip=10.0.163.141)\u001b[0m HPU available: False, using: 0 HPUs\n",
"\u001b[36m(RayTrainWorker pid=4246, ip=10.0.191.43)\u001b[0m /home/ray/anaconda3/lib/python3.10/site-packages/huggingface_hub/file_download.py:795: FutureWarning: `resume_download` is deprecated and will be removed in version 1.0.0. Downloads always resume when possible. If you want to force a new download, use `force_download=True`.\u001b[32m [repeated 15x across cluster]\u001b[0m\n",
"\u001b[36m(RayTrainWorker pid=4246, ip=10.0.191.43)\u001b[0m warnings.warn(\u001b[32m [repeated 15x across cluster]\u001b[0m\n",
"\u001b[36m(RayTrainWorker pid=4088, ip=10.0.163.141)\u001b[0m 2025-09-30 14:20:07.970624: I tensorflow/core/util/port.cc:113] oneDNN custom operations are on. You may see slightly different numerical results due to floating-point round-off errors from different computation orders. To turn them off, set the environment variable `TF_ENABLE_ONEDNN_OPTS=0`.\n",
"\u001b[36m(RayTrainWorker pid=4088, ip=10.0.163.141)\u001b[0m 2025-09-30 14:20:08.208262: E external/local_xla/xla/stream_executor/cuda/cuda_dnn.cc:9261] Unable to register cuDNN factory: Attempting to register factory for plugin cuDNN when one has already been registered\n",
"\u001b[36m(RayTrainWorker pid=4088, ip=10.0.163.141)\u001b[0m 2025-09-30 14:20:08.208291: E external/local_xla/xla/stream_executor/cuda/cuda_fft.cc:607] Unable to register cuFFT factory: Attempting to register factory for plugin cuFFT when one has already been registered\n",
"\u001b[36m(RayTrainWorker pid=4088, ip=10.0.163.141)\u001b[0m 2025-09-30 14:20:08.231782: E external/local_xla/xla/stream_executor/cuda/cuda_blas.cc:1515] Unable to register cuBLAS factory: Attempting to register factory for plugin cuBLAS when one has already been registered\n",
"\u001b[36m(RayTrainWorker pid=4088, ip=10.0.163.141)\u001b[0m 2025-09-30 14:20:08.277889: I tensorflow/core/platform/cpu_feature_guard.cc:182] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations.\n",
"\u001b[36m(RayTrainWorker pid=4088, ip=10.0.163.141)\u001b[0m To enable the following instructions: AVX2 AVX512F AVX512_VNNI FMA, in other operations, rebuild TensorFlow with the appropriate compiler flags.\n",
"\u001b[36m(SplitCoordinator pid=7661)\u001b[0m /home/ray/anaconda3/lib/python3.10/site-packages/transformers/utils/generic.py:441: UserWarning: torch.utils._pytree._register_pytree_node is deprecated. Please use torch.utils._pytree.register_pytree_node instead.\n",
"\u001b[36m(SplitCoordinator pid=7661)\u001b[0m _torch_pytree._register_pytree_node(\n",
"\u001b[36m(RayTrainWorker pid=4088, ip=10.0.163.141)\u001b[0m 2025-09-30 14:20:10.134645: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Could not find TensorRT\n",
"\u001b[36m(RayTrainWorker pid=4088, ip=10.0.163.141)\u001b[0m LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0,1,2,3]\n",
"\u001b[36m(RayTrainWorker pid=4088, ip=10.0.163.141)\u001b[0m FullyShardedDataParallel(\n",
"\u001b[36m(RayTrainWorker pid=4088, ip=10.0.163.141)\u001b[0m (_fsdp_wrapped_module): DollyV2Model(\n",
"\u001b[36m(RayTrainWorker pid=4088, ip=10.0.163.141)\u001b[0m (model): GPTNeoXForCausalLM(\n",
"\u001b[36m(RayTrainWorker pid=4088, ip=10.0.163.141)\u001b[0m (gpt_neox): GPTNeoXModel(\n",
"\u001b[36m(RayTrainWorker pid=4088, ip=10.0.163.141)\u001b[0m (embed_in): Embedding(50280, 4096)\n",
"\u001b[36m(RayTrainWorker pid=4088, ip=10.0.163.141)\u001b[0m (emb_dropout): Dropout(p=0.0, inplace=False)\n",
"\u001b[36m(RayTrainWorker pid=4088, ip=10.0.163.141)\u001b[0m (layers): ModuleList(\n",
"\u001b[36m(RayTrainWorker pid=4088, ip=10.0.163.141)\u001b[0m (0-31): 32 x FullyShardedDataParallel(\n",
"\u001b[36m(RayTrainWorker pid=4088, ip=10.0.163.141)\u001b[0m (_fsdp_wrapped_module): CheckpointWrapper(\n",
"\u001b[36m(RayTrainWorker pid=4088, ip=10.0.163.141)\u001b[0m (_checkpoint_wrapped_module): GPTNeoXLayer(\n",
"\u001b[36m(RayTrainWorker pid=4088, ip=10.0.163.141)\u001b[0m (input_layernorm): LayerNorm((4096,), eps=1e-05, elementwise_affine=True)\n",
"\u001b[36m(RayTrainWorker pid=4088, ip=10.0.163.141)\u001b[0m (post_attention_layernorm): LayerNorm((4096,), eps=1e-05, elementwise_affine=True)\n",
"\u001b[36m(RayTrainWorker pid=4088, ip=10.0.163.141)\u001b[0m (post_attention_dropout): Dropout(p=0.0, inplace=False)\n",
"\u001b[36m(RayTrainWorker pid=4088, ip=10.0.163.141)\u001b[0m (post_mlp_dropout): Dropout(p=0.0, inplace=False)\n",
"\u001b[36m(RayTrainWorker pid=4088, ip=10.0.163.141)\u001b[0m (attention): GPTNeoXAttention(\n",
"\u001b[36m(RayTrainWorker pid=4088, ip=10.0.163.141)\u001b[0m (rotary_emb): GPTNeoXRotaryEmbedding()\n",
"\u001b[36m(RayTrainWorker pid=4088, ip=10.0.163.141)\u001b[0m (query_key_value): Linear(in_features=4096, out_features=12288, bias=True)\n",
"\u001b[36m(RayTrainWorker pid=4088, ip=10.0.163.141)\u001b[0m (dense): Linear(in_features=4096, out_features=4096, bias=True)\n",
"\u001b[36m(RayTrainWorker pid=4088, ip=10.0.163.141)\u001b[0m (attention_dropout): Dropout(p=0.0, inplace=False)\n",
"\u001b[36m(RayTrainWorker pid=4088, ip=10.0.163.141)\u001b[0m )\n",
"\u001b[36m(RayTrainWorker pid=4088, ip=10.0.163.141)\u001b[0m (mlp): GPTNeoXMLP(\n",
"\u001b[36m(RayTrainWorker pid=4088, ip=10.0.163.141)\u001b[0m (dense_h_to_4h): Linear(in_features=4096, out_features=16384, bias=True)\n",
"\u001b[36m(RayTrainWorker pid=4088, ip=10.0.163.141)\u001b[0m (dense_4h_to_h): Linear(in_features=16384, out_features=4096, bias=True)\n",
"\u001b[36m(RayTrainWorker pid=4088, ip=10.0.163.141)\u001b[0m (act): GELUActivation()\n",
"\u001b[36m(RayTrainWorker pid=4088, ip=10.0.163.141)\u001b[0m )\n",
"\u001b[36m(RayTrainWorker pid=4088, ip=10.0.163.141)\u001b[0m )\n",
"\u001b[36m(RayTrainWorker pid=4088, ip=10.0.163.141)\u001b[0m )\n",
"\u001b[36m(RayTrainWorker pid=4088, ip=10.0.163.141)\u001b[0m )\n",
"\u001b[36m(RayTrainWorker pid=4088, ip=10.0.163.141)\u001b[0m )\n",
"\u001b[36m(RayTrainWorker pid=4088, ip=10.0.163.141)\u001b[0m (final_layer_norm): LayerNorm((4096,), eps=1e-05, elementwise_affine=True)\n",
"\u001b[36m(RayTrainWorker pid=4088, ip=10.0.163.141)\u001b[0m )\n",
"\u001b[36m(RayTrainWorker pid=4088, ip=10.0.163.141)\u001b[0m (embed_out): Linear(in_features=4096, out_features=50280, bias=False)\n",
"\u001b[36m(RayTrainWorker pid=4088, ip=10.0.163.141)\u001b[0m )\n",
"\u001b[36m(RayTrainWorker pid=4088, ip=10.0.163.141)\u001b[0m )\n",
"\u001b[36m(RayTrainWorker pid=4088, ip=10.0.163.141)\u001b[0m )\n",
"\u001b[36m(RayTrainWorker pid=4088, ip=10.0.163.141)\u001b[0m /tmp/ray/session_2025-09-30_14-10-46_627006_2395/runtime_resources/pip/72a6e451f55d87eb50ebbf5bc30a4a57ed513d34/virtualenv/lib/python3.10/site-packages/lightning/pytorch/utilities/model_summary/model_summary.py:231: Precision 16-mixed is not supported by the model summary. Estimated model size in MB will not be accurate. Using 32 bits instead.\n",
"\u001b[36m(RayTrainWorker pid=4088, ip=10.0.163.141)\u001b[0m \n",
"\u001b[36m(RayTrainWorker pid=4088, ip=10.0.163.141)\u001b[0m | Name | Type | Params | Mode\n",
"\u001b[36m(RayTrainWorker pid=4088, ip=10.0.163.141)\u001b[0m ----------------------------------------------------\n",
"\u001b[36m(RayTrainWorker pid=4088, ip=10.0.163.141)\u001b[0m 0 | model | GPTNeoXForCausalLM | 428 M | eval\n",
"\u001b[36m(RayTrainWorker pid=4088, ip=10.0.163.141)\u001b[0m ----------------------------------------------------\n",
"\u001b[36m(RayTrainWorker pid=4088, ip=10.0.163.141)\u001b[0m 428 M Trainable params\n",
"\u001b[36m(RayTrainWorker pid=4088, ip=10.0.163.141)\u001b[0m 0 Non-trainable params\n",
"\u001b[36m(RayTrainWorker pid=4088, ip=10.0.163.141)\u001b[0m 428 M Total params\n",
"\u001b[36m(RayTrainWorker pid=4088, ip=10.0.163.141)\u001b[0m 1,714.014 Total estimated model params size (MB)\n",
"\u001b[36m(RayTrainWorker pid=4088, ip=10.0.163.141)\u001b[0m 64 Modules in train mode\n",
"\u001b[36m(RayTrainWorker pid=4088, ip=10.0.163.141)\u001b[0m 455 Modules in eval mode\n",
"\u001b[36m(RayTrainWorker pid=4246, ip=10.0.191.43)\u001b[0m LOCAL_RANK: 2 - CUDA_VISIBLE_DEVICES: [0,1,2,3]\u001b[32m [repeated 15x across cluster]\u001b[0m\n",
"\u001b[36m(SplitCoordinator pid=7661)\u001b[0m Registered dataset logger for dataset train_12_0\n",
"\u001b[36m(SplitCoordinator pid=7661)\u001b[0m Starting execution of Dataset train_12_0. Full logs are in /tmp/ray/session_2025-09-30_14-10-46_627006_2395/logs/ray-data\n",
"\u001b[36m(SplitCoordinator pid=7661)\u001b[0m Execution plan of Dataset train_12_0: InputDataBuffer[Input] -> TaskPoolMapOperator[MapBatches(split_text)->MapBatches(tokenize)] -> LimitOperator[limit=400] -> OutputSplitter[split(16, equal=True)]\n",
"\u001b[36m(SplitCoordinator pid=7661)\u001b[0m ⚠️ Ray's object store is configured to use only 28.7% of available memory (229.3GiB out of 800.0GiB total). For optimal Ray Data performance, we recommend setting the object store to at least 50% of available memory. You can do this by setting the 'object_store_memory' parameter when calling ray.init() or by setting the RAY_DEFAULT_OBJECT_STORE_MEMORY_PROPORTION environment variable.\n"
]
},
{
"data": {
"application/vnd.jupyter.widget-view+json": {
"model_id": "ab803bffe2224e6591bc452cac07f2a8",
"version_major": 2,
"version_minor": 0
},
"text/plain": [
"(pid=7661) Running 0: 0.00 row [00:00, ? row/s]"
]
},
"metadata": {},
"output_type": "display_data"
},
{
"data": {
"application/vnd.jupyter.widget-view+json": {
"model_id": "8d647925625f49c19e16fddce0fab359",
"version_major": 2,
"version_minor": 0
},
"text/plain": [
"(pid=7661) - MapBatches(split_text)->MapBatches(tokenize) 1: 0.00 row [00:00, ? row/s]"
]
},
"metadata": {},
"output_type": "display_data"
},
{
"data": {
"application/vnd.jupyter.widget-view+json": {
"model_id": "30978a8d911a44dcb168b0d5a386a42c",
"version_major": 2,
"version_minor": 0
},
"text/plain": [
"(pid=7661) - limit=400 2: 0.00 row [00:00, ? row/s]"
]
},
"metadata": {},
"output_type": "display_data"
},
{
"data": {
"application/vnd.jupyter.widget-view+json": {
"model_id": "432c8cb6d3e84b749d9341ff104bb25c",
"version_major": 2,
"version_minor": 0
},
"text/plain": [
"(pid=7661) - split(16, equal=True) 3: 0.00 row [00:00, ? row/s]"
]
},
"metadata": {},
"output_type": "display_data"
},
{
"name": "stderr",
"output_type": "stream",
"text": [
"\u001b[36m(MapBatches(split_text)->MapBatches(tokenize) pid=4089, ip=10.0.191.43)\u001b[0m /home/ray/anaconda3/lib/python3.10/site-packages/transformers/utils/generic.py:441: UserWarning: torch.utils._pytree._register_pytree_node is deprecated. Please use torch.utils._pytree.register_pytree_node instead.\n",
"\u001b[36m(MapBatches(split_text)->MapBatches(tokenize) pid=4089, ip=10.0.191.43)\u001b[0m _torch_pytree._register_pytree_node(\n",
"\u001b[36m(MapBatches(split_text)->MapBatches(tokenize) pid=4089, ip=10.0.191.43)\u001b[0m /home/ray/anaconda3/lib/python3.10/site-packages/huggingface_hub/file_download.py:795: FutureWarning: `resume_download` is deprecated and will be removed in version 1.0.0. Downloads always resume when possible. If you want to force a new download, use `force_download=True`.\n",
"\u001b[36m(MapBatches(split_text)->MapBatches(tokenize) pid=4089, ip=10.0.191.43)\u001b[0m warnings.warn(\n",
"\u001b[36m(MapBatches(split_text)->MapBatches(tokenize) pid=4089, ip=10.0.191.43)\u001b[0m Special tokens have been added in the vocabulary, make sure the associated word embeddings are fine-tuned or trained.\n",
"\u001b[36m(SplitCoordinator pid=7661)\u001b[0m /home/ray/anaconda3/lib/python3.10/site-packages/ray/anyscale/data/_internal/cluster_autoscaler/productivity_calculator.py:174: RuntimeWarning: invalid value encountered in divide\n",
"\u001b[36m(SplitCoordinator pid=7661)\u001b[0m gpu_fraction_per_op = (optimal_num_tasks_per_op * num_gpus_per_op) / np.sum(\n",
"\u001b[36m(SplitCoordinator pid=7661)\u001b[0m ✔️ Dataset train_12_0 execution finished in 5.10 seconds\n",
"\u001b[36m(RayTrainWorker pid=4088, ip=10.0.163.141)\u001b[0m /tmp/ray/session_2025-09-30_14-10-46_627006_2395/runtime_resources/pip/72a6e451f55d87eb50ebbf5bc30a4a57ed513d34/virtualenv/lib/python3.10/site-packages/lightning/pytorch/loops/fit_loop.py:527: Found 455 module(s) in eval mode at the start of training. This may lead to unexpected behavior during training. If this is intentional, you can ignore this warning.\n"
]
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"Epoch 0: | | 0/? [00:00<?, ?it/s] 141)\u001b[0m \n",
"Epoch 0: | | 1/? [00:20<00:00, 0.05it/s, v_num=0, train_loss=14.60]\n",
"Epoch 0: | | 2/? [00:36<00:00, 0.06it/s, v_num=0, train_loss=13.90]\n",
"Epoch 0: | | 3/? [00:52<00:00, 0.06it/s, v_num=0, train_loss=14.80]\n",
"Epoch 0: | | 4/? [01:09<00:00, 0.06it/s, v_num=0, train_loss=14.10]\n",
"Epoch 0: | | 5/? [01:27<00:00, 0.06it/s, v_num=0, train_loss=14.90]\n"
]
},
{
"name": "stderr",
"output_type": "stream",
"text": [
"\u001b[36m(RayTrainWorker pid=4090, ip=10.0.191.43)\u001b[0m Checkpoint successfully created at: Checkpoint(filesystem=local, path=/mnt/cluster_storage/finetune_dolly-v2-7b-trial1/checkpoint_2025-09-30_14-23-45.998076)\n",
"\u001b[36m(RayTrainWorker pid=4090, ip=10.0.191.43)\u001b[0m Reporting training result 1: TrainingResult(checkpoint=Checkpoint(filesystem=local, path=/mnt/cluster_storage/finetune_dolly-v2-7b-trial1/checkpoint_2025-09-30_14-23-45.998076), metrics={'train_loss': 14.9804105758667, 'epoch': 0, 'step': 5})\n",
"\u001b[36m(RayTrainWorker pid=4088, ip=10.0.163.141)\u001b[0m Checkpoint successfully created at: Checkpoint(filesystem=local, path=/mnt/cluster_storage/finetune_dolly-v2-7b-trial1/checkpoint_2025-09-30_14-23-45.998076)\u001b[32m [repeated 12x across cluster]\u001b[0m\n",
"\u001b[36m(RayTrainWorker pid=4088, ip=10.0.163.141)\u001b[0m Reporting training result 1: TrainingResult(checkpoint=Checkpoint(filesystem=local, path=/mnt/cluster_storage/finetune_dolly-v2-7b-trial1/checkpoint_2025-09-30_14-23-45.998076), metrics={'train_loss': 14.879826545715332, 'epoch': 0, 'step': 5})\u001b[32m [repeated 12x across cluster]\u001b[0m\n",
"\u001b[36m(RayTrainWorker pid=4088, ip=10.0.163.141)\u001b[0m `Trainer.fit` stopped: `max_epochs=1` reached.\n"
]
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"Epoch 0: | | 5/? [06:07<00:00, 0.01it/s, v_num=0, train_loss=14.90]\n"
]
},
{
"data": {
"text/plain": [
"Result(metrics={'train_loss': 14.879826545715332, 'epoch': 0, 'step': 5}, checkpoint=Checkpoint(filesystem=local, path=/mnt/cluster_storage/finetune_dolly-v2-7b-trial1/checkpoint_2025-09-30_14-23-45.998076), error=None, path='/mnt/cluster_storage/finetune_dolly-v2-7b-trial1', metrics_dataframe= train_loss epoch step\n",
"0 14.879827 0 5, best_checkpoints=[(Checkpoint(filesystem=local, path=/mnt/cluster_storage/finetune_dolly-v2-7b-trial1/checkpoint_2025-09-30_14-23-45.998076), {'train_loss': 14.879826545715332, 'epoch': 0, 'step': 5})], _storage_filesystem=<pyarrow._fs.LocalFileSystem object at 0x737af20de5b0>)"
]
},
"execution_count": 16,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"from ray.train.torch import TorchTrainer\n",
"from ray.train import RunConfig, ScalingConfig, CheckpointConfig\n",
"\n",
"# Save Ray Train checkpoints according to the performance on validation set\n",
"run_config = RunConfig(\n",
" name=\"finetune_dolly-v2-7b-trial1\",\n",
" storage_path=storage_path,\n",
" checkpoint_config=CheckpointConfig(num_to_keep=1),\n",
")\n",
"\n",
"# Scale the FSDP training workload across 16 GPUs\n",
"# You can change this config based on your compute resources.\n",
"scaling_config = ScalingConfig(\n",
" num_workers=num_workers, use_gpu=True\n",
")\n",
"\n",
"# Configuration to pass into train_func\n",
"train_config = {\n",
" \"lr\": 2e-5,\n",
" \"eps\": 1e-8,\n",
" \"strategy\": fsdp_strategy,\n",
" \"batch_size_per_worker\": batch_size_per_worker\n",
"}\n",
"\n",
"# Define a TorchTrainer and launch you training workload\n",
"ray_trainer = TorchTrainer(\n",
" train_func,\n",
" train_loop_config=train_config,\n",
" run_config=run_config,\n",
" scaling_config=scaling_config,\n",
" datasets={\"train\": train_ds},\n",
")\n",
"result = ray_trainer.fit()\n",
"\n",
"result\n"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"We finished training in 2877s. The price for an on-demand g4dn.4xlarge instance is `$1.204/hour`, while a g4dn.8xlarge instance costs `$2.176/hour`. The total cost would be `($1.204 * 15 + $2.176) * 2877 / 3600 = $16.17`."
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Text-generation with HuggingFace Pipeline\n",
"\n",
"We can use the [HuggingFace Pipeline](https://huggingface.co/docs/transformers/main_classes/pipelines) to generate predictions from our fine-tuned model. Let's input some prompts and see if our tuned Dolly can speak like Shakespeare:"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"tags": []
},
"outputs": [],
"source": [
"import os\n",
"from transformers import pipeline\n",
"\n",
"@ray.remote(num_gpus=1)\n",
"def generate_tokens():\n",
" tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, padding_side=\"right\")\n",
"\n",
" ckpt_path = os.path.join(result.checkpoint.path, \"checkpoint.ckpt\")\n",
"\n",
" dolly = DollyV2Model.load_from_checkpoint(ckpt_path, map_location=torch.device(\"cpu\"))\n",
"\n",
" nlp_pipeline = pipeline(\n",
" task=\"text-generation\", \n",
" model=dolly.model, \n",
" tokenizer=tokenizer, \n",
" device_map=\"auto\"\n",
" )\n",
"\n",
" tokens = []\n",
" for prompt in [\"This is\", \"I am\", \"Once more\"]:\n",
" tokens.append(nlp_pipeline(prompt, max_new_tokens=20, do_sample=True, pad_token_id=tokenizer.eos_token_id))\n",
"\n",
" return tokens\n",
"\n",
"ref = generate_tokens.remote()\n",
"output = ray.get(ref)"
]
},
{
"cell_type": "code",
"execution_count": 23,
"metadata": {
"tags": []
},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"[{'generated_text': \"This is more like it:\\n\\nIt's just a guess, but maybe the extra processing power of Intel\"}]\n",
"[{'generated_text': \"I am the biggest fan of your wife's writing, and this novella was fantastic. So interesting to see\"}]\n",
"[{'generated_text': 'Once more I wish I could make sense of it.\" \"My friend, you can leave all this behind you'}]\n"
]
}
],
"source": [
"for generated_tokens in output:\n",
" print(generated_tokens)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"References:\n",
"- [PyTorch FSDP Tutorial](https://www.youtube.com/watch?v=8_k76AHu__s&list=PL_lsbAsL_o2BT6aerEKgIoufVD_fodnuT)\n",
"- [Getting Started with Fully Sharded Data Parallel(FSDP)](https://pytorch.org/tutorials/intermediate/FSDP_tutorial.html#:~:text=FSDP%20is%20a%20type%20of,sizes%20for%20our%20training%20job.)\n",
"- [Fully Sharded Data Parallel: faster AI training with fewer GPUs](https://engineering.fb.com/2021/07/15/open-source/fsdp/)\n",
"- [Hugging Face: dolly-v2-7b Model Card](https://huggingface.co/databricks/dolly-v2-7b)\n",
"- [Hugging Face: Handling big models for inference](https://huggingface.co/docs/accelerate/usage_guides/big_modeling)"
]
}
],
"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.18"
},
"orphan": true
},
"nbformat": 4,
"nbformat_minor": 4
}
doc/source/train/examples/lightning/lightning_cola_advanced.ipynb
+1186
View File
File diff suppressed because it is too large Load Diff
doc/source/train/examples/lightning/lightning_mnist_example.ipynb
+947
View File
@@ -0,0 +1,947 @@
{
"cells": [
{
"cell_type": "markdown",
"metadata": {},
"source": [
"# Train a Pytorch Lightning Image Classifier\n",
"\n",
"<a id=\"try-anyscale-quickstart-lightning_mnist_example\" href=\"https://console.anyscale.com/register/ha?render_flow=ray&utm_source=ray_docs&utm_medium=docs&utm_campaign=lightning_mnist_example\">\n",
" <img src=\"../../../_static/img/run-on-anyscale.svg\" alt=\"try-anyscale-quickstart\">\n",
"</a>\n",
"<br></br>\n",
"\n",
"This example introduces how to train a Pytorch Lightning Module using Ray Train {class}`TorchTrainer <ray.train.torch.TorchTrainer>`. It demonstrates how to train a basic neural network on the MNIST dataset with distributed data parallelism.\n"
]
},
{
"cell_type": "code",
"execution_count": 1,
"metadata": {},
"outputs": [],
"source": [
"!pip install \"torchmetrics>=0.9\" \"lightning>=2.0\""
]
},
{
"cell_type": "code",
"execution_count": 2,
"metadata": {},
"outputs": [],
"source": [
"import os\n",
"import numpy as np\n",
"import random\n",
"import torch\n",
"import torch.nn as nn\n",
"import torch.nn.functional as F\n",
"from filelock import FileLock\n",
"from torch.utils.data import DataLoader, random_split, Subset\n",
"from torchmetrics import Accuracy\n",
"from torchvision.datasets import MNIST\n",
"from torchvision import transforms\n",
"\n",
"try:\n",
" import lightning.pytorch as pl\n",
" from lightning.pytorch.loggers import CSVLogger\n",
"except ModuleNotFoundError:\n",
" import pytorch_lightning as pl\n",
" from pytorch_lightning.loggers import CSVLogger"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Prepare a dataset and module\n",
"\n",
"The Pytorch Lightning Trainer takes either `torch.utils.data.DataLoader` or `pl.LightningDataModule` as data inputs. You can continue using them without any changes with Ray Train. "
]
},
{
"cell_type": "code",
"execution_count": 4,
"metadata": {},
"outputs": [],
"source": [
"class MNISTDataModule(pl.LightningDataModule):\n",
" def __init__(self, batch_size=100):\n",
" super().__init__()\n",
" self.data_dir = os.getcwd()\n",
" self.batch_size = batch_size\n",
" self.transform = transforms.Compose(\n",
" [transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))]\n",
" )\n",
"\n",
" def setup(self, stage=None):\n",
" with FileLock(f\"{self.data_dir}.lock\"):\n",
" mnist = MNIST(\n",
" self.data_dir, train=True, download=True, transform=self.transform\n",
" )\n",
"\n",
" # Split data into train and val sets\n",
" self.mnist_train, self.mnist_val = random_split(mnist, [55000, 5000])\n",
"\n",
" def train_dataloader(self):\n",
" return DataLoader(self.mnist_train, batch_size=self.batch_size, num_workers=4)\n",
"\n",
" def val_dataloader(self):\n",
" return DataLoader(self.mnist_val, batch_size=self.batch_size, num_workers=4)\n",
"\n",
" def test_dataloader(self):\n",
" with FileLock(f\"{self.data_dir}.lock\"):\n",
" self.mnist_test = MNIST(\n",
" self.data_dir, train=False, download=True, transform=self.transform\n",
" )\n",
" return DataLoader(self.mnist_test, batch_size=self.batch_size, num_workers=4)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Next, define a simple multi-layer perception as the subclass of `pl.LightningModule`."
]
},
{
"cell_type": "code",
"execution_count": 5,
"metadata": {},
"outputs": [],
"source": [
"class MNISTClassifier(pl.LightningModule):\n",
" def __init__(self, lr=1e-3, feature_dim=128):\n",
" torch.manual_seed(421)\n",
" super(MNISTClassifier, self).__init__()\n",
" self.save_hyperparameters()\n",
"\n",
" self.linear_relu_stack = nn.Sequential(\n",
" nn.Linear(28 * 28, feature_dim),\n",
" nn.ReLU(),\n",
" nn.Linear(feature_dim, 10),\n",
" nn.ReLU(),\n",
" )\n",
" self.lr = lr\n",
" self.accuracy = Accuracy(task=\"multiclass\", num_classes=10, top_k=1)\n",
" self.eval_loss = []\n",
" self.eval_accuracy = []\n",
" self.test_accuracy = []\n",
" pl.seed_everything(888)\n",
"\n",
" def forward(self, x):\n",
" x = x.view(-1, 28 * 28)\n",
" x = self.linear_relu_stack(x)\n",
" return x\n",
"\n",
" def training_step(self, batch, batch_idx):\n",
" x, y = batch\n",
" y_hat = self(x)\n",
" loss = torch.nn.functional.cross_entropy(y_hat, y)\n",
" self.log(\"train_loss\", loss)\n",
" return loss\n",
"\n",
" def validation_step(self, val_batch, batch_idx):\n",
" loss, acc = self._shared_eval(val_batch)\n",
" self.log(\"val_accuracy\", acc)\n",
" self.eval_loss.append(loss)\n",
" self.eval_accuracy.append(acc)\n",
" return {\"val_loss\": loss, \"val_accuracy\": acc}\n",
"\n",
" def test_step(self, test_batch, batch_idx):\n",
" loss, acc = self._shared_eval(test_batch)\n",
" self.test_accuracy.append(acc)\n",
" self.log(\"test_accuracy\", acc, sync_dist=True, on_epoch=True)\n",
" return {\"test_loss\": loss, \"test_accuracy\": acc}\n",
"\n",
" def _shared_eval(self, batch):\n",
" x, y = batch\n",
" logits = self.forward(x)\n",
" loss = F.nll_loss(logits, y)\n",
" acc = self.accuracy(logits, y)\n",
" return loss, acc\n",
"\n",
" def on_validation_epoch_end(self):\n",
" avg_loss = torch.stack(self.eval_loss).mean()\n",
" avg_acc = torch.stack(self.eval_accuracy).mean()\n",
" self.log(\"val_loss\", avg_loss, sync_dist=True)\n",
" self.log(\"val_accuracy\", avg_acc, sync_dist=True)\n",
" self.eval_loss.clear()\n",
" self.eval_accuracy.clear()\n",
" \n",
" def configure_optimizers(self):\n",
" optimizer = torch.optim.Adam(self.parameters(), lr=self.lr)\n",
" return optimizer"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"You don't need to modify the definition of the PyTorch Lightning model or datamodule."
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Define a training function\n",
"\n",
"This code defines a {ref}`training function <train-overview-training-function>` for each worker. Comparing the training function with the original PyTorch Lightning code, notice three main differences:\n",
"\n",
"- Distributed strategy: Use {class}`RayDDPStrategy <ray.train.lightning.RayDDPStrategy>`.\n",
"- Cluster environment: Use {class}`RayLightningEnvironment <ray.train.lightning.RayLightningEnvironment>`.\n",
"- Parallel devices: Always set to `devices=\"auto\"` to use all available devices configured by ``TorchTrainer``.\n",
"\n",
"See {ref}`Getting Started with PyTorch Lightning <train-pytorch-lightning>` for more information.\n",
"\n",
"\n",
"For checkpoint reporting, Ray Train provides a minimal {class}`RayTrainReportCallback <ray.train.lightning.RayTrainReportCallback>` class that reports metrics and checkpoints at the end of each train epoch. For more complex checkpoint logic, implement custom callbacks. See {ref}`Saving and Loading Checkpoint <train-checkpointing>`."
]
},
{
"cell_type": "code",
"execution_count": 6,
"metadata": {},
"outputs": [],
"source": [
"use_gpu = True # Set to False if you want to run without GPUs\n",
"num_workers = 4"
]
},
{
"cell_type": "code",
"execution_count": 7,
"metadata": {},
"outputs": [],
"source": [
"try:\n",
" import lightning.pytorch as pl\n",
"except ModuleNotFoundError:\n",
" import pytorch_lightning as pl\n",
"\n",
"from ray.train import RunConfig, ScalingConfig, CheckpointConfig\n",
"from ray.train.torch import TorchTrainer\n",
"from ray.train.lightning import (\n",
" RayDDPStrategy,\n",
" RayLightningEnvironment,\n",
" RayTrainReportCallback,\n",
" prepare_trainer,\n",
")\n",
"\n",
"def train_func_per_worker():\n",
" model = MNISTClassifier(lr=1e-3, feature_dim=128)\n",
" datamodule = MNISTDataModule(batch_size=128)\n",
"\n",
" trainer = pl.Trainer(\n",
" devices=\"auto\",\n",
" strategy=RayDDPStrategy(),\n",
" plugins=[RayLightningEnvironment()],\n",
" callbacks=[RayTrainReportCallback()],\n",
" max_epochs=10,\n",
" accelerator=\"gpu\" if use_gpu else \"cpu\",\n",
" log_every_n_steps=100,\n",
" logger=CSVLogger(\"logs\"),\n",
" )\n",
" \n",
" trainer = prepare_trainer(trainer)\n",
" \n",
" # Train model\n",
" trainer.fit(model, datamodule=datamodule)\n",
"\n",
" # Evaluation on the test dataset\n",
" trainer.test(model, datamodule=datamodule)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Now put everything together:"
]
},
{
"cell_type": "code",
"execution_count": 8,
"metadata": {},
"outputs": [],
"source": [
"scaling_config = ScalingConfig(num_workers=num_workers, use_gpu=use_gpu)\n",
"\n",
"run_config = RunConfig(\n",
" name=\"ptl-mnist-example\",\n",
" storage_path=\"/tmp/ray_results\",\n",
" checkpoint_config=CheckpointConfig(\n",
" num_to_keep=3,\n",
" checkpoint_score_attribute=\"val_accuracy\",\n",
" checkpoint_score_order=\"max\",\n",
" ),\n",
")\n",
"\n",
"trainer = TorchTrainer(\n",
" train_func_per_worker,\n",
" scaling_config=scaling_config,\n",
" run_config=run_config,\n",
")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Now fit your trainer:"
]
},
{
"cell_type": "code",
"execution_count": 9,
"metadata": {},
"outputs": [
{
"data": {
"text/html": [
"<div class=\"tuneStatus\">\n",
" <div style=\"display: flex;flex-direction: row\">\n",
" <div style=\"display: flex;flex-direction: column;\">\n",
" <h3>Tune Status</h3>\n",
" <table>\n",
"<tbody>\n",
"<tr><td>Current time:</td><td>2023-08-07 23:41:11</td></tr>\n",
"<tr><td>Running for: </td><td>00:00:39.80 </td></tr>\n",
"<tr><td>Memory: </td><td>24.2/186.6 GiB </td></tr>\n",
"</tbody>\n",
"</table>\n",
" </div>\n",
" <div class=\"vDivider\"></div>\n",
" <div class=\"systemInfo\">\n",
" <h3>System Info</h3>\n",
" Using FIFO scheduling algorithm.<br>Logical resource usage: 1.0/48 CPUs, 4.0/4 GPUs\n",
" </div>\n",
" \n",
" </div>\n",
" <div class=\"hDivider\"></div>\n",
" <div class=\"trialStatus\">\n",
" <h3>Trial Status</h3>\n",
" <table>\n",
"<thead>\n",
"<tr><th>Trial name </th><th>status </th><th>loc </th><th style=\"text-align: right;\"> iter</th><th style=\"text-align: right;\"> total time (s)</th><th style=\"text-align: right;\"> train_loss</th><th style=\"text-align: right;\"> val_accuracy</th><th style=\"text-align: right;\"> val_loss</th></tr>\n",
"</thead>\n",
"<tbody>\n",
"<tr><td>TorchTrainer_78346_00000</td><td>TERMINATED</td><td>10.0.6.244:120026</td><td style=\"text-align: right;\"> 10</td><td style=\"text-align: right;\"> 29.0221</td><td style=\"text-align: right;\"> 0.0315938</td><td style=\"text-align: right;\"> 0.970002</td><td style=\"text-align: right;\"> -12.3466</td></tr>\n",
"</tbody>\n",
"</table>\n",
" </div>\n",
"</div>\n",
"<style>\n",
".tuneStatus {\n",
" color: var(--jp-ui-font-color1);\n",
"}\n",
".tuneStatus .systemInfo {\n",
" display: flex;\n",
" flex-direction: column;\n",
"}\n",
".tuneStatus td {\n",
" white-space: nowrap;\n",
"}\n",
".tuneStatus .trialStatus {\n",
" display: flex;\n",
" flex-direction: column;\n",
"}\n",
".tuneStatus h3 {\n",
" font-weight: bold;\n",
"}\n",
".tuneStatus .hDivider {\n",
" border-bottom-width: var(--jp-border-width);\n",
" border-bottom-color: var(--jp-border-color0);\n",
" border-bottom-style: solid;\n",
"}\n",
".tuneStatus .vDivider {\n",
" border-left-width: var(--jp-border-width);\n",
" border-left-color: var(--jp-border-color0);\n",
" border-left-style: solid;\n",
" margin: 0.5em 1em 0.5em 1em;\n",
"}\n",
"</style>\n"
],
"text/plain": [
"<IPython.core.display.HTML object>"
]
},
"metadata": {},
"output_type": "display_data"
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"\u001b[2m\u001b[36m(TorchTrainer pid=120026)\u001b[0m Starting distributed worker processes: ['120176 (10.0.6.244)', '120177 (10.0.6.244)', '120178 (10.0.6.244)', '120179 (10.0.6.244)']\n",
"\u001b[2m\u001b[36m(RayTrainWorker pid=120176)\u001b[0m Setting up process group for: env:// [rank=0, world_size=4]\n",
"\u001b[2m\u001b[36m(RayTrainWorker pid=120176)\u001b[0m [rank: 0] Global seed set to 888\n",
"\u001b[2m\u001b[36m(RayTrainWorker pid=120176)\u001b[0m GPU available: True (cuda), used: True\n",
"\u001b[2m\u001b[36m(RayTrainWorker pid=120176)\u001b[0m TPU available: False, using: 0 TPU cores\n",
"\u001b[2m\u001b[36m(RayTrainWorker pid=120176)\u001b[0m IPU available: False, using: 0 IPUs\n",
"\u001b[2m\u001b[36m(RayTrainWorker pid=120176)\u001b[0m HPU available: False, using: 0 HPUs\n"
]
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"\u001b[2m\u001b[36m(RayTrainWorker pid=120178)\u001b[0m Downloading http://yann.lecun.com/exdb/mnist/train-images-idx3-ubyte.gz\n",
"\u001b[2m\u001b[36m(RayTrainWorker pid=120178)\u001b[0m Downloading http://yann.lecun.com/exdb/mnist/train-images-idx3-ubyte.gz to /tmp/ray_results/ptl-mnist-example/TorchTrainer_78346_00000_0_2023-08-07_23-40-31/rank_2/MNIST/raw/train-images-idx3-ubyte.gz\n"
]
},
{
"name": "stdout",
"output_type": "stream",
"text": [
" 0%| | 0/9912422 [00:00<?, ?it/s]\n",
"100%|██████████| 9912422/9912422 [00:00<00:00, 94562894.32it/s]\n",
" 9%|▉ | 917504/9912422 [00:00<00:00, 9166590.91it/s]\n",
"100%|██████████| 9912422/9912422 [00:00<00:00, 115619443.32it/s]\n"
]
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"\u001b[2m\u001b[36m(RayTrainWorker pid=120179)\u001b[0m Extracting /tmp/ray_results/ptl-mnist-example/TorchTrainer_78346_00000_0_2023-08-07_23-40-31/rank_3/MNIST/raw/train-images-idx3-ubyte.gz to /tmp/ray_results/ptl-mnist-example/TorchTrainer_78346_00000_0_2023-08-07_23-40-31/rank_3/MNIST/raw\n",
"\u001b[2m\u001b[36m(RayTrainWorker pid=120176)\u001b[0m \n"
]
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"\u001b[2m\u001b[36m(RayTrainWorker pid=120177)\u001b[0m Missing logger folder: logs/lightning_logs\n",
"\u001b[2m\u001b[36m(RayTrainWorker pid=120176)\u001b[0m LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0,1,2,3]\n",
"\u001b[2m\u001b[36m(RayTrainWorker pid=120176)\u001b[0m | Name | Type | Params\n",
"\u001b[2m\u001b[36m(RayTrainWorker pid=120176)\u001b[0m ---------------------------------------------------------\n",
"\u001b[2m\u001b[36m(RayTrainWorker pid=120176)\u001b[0m 0 | linear_relu_stack | Sequential | 101 K \n",
"\u001b[2m\u001b[36m(RayTrainWorker pid=120176)\u001b[0m 1 | accuracy | MulticlassAccuracy | 0 \n",
"\u001b[2m\u001b[36m(RayTrainWorker pid=120176)\u001b[0m ---------------------------------------------------------\n",
"\u001b[2m\u001b[36m(RayTrainWorker pid=120176)\u001b[0m 101 K Trainable params\n",
"\u001b[2m\u001b[36m(RayTrainWorker pid=120176)\u001b[0m 0 Non-trainable params\n",
"\u001b[2m\u001b[36m(RayTrainWorker pid=120176)\u001b[0m 101 K Total params\n",
"\u001b[2m\u001b[36m(RayTrainWorker pid=120176)\u001b[0m 0.407 Total estimated model params size (MB)\n"
]
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"Sanity Checking: 0it [00:00, ?it/s])\u001b[0m \n",
"Sanity Checking DataLoader 0: 0%| | 0/2 [00:00<?, ?it/s]\n",
"Sanity Checking DataLoader 0: 100%|██████████| 2/2 [00:00<00:00, 2.69it/s]\n"
]
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"\u001b[2m\u001b[36m(RayTrainWorker pid=120176)\u001b[0m /mnt/cluster_storage/pypi/lib/python3.9/site-packages/pytorch_lightning/trainer/connectors/logger_connector/result.py:432: PossibleUserWarning: It is recommended to use `self.log('val_accuracy', ..., sync_dist=True)` when logging on epoch level in distributed setting to accumulate the metric across devices.\n",
"\u001b[2m\u001b[36m(RayTrainWorker pid=120176)\u001b[0m warning_cache.warn(\n",
"\u001b[2m\u001b[36m(RayTrainWorker pid=120179)\u001b[0m [rank: 3] Global seed set to 888\u001b[32m [repeated 7x across cluster] (Ray deduplicates logs by default. Set RAY_DEDUP_LOGS=0 to disable log deduplication, or see https://docs.ray.io/en/master/ray-observability/ray-logging.html#log-deduplication for more options.)\u001b[0m\n"
]
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"Epoch 0: 0%| | 0/108 [00:00<?, ?it/s] \n",
"Epoch 0: 12%|█▏ | 13/108 [00:00<00:02, 39.35it/s, v_num=0]\n",
"Epoch 0: 25%|██▌ | 27/108 [00:00<00:01, 59.26it/s, v_num=0]\n",
"Epoch 0: 26%|██▌ | 28/108 [00:00<00:01, 61.03it/s, v_num=0]\n",
"Epoch 0: 27%|██▋ | 29/108 [00:00<00:01, 62.76it/s, v_num=0]\n",
"Epoch 0: 42%|████▏ | 45/108 [00:00<00:00, 81.02it/s, v_num=0]\n",
"Epoch 0: 53%|█████▎ | 57/108 [00:00<00:00, 86.01it/s, v_num=0]\n",
"Epoch 0: 64%|██████▍ | 69/108 [00:00<00:00, 88.63it/s, v_num=0]\n",
"Epoch 0: 81%|████████ | 87/108 [00:00<00:00, 98.04it/s, v_num=0]\n",
"Epoch 0: 81%|████████▏ | 88/108 [00:00<00:00, 98.69it/s, v_num=0]\n",
"Epoch 0: 82%|████████▏ | 89/108 [00:00<00:00, 99.34it/s, v_num=0]\n",
"Epoch 0: 96%|█████████▋| 104/108 [00:00<00:00, 104.14it/s, v_num=0]\n",
"Epoch 0: 97%|█████████▋| 105/108 [00:01<00:00, 104.71it/s, v_num=0]\n",
"Epoch 0: 98%|█████████▊| 106/108 [00:01<00:00, 105.22it/s, v_num=0]\n",
"Epoch 0: 100%|██████████| 108/108 [00:01<00:00, 105.79it/s, v_num=0]\n",
"Validation: 0it [00:00, ?it/s]\u001b[A76)\u001b[0m \n",
"Validation: 0%| | 0/10 [00:00<?, ?it/s]\u001b[A\n",
"Validation DataLoader 0: 0%| | 0/10 [00:00<?, ?it/s]\u001b[A\n",
"Validation DataLoader 0: 10%|█ | 1/10 [00:00<00:00, 171.69it/s]\u001b[A\n",
"Validation DataLoader 0: 20%|██ | 2/10 [00:00<00:00, 200.99it/s]\u001b[A\n",
"Validation DataLoader 0: 30%|███ | 3/10 [00:00<00:00, 221.66it/s]\u001b[A\n",
"Validation DataLoader 0: 40%|████ | 4/10 [00:00<00:00, 215.50it/s]\u001b[A\n",
"Validation DataLoader 0: 50%|█████ | 5/10 [00:00<00:00, 194.14it/s]\u001b[A\n",
"Validation DataLoader 0: 60%|██████ | 6/10 [00:00<00:00, 205.63it/s]\u001b[A\n",
"Validation DataLoader 0: 70%|███████ | 7/10 [00:00<00:00, 215.27it/s]\u001b[A\n",
"Validation DataLoader 0: 80%|████████ | 8/10 [00:00<00:00, 216.26it/s]\u001b[A\n",
"Validation DataLoader 0: 90%|█████████ | 9/10 [00:00<00:00, 198.67it/s]\u001b[A\n",
"Validation DataLoader 0: 100%|██████████| 10/10 [00:00<00:00, 205.79it/s]\u001b[A\n",
"Epoch 0: 100%|██████████| 108/108 [00:01<00:00, 79.84it/s, v_num=0] \n",
" \u001b[A\n",
"Epoch 1: 0%| | 0/108 [00:00<?, ?it/s, v_num=0] \n",
"Epoch 1: 11%|█ | 12/108 [00:00<00:02, 32.36it/s, v_num=0]\n",
"Epoch 1: 23%|██▎ | 25/108 [00:00<00:01, 50.16it/s, v_num=0]\n",
"Epoch 1: 37%|███▋ | 40/108 [00:00<00:01, 65.95it/s, v_num=0]\n",
"Epoch 1: 38%|███▊ | 41/108 [00:00<00:00, 67.05it/s, v_num=0]\n",
"Epoch 1: 50%|█████ | 54/108 [00:00<00:00, 75.52it/s, v_num=0]\n",
"Epoch 1: 51%|█████ | 55/108 [00:00<00:00, 76.40it/s, v_num=0]\n",
"Epoch 1: 62%|██████▏ | 67/108 [00:00<00:00, 81.72it/s, v_num=0]\n",
"\u001b[2m\u001b[36m(RayTrainWorker pid=120178)\u001b[0m Downloading http://yann.lecun.com/exdb/mnist/t10k-labels-idx1-ubyte.gz\u001b[32m [repeated 15x across cluster]\u001b[0m\n",
"Epoch 1: 77%|███████▋ | 83/108 [00:00<00:00, 89.48it/s, v_num=0]\n",
"\u001b[2m\u001b[36m(RayTrainWorker pid=120178)\u001b[0m Downloading http://yann.lecun.com/exdb/mnist/t10k-labels-idx1-ubyte.gz to /tmp/ray_results/ptl-mnist-example/TorchTrainer_78346_00000_0_2023-08-07_23-40-31/rank_2/MNIST/raw/t10k-labels-idx1-ubyte.gz\u001b[32m [repeated 15x across cluster]\u001b[0m\n",
"Epoch 1: 78%|███████▊ | 84/108 [00:00<00:00, 89.21it/s, v_num=0]\n",
"Epoch 1: 91%|█████████ | 98/108 [00:01<00:00, 93.27it/s, v_num=0]\n",
"\u001b[2m\u001b[36m(RayTrainWorker pid=120178)\u001b[0m Extracting /tmp/ray_results/ptl-mnist-example/TorchTrainer_78346_00000_0_2023-08-07_23-40-31/rank_2/MNIST/raw/t10k-labels-idx1-ubyte.gz to /tmp/ray_results/ptl-mnist-example/TorchTrainer_78346_00000_0_2023-08-07_23-40-31/rank_2/MNIST/raw\u001b[32m [repeated 15x across cluster]\u001b[0m\n",
"Epoch 1: 92%|█████████▏| 99/108 [00:01<00:00, 93.94it/s, v_num=0]\n",
"Epoch 1: 93%|█████████▎| 100/108 [00:01<00:00, 94.57it/s, v_num=0]\n",
"Epoch 1: 100%|██████████| 108/108 [00:01<00:00, 98.06it/s, v_num=0]\n",
"Validation: 0it [00:00, ?it/s]\u001b[A76)\u001b[0m \n",
"Validation: 0%| | 0/10 [00:00<?, ?it/s]\u001b[A\n",
"Validation DataLoader 0: 0%| | 0/10 [00:00<?, ?it/s]\u001b[A\n",
"Validation DataLoader 0: 10%|█ | 1/10 [00:00<00:00, 320.27it/s]\u001b[A\n",
"Validation DataLoader 0: 20%|██ | 2/10 [00:00<00:00, 291.99it/s]\u001b[A\n",
"\u001b[2m\u001b[36m(RayTrainWorker pid=120176)\u001b[0m \u001b[32m [repeated 19x across cluster]\u001b[0m\n",
"Validation DataLoader 0: 30%|███ | 3/10 [00:00<00:00, 291.61it/s]\u001b[A\n",
"Validation DataLoader 0: 40%|████ | 4/10 [00:00<00:00, 268.90it/s]\u001b[A\n",
"Validation DataLoader 0: 50%|█████ | 5/10 [00:00<00:00, 290.07it/s]\u001b[A\n",
"Validation DataLoader 0: 60%|██████ | 6/10 [00:00<00:00, 293.52it/s]\u001b[A\n",
"Validation DataLoader 0: 70%|███████ | 7/10 [00:00<00:00, 299.70it/s]\u001b[A\n",
"Validation DataLoader 0: 80%|████████ | 8/10 [00:00<00:00, 304.80it/s]\u001b[A\n",
"Validation DataLoader 0: 90%|█████████ | 9/10 [00:00<00:00, 310.16it/s]\u001b[A\n",
"Validation DataLoader 0: 100%|██████████| 10/10 [00:00<00:00, 303.63it/s]\u001b[A\n",
"Epoch 1: 100%|██████████| 108/108 [00:01<00:00, 76.12it/s, v_num=0]\n",
" \u001b[A\n",
"Epoch 2: 0%| | 0/108 [00:00<?, ?it/s, v_num=0] \n",
"Epoch 2: 7%|▋ | 8/108 [00:00<00:03, 25.23it/s, v_num=0]\n",
"Epoch 2: 16%|█▌ | 17/108 [00:00<00:02, 39.73it/s, v_num=0]\n",
"Epoch 2: 17%|█▋ | 18/108 [00:00<00:02, 41.60it/s, v_num=0]\n",
"Epoch 2: 18%|█▊ | 19/108 [00:00<00:02, 43.49it/s, v_num=0]\n",
"Epoch 2: 18%|█▊ | 19/108 [00:00<00:02, 43.46it/s, v_num=0]\n",
"Epoch 2: 19%|█▊ | 20/108 [00:00<00:01, 45.27it/s, v_num=0]\n",
"Epoch 2: 27%|██▋ | 29/108 [00:00<00:01, 53.08it/s, v_num=0]\n",
"Epoch 2: 42%|████▏ | 45/108 [00:00<00:00, 69.12it/s, v_num=0]\n",
"Epoch 2: 43%|████▎ | 46/108 [00:00<00:00, 70.31it/s, v_num=0]\n",
"Epoch 2: 44%|████▎ | 47/108 [00:00<00:00, 71.51it/s, v_num=0]\n",
"Epoch 2: 44%|████▍ | 48/108 [00:00<00:00, 72.71it/s, v_num=0]\n",
"Epoch 2: 59%|█████▉ | 64/108 [00:00<00:00, 83.97it/s, v_num=0]\n",
"Epoch 2: 75%|███████▌ | 81/108 [00:00<00:00, 93.77it/s, v_num=0]\n",
"Epoch 2: 90%|████████▉ | 97/108 [00:00<00:00, 99.35it/s, v_num=0]\n",
"Epoch 2: 100%|██████████| 108/108 [00:01<00:00, 101.71it/s, v_num=0]\n",
"Validation: 0it [00:00, ?it/s]\u001b[A76)\u001b[0m \n",
"Validation: 0%| | 0/10 [00:00<?, ?it/s]\u001b[A\n",
"Validation DataLoader 0: 0%| | 0/10 [00:00<?, ?it/s]\u001b[A\n",
"Validation DataLoader 0: 10%|█ | 1/10 [00:00<00:00, 212.13it/s]\u001b[A\n",
"Validation DataLoader 0: 20%|██ | 2/10 [00:00<00:00, 184.45it/s]\u001b[A\n",
"Validation DataLoader 0: 30%|███ | 3/10 [00:00<00:00, 228.42it/s]\u001b[A\n",
"Validation DataLoader 0: 40%|████ | 4/10 [00:00<00:00, 225.00it/s]\u001b[A\n",
"Validation DataLoader 0: 50%|█████ | 5/10 [00:00<00:00, 250.65it/s]\u001b[A\n",
"Validation DataLoader 0: 60%|██████ | 6/10 [00:00<00:00, 251.36it/s]\u001b[A\n",
"Validation DataLoader 0: 70%|███████ | 7/10 [00:00<00:00, 268.85it/s]\u001b[A\n",
"Validation DataLoader 0: 80%|████████ | 8/10 [00:00<00:00, 256.15it/s]\u001b[A\n",
"Validation DataLoader 0: 90%|█████████ | 9/10 [00:00<00:00, 269.87it/s]\u001b[A\n",
"Epoch 2: 100%|██████████| 108/108 [00:01<00:00, 77.52it/s, v_num=0] it/s]\u001b[A\n",
" \u001b[A\n",
"Epoch 3: 0%| | 0/108 [00:00<?, ?it/s, v_num=0] \n",
"Epoch 3: 8%|▊ | 9/108 [00:00<00:03, 25.68it/s, v_num=0]\n",
"Epoch 3: 9%|▉ | 10/108 [00:00<00:03, 28.26it/s, v_num=0]\n",
"Epoch 3: 20%|██ | 22/108 [00:00<00:01, 48.10it/s, v_num=0]\n",
"Epoch 3: 21%|██▏ | 23/108 [00:00<00:01, 49.73it/s, v_num=0]\n",
"Epoch 3: 22%|██▏ | 24/108 [00:00<00:01, 51.34it/s, v_num=0]\n",
"Epoch 3: 23%|██▎ | 25/108 [00:00<00:01, 52.98it/s, v_num=0]\n",
"Epoch 3: 37%|███▋ | 40/108 [00:00<00:00, 69.67it/s, v_num=0]\n",
"Epoch 3: 51%|█████ | 55/108 [00:00<00:00, 80.93it/s, v_num=0]\n",
"Epoch 3: 64%|██████▍ | 69/108 [00:00<00:00, 87.15it/s, v_num=0]\n",
"Epoch 3: 65%|██████▍ | 70/108 [00:00<00:00, 88.04it/s, v_num=0]\n",
"Epoch 3: 66%|██████▌ | 71/108 [00:00<00:00, 88.92it/s, v_num=0]\n",
"Epoch 3: 77%|███████▋ | 83/108 [00:00<00:00, 92.62it/s, v_num=0]\n",
"Epoch 3: 86%|████████▌ | 93/108 [00:01<00:00, 92.33it/s, v_num=0]\n",
"Epoch 3: 87%|████████▋ | 94/108 [00:01<00:00, 92.93it/s, v_num=0]\n",
"Epoch 3: 88%|████████▊ | 95/108 [00:01<00:00, 93.61it/s, v_num=0]\n",
"Epoch 3: 100%|██████████| 108/108 [00:01<00:00, 97.43it/s, v_num=0]\n",
"Validation: 0it [00:00, ?it/s]\u001b[A76)\u001b[0m \n",
"Validation: 0%| | 0/10 [00:00<?, ?it/s]\u001b[A\n",
"Validation DataLoader 0: 0%| | 0/10 [00:00<?, ?it/s]\u001b[A\n",
"Validation DataLoader 0: 10%|█ | 1/10 [00:00<00:00, 308.50it/s]\u001b[A\n",
"Validation DataLoader 0: 20%|██ | 2/10 [00:00<00:00, 344.87it/s]\u001b[A\n",
"Validation DataLoader 0: 30%|███ | 3/10 [00:00<00:00, 375.98it/s]\u001b[A\n",
"Validation DataLoader 0: 40%|████ | 4/10 [00:00<00:00, 335.26it/s]\u001b[A\n",
"Validation DataLoader 0: 50%|█████ | 5/10 [00:00<00:00, 327.34it/s]\u001b[A\n",
"Validation DataLoader 0: 60%|██████ | 6/10 [00:00<00:00, 317.66it/s]\u001b[A\n",
"Validation DataLoader 0: 70%|███████ | 7/10 [00:00<00:00, 332.79it/s]\u001b[A\n",
"Validation DataLoader 0: 80%|████████ | 8/10 [00:00<00:00, 188.14it/s]\u001b[A\n",
"Validation DataLoader 0: 90%|█████████ | 9/10 [00:00<00:00, 201.21it/s]\u001b[A\n",
"Epoch 3: 100%|██████████| 108/108 [00:01<00:00, 75.94it/s, v_num=0]6it/s]\u001b[A\n",
" \u001b[A\n",
"Epoch 4: 0%| | 0/108 [00:00<?, ?it/s, v_num=0] \n",
"Epoch 4: 10%|█ | 11/108 [00:00<00:03, 30.09it/s, v_num=0]\n",
"Epoch 4: 20%|██ | 22/108 [00:00<00:01, 46.96it/s, v_num=0]\n",
"Epoch 4: 21%|██▏ | 23/108 [00:00<00:01, 47.88it/s, v_num=0]\n",
"Epoch 4: 35%|███▌ | 38/108 [00:00<00:01, 65.26it/s, v_num=0]\n",
"Epoch 4: 36%|███▌ | 39/108 [00:00<00:01, 65.73it/s, v_num=0]\n",
"Epoch 4: 53%|█████▎ | 57/108 [00:00<00:00, 81.51it/s, v_num=0]\n",
"Epoch 4: 54%|█████▎ | 58/108 [00:00<00:00, 82.56it/s, v_num=0]\n",
"Epoch 4: 68%|██████▊ | 73/108 [00:00<00:00, 89.69it/s, v_num=0]\n",
"Epoch 4: 69%|██████▊ | 74/108 [00:00<00:00, 90.53it/s, v_num=0]\n",
"Epoch 4: 83%|████████▎ | 90/108 [00:00<00:00, 98.32it/s, v_num=0]\n",
"Epoch 4: 98%|█████████▊| 106/108 [00:01<00:00, 103.12it/s, v_num=0]\n",
"Epoch 4: 100%|██████████| 108/108 [00:01<00:00, 103.78it/s, v_num=0]\n",
"Validation: 0it [00:00, ?it/s]\u001b[A76)\u001b[0m \n",
"Validation: 0%| | 0/10 [00:00<?, ?it/s]\u001b[A\n",
"Validation DataLoader 0: 0%| | 0/10 [00:00<?, ?it/s]\u001b[A\n",
"Validation DataLoader 0: 10%|█ | 1/10 [00:00<00:00, 268.49it/s]\u001b[A\n",
"Validation DataLoader 0: 20%|██ | 2/10 [00:00<00:00, 298.62it/s]\u001b[A\n",
"Validation DataLoader 0: 30%|███ | 3/10 [00:00<00:00, 282.88it/s]\u001b[A\n",
"Validation DataLoader 0: 40%|████ | 4/10 [00:00<00:00, 256.50it/s]\u001b[A\n",
"Validation DataLoader 0: 50%|█████ | 5/10 [00:00<00:00, 276.28it/s]\u001b[A\n",
"Validation DataLoader 0: 60%|██████ | 6/10 [00:00<00:00, 268.05it/s]\u001b[A\n",
"Validation DataLoader 0: 70%|███████ | 7/10 [00:00<00:00, 276.18it/s]\u001b[A\n",
"Validation DataLoader 0: 80%|████████ | 8/10 [00:00<00:00, 290.08it/s]\u001b[A\n",
"Validation DataLoader 0: 90%|█████████ | 9/10 [00:00<00:00, 261.92it/s]\u001b[A\n",
"Validation DataLoader 0: 100%|██████████| 10/10 [00:00<00:00, 274.00it/s]\u001b[A\n",
"Epoch 4: 100%|██████████| 108/108 [00:01<00:00, 78.54it/s, v_num=0] \n",
" \u001b[A\n",
"Epoch 5: 0%| | 0/108 [00:00<?, ?it/s, v_num=0] \n",
"Epoch 5: 5%|▍ | 5/108 [00:00<00:06, 15.52it/s, v_num=0]\n",
"\u001b[2m\u001b[36m(RayTrainWorker pid=120176)\u001b[0m \u001b[32m [repeated 9x across cluster]\u001b[0m\n",
"Epoch 5: 17%|█▋ | 18/108 [00:00<00:02, 42.53it/s, v_num=0]\n",
"Epoch 5: 26%|██▌ | 28/108 [00:00<00:01, 52.36it/s, v_num=0]\n",
"Epoch 5: 27%|██▋ | 29/108 [00:00<00:01, 53.91it/s, v_num=0]\n",
"Epoch 5: 28%|██▊ | 30/108 [00:00<00:01, 55.45it/s, v_num=0]\n",
"Epoch 5: 29%|██▊ | 31/108 [00:00<00:01, 56.96it/s, v_num=0]\n",
"Epoch 5: 37%|███▋ | 40/108 [00:00<00:01, 61.48it/s, v_num=0]\n",
"Epoch 5: 38%|███▊ | 41/108 [00:00<00:01, 62.61it/s, v_num=0]\n",
"Epoch 5: 39%|███▉ | 42/108 [00:00<00:01, 63.79it/s, v_num=0]\n",
"Epoch 5: 40%|███▉ | 43/108 [00:00<00:01, 64.96it/s, v_num=0]\n",
"Epoch 5: 48%|████▊ | 52/108 [00:00<00:00, 67.96it/s, v_num=0]\n",
"Epoch 5: 49%|████▉ | 53/108 [00:00<00:00, 68.88it/s, v_num=0]\n",
"Epoch 5: 50%|█████ | 54/108 [00:00<00:00, 69.77it/s, v_num=0]\n",
"Epoch 5: 62%|██████▏ | 67/108 [00:00<00:00, 76.43it/s, v_num=0]\n",
"Epoch 5: 78%|███████▊ | 84/108 [00:00<00:00, 85.56it/s, v_num=0]\n",
"Epoch 5: 79%|███████▊ | 85/108 [00:00<00:00, 86.17it/s, v_num=0]\n",
"Epoch 5: 93%|█████████▎| 100/108 [00:01<00:00, 92.27it/s, v_num=0]\n",
"Epoch 5: 100%|██████████| 108/108 [00:01<00:00, 94.81it/s, v_num=0]\n",
"Validation: 0it [00:00, ?it/s]\u001b[A76)\u001b[0m \n",
"Validation: 0%| | 0/10 [00:00<?, ?it/s]\u001b[A\n",
"Validation DataLoader 0: 0%| | 0/10 [00:00<?, ?it/s]\u001b[A\n",
"Validation DataLoader 0: 10%|█ | 1/10 [00:00<00:00, 255.91it/s]\u001b[A\n",
"Validation DataLoader 0: 20%|██ | 2/10 [00:00<00:00, 206.50it/s]\u001b[A\n",
"Validation DataLoader 0: 30%|███ | 3/10 [00:00<00:00, 214.91it/s]\u001b[A\n",
"Validation DataLoader 0: 40%|████ | 4/10 [00:00<00:00, 236.35it/s]\u001b[A\n",
"Validation DataLoader 0: 50%|█████ | 5/10 [00:00<00:00, 240.05it/s]\u001b[A\n",
"Validation DataLoader 0: 60%|██████ | 6/10 [00:00<00:00, 234.90it/s]\u001b[A\n",
"Validation DataLoader 0: 70%|███████ | 7/10 [00:00<00:00, 240.78it/s]\u001b[A\n",
"Validation DataLoader 0: 80%|████████ | 8/10 [00:00<00:00, 252.00it/s]\u001b[A\n",
"Validation DataLoader 0: 90%|█████████ | 9/10 [00:00<00:00, 255.18it/s]\u001b[A\n",
"Epoch 5: 100%|██████████| 108/108 [00:01<00:00, 72.36it/s, v_num=0]1it/s]\u001b[A\n",
" \u001b[A\n",
"Epoch 6: 0%| | 0/108 [00:00<?, ?it/s, v_num=0] \n",
"Epoch 6: 15%|█▍ | 16/108 [00:00<00:02, 44.27it/s, v_num=0]\n",
"Epoch 6: 25%|██▌ | 27/108 [00:00<00:01, 57.21it/s, v_num=0]\n",
"Epoch 6: 38%|███▊ | 41/108 [00:00<00:00, 70.86it/s, v_num=0]\n",
"Epoch 6: 39%|███▉ | 42/108 [00:00<00:00, 71.82it/s, v_num=0]\n",
"Epoch 6: 55%|█████▍ | 59/108 [00:00<00:00, 85.97it/s, v_num=0]\n",
"Epoch 6: 68%|██████▊ | 73/108 [00:00<00:00, 91.53it/s, v_num=0]\n",
"Epoch 6: 81%|████████ | 87/108 [00:00<00:00, 96.88it/s, v_num=0]\n",
"Epoch 6: 92%|█████████▏| 99/108 [00:00<00:00, 99.33it/s, v_num=0]\n",
"Epoch 6: 93%|█████████▎| 100/108 [00:01<00:00, 98.66it/s, v_num=0]\n",
"Epoch 6: 94%|█████████▎| 101/108 [00:01<00:00, 99.34it/s, v_num=0]\n",
"Epoch 6: 100%|██████████| 108/108 [00:01<00:00, 102.79it/s, v_num=0]\n",
"Validation: 0it [00:00, ?it/s]\u001b[A76)\u001b[0m \n",
"Validation: 0%| | 0/10 [00:00<?, ?it/s]\u001b[A\n",
"Validation DataLoader 0: 0%| | 0/10 [00:00<?, ?it/s]\u001b[A\n",
"Validation DataLoader 0: 10%|█ | 1/10 [00:00<00:00, 197.51it/s]\u001b[A\n",
"Validation DataLoader 0: 20%|██ | 2/10 [00:00<00:00, 143.68it/s]\u001b[A\n",
"Validation DataLoader 0: 30%|███ | 3/10 [00:00<00:00, 156.17it/s]\u001b[A\n",
"Validation DataLoader 0: 40%|████ | 4/10 [00:00<00:00, 180.52it/s]\u001b[A\n",
"Validation DataLoader 0: 50%|█████ | 5/10 [00:00<00:00, 205.25it/s]\u001b[A\n",
"Validation DataLoader 0: 60%|██████ | 6/10 [00:00<00:00, 212.20it/s]\u001b[A\n",
"Validation DataLoader 0: 70%|███████ | 7/10 [00:00<00:00, 195.64it/s]\u001b[A\n",
"Validation DataLoader 0: 80%|████████ | 8/10 [00:00<00:00, 211.21it/s]\u001b[A\n",
"Validation DataLoader 0: 90%|█████████ | 9/10 [00:00<00:00, 225.13it/s]\u001b[A\n",
"Epoch 6: 100%|██████████| 108/108 [00:01<00:00, 76.04it/s, v_num=0] it/s]\u001b[A\n",
" \u001b[A\n",
"Epoch 7: 0%| | 0/108 [00:00<?, ?it/s, v_num=0] \n",
"Epoch 7: 11%|█ | 12/108 [00:00<00:02, 33.31it/s, v_num=0]\n",
"Epoch 7: 20%|██ | 22/108 [00:00<00:01, 46.90it/s, v_num=0]\n",
"Epoch 7: 22%|██▏ | 24/108 [00:00<00:01, 50.49it/s, v_num=0]\n",
"Epoch 7: 31%|███▏ | 34/108 [00:00<00:01, 58.20it/s, v_num=0]\n",
"Epoch 7: 32%|███▏ | 35/108 [00:00<00:01, 59.59it/s, v_num=0]\n",
"Epoch 7: 33%|███▎ | 36/108 [00:00<00:01, 60.97it/s, v_num=0]\n",
"Epoch 7: 48%|████▊ | 52/108 [00:00<00:00, 74.69it/s, v_num=0]\n",
"Epoch 7: 64%|██████▍ | 69/108 [00:00<00:00, 85.96it/s, v_num=0]\n",
"Epoch 7: 80%|███████▉ | 86/108 [00:00<00:00, 94.41it/s, v_num=0]\n",
"Epoch 7: 81%|████████▏ | 88/108 [00:00<00:00, 95.91it/s, v_num=0]\n",
"Epoch 7: 97%|█████████▋| 105/108 [00:01<00:00, 102.61it/s, v_num=0]\n",
"Epoch 7: 100%|██████████| 108/108 [00:01<00:00, 103.00it/s, v_num=0]\n",
"Validation: 0it [00:00, ?it/s]\u001b[A76)\u001b[0m \n",
"Validation: 0%| | 0/10 [00:00<?, ?it/s]\u001b[A\n",
"Validation DataLoader 0: 0%| | 0/10 [00:00<?, ?it/s]\u001b[A\n",
"Validation DataLoader 0: 10%|█ | 1/10 [00:00<00:00, 215.46it/s]\u001b[A\n",
"Validation DataLoader 0: 20%|██ | 2/10 [00:00<00:00, 246.46it/s]\u001b[A\n",
"Validation DataLoader 0: 30%|███ | 3/10 [00:00<00:00, 264.39it/s]\u001b[A\n",
"Validation DataLoader 0: 40%|████ | 4/10 [00:00<00:00, 256.84it/s]\u001b[A\n",
"Validation DataLoader 0: 50%|█████ | 5/10 [00:00<00:00, 218.46it/s]\u001b[A\n",
"Validation DataLoader 0: 60%|██████ | 6/10 [00:00<00:00, 230.90it/s]\u001b[A\n",
"Validation DataLoader 0: 70%|███████ | 7/10 [00:00<00:00, 243.53it/s]\u001b[A\n",
"Validation DataLoader 0: 80%|████████ | 8/10 [00:00<00:00, 253.83it/s]\u001b[A\n",
"Validation DataLoader 0: 90%|█████████ | 9/10 [00:00<00:00, 249.22it/s]\u001b[A\n",
"Epoch 7: 100%|██████████| 108/108 [00:01<00:00, 78.36it/s, v_num=0] it/s]\u001b[A\n",
" \u001b[A\n",
"Epoch 8: 0%| | 0/108 [00:00<?, ?it/s, v_num=0] \n",
"Epoch 8: 7%|▋ | 8/108 [00:00<00:03, 25.72it/s, v_num=0]\n",
"Epoch 8: 19%|█▊ | 20/108 [00:00<00:01, 47.54it/s, v_num=0]\n",
"Epoch 8: 31%|███ | 33/108 [00:00<00:01, 62.61it/s, v_num=0]\n",
"\u001b[2m\u001b[36m(RayTrainWorker pid=120176)\u001b[0m \u001b[32m [repeated 7x across cluster]\u001b[0m\n",
"Epoch 8: 44%|████▎ | 47/108 [00:00<00:00, 74.45it/s, v_num=0]\n",
"Epoch 8: 55%|█████▍ | 59/108 [00:00<00:00, 81.84it/s, v_num=0]\n",
"Epoch 8: 56%|█████▌ | 60/108 [00:00<00:00, 80.82it/s, v_num=0]\n",
"Epoch 8: 57%|█████▋ | 62/108 [00:00<00:00, 82.74it/s, v_num=0]\n",
"Epoch 8: 58%|█████▊ | 63/108 [00:00<00:00, 83.69it/s, v_num=0]\n",
"Epoch 8: 70%|███████ | 76/108 [00:00<00:00, 88.60it/s, v_num=0]\n",
"Epoch 8: 85%|████████▌ | 92/108 [00:00<00:00, 96.53it/s, v_num=0]\n",
"Epoch 8: 86%|████████▌ | 93/108 [00:00<00:00, 96.21it/s, v_num=0]\n",
"Epoch 8: 87%|████████▋ | 94/108 [00:00<00:00, 96.72it/s, v_num=0]\n",
"Epoch 8: 88%|████████▊ | 95/108 [00:00<00:00, 97.32it/s, v_num=0]\n",
"Epoch 8: 89%|████████▉ | 96/108 [00:00<00:00, 98.03it/s, v_num=0]\n",
"Epoch 8: 100%|██████████| 108/108 [00:01<00:00, 102.15it/s, v_num=0]\n",
"Validation: 0it [00:00, ?it/s]\u001b[A76)\u001b[0m \n",
"Validation: 0%| | 0/10 [00:00<?, ?it/s]\u001b[A\n",
"Validation DataLoader 0: 0%| | 0/10 [00:00<?, ?it/s]\u001b[A\n",
"Validation DataLoader 0: 10%|█ | 1/10 [00:00<00:00, 228.96it/s]\u001b[A\n",
"Validation DataLoader 0: 20%|██ | 2/10 [00:00<00:00, 220.63it/s]\u001b[A\n",
"Validation DataLoader 0: 30%|███ | 3/10 [00:00<00:00, 220.41it/s]\u001b[A\n",
"Validation DataLoader 0: 40%|████ | 4/10 [00:00<00:00, 208.74it/s]\u001b[A\n",
"Validation DataLoader 0: 50%|█████ | 5/10 [00:00<00:00, 221.74it/s]\u001b[A\n",
"Validation DataLoader 0: 60%|██████ | 6/10 [00:00<00:00, 243.64it/s]\u001b[A\n",
"Validation DataLoader 0: 70%|███████ | 7/10 [00:00<00:00, 253.60it/s]\u001b[A\n",
"Validation DataLoader 0: 80%|████████ | 8/10 [00:00<00:00, 254.93it/s]\u001b[A\n",
"Validation DataLoader 0: 90%|█████████ | 9/10 [00:00<00:00, 207.23it/s]\u001b[A\n",
"Epoch 8: 100%|██████████| 108/108 [00:01<00:00, 78.28it/s, v_num=0] it/s]\u001b[A\n",
" \u001b[A\n",
"Epoch 9: 0%| | 0/108 [00:00<?, ?it/s, v_num=0] \n",
"Epoch 9: 11%|█ | 12/108 [00:00<00:02, 33.03it/s, v_num=0]\n",
"Epoch 9: 21%|██▏ | 23/108 [00:00<00:01, 48.82it/s, v_num=0]\n",
"Epoch 9: 31%|███ | 33/108 [00:00<00:01, 58.62it/s, v_num=0]\n",
"Epoch 9: 31%|███▏ | 34/108 [00:00<00:01, 58.61it/s, v_num=0]\n",
"Epoch 9: 32%|███▏ | 35/108 [00:00<00:01, 59.89it/s, v_num=0]\n",
"Epoch 9: 33%|███▎ | 36/108 [00:00<00:01, 61.11it/s, v_num=0]\n",
"Epoch 9: 46%|████▋ | 50/108 [00:00<00:00, 71.95it/s, v_num=0]\n",
"Epoch 9: 61%|██████ | 66/108 [00:00<00:00, 82.62it/s, v_num=0]\n",
"Epoch 9: 70%|███████ | 76/108 [00:00<00:00, 83.77it/s, v_num=0]\n",
"Epoch 9: 71%|███████▏ | 77/108 [00:00<00:00, 84.54it/s, v_num=0]\n",
"Epoch 9: 72%|███████▏ | 78/108 [00:00<00:00, 85.33it/s, v_num=0]\n",
"Epoch 9: 88%|████████▊ | 95/108 [00:01<00:00, 93.18it/s, v_num=0]\n",
"Epoch 9: 100%|██████████| 108/108 [00:01<00:00, 98.27it/s, v_num=0]\n",
"Validation: 0it [00:00, ?it/s]\u001b[A76)\u001b[0m \n",
"Validation: 0%| | 0/10 [00:00<?, ?it/s]\u001b[A\n",
"Validation DataLoader 0: 0%| | 0/10 [00:00<?, ?it/s]\u001b[A\n",
"Validation DataLoader 0: 10%|█ | 1/10 [00:00<00:00, 305.42it/s]\u001b[A\n",
"Validation DataLoader 0: 20%|██ | 2/10 [00:00<00:00, 337.39it/s]\u001b[A\n",
"Validation DataLoader 0: 30%|███ | 3/10 [00:00<00:00, 368.65it/s]\u001b[A\n",
"Validation DataLoader 0: 40%|████ | 4/10 [00:00<00:00, 361.22it/s]\u001b[A\n",
"Validation DataLoader 0: 50%|█████ | 5/10 [00:00<00:00, 250.96it/s]\u001b[A\n",
"Validation DataLoader 0: 60%|██████ | 6/10 [00:00<00:00, 271.98it/s]\u001b[A\n",
"Validation DataLoader 0: 70%|███████ | 7/10 [00:00<00:00, 289.64it/s]\u001b[A\n",
"Validation DataLoader 0: 80%|████████ | 8/10 [00:00<00:00, 304.16it/s]\u001b[A\n",
"Validation DataLoader 0: 90%|█████████ | 9/10 [00:00<00:00, 184.87it/s]\u001b[A\n",
"Validation DataLoader 0: 100%|██████████| 10/10 [00:00<00:00, 196.99it/s]\u001b[A\n",
"Epoch 9: 100%|██████████| 108/108 [00:01<00:00, 74.63it/s, v_num=0]\n",
" \u001b[A\n"
]
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"\u001b[2m\u001b[36m(RayTrainWorker pid=120176)\u001b[0m `Trainer.fit` stopped: `max_epochs=10` reached.\n",
"100%|██████████| 4542/4542 [00:00<00:00, 48474627.91it/s]\u001b[32m [repeated 14x across cluster]\u001b[0m\n",
"100%|██████████| 9912422/9912422 [00:00<00:00, 90032420.31it/s]\u001b[32m [repeated 2x across cluster]\u001b[0m\n",
"\u001b[2m\u001b[36m(RayTrainWorker pid=120176)\u001b[0m \u001b[32m [repeated 5x across cluster]\u001b[0m\n",
"\u001b[2m\u001b[36m(RayTrainWorker pid=120178)\u001b[0m Missing logger folder: logs/lightning_logs\u001b[32m [repeated 2x across cluster]\u001b[0m\n",
"\u001b[2m\u001b[36m(RayTrainWorker pid=120179)\u001b[0m LOCAL_RANK: 3 - CUDA_VISIBLE_DEVICES: [0,1,2,3]\u001b[32m [repeated 3x across cluster]\u001b[0m\n",
"\u001b[2m\u001b[36m(RayTrainWorker pid=120176)\u001b[0m [rank: 0] Global seed set to 888\n"
]
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"Epoch 9: 100%|██████████| 108/108 [00:01<00:00, 66.61it/s, v_num=0]\n"
]
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"\u001b[2m\u001b[36m(RayTrainWorker pid=120176)\u001b[0m /mnt/cluster_storage/pypi/lib/python3.9/site-packages/pytorch_lightning/trainer/connectors/data_connector.py:225: PossibleUserWarning: Using `DistributedSampler` with the dataloaders. During `trainer.test()`, it is recommended to use `Trainer(devices=1, num_nodes=1)` to ensure each sample/batch gets evaluated exactly once. Otherwise, multi-device settings use `DistributedSampler` that replicates some samples to make sure all devices have same batch size in case of uneven inputs.\n",
"\u001b[2m\u001b[36m(RayTrainWorker pid=120176)\u001b[0m rank_zero_warn(\n"
]
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"Testing DataLoader 0: 25%|██▌ | 5/20 [00:00<00:00, 146.57it/s]\n",
"Testing DataLoader 0: 100%|██████████| 20/20 [00:00<00:00, 163.98it/s]\n",
"Testing DataLoader 0: 100%|██████████| 20/20 [00:00<00:00, 125.34it/s]\n",
"\u001b[2m\u001b[36m(RayTrainWorker pid=120176)\u001b[0m ┏━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━┓\n",
"\u001b[2m\u001b[36m(RayTrainWorker pid=120176)\u001b[0m ┃ Test metric ┃ DataLoader 0 ┃\n",
"\u001b[2m\u001b[36m(RayTrainWorker pid=120176)\u001b[0m ┡━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━┩\n",
"\u001b[2m\u001b[36m(RayTrainWorker pid=120176)\u001b[0m │ test_accuracy │ 0.9740999937057495 │\n",
"\u001b[2m\u001b[36m(RayTrainWorker pid=120176)\u001b[0m └───────────────────────────┴───────────────────────────┘\n"
]
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"2023-08-07 23:41:11,072\tINFO tune.py:1145 -- Total run time: 39.92 seconds (39.80 seconds for the tuning loop).\n"
]
}
],
"source": [
"result = trainer.fit()"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Check training results and checkpoints"
]
},
{
"cell_type": "code",
"execution_count": 10,
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"Result(\n",
" metrics={'train_loss': 0.03159375861287117, 'val_accuracy': 0.9700015783309937, 'val_loss': -12.346583366394043, 'epoch': 9, 'step': 1080},\n",
" path='/tmp/ray_results/ptl-mnist-example/TorchTrainer_78346_00000_0_2023-08-07_23-40-31',\n",
" checkpoint=LegacyTorchCheckpoint(local_path=/tmp/ray_results/ptl-mnist-example/TorchTrainer_78346_00000_0_2023-08-07_23-40-31/checkpoint_000009)\n",
")"
]
},
"execution_count": null,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"result"
]
},
{
"cell_type": "code",
"execution_count": 11,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Validation Accuracy: 0.9700015783309937\n",
"Trial Directory: /tmp/ray_results/ptl-mnist-example/TorchTrainer_78346_00000_0_2023-08-07_23-40-31\n",
"['checkpoint_000007', 'checkpoint_000008', 'checkpoint_000009', 'events.out.tfevents.1691476838.ip-10-0-6-244', 'params.json', 'params.pkl', 'progress.csv', 'rank_0', 'rank_0.lock', 'rank_1', 'rank_1.lock', 'rank_2', 'rank_2.lock', 'rank_3', 'rank_3.lock', 'result.json']\n"
]
}
],
"source": [
"print(\"Validation Accuracy: \", result.metrics[\"val_accuracy\"])\n",
"print(\"Trial Directory: \", result.path)\n",
"print(sorted(os.listdir(result.path)))"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Ray Train saved three checkpoints(`checkpoint_000007`, `checkpoint_000008`, `checkpoint_000009`) in the trial directory. The following code retrieves the latest checkpoint from the fit results and loads it back into the model.\n",
"\n",
"If you lost the in-memory result object, you can restore the model from the checkpoint file. The checkpoint path is: `/tmp/ray_results/ptl-mnist-example/TorchTrainer_eb925_00000_0_2023-08-07_23-15-06/checkpoint_000009/checkpoint.ckpt`."
]
},
{
"cell_type": "code",
"execution_count": 12,
"metadata": {
"tags": []
},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Global seed set to 888\n"
]
},
{
"data": {
"text/plain": [
"MNISTClassifier(\n",
" (linear_relu_stack): Sequential(\n",
" (0): Linear(in_features=784, out_features=128, bias=True)\n",
" (1): ReLU()\n",
" (2): Linear(in_features=128, out_features=10, bias=True)\n",
" (3): ReLU()\n",
" )\n",
" (accuracy): MulticlassAccuracy()\n",
")"
]
},
"execution_count": null,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"checkpoint = result.checkpoint\n",
"\n",
"with checkpoint.as_directory() as ckpt_dir:\n",
" best_model = MNISTClassifier.load_from_checkpoint(f\"{ckpt_dir}/checkpoint.ckpt\")\n",
"\n",
"best_model"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## See also\n",
"\n",
"* {ref}`Getting Started with PyTorch Lightning <train-pytorch-lightning>` for a tutorial on using Ray Train and PyTorch Lightning \n",
"\n",
"* {doc}`Ray Train Examples <../../examples>` for more use cases"
]
}
],
"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"
},
"orphan": true,
"vscode": {
"interpreter": {
"hash": "a8c1140d108077f4faeb76b2438f85e4ed675f93d004359552883616a1acd54c"
}
}
},
"nbformat": 4,
"nbformat_minor": 4
}
doc/source/train/examples/lightning/vicuna_13b_lightning_deepspeed_finetune.ipynb
+1272
View File
File diff suppressed because it is too large Load Diff
doc/source/train/examples/pytorch/BUILD.bazel
+43
View File
@@ -0,0 +1,43 @@
load("//bazel:python.bzl", "py_test_run_all_notebooks")
filegroup(
name = "train_pytorch_examples",
srcs = glob(["*.ipynb"]),
visibility = ["//doc:__subpackages__"],
)
# GPU Tests
py_test_run_all_notebooks(
size = "large",
include = ["*.ipynb"],
data = ["//doc/source/train/examples/pytorch:train_pytorch_examples"],
exclude = ["convert_existing_pytorch_code_to_ray_train.ipynb"], # CPU test
tags = [
"exclusive",
"gpu",
"ray_air",
"team:ml",
],
)
# CPU Tests
py_test_run_all_notebooks(
size = "large",
include = ["convert_existing_pytorch_code_to_ray_train.ipynb"],
data = ["//doc/source/train/examples/pytorch:train_pytorch_examples"],
exclude = [],
tags = [
"exclusive",
"ray_air",
"team:ml",
],
)
filegroup(
name = "pytorch_examples_ci_configs",
srcs = glob([
"**/ci/aws.yaml",
"**/ci/gce.yaml",
]),
visibility = ["//doc/source/train/examples:__pkg__"],
)
doc/source/train/examples/pytorch/convert_existing_pytorch_code_to_ray_train.ipynb
+1030
View File
File diff suppressed because it is too large Load Diff
doc/source/train/examples/pytorch/dreambooth_finetuning.rst
+304
View File
@@ -0,0 +1,304 @@
:orphan:
Fine-tuning of Stable Diffusion with DreamBooth and Ray Train
=============================================================
.. raw:: html
<a id="try-anyscale-quickstart-dreambooth_finetuning" target="_blank" href="https://console.anyscale.com/register/ha?render_flow=ray&utm_source=ray_docs&utm_medium=docs&utm_campaign=dreambooth_finetuning">
<img src="../../../_static/img/run-on-anyscale.svg" alt="Run on Anyscale" />
<br/><br/>
</a>
This is an intermediate example that shows how to do DreamBooth fine-tuning of a Stable Diffusion model using Ray Train.
It demonstrates how to use :ref:`Ray Data <data>` with PyTorch Lightning in Ray Train.
See the original `DreamBooth project homepage <https://dreambooth.github.io/>`_ for more details on what this fine-tuning method achieves.
.. image:: https://dreambooth.github.io/DreamBooth_files/high_level.png
:target: https://dreambooth.github.io
:alt: DreamBooth fine-tuning overview
This example builds on `this Hugging Face 🤗 tutorial <https://huggingface.co/docs/diffusers/training/dreambooth>`_.
See the Hugging Face tutorial for useful explanations and suggestions on hyperparameters.
**Adapting this example to Ray Train allows you to easily scale up the fine-tuning to an arbitrary number of distributed training workers.**
**Compute requirements:**
* Because of the large model sizes, you need a machine with at least 1 A10G GPU.
* Each training worker uses 1 GPU. You can use multiple GPUs or workers to leverage data-parallel training to speed up training time.
This example fine-tunes both the ``text_encoder`` and ``unet`` models used in the stable diffusion process, with respect to a prior preserving loss.
.. image:: /templates/05_dreambooth_finetuning/dreambooth/images/dreambooth_example.png
:alt: DreamBooth overview
Find the full code repository at `https://github.com/ray-project/ray/tree/master/doc/source/templates/05_dreambooth_finetuning <https://github.com/ray-project/ray/tree/master/doc/source/templates/05_dreambooth_finetuning>`_
How it works
------------
This example uses Ray Data for data loading and Ray Train for distributed training.
Data loading
^^^^^^^^^^^^
.. note::
Find the latest version of the code at `dataset.py <https://github.com/ray-project/ray/blob/master/doc/source/templates/05_dreambooth_finetuning/dreambooth/dataset.py>`_
The latest version might differ slightly from the code presented here.
Use Ray Data for data loading. The code has three interesting parts.
First, load two datasets using :func:`ray.data.read_images`:
.. literalinclude:: /templates/05_dreambooth_finetuning/dreambooth/dataset.py
:language: python
:start-at: instance_dataset = read
:end-at: class_dataset = read
:dedent: 4
Then, tokenize the prompt that generated these images:
.. literalinclude:: /templates/05_dreambooth_finetuning/dreambooth/dataset.py
:language: python
:start-at: tokenizer = AutoTokenizer
:end-at: instance_prompt_ids = _tokenize
:dedent: 4
And lastly, apply a ``torchvision`` preprocessing pipeline to the images:
.. literalinclude:: /templates/05_dreambooth_finetuning/dreambooth/dataset.py
:language: python
:start-after: START: image preprocessing
:end-before: END: image preprocessing
:dedent: 4
Apply all three parts in a final step:
.. literalinclude:: /templates/05_dreambooth_finetuning/dreambooth/dataset.py
:language: python
:start-after: START: Apply preprocessing
:end-before: END: Apply preprocessing
:dedent: 4
Distributed training
^^^^^^^^^^^^^^^^^^^^
.. note::
Find the latest version of the code at `train.py <https://github.com/ray-project/ray/blob/master/doc/source/templates/05_dreambooth_finetuning/dreambooth/train.py>`_
The latest version might differ slightly from the code presented here.
The central part of the training code is the :ref:`training function <train-overview-training-function>`. This function accepts a configuration dict that contains the hyperparameters. It then defines a regular PyTorch training loop.
You interact with the Ray Train API in only a few locations, which follow in-line comments in the snippet below.
Remember that you want to do data-parallel training for all the models.
#. Load the data shard for each worker with `session.get_dataset_shard("train")``
#. Iterate over the dataset with `train_dataset.iter_torch_batches()``
#. Report results to Ray Train with `session.report(results)``
The code is compacted for brevity. The `full code <https://github.com/ray-project/ray/blob/master/doc/source/templates/05_dreambooth_finetuning/dreambooth/train.py>`_ is more thoroughly annotated.
.. literalinclude:: /templates/05_dreambooth_finetuning/dreambooth/train.py
:language: python
:start-at: def train_fn(config)
:end-before: END: Training loop
You can then run this training function with Ray Train's TorchTrainer:
.. literalinclude:: /templates/05_dreambooth_finetuning/dreambooth/train.py
:language: python
:start-at: args = train_arguments
:end-at: trainer.fit()
:dedent: 4
Configure the scale
^^^^^^^^^^^^^^^^^^^
In the TorchTrainer, you can easily configure the scale.
The preceding example uses the ``num_workers`` argument to specify the number
of workers. This argument defaults to 2 workers with 1 GPU each, totalling to 2 GPUs.
To run the example on 4 GPUs, set the number of workers to 4 using ``--num-workers=4``.
Or you can change the scaling config directly:
.. code-block:: diff
scaling_config=ScalingConfig(
use_gpu=True,
- num_workers=args.num_workers,
+ num_workers=4,
)
If you're running multi-node training, make sure that all nodes have access to a shared
storage like NFS or EFS. In the following example script, you can adjust the location with the
``DATA_PREFIX`` environment variable.
Training throughput
~~~~~~~~~~~~~~~~~~~
Compare throughput of the preceding training runs that used 1, 2, and 4 workers or GPUs.
Consider the following setup:
* 1 GCE g2-standard-48-nvidia-l4-4 instance with 4 GPUs
* Model as configured below
* Data from this example
* 200 regularization images
* Training for 4 epochs (local batch size = 2)
* 3 runs per configuration
You expect that the training time should benefit from scale and decreases when running with
more workers and GPUs.
.. image:: /templates/05_dreambooth_finetuning/dreambooth/images/dreambooth_training.png
:alt: DreamBooth training times
.. list-table::
:header-rows: 1
* - Number of workers/GPUs
- Training time (seconds)
* - 1
- 802.14
* - 2
- 487.82
* - 4
- 313.25
While the training time decreases linearly with the amount of workers/GPUs, you can observe some penalty.
Specifically, with double the amount of workers you don't get half of the training time.
This penalty is most likely due to additional communication between processes and the transfer of large model
weights. You are also only training with a batch size of one because of the GPU memory limitation. On larger
GPUs with higher batch sizes you would expect a greater benefit from scaling out.
Run the example
---------------
First, download the pre-trained Stable Diffusion model as a starting point.
Then train this model with a few images of a subject.
To achieve this, choose a non-word as an identifier, such as ``unqtkn``. When fine-tuning the model with this subject, you teach the model that the prompt is ``A photo of a unqtkn <class>``.
After fine-tuning you can run inference with this specific prompt.
For instance: ``A photo of a unqtkn <class>`` creates an image of the subject.
Similarly, ``A photo of a unqtkn <class> at the beach`` creates an image of the subject at the beach.
Step 0: Preparation
^^^^^^^^^^^^^^^^^^^
Clone the Ray repository, go to the example directory, and install dependencies.
.. code-block:: bash
git clone https://github.com/ray-project/ray.git
cd doc/source/templates/05_dreambooth_finetuning
pip install -Ur dreambooth/requirements.txt
Prepare some directories and environment variables.
.. literalinclude:: /templates/05_dreambooth_finetuning/dreambooth_run.sh
:language: bash
:start-after: __preparation_start__
:end-before: __preparation_end__
Step 1: Download the pre-trained model
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
Download and cache a pre-trained Stable Diffusion model locally.
.. literalinclude:: /templates/05_dreambooth_finetuning/dreambooth_run.sh
:language: bash
:start-after: __cache_model_start__
:end-before: __cache_model_end__
You can access the downloaded model checkpoint at the ``$ORIG_MODEL_PATH``.
Step 2: Supply images of your subject
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
Use one of the sample datasets, like `dog` or `lego car`, or provide your own directory
of images, and specify the directory with the ``$INSTANCE_DIR`` environment variable.
Then, copy these images to ``$IMAGES_OWN_DIR``.
.. literalinclude:: /templates/05_dreambooth_finetuning/dreambooth_run.sh
:language: bash
:start-after: __supply_own_images_start__
:end-before: __supply_own_images_end__
The ``$CLASS_NAME`` should be the general category of your subject.
The images produced by the prompt ``photo of a unqtkn <class>`` should be diverse images
that are different enough from the subject in order for generated images to clearly
show the effect of fine-tuning.
Step 3: Create the regularization images
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
Create a regularization image set for a class of subjects using the pre-trained
Stable Diffusion model. This regularization set ensures that
the model still produces decent images for random images of the same class,
rather than just optimize for producing good images of the subject.
.. literalinclude:: /templates/05_dreambooth_finetuning/dreambooth_run.sh
:language: bash
:start-after: Step 3: START
:end-before: Step 3: END
Use Ray Data to do batch inference with 4 workers, to generate more images in parallel.
Step 4: Fine-tune the model
^^^^^^^^^^^^^^^^^^^^^^^^^^^
Save a few, like 4 to 5, images of the subject being fine-tuned
in a local directory. Then launch the training job with:
.. literalinclude:: /templates/05_dreambooth_finetuning/dreambooth_run.sh
:language: bash
:start-after: Step 4: START
:end-before: Step 4: END
Step 5: Generate images of the subject
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
Try your model with the same command line as Step 2, but point
to your own model this time.
.. literalinclude:: /templates/05_dreambooth_finetuning/dreambooth_run.sh
:language: bash
:start-after: Step 5: START
:end-before: Step 5: END
Next, try replacing the prompt with something more interesting.
For example, for the dog subject, you can try:
- "photo of a unqtkn dog in a bucket"
- "photo of a unqtkn dog sleeping"
- "photo of a unqtkn dog in a doghouse"
See also
--------
* :doc:`Ray Train Examples <../../examples>` for more use cases
* :ref:`Ray Train User Guides <train-user-guides>` for how-to guides
doc/source/train/examples/pytorch/pytorch_resnet_finetune.ipynb
+519
View File
@@ -0,0 +1,519 @@
{
"cells": [
{
"cell_type": "markdown",
"metadata": {},
"source": [
"# Finetuning a Pytorch Image Classifier with Ray Train\n",
"\n",
"<a id=\"try-anyscale-quickstart-pytorch_resnet_finetune\" href=\"https://console.anyscale.com/register/ha?render_flow=ray&utm_source=ray_docs&utm_medium=docs&utm_campaign=pytorch_resnet_finetune\">\n",
" <img src=\"../../../_static/img/run-on-anyscale.svg\" alt=\"try-anyscale-quickstart\">\n",
"</a>\n",
"<br></br>\n",
"\n",
"This example fine-tunes a pre-trained ResNet model with Ray Train. \n",
"\n",
"For this example, the network architecture consists of the intermediate layer output of a pre-trained ResNet model, which feeds into a randomly initialized linear layer that outputs classification logits for our new task.\n",
"\n",
"\n"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Load and preprocess finetuning dataset\n",
"This example is adapted from Pytorch's [Transfer Learning for Computer Vision](https://pytorch.org/tutorials/beginner/transfer_learning_tutorial.html) tutorial.\n",
"We will use *hymenoptera_data* as the finetuning dataset, which contains two classes (bees and ants) and 397 total images (across training and validation). This is a quite small dataset and used only for demonstration purposes. "
]
},
{
"cell_type": "code",
"execution_count": 15,
"metadata": {
"tags": [
"remove-cell"
]
},
"outputs": [],
"source": [
"# To run full example, set SMOKE_TEST as False\n",
"SMOKE_TEST = True\n"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"The dataset is publicly available [here](https://www.kaggle.com/datasets/ajayrana/hymenoptera-data). Note that it is structured with directory names as the labels. Use `torchvision.datasets.ImageFolder()` to load the images and their corresponding labels."
]
},
{
"cell_type": "code",
"execution_count": 2,
"metadata": {},
"outputs": [],
"source": [
"import os\n",
"import torch\n",
"import torch.nn as nn\n",
"import torch.optim as optim\n",
"from torch.utils.data import DataLoader\n",
"from torchvision import datasets, models, transforms\n",
"import numpy as np\n",
"\n",
"# Data augmentation and normalization for training\n",
"# Just normalization for validation\n",
"data_transforms = {\n",
" \"train\": transforms.Compose(\n",
" [\n",
" transforms.RandomResizedCrop(224),\n",
" transforms.RandomHorizontalFlip(),\n",
" transforms.ToTensor(),\n",
" transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),\n",
" ]\n",
" ),\n",
" \"val\": transforms.Compose(\n",
" [\n",
" transforms.Resize(224),\n",
" transforms.CenterCrop(224),\n",
" transforms.ToTensor(),\n",
" transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),\n",
" ]\n",
" ),\n",
"}\n",
"\n",
"def download_datasets():\n",
" os.system(\n",
" \"wget https://download.pytorch.org/tutorial/hymenoptera_data.zip >/dev/null 2>&1\"\n",
" )\n",
" os.system(\"unzip hymenoptera_data.zip >/dev/null 2>&1\")\n",
"\n",
"# Download and build torch datasets\n",
"def build_datasets():\n",
" torch_datasets = {}\n",
" for split in [\"train\", \"val\"]:\n",
" torch_datasets[split] = datasets.ImageFolder(\n",
" os.path.join(\"./hymenoptera_data\", split), data_transforms[split]\n",
" )\n",
" return torch_datasets\n"
]
},
{
"cell_type": "code",
"execution_count": 16,
"metadata": {
"tags": [
"remove-cell"
]
},
"outputs": [],
"source": [
"if SMOKE_TEST:\n",
" from torch.utils.data import Subset\n",
"\n",
" def build_datasets():\n",
" torch_datasets = {}\n",
" for split in [\"train\", \"val\"]:\n",
" torch_datasets[split] = datasets.ImageFolder(\n",
" os.path.join(\"./hymenoptera_data\", split), data_transforms[split]\n",
" )\n",
" \n",
" # Only take a subset for smoke test\n",
" for split in [\"train\", \"val\"]:\n",
" indices = list(range(100))\n",
" torch_datasets[split] = Subset(torch_datasets[split], indices)\n",
" return torch_datasets\n"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Initialize Model and Fine-tuning configs"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Next, let's define the training configuration that will be passed into the training loop function later."
]
},
{
"cell_type": "code",
"execution_count": 4,
"metadata": {},
"outputs": [],
"source": [
"train_loop_config = {\n",
" \"input_size\": 224, # Input image size (224 x 224)\n",
" \"batch_size\": 32, # Batch size for training\n",
" \"num_epochs\": 10, # Number of epochs to train for\n",
" \"lr\": 0.001, # Learning Rate\n",
" \"momentum\": 0.9, # SGD optimizer momentum\n",
"}\n"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Next, let's define our model. You can either create a model from pre-trained weights or reload the model checkpoint from a previous run."
]
},
{
"cell_type": "code",
"execution_count": 5,
"metadata": {},
"outputs": [],
"source": [
"import os\n",
"import torch\n",
"from ray.train import Checkpoint\n",
"\n",
"# Option 1: Initialize model with pretrained weights\n",
"def initialize_model():\n",
" # Load pretrained model params\n",
" model = models.resnet50(pretrained=True)\n",
"\n",
" # Replace the original classifier with a new Linear layer\n",
" num_features = model.fc.in_features\n",
" model.fc = nn.Linear(num_features, 2)\n",
"\n",
" # Ensure all params get updated during finetuning\n",
" for param in model.parameters():\n",
" param.requires_grad = True\n",
" return model\n",
"\n",
"\n",
"# Option 2: Initialize model with an Train checkpoint\n",
"# Replace this with your own uri\n",
"CHECKPOINT_FROM_S3 = Checkpoint(\n",
" path=\"s3://air-example-data/finetune-resnet-checkpoint/TorchTrainer_4f69f_00000_0_2023-02-14_14-04-09/checkpoint_000001/\"\n",
")\n",
"\n",
"\n",
"def initialize_model_from_checkpoint(checkpoint: Checkpoint):\n",
" with checkpoint.as_directory() as tmpdir:\n",
" state_dict = torch.load(os.path.join(tmpdir, \"checkpoint.pt\"))\n",
" resnet50 = initialize_model()\n",
" resnet50.load_state_dict(state_dict[\"model\"])\n",
" return resnet50\n"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Define the Training Loop\n",
"\n",
"The `train_loop_per_worker` function defines the fine-tuning procedure for each worker.\n",
"\n",
"**1. Prepare dataloaders for each worker**:\n",
"- This tutorial assumes you are using PyTorch's native `torch.utils.data.Dataset` for data input. {meth}`train.torch.prepare_data_loader() <ray.train.torch.prepare_data_loader>` prepares your dataLoader for distributed execution. You can also use Ray Data for more efficient preprocessing. For more details on using Ray Data for images, see the {doc}`Working with Images </data/working-with-images>` Ray Data user guide.\n",
"\n",
"**2. Prepare your model**:\n",
"- {meth}`train.torch.prepare_model() <ray.train.torch.prepare_model>` prepares the model for distributed training. Under the hood, it converts your torch model to `DistributedDataParallel` model, which synchronize its weights across all workers.\n",
"\n",
"**3. Report metrics and checkpoint**:\n",
"- {meth}`train.report() <ray.train.report>` will report metrics and checkpoints to Ray Train.\n",
"- Saving checkpoints through {meth}`train.report(metrics, checkpoint=...) <ray.train.report>` will automatically [upload checkpoints to cloud storage](tune-cloud-checkpointing) (if configured), and allow you to easily enable Ray Train worker fault tolerance in the future."
]
},
{
"cell_type": "code",
"execution_count": 6,
"metadata": {},
"outputs": [],
"source": [
"import os\n",
"from tempfile import TemporaryDirectory\n",
"\n",
"import ray.train as train\n",
"from ray.train import Checkpoint\n",
"\n",
"\n",
"\n",
"def evaluate(logits, labels):\n",
" _, preds = torch.max(logits, 1)\n",
" corrects = torch.sum(preds == labels).item()\n",
" return corrects\n",
"\n",
"\n",
"def train_loop_per_worker(configs):\n",
" import warnings\n",
"\n",
" warnings.filterwarnings(\"ignore\")\n",
"\n",
" # Calculate the batch size for a single worker\n",
" worker_batch_size = configs[\"batch_size\"] // train.get_context().get_world_size()\n",
"\n",
" # Download dataset once on local rank 0 worker\n",
" if train.get_context().get_local_rank() == 0:\n",
" download_datasets()\n",
" torch.distributed.barrier()\n",
"\n",
" # Build datasets on each worker\n",
" torch_datasets = build_datasets()\n",
"\n",
" # Prepare dataloader for each worker\n",
" dataloaders = dict()\n",
" dataloaders[\"train\"] = DataLoader(\n",
" torch_datasets[\"train\"], batch_size=worker_batch_size, shuffle=True\n",
" )\n",
" dataloaders[\"val\"] = DataLoader(\n",
" torch_datasets[\"val\"], batch_size=worker_batch_size, shuffle=False\n",
" )\n",
"\n",
" # Distribute\n",
" dataloaders[\"train\"] = train.torch.prepare_data_loader(dataloaders[\"train\"])\n",
" dataloaders[\"val\"] = train.torch.prepare_data_loader(dataloaders[\"val\"])\n",
"\n",
" device = train.torch.get_device()\n",
"\n",
" # Prepare DDP Model, optimizer, and loss function\n",
" model = initialize_model()\n",
" model = train.torch.prepare_model(model)\n",
"\n",
" optimizer = optim.SGD(\n",
" model.parameters(), lr=configs[\"lr\"], momentum=configs[\"momentum\"]\n",
" )\n",
" criterion = nn.CrossEntropyLoss()\n",
"\n",
" # Start training loops\n",
" for epoch in range(configs[\"num_epochs\"]):\n",
" # Each epoch has a training and validation phase\n",
" for phase in [\"train\", \"val\"]:\n",
" if phase == \"train\":\n",
" model.train() # Set model to training mode\n",
" else:\n",
" model.eval() # Set model to evaluate mode\n",
"\n",
" running_loss = 0.0\n",
" running_corrects = 0\n",
"\n",
" if train.get_context().get_world_size() > 1:\n",
" dataloaders[phase].sampler.set_epoch(epoch)\n",
"\n",
" for inputs, labels in dataloaders[phase]:\n",
" inputs = inputs.to(device)\n",
" labels = labels.to(device)\n",
"\n",
" # zero the parameter gradients\n",
" optimizer.zero_grad()\n",
"\n",
" # forward\n",
" with torch.set_grad_enabled(phase == \"train\"):\n",
" # Get model outputs and calculate loss\n",
" outputs = model(inputs)\n",
" loss = criterion(outputs, labels)\n",
"\n",
" # backward + optimize only if in training phase\n",
" if phase == \"train\":\n",
" loss.backward()\n",
" optimizer.step()\n",
"\n",
" # calculate statistics\n",
" running_loss += loss.item() * inputs.size(0)\n",
" running_corrects += evaluate(outputs, labels)\n",
"\n",
" size = len(torch_datasets[phase]) // train.get_context().get_world_size()\n",
" epoch_loss = running_loss / size\n",
" epoch_acc = running_corrects / size\n",
"\n",
" if train.get_context().get_world_rank() == 0:\n",
" print(\n",
" \"Epoch {}-{} Loss: {:.4f} Acc: {:.4f}\".format(\n",
" epoch, phase, epoch_loss, epoch_acc\n",
" )\n",
" )\n",
"\n",
" # Report metrics and checkpoint every epoch\n",
" if phase == \"val\":\n",
" with TemporaryDirectory() as tmpdir:\n",
" state_dict = {\n",
" \"epoch\": epoch,\n",
" \"model\": model.module.state_dict(),\n",
" \"optimizer_state_dict\": optimizer.state_dict(),\n",
" }\n",
" torch.save(state_dict, os.path.join(tmpdir, \"checkpoint.pt\"))\n",
" train.report(\n",
" metrics={\"loss\": epoch_loss, \"acc\": epoch_acc},\n",
" checkpoint=Checkpoint.from_directory(tmpdir),\n",
" )\n"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Next, setup the TorchTrainer:"
]
},
{
"cell_type": "code",
"execution_count": 7,
"metadata": {},
"outputs": [],
"source": [
"from ray.train.torch import TorchTrainer\n",
"from ray.train import ScalingConfig, RunConfig, CheckpointConfig\n",
"\n",
"# Scale out model training across 4 GPUs.\n",
"scaling_config = ScalingConfig(\n",
" num_workers=4, use_gpu=True, resources_per_worker={\"CPU\": 1, \"GPU\": 1}\n",
")\n",
"\n",
"# Save the latest checkpoint\n",
"checkpoint_config = CheckpointConfig(num_to_keep=1)\n",
"\n",
"# Set experiment name and checkpoint configs\n",
"run_config = RunConfig(\n",
" name=\"finetune-resnet\",\n",
" storage_path=\"/tmp/ray_results\",\n",
" checkpoint_config=checkpoint_config,\n",
")\n"
]
},
{
"cell_type": "code",
"execution_count": 8,
"metadata": {
"tags": [
"remove-cell"
]
},
"outputs": [],
"source": [
"if SMOKE_TEST:\n",
" scaling_config = ScalingConfig(\n",
" num_workers=8, use_gpu=False, resources_per_worker={\"CPU\": 1}\n",
" )\n",
" train_loop_config[\"num_epochs\"] = 1\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"trainer = TorchTrainer(\n",
" train_loop_per_worker=train_loop_per_worker,\n",
" train_loop_config=train_loop_config,\n",
" scaling_config=scaling_config,\n",
" run_config=run_config,\n",
")\n",
"\n",
"result = trainer.fit()\n",
"print(result)\n"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Load the checkpoint for prediction:\n",
"\n",
" \n",
" The metadata and checkpoints have already been saved in the `storage_path` specified in TorchTrainer:"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"We now need to load the trained model and evaluate it on test data. The best model parameters have been saved in `log_dir`. We can load the resulting checkpoint from our fine-tuning run using the previously defined `initialize_model_from_checkpoint()` function."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"model = initialize_model_from_checkpoint(result.checkpoint)\n",
"device = torch.device(\"cuda\")\n"
]
},
{
"cell_type": "code",
"execution_count": 12,
"metadata": {
"tags": [
"remove-cell"
]
},
"outputs": [],
"source": [
"if SMOKE_TEST:\n",
" device = torch.device(\"cpu\")\n"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Finally, define a simple evaluation loop and check the performance of the checkpoint model."
]
},
{
"cell_type": "code",
"execution_count": 14,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Accuracy: 0.934640522875817\n"
]
}
],
"source": [
"model = model.to(device)\n",
"model.eval()\n",
"\n",
"download_datasets()\n",
"torch_datasets = build_datasets()\n",
"dataloader = DataLoader(torch_datasets[\"val\"], batch_size=32, num_workers=4)\n",
"corrects = 0\n",
"for inputs, labels in dataloader:\n",
" inputs = inputs.to(device)\n",
" labels = labels.to(device)\n",
" preds = model(inputs)\n",
" corrects += evaluate(preds, labels)\n",
"\n",
"print(\"Accuracy: \", corrects / len(dataloader.dataset))\n"
]
}
],
"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.8.13"
},
"orphan": true,
"vscode": {
"interpreter": {
"hash": "a8c1140d108077f4faeb76b2438f85e4ed675f93d004359552883616a1acd54c"
}
}
},
"nbformat": 4,
"nbformat_minor": 4
}
doc/source/train/examples/pytorch/torch_detection.ipynb
+1393
View File
File diff suppressed because one or more lines are too long
doc/source/train/examples/pytorch/torch_fashion_mnist_example.rst
+27
View File
@@ -0,0 +1,27 @@
:orphan:
.. _train-pytorch-fashion-mnist:
Train a PyTorch model on Fashion MNIST
======================================
.. raw:: html
<a id="try-anyscale-quickstart-torch_fashion_mnist_example" target="_blank" href="https://console.anyscale.com/register/ha?render_flow=ray&utm_source=ray_docs&utm_medium=docs&utm_campaign=torch_fashion_mnist_example">
<img src="../../../_static/img/run-on-anyscale.svg" alt="Run on Anyscale" />
<br/><br/>
</a>
This example runs distributed training of a PyTorch model on Fashion MNIST with Ray Train.
Code example
------------
.. literalinclude:: /../../python/ray/train/examples/pytorch/torch_fashion_mnist_example.py
See also
--------
* :ref:`Get Started with PyTorch <train-pytorch>` for a tutorial on using Ray Train and PyTorch
* :doc:`Ray Train Examples <../../examples>` for more use cases
doc/source/train/examples/pytorch/torch_regression_example.rst
+14
View File
@@ -0,0 +1,14 @@
:orphan:
torch_regression_example
========================
.. raw:: html
<a id="try-anyscale-quickstart-torch_regression_example" target="_blank" href="https://console.anyscale.com/register/ha?render_flow=ray&utm_source=ray_docs&utm_medium=docs&utm_campaign=torch_regression_example">
<img src="../../../_static/img/run-on-anyscale.svg" alt="Run on Anyscale" />
<br/><br/>
</a>
.. literalinclude:: /../../python/ray/train/examples/pytorch/torch_regression_example.py
doc/source/train/examples/tf/tensorflow_mnist_example.rst
+26
View File
@@ -0,0 +1,26 @@
:orphan:
Training with TensorFlow and Ray Train
======================================
.. raw:: html
<a id="try-anyscale-quickstart-tensorflow_mnist_example" target="_blank" href="https://console.anyscale.com/register/ha?render_flow=ray&utm_source=ray_docs&utm_medium=docs&utm_campaign=tensorflow_mnist_example">
<img src="../../../_static/img/run-on-anyscale.svg" alt="Run on Anyscale" />
<br/><br/>
</a>
This basic example runs distributed training of a TensorFlow model on MNIST with Ray Train.
Code example
------------
.. literalinclude:: /../../python/ray/train/examples/tf/tensorflow_mnist_example.py
See also
--------
* :doc:`Ray Train Examples <../../examples>` for more use cases.
* :ref:`Distributed Tensorflow & Keras <train-tensorflow-overview>` for a tutorial.
doc/source/train/examples/tf/tensorflow_regression_example.rst
+13
View File
@@ -0,0 +1,13 @@
:orphan:
tensorflow_regression_example
=============================
.. raw:: html
<a id="try-anyscale-quickstart-tensorflow_regression_example" target="_blank" href="https://console.anyscale.com/register/ha?render_flow=ray&utm_source=ray_docs&utm_medium=docs&utm_campaign=tensorflow_regression_example">
<img src="../../../_static/img/run-on-anyscale.svg" alt="Run on Anyscale" />
<br/><br/>
</a>
.. literalinclude:: /../../python/ray/train/examples/tf/tensorflow_regression_example.py
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