In addition to Ray Data, you can continue to use framework-native data utilities with Ray Train, such as PyTorch Dataset, Hugging Face Dataset, and Lightning DataModule.

pip install -U "ray[data,train]"

.. tab-item:: PyTorch

    .. code-block:: python
        :emphasize-lines: 14,21,29,33-35,53

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

        # Set this to True to use GPU.
        # If False, do CPU training instead of GPU training.
        use_gpu = False

        # Step 1: Create a Ray Dataset from in-memory Python lists.
        # You can also create a Ray Dataset from many other sources and file
        # formats.
        train_dataset = ray.data.from_items([{"x": [x], "y": [2 * x]} for x in range(200)])

        # Step 2: Preprocess your Ray Dataset.
        def increment(batch):
            batch["y"] = batch["y"] + 1
            return batch

        train_dataset = train_dataset.map_batches(increment)


        def train_func():
            batch_size = 16

            # Step 4: Access the dataset shard for the training worker via
            # ``get_dataset_shard``.
            train_data_shard = train.get_dataset_shard("train")
            # `iter_torch_batches` returns an iterable object that
            # yield tensor batches. Ray Data automatically moves the Tensor batches
            # to GPU if you enable GPU training.
            train_dataloader = train_data_shard.iter_torch_batches(
                batch_size=batch_size, dtypes=torch.float32
            )

            for epoch_idx in range(1):
                for batch in train_dataloader:
                    inputs, labels = batch["x"], batch["y"]
                    assert type(inputs) == torch.Tensor
                    assert type(labels) == torch.Tensor
                    assert inputs.shape[0] == batch_size
                    assert labels.shape[0] == batch_size
                    # Only check one batch for demo purposes.
                    # Replace the above with your actual model training code.
                    break

        # Step 3: Create a TorchTrainer. Specify the number of training workers and
        # pass in your Ray Dataset.
        # The Ray Dataset is automatically split across all training workers.
        trainer = TorchTrainer(
            train_func,
            datasets={"train": train_dataset},
            scaling_config=ScalingConfig(num_workers=2, use_gpu=use_gpu)
        )
        result = trainer.fit()

.. tab-item:: PyTorch Lightning

    .. code-block:: python
        :emphasize-lines: 4-5,10-11,14-15,26-27,33

        from ray import train

        # Create the train and validation datasets.
        train_data = ray.data.read_csv("./train.csv")
        val_data = ray.data.read_csv("./validation.csv")

        def train_func_per_worker():
            # Access Ray datsets in your train_func via ``get_dataset_shard``.
            # Ray Data shards all datasets across workers by default.
            train_ds = train.get_dataset_shard("train")
            val_ds = train.get_dataset_shard("validation")

            # Create Ray dataset iterables via ``iter_torch_batches``.
            train_dataloader = train_ds.iter_torch_batches(batch_size=16)
            val_dataloader = val_ds.iter_torch_batches(batch_size=16)

            ...

            trainer = pl.Trainer(
                # ...
            )

            # Feed the Ray dataset iterables to ``pl.Trainer.fit``.
            trainer.fit(
                model,
                train_dataloaders=train_dataloader,
                val_dataloaders=val_dataloader
            )

        trainer = TorchTrainer(
            train_func,
            # You can pass in multiple datasets to the Trainer.
            datasets={"train": train_data, "validation": val_data},
            scaling_config=ScalingConfig(num_workers=4),
        )
        trainer.fit()

