Get Started with Hugging Face Transformers#

This tutorial walks through the process of converting an existing Hugging Face Transformers script to use Ray Train.

Learn how to:

  1. Configure a training function to report metrics and save checkpoints.

  2. Configure scaling and CPU or GPU resource requirements for your training job.

  3. Launch your distributed training job with a TorchTrainer.


For reference, the final code follows:

from ray.train.torch import TorchTrainer
from ray.train import ScalingConfig

def train_func(config):
    # Your Transformers training code here.

scaling_config = ScalingConfig(num_workers=2, use_gpu=True)
trainer = TorchTrainer(train_func, scaling_config=scaling_config)
result =
  1. train_func is the Python code that executes on each distributed training worker.

  2. ScalingConfig defines the number of distributed training workers and whether to use GPUs.

  3. TorchTrainer launches the distributed training job.

Compare a Hugging Face Transformers training script with and without Ray Train.

# Adapted from Hugging Face tutorial:

import numpy as np
import evaluate
from datasets import load_dataset
from transformers import (

# Datasets
dataset = load_dataset("yelp_review_full")
tokenizer = AutoTokenizer.from_pretrained("bert-base-cased")

def tokenize_function(examples):
    return tokenizer(examples["text"], padding="max_length", truncation=True)

small_train_dataset = dataset["train"].select(range(1000)).map(tokenize_function, batched=True)
small_eval_dataset = dataset["test"].select(range(1000)).map(tokenize_function, batched=True)

# Model
model = AutoModelForSequenceClassification.from_pretrained(
    "bert-base-cased", num_labels=5

# Metrics
metric = evaluate.load("accuracy")

def compute_metrics(eval_pred):
    logits, labels = eval_pred
    predictions = np.argmax(logits, axis=-1)
    return metric.compute(predictions=predictions, references=labels)

# Hugging Face Trainer
training_args = TrainingArguments(
    output_dir="test_trainer", evaluation_strategy="epoch", report_to="none"

trainer = Trainer(

# Start Training

Set up a training function#

First, update your training code to support distributed training. You can begin by wrapping your code in a training function:

def train_func(config):
    # Your Transformers training code here.

This function executes on each distributed training worker. Ray Train sets up the distributed process group on each worker before entering this function.

Put all the logic into this function, including dataset construction and preprocessing, model initialization, transformers trainer definition and more.


If you are using Hugging Face Datasets or Evaluate, make sure to call datasets.load_dataset and evaluate.load inside the training function. Don’t pass the loaded datasets and metrics from outside of the training function, because it might cause serialization errors while transferring the objects to the workers.

Report checkpoints and metrics#

To persist your checkpoints and monitor training progress, add a ray.train.huggingface.transformers.RayTrainReportCallback utility callback to your Trainer.

 import transformers
 from ray.train.huggingface.transformers import RayTrainReportCallback

 def train_func(config):
     trainer = transformers.Trainer(...)
+    trainer.add_callback(RayTrainReportCallback())

Reporting metrics and checkpoints to Ray Train ensures that you can use Ray Tune and fault-tolerant training. Note that the ray.train.huggingface.transformers.RayTrainReportCallback only provides a simple implementation, and you can further customize it.

Prepare a Transformers Trainer#

Finally, pass your Transformers Trainer into prepare_trainer() to validate your configurations and enable Ray Data Integration.

 import transformers
 import ray.train.huggingface.transformers

 def train_func(config):
     trainer = transformers.Trainer(...)
+    trainer = ray.train.huggingface.transformers.prepare_trainer(trainer)

Configure scale and GPUs#

Outside of your training function, create a ScalingConfig object to configure:

  1. num_workers - The number of distributed training worker processes.

  2. use_gpu - Whether each worker should use a GPU (or CPU).

from ray.train import ScalingConfig
scaling_config = ScalingConfig(num_workers=2, use_gpu=True)

For more details, see Configuring Scale and GPUs.

Launch a training job#

Tying this all together, you can now launch a distributed training job with a TorchTrainer.

from ray.train.torch import TorchTrainer

trainer = TorchTrainer(train_func, scaling_config=scaling_config)
result =

Refer to Run Configuration in Train (RunConfig) for more configuration options for TorchTrainer.

Access training results#

After training completes, a Result object is returned which contains information about the training run, including the metrics and checkpoints reported during training.

result.metrics     # The metrics reported during training.
result.checkpoint  # The latest checkpoint reported during training.
result.path     # The path where logs are stored.
result.error       # The exception that was raised, if training failed.

Next steps#

After you have converted your Hugging Face Transformers training script to use Ray Train:

  • See User Guides to learn more about how to perform specific tasks.

  • Browse the Examples for end-to-end examples of how to use Ray Train.

  • Dive into the API Reference for more details on the classes and methods used in this tutorial.

TransformersTrainer Migration Guide#

Ray 2.1 introduced the TransformersTrainer, which exposes a trainer_init_per_worker interface to define transformers.Trainer, then runs a pre-defined training function in a black box.

Ray 2.7 introduced the newly unified TorchTrainer API, which offers enhanced transparency, flexibility, and simplicity. This API aligns more with standard Hugging Face Transformers scripts, ensuring that you have better control over your native Transformers training code.

import transformers
from transformers import AutoConfig, AutoModelForCausalLM
from datasets import load_dataset

import ray
from ray.train.huggingface import TransformersTrainer
from ray.train import ScalingConfig

# Dataset
def preprocess(examples):

hf_datasets = load_dataset("wikitext", "wikitext-2-raw-v1")
processed_ds =, ...)

ray_train_ds =["train"])
ray_eval_ds =["validation"])

# Define the Trainer generation function
def trainer_init_per_worker(train_dataset, eval_dataset, **config):
    MODEL_NAME = "gpt2"
    model_config = AutoConfig.from_pretrained(MODEL_NAME)
    model = AutoModelForCausalLM.from_config(model_config)
    args = transformers.TrainingArguments(
    return transformers.Trainer(

# Build a Ray TransformersTrainer
scaling_config = ScalingConfig(num_workers=4, use_gpu=True)
ray_trainer = TransformersTrainer(
    datasets={"train": ray_train_ds, "evaluation": ray_eval_ds},
result =