.. tab-item:: HuggingFace Transformers

    .. code-block:: python
        :emphasize-lines: 7-9,14-15,18-19,25,31-32,42

        import ray
        import ray.train
        from huggingface_hub import HfFileSystem

        ...

        # Create the train and evaluation datasets using HfFileSystem.
        fs = HfFileSystem()
        train_data = ray.data.read_parquet("hf://datasets/your-dataset/train/", filesystem=fs)
        eval_data = ray.data.read_parquet("hf://datasets/your-dataset/validation/", filesystem=fs)

        def train_func():
            # Access Ray datsets in your train_func via ``get_dataset_shard``.
            # Ray Data shards all datasets across workers by default.
            train_ds = ray.train.get_dataset_shard("train")
            eval_ds = ray.train.get_dataset_shard("evaluation")

            # Create Ray dataset iterables via ``iter_torch_batches``.
            train_iterable_ds = train_ds.iter_torch_batches(batch_size=16)
            eval_iterable_ds = eval_ds.iter_torch_batches(batch_size=16)

            ...

            args = transformers.TrainingArguments(
                ...,
                max_steps=max_steps # Required for iterable datasets
            )

            trainer = transformers.Trainer(
                ...,
                model=model,
                train_dataset=train_iterable_ds,
                eval_dataset=eval_iterable_ds,
            )

            # Prepare your Transformers Trainer
            trainer = ray.train.huggingface.transformers.prepare_trainer(trainer)
            trainer.train()

        trainer = TorchTrainer(
            train_func,
            # You can pass in multiple datasets to the Trainer.
            datasets={"train": train_data, "evaluation": val_data},
            scaling_config=ScalingConfig(num_workers=4, use_gpu=True),
        )
        trainer.fit()

Ray Datasets can be created from many different data sources and formats. For more details, see :ref:`Loading Data <loading_data>`.

.. _train-datasets-preprocess:

Preprocessing data

Your preprocessed datasets can be passed into a Ray Train Trainer (e.g. :class:`~ray.train.torch.TorchTrainer`) through the ``datasets`` argument.

The datasets passed into the Trainer's ``datasets`` can be accessed inside of the ``train_loop_per_worker`` run on each distributed training worker by calling :meth:`ray.train.get_dataset_shard`.

Ray Data splits all datasets across the training workers by default. :meth:`~ray.train.get_dataset_shard` returns ``1/n`` of the dataset, where ``n`` is the number of training workers.

Ray Data does data splitting in a streaming fashion on the fly.

.. note::

    Be aware that because Ray Data splits the evaluation dataset, you have to aggregate the evaluation results across workers.
    You might consider using `TorchMetrics <https://torchmetrics.readthedocs.io/en/latest/>`_ (:doc:`example <../examples/deepspeed/deepspeed_example>`) or
    utilities available in other frameworks that you can explore.

This behavior can be overwritten by passing in the ``dataset_config`` argument. For more information on configuring splitting logic, see :ref:`Splitting datasets <train-datasets-split>`.

.. _train-datasets-consume:

Consuming data
~~~~~~~~~~~~~~

Inside the ``train_loop_per_worker``, each worker can access its shard of the dataset via :meth:`ray.train.get_dataset_shard`.

This data can be consumed in a variety of ways:

- To create a generic Iterable of batches, you can call :meth:`~ray.data.DataIterator.iter_batches`.
- To create a replacement for a PyTorch DataLoader, you can call :meth:`~ray.data.DataIterator.iter_torch_batches`.

For more details on how to iterate over your data, see :ref:`Iterating over data <iterating-over-data>`.

.. _train-datasets-pytorch:

Starting with PyTorch data
--------------------------

Some frameworks provide their own dataset and data loading utilities. For example:

- **PyTorch:** `Dataset & DataLoader <https://pytorch.org/tutorials/beginner/basics/data_tutorial.html>`_
- **Hugging Face:** `Dataset <https://huggingface.co/docs/datasets/index>`_
- **PyTorch Lightning:** `LightningDataModule <https://lightning.ai/docs/pytorch/stable/data/datamodule.html>`_

You can still use these framework data utilities directly with Ray Train.

At a high level, you can compare these concepts as follows:

.. list-table::
   :header-rows: 1

   * - PyTorch API
     - HuggingFace API
     - Ray Data API
   * - `torch.utils.data.Dataset <https://pytorch.org/docs/stable/data.html#torch.utils.data.Dataset>`_
     - `datasets.Dataset <https://huggingface.co/docs/datasets/main/en/package_reference/main_classes#datasets.Dataset>`_
     - :class:`ray.data.Dataset`
   * - `torch.utils.data.DataLoader <https://pytorch.org/docs/stable/data.html#torch.utils.data.DataLoader>`_
     - n/a
     - :meth:`ray.data.Dataset.iter_torch_batches`

For more details, see the following sections for each framework:

.. tab-set::

    .. tab-item:: PyTorch DataLoader

        **Option 1 (with Ray Data):**

        1. Convert your PyTorch Dataset to a Ray Dataset.
        2. Pass the Ray Dataset into the TorchTrainer via  ``datasets`` argument.
        3. Inside your ``train_loop_per_worker``, you can access the dataset via :meth:`ray.train.get_dataset_shard`.
        4. Create a dataset iterable via :meth:`ray.data.DataIterator.iter_torch_batches`.

        For more details, see the :ref:`Migrating from PyTorch Datasets and DataLoaders <migrate_pytorch>`.

        **Option 2 (without Ray Data):**

        1. Instantiate the Torch Dataset and DataLoader directly in the ``train_loop_per_worker``.
        2. Use the :meth:`ray.train.torch.prepare_data_loader` utility to set up the DataLoader for distributed training.

    .. tab-item:: LightningDataModule

        The ``LightningDataModule`` is created with PyTorch ``Dataset``\s and ``DataLoader``\s. You can apply the same logic here.

    .. tab-item:: Hugging Face Dataset

        **Option 1 (with Ray Data):**

        1. Convert your Hugging Face Dataset to a Ray Dataset. For instructions, see :ref:`Ray Data for Hugging Face <loading_datasets_from_ml_libraries>`.
        2. Pass the Ray Dataset into the TorchTrainer via the ``datasets`` argument.
        3. Inside your ``train_loop_per_worker``, access the sharded dataset via :meth:`ray.train.get_dataset_shard`.
        4. Create a iterable dataset via :meth:`ray.data.DataIterator.iter_torch_batches`.
        5. Pass the iterable dataset while initializing ``transformers.Trainer``.
        6. Wrap your transformers trainer with the :meth:`ray.train.huggingface.transformers.prepare_trainer` utility.

        **Option 2 (without Ray Data):**

        1. Instantiate the Hugging Face Dataset directly in the ``train_loop_per_worker``.
        2. Pass the Hugging Face Dataset into ``transformers.Trainer`` during initialization.

.. tip::

    When using Torch or Hugging Face Datasets directly without Ray Data, make sure to instantiate your Dataset *inside* the ``train_loop_per_worker``.
    Instantiating the Dataset outside of the ``train_loop_per_worker`` and passing it in via global scope
    can cause errors due to the Dataset not being serializable.

.. note::

    When using PyTorch DataLoader with more than 1 worker, you should set the
    process start method to be `forkserver` or `spawn`.
    :ref:`Forking Ray Actors and Tasks is an anti-pattern <forking-ray-processes-antipattern>` that
    can lead to unexpected issues such as deadlocks.

    .. code-block:: python

        data_loader = DataLoader(
            dataset,
            num_workers=2, 
            multiprocessing_context=multiprocessing.get_context("forkserver"), 
            ...
        )

.. _train-datasets-split:

Splitting datasets
------------------
By default, Ray Train splits all datasets across workers using :meth:`Dataset.streaming_split <ray.data.Dataset.streaming_split>`. Each worker sees a disjoint subset of the data, instead of iterating over the entire dataset.

If want to customize which datasets are split, pass in a :class:`DataConfig <ray.train.DataConfig>` to the Trainer constructor.

For example, to split only the training dataset, do the following:

.. testcode::

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

    ds = ray.data.read_text(
        "s3://anonymous@ray-example-data/sms_spam_collection_subset.txt"
    )
    train_ds, val_ds = ds.train_test_split(0.3)

    def train_loop_per_worker():
        # Get the sharded training dataset
        train_ds = train.get_dataset_shard("train")
        for _ in range(2):
            for batch in train_ds.iter_batches(batch_size=128):
                print("Do some training on batch", batch)

        # Get the unsharded full validation dataset
        val_ds = train.get_dataset_shard("val")
        for _ in range(2):
            for batch in val_ds.iter_batches(batch_size=128):
                print("Do some evaluation on batch", batch)

    my_trainer = TorchTrainer(
        train_loop_per_worker,
        scaling_config=ScalingConfig(num_workers=2),
        datasets={"train": train_ds, "val": val_ds},
        dataset_config=ray.train.DataConfig(
            datasets_to_split=["train"],
        ),
    )
    my_trainer.fit()


Full customization (advanced)

# Note that this example class is doing the same thing as the basic DataConfig
# implementation included with Ray Train.
from typing import Optional, Dict, List

import ray
from ray import train
from ray.train.torch import TorchTrainer
from ray.train import DataConfig, ScalingConfig
from ray.data import Dataset, DataIterator, NodeIdStr
from ray.actor import ActorHandle

ds = ray.data.read_text(
    "s3://anonymous@ray-example-data/sms_spam_collection_subset.txt"
)

def train_loop_per_worker():
    # Get an iterator to the dataset we passed in below.
    it = train.get_dataset_shard("train")
    for _ in range(2):
        for batch in it.iter_batches(batch_size=128):
            print("Do some training on batch", batch)


class MyCustomDataConfig(DataConfig):
    def configure(
        self,
        datasets: Dict[str, Dataset],
        world_size: int,
        worker_handles: Optional[List[ActorHandle]],
        worker_node_ids: Optional[List[NodeIdStr]],
        **kwargs,
    ) -> List[Dict[str, DataIterator]]:
        assert len(datasets) == 1, "This example only handles the simple case"

        # Configure Ray Data for ingest.
        ctx = ray.data.DataContext.get_current()
        ctx.execution_options = DataConfig.default_ingest_options()

        # Split the stream into shards.
        iterator_shards = datasets["train"].streaming_split(
            world_size, equal=True, locality_hints=worker_node_ids
        )

        # Return the assigned iterators for each worker.
        return [{"train": it} for it in iterator_shards]


my_trainer = TorchTrainer(
    train_loop_per_worker,
    scaling_config=ScalingConfig(num_workers=2),
    datasets={"train": ds},
    dataset_config=MyCustomDataConfig(),
)
my_trainer.fit()

import ray

# Preserve ordering in Ray Datasets for reproducibility.
ctx = ray.data.DataContext.get_current()
ctx.execution_options.preserve_order = True

ds = ray.data.read_text(
    "s3://anonymous@ray-example-data/sms_spam_collection_subset.txt"
)

import base64
import numpy as np
from tempfile import TemporaryDirectory

import ray
from ray import train
from ray.train import Checkpoint, ScalingConfig
from ray.train.torch import TorchTrainer
from ray.data.preprocessors import Concatenator, StandardScaler

dataset = ray.data.read_csv("s3://anonymous@air-example-data/breast_cancer.csv")

# Create preprocessors to scale some columns and concatenate the results.
scaler = StandardScaler(columns=["mean radius", "mean texture"])
columns_to_concatenate = dataset.columns()
columns_to_concatenate.remove("target")
concatenator = Concatenator(columns=columns_to_concatenate, dtype=np.float32)

# Compute dataset statistics and get transformed datasets. Note that the
# fit call is executed immediately, but the transformation is lazy.
dataset = scaler.fit_transform(dataset)
dataset = concatenator.fit_transform(dataset)

def train_loop_per_worker():
    context = train.get_context()
    print(context.get_metadata())  # prints {"preprocessor_pkl": ...}

    # Get an iterator to the dataset we passed in below.
    it = train.get_dataset_shard("train")
    for _ in range(2):
        # Prefetch 10 batches at a time.
        for batch in it.iter_batches(batch_size=128, prefetch_batches=10):
            print("Do some training on batch", batch)

    # Save a checkpoint.
    with TemporaryDirectory() as temp_dir:
        train.report(
            {"score": 2.0},
            checkpoint=Checkpoint.from_directory(temp_dir),
        )

# Serialize the preprocessor. Since serialize() returns bytes,
# convert to base64 string for JSON compatibility.
serialized_preprocessor = base64.b64encode(scaler.serialize()).decode("ascii")

my_trainer = TorchTrainer(
    train_loop_per_worker,
    scaling_config=ScalingConfig(num_workers=2),
    datasets={"train": dataset},
    metadata={"preprocessor_pkl": serialized_preprocessor},
)

# Get the fitted preprocessor back from the result metadata.
metadata = my_trainer.fit().checkpoint.get_metadata()
# Decode from base64 before deserializing
serialized_data = base64.b64decode(metadata["preprocessor_pkl"])
print(StandardScaler.deserialize(serialized_data))

While iterating over a dataset for training, you can increase ``prefetch_batches`` in :meth:`iter_batches <ray.data.DataIterator.iter_batches>` or :meth:`iter_torch_batches <ray.data.DataIterator.iter_torch_batches>` to further increase performance. While training on the current batch, this approach launches background threads to fetch and process the next ``N`` batches.

This approach can help if training is bottlenecked on cross-node data transfer or on last-mile preprocessing such as converting batches to tensors or executing ``collate_fn``. However, increasing ``prefetch_batches`` leads to more data that needs to be held in heap memory. By default, ``prefetch_batches`` is set to 1.

For example, the following code prefetches 10 batches at a time for each training worker:

.. testcode::

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

    ds = ray.data.read_text(
        "s3://anonymous@ray-example-data/sms_spam_collection_subset.txt"
    )

    def train_loop_per_worker():
        # Get an iterator to the dataset we passed in below.
        it = train.get_dataset_shard("train")
        for _ in range(2):
            # Prefetch 10 batches at a time.
            for batch in it.iter_batches(batch_size=128, prefetch_batches=10):
                print("Do some training on batch", batch)

    my_trainer = TorchTrainer(
        train_loop_per_worker,
        scaling_config=ScalingConfig(num_workers=2),
        datasets={"train": ds},
    )
    my_trainer.fit()

Avoid heavy transformation in collate_fn

If your preprocessed Dataset is small enough to fit in Ray object store memory (by default this is 30% of total cluster RAM), *materialize* the preprocessed dataset in Ray's built-in object store, by calling :meth:`materialize() <ray.data.Dataset.materialize>` on the preprocessed dataset. This method tells Ray Data to compute the entire preprocessed and pin it in the Ray object store memory. As a result, when iterating over the dataset repeatedly, the preprocessing operations do not need to be re-run. However, if the preprocessed data is too large to fit into Ray object store memory, this approach will greatly decreases performance as data needs to be spilled to and read back from disk.

Transformations that you want to run per-epoch, such as randomization, should go after the materialize call.

.. testcode::

    from typing import Dict
    import numpy as np
    import ray

    # Load the data.
    train_ds = ray.data.read_parquet("s3://anonymous@ray-example-data/iris.parquet")

    # Define a preprocessing function.
    def normalize_length(batch: Dict[str, np.ndarray]) -> Dict[str, np.ndarray]:
        new_col = batch["sepal.length"] / np.max(batch["sepal.length"])
        batch["normalized.sepal.length"] = new_col
        del batch["sepal.length"]
        return batch

    # Preprocess the data. Transformations that are made before the materialize call
    # below are only run once.
    train_ds = train_ds.map_batches(normalize_length)

    # Materialize the dataset in object store memory.
    # Only do this if train_ds is small enough to fit in object store memory.
    train_ds = train_ds.materialize()

    # Dummy augmentation transform.
    def augment_data(batch):
        return batch

    # Add per-epoch preprocessing. Transformations that you want to run per-epoch, such
    # as data augmentation or randomization, should go after the materialize call.
    train_ds = train_ds.map_batches(augment_data)

    # Pass train_ds to the Trainer


Adding CPU-only nodes to your cluster