chore: import upstream snapshot with attribution

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
@@ -0,0 +1,20 @@
:orphan:
.. _train-collate-utils:
Collate Utilities
=================
.. literalinclude:: ../doc_code/collate_utils.py
:language: python
.. _random-text-generator:
Random Text Generator
=====================
The following helper functions generate random text samples with labels:
.. literalinclude:: ../doc_code/random_text_generator.py
:language: python
@@ -0,0 +1,376 @@
.. _train-validating-checkpoints:
Validating checkpoints asynchronously
=====================================
During training, you may want to validate the model periodically to monitor training progress.
The standard way to do this is to periodically switch between training and validation within
the training loop. Instead, Ray Train allows you to asynchronously validate the model in a
separate Ray task, which does the following:
* Runs validation in parallel without blocking the training loop
* Runs validation on different, potentially cheaper hardware than training, since validation
doesn't require optimizer states or gradients and can use 2-4x less GPU memory
* Leverages :ref:`autoscaling <vms-autoscaling>` to launch user-specified machines only for the duration of the validation
* Lets training continue immediately after saving a checkpoint with partial metrics (for example, loss)
and then receives validation metrics (for example, accuracy) as soon as they are available. If the initial
and validated metrics share the same key, the validated metrics overwrite the initial metrics.
When to use async validation
----------------------------
Asynchronous validation is preferable to alternating between training and validation within the
same training loop in the following scenarios:
* **Validation takes a large percentage of total training time.** If validation is a significant
fraction of your end-to-end training time, running it asynchronously can substantially reduce
wall clock time by overlapping validation with training.
* **Cheaper GPUs are available for validation.** Validation doesn't require optimizer states or
gradients, so it can use 2-4x less GPU memory than training. If you have a pool of cheaper GPUs
or an autoscaling setup that can provision them, async validation lets you run validation on
those cheaper machines instead of occupying your expensive training GPUs.
* **Training throughput stops scaling linearly with more workers.** As worker count increases,
allreduce overhead grows and limits training speed, so doubling workers no longer doubles
throughput. Validation, however, scales more linearly since it requires no gradient synchronization.
Asynchronous validation can therefore utilize otherwise idle cluster capacity without impacting
training.
The best way to know if async validation helps your workload is to try it. Converting is
straightforward (see the tutorial below), so you can run both approaches and compare.
Tutorial
--------
First, define a ``validation_fn`` that takes a :class:`ray.train.Checkpoint` to validate
and any number of json-serializable keyword arguments. This function should return a dictionary
of metrics from that validation.
The following is a simple example for teaching purposes only. It is impractical
because the validation task always runs on cpu; for a more realistic example, see
:ref:`train-distributed-validate-fn`.
.. literalinclude:: ../doc_code/asynchronous_validation.py
:language: python
:start-after: __validation_fn_simple_start__
:end-before: __validation_fn_simple_end__
.. note::
In this example, the validation dataset is a ray.data.Dataset object, which is not
json-serializable. We therefore include it with the validation_fn closure instead of passing
it as a keyword argument.
.. warning::
Don't pass large objects to the ``validation_fn`` because Ray Train runs it as a Ray task and
serializes all captured variables. Instead, package large objects in the ``Checkpoint`` and
access them from shared storage later as explained in :ref:`train-checkpointing`.
Next, register your ``validation_fn`` with your trainer by settings its ``validation_config`` argument to a
:class:`~ray.train.v2.api.report_config.ValidationConfig` object that contains your ``validation_fn``
and any default keyword arguments you want to pass to your ``validation_fn``.
Next, within your rank 0 worker's training loop, call :func:`ray.train.report` with ``validation``
set to True, which will call your ``validation_fn`` with the default keyword arguments you passed to the trainer.
Alternatively, you can set ``validation`` to a :class:`~ray.train.v2.api.report_config.ValidationTaskConfig` object
that contains keyword arguments that will override matching keyword arguments you passed to the trainer. If
``validation`` is False, Ray Train will not run validation.
.. literalinclude:: ../doc_code/asynchronous_validation.py
:language: python
:start-after: __validation_fn_report_start__
:end-before: __validation_fn_report_end__
Finally, after training is done, you can access your checkpoints and their associated metrics with the
:class:`ray.train.Result` object. See :ref:`train-inspect-results` for more details.
.. _train-distributed-validate-fn:
Write a distributed validation function
---------------------------------------
The ``validation_fn`` above runs in a single Ray task, but you can improve its performance by spawning
even more Ray tasks or actors. The Ray team recommends doing this with one of the following approaches:
* Creating a :class:`ray.train.torch.TorchTrainer` that only does validation, not training.
* Using :func:`ray.data.Dataset.map_batches` to calculate metrics on a validation set.
Choose an approach
~~~~~~~~~~~~~~~~~~
You should use ``TorchTrainer`` if:
* You want to keep your existing validation logic and avoid migrating to Ray Data.
The training function API lets you fully customize the validation loop to match your current setup.
* Your validation code depends on running within a Torch process group — for example, your
metric aggregation logic uses collective communication calls, or your model parallelism
setup requires cross-GPU communication during the forward pass.
* You want a more consistent training and validation experience. The ``map_batches`` approach involves
running multiple Ray Data Datasets in a single ray cluster; we are currently working on better support
for this.
You should use ``map_batches`` if:
* You care about validation performance. Preliminary benchmarks show that ``map_batches`` is
faster.
* You prefer Ray Datas native metric aggregation APIs over PyTorch, where you must implement
aggregation manually using low-level collective operations or rely on third-party libraries
such as `torchmetrics <https://lightning.ai/docs/torchmetrics/stable>`_.
Example: validation with Ray Train TorchTrainer
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
Here is a ``validation_fn`` that uses a ``TorchTrainer`` to calculate average cross entropy
loss on a validation set. Note the following about this example:
* ``TorchTrainer`` is typically used for training, but you can use it for validation like in this
example allowing different resource requirements for training and validation, for example,
A100 for training and A10G for validation.
* The validation train function returns its metrics directly from worker 0 rather than calling
``ray.train.report`` which is accessible via ``result.return_value``. These values can't be torch
tensors and must be python based like ``ray.train.report``.
.. literalinclude:: ../doc_code/asynchronous_validation.py
:language: python
:start-after: __validation_fn_torch_trainer_start__
:end-before: __validation_fn_torch_trainer_end__
Example: validation with Ray Data map_batches
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
The following is a ``validation_fn`` that uses :func:`ray.data.Dataset.map_batches` to
calculate average accuracy on a validation set. To learn more about how to use
``map_batches`` for batch inference, see :ref:`batch_inference_home`.
.. literalinclude:: ../doc_code/asynchronous_validation.py
:language: python
:start-after: __validation_fn_map_batches_start__
:end-before: __validation_fn_map_batches_end__
Isolating training and validation with subclusters
---------------------------------------------------
When training and validation run concurrently on the same Ray cluster,
they compete for the same nodes by default. To give each phase its own
slice of the cluster — for example, A100s for training and A10Gs for
validation — label your worker pools with a ``ray-subcluster`` value and
pin each Dataset to its subcluster. See :ref:`data_concurrent_execution`
for the background and compute-config setup.
The pattern differs slightly between the ``TorchTrainer`` validation_fn
and the ``map_batches`` validation_fn, because only the former goes
through ``ray.train.DataConfig``.
**TorchTrainer validation_fn.** Set the validation Dataset's selector
through the sub-trainer's ``dataset_config``:
.. code-block:: python
from ray.data import ExecutionOptions
def validation_fn(checkpoint, ...) -> dict:
trainer = ray.train.torch.TorchTrainer(
...,
datasets={"validation": validation_dataset},
dataset_config=ray.train.DataConfig(
execution_options={
"validation": ExecutionOptions(
label_selector={"ray-subcluster": "validation"}
),
},
),
)
...
**map_batches validation_fn.** The ``map_batches`` path doesn't take a
``DataConfig``. Construct ``validation_dataset`` under a
``DataContext.current()`` block so the selector is baked into the
Dataset at construction — every downstream operator inherits it:
.. code-block:: python
ctx = ray.data.DataContext.get_current().copy()
ctx.execution_options.label_selector = {"ray-subcluster": "validation"}
with ray.data.DataContext.current(ctx):
validation_dataset = ray.data.read_parquet(...)
def validation_fn(checkpoint) -> dict:
eval_res = validation_dataset.map_batches(...)
...
**Training-side configuration.** A Train pipeline needs the selector
specified in two places — they cover different phases and are not
redundant:
1. **At Dataset construction**, via the ``DataContext.current()`` context
manager, so construction-time tasks (parquet schema inference, file
listing) land on training nodes.
2. **In the trainer's** ``dataset_config``, because Train wholesale
replaces ``ds.context.execution_options`` with ``DataConfig``'s
per-dataset entry at training start. Anything not restated in
``DataConfig.execution_options````label_selector`` included — is
dropped, so per-worker ingest would lose its pinning.
.. code-block:: python
from ray.data import ExecutionOptions
def run_trainer() -> ray.train.Result:
# (1) Pin construction-time tasks.
ctx = ray.data.DataContext.get_current().copy()
ctx.execution_options.label_selector = {"ray-subcluster": "training"}
with ray.data.DataContext.current(ctx):
train_dataset = ray.data.read_parquet(...)
# (2) Pin per-worker ingest — Train replaces ds.context options
# wholesale, so the selector must be restated here.
trainer = ray.train.torch.TorchTrainer(
...,
datasets={"train": train_dataset},
dataset_config=ray.train.DataConfig(
datasets_to_split=["train"],
execution_options={
"train": ExecutionOptions(
label_selector={"ray-subcluster": "training"}
),
},
),
)
...
.. note::
For *interleaved* validation — where you reuse the training workers
to validate on a separate "validation" Dataset inside the same
``TorchTrainer`` — pass both Datasets to ``datasets={...}`` and give
both an entry in ``DataConfig.execution_options`` so they're each
scoped to their own subcluster:
.. code-block:: python
from ray.data import ExecutionOptions
dataset_config = ray.train.DataConfig(
datasets_to_split=["train", "validation"],
execution_options={
"train": ExecutionOptions(
label_selector={"ray-subcluster": "training"}
),
"validation": ExecutionOptions(
label_selector={"ray-subcluster": "validation"}
),
},
)
Tuning asynchronous validation
------------------------------
Overlapping validation and training
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
Asynchronous validation is most beneficial when training and validation fully overlap. If one
finishes before the other, some workers sit idle. :ref:`Autoscaling <vms-autoscaling>` lets you
spin up workers only for the duration of validation, which mitigates this but doesn't fully
eliminate the gap.
You can tune the following knobs to overlap validation and training as closely as possible:
* **Number of workers**: Tune the number of validation workers relative to training workers so that
the two phases overlap as closely as possible.
* **Batch size**: A larger batch size typically improves throughput, but it can negatively impact
training convergence and may lead to out-of-memory (OOM) errors.
* **Validation frequency**: Choose a validation cadence and dataset size that balance overlap with
training. Validating too frequently or over too many rows can create a long validation tail.
Also note that breaking early from a Ray Data iterator may lead to resource leaks - this will be
fixed in a future release.
Ray Data production vs consumption
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
See :ref:`balancing-data-production-consumption` for tips on balancing data production and consumption rates.
Checkpoint metrics lifecycle
-----------------------------
During the training loop the following happens to your checkpoints and metrics :
1. You report a checkpoint with some initial metrics, such as training loss, as well as a
:class:`~ray.train.v2.api.report_config.ValidationTaskConfig` object that contains the keyword
arguments to pass to the ``validation_fn``.
2. Ray Train asynchronously runs your ``validation_fn`` with that checkpoint and configuration.
3. When that validation task completes, Ray Train associates the metrics returned by your ``validation_fn``
with that checkpoint.
4. After training is done, you can access your checkpoints and their associated metrics with the
:class:`ray.train.Result` object. See :ref:`train-inspect-results` for more details.
.. figure:: ../images/checkpoint_metrics_lifecycle.png
How Ray Train populates checkpoint metrics during training and how you access them after training.
Experiment tracking
-------------------
In normal :ref:`experiment tracking with Ray Train <train-experiment-tracking-native>`,
you handle creating, logging to, and finishing the experiment tracking run from
the rank 0 training worker. However, asynchronous validation complicates this because
validation metrics are computed outside of the training worker, in a separate
Ray task.
Most modern experiment tracking configurations (for example,
`W&B distributed training <https://docs.wandb.ai/models/track/log/distributed-training#track-all-processes-to-a-single-run>`_)
support writing to the same run from different threads or processes. Other configurations,
such as the `MLflow fluent API <https://mlflow.org/docs/latest/api_reference/python_api/mlflow.html>`_, may not.
Writing to the same run
~~~~~~~~~~~~~~~~~~~~~~~
If your experiment tracking library supports writing to the same run from different
processes, the rank 0 training worker can start the run and the validation task can
join it and log validation metrics directly.
.. tab-set::
.. tab-item:: W&B
.. literalinclude:: ../doc_code/asynchronous_validation.py
:language: python
:start-after: __exp_tracking_same_run_wandb_start__
:end-before: __exp_tracking_same_run_wandb_end__
.. tab-item:: MLflow (non-fluent)
.. literalinclude:: ../doc_code/asynchronous_validation.py
:language: python
:start-after: __exp_tracking_same_run_mlflow_start__
:end-before: __exp_tracking_same_run_mlflow_end__
Reliability
~~~~~~~~~~~
If experiment tracking logging fails (for example, due to a transient network error),
you have two options for retrying:
1. **Wrap your logging calls in a try/except block** within the ``validation_fn`` and
retry the logging manually with your experiment tracker's API.
2. **Use** :func:`ray.train.get_all_reported_checkpoints` **periodically during training** to
retrieve all reported checkpoints and their associated metrics, then re-log any missing
entries to your experiment tracker.
Writing to different runs
~~~~~~~~~~~~~~~~~~~~~~~~~
If your experiment tracking library does not support writing to the same run from different
processes, the validation task must start a new run each time it logs validation metrics.
Many tracking libraries provide ways to group related runs together so that training and
validation runs are still associated.
.. tab-set::
.. tab-item:: W&B
Use `W&B run grouping <https://docs.wandb.ai/models/runs/grouping>`_ to group
the training run and validation runs together.
.. tab-item:: MLflow
Use `MLflow parent and child runs <https://mlflow.org/docs/latest/ml/traditional-ml/tutorials/hyperparameter-tuning/part1-child-runs/#adapting-for-parent-and-child-runs>`_
to group the training run and validation runs together.
@@ -0,0 +1,470 @@
.. _train-checkpointing:
Saving and Loading Checkpoints
==============================
Ray Train provides a way to snapshot training progress with :class:`Checkpoints <ray.train.Checkpoint>`.
This is useful for:
1. **Storing the best-performing model weights:** Save your model to persistent storage, and use it for downstream serving or inference.
2. **Fault tolerance:** Handle worker process and node failures in a long-running training job and leverage pre-emptible machines.
3. **Distributed checkpointing:** Ray Train checkpointing can be used to
:ref:`upload model shards from multiple workers in parallel. <train-distributed-checkpointing>`
.. _train-dl-saving-checkpoints:
Saving checkpoints during training
----------------------------------
The :class:`Checkpoint <ray.train.Checkpoint>` is a lightweight interface provided
by Ray Train that represents a *directory* that exists on local or remote storage.
For example, a checkpoint could point to a directory in cloud storage:
``s3://my-bucket/my-checkpoint-dir``.
A locally available checkpoint points to a location on the local filesystem:
``/tmp/my-checkpoint-dir``.
Here's how you save a checkpoint in the training loop:
1. Write your model checkpoint to a local directory.
- Since a :class:`Checkpoint <ray.train.Checkpoint>` just points to a directory, the contents are completely up to you.
- This means that you can use any serialization format you want.
- This makes it **easy to use familiar checkpoint utilities provided by training frameworks**, such as
``torch.save``, ``pl.Trainer.save_checkpoint``, Accelerate's ``accelerator.save_model``,
Transformers' ``save_pretrained``, ``tf.keras.Model.save``, etc.
2. Create a :class:`Checkpoint <ray.train.Checkpoint>` from the directory using :meth:`Checkpoint.from_directory <ray.train.Checkpoint.from_directory>`.
3. Report the checkpoint to Ray Train using :func:`ray.train.report(metrics, checkpoint=...) <ray.train.report>`.
- The metrics reported alongside the checkpoint are used to :ref:`keep track of the best-performing checkpoints <train-dl-configure-checkpoints>`.
- This will **upload the checkpoint to persistent storage** if configured. See :ref:`persistent-storage-guide`.
.. figure:: ../images/checkpoint_lifecycle.png
The lifecycle of a :class:`~ray.train.Checkpoint`, from being saved locally
to disk to being uploaded to persistent storage via ``train.report``.
As shown in the figure above, the best practice for saving checkpoints is to
first dump the checkpoint to a local temporary directory. Then, the call to ``train.report``
uploads the checkpoint to its final persistent storage location.
Then, the local temporary directory can be safely cleaned up to free up disk space
(e.g., from exiting the ``tempfile.TemporaryDirectory`` context).
.. tip::
In standard DDP training, where each worker has a copy of the full-model, you should
only save and report a checkpoint from a single worker to prevent redundant uploads.
This typically looks like:
.. literalinclude:: ../doc_code/checkpoints.py
:language: python
:start-after: __checkpoint_from_single_worker_start__
:end-before: __checkpoint_from_single_worker_end__
If using parallel training strategies such as DeepSpeed Zero and FSDP, where
each worker only has a shard of the full training state, you can save and report a checkpoint
from each worker. See :ref:`train-distributed-checkpointing` for an example.
Here are a few examples of saving checkpoints with different training frameworks:
.. tab-set::
.. tab-item:: Native PyTorch
.. literalinclude:: ../doc_code/checkpoints.py
:language: python
:start-after: __pytorch_save_start__
:end-before: __pytorch_save_end__
.. tip::
You most likely want to unwrap the DDP model before saving it to a checkpoint.
``model.module.state_dict()`` is the state dict without each key having a ``"module."`` prefix.
.. tab-item:: PyTorch Lightning
Ray Train leverages PyTorch Lightning's ``Callback`` interface to report metrics
and checkpoints. We provide a simple callback implementation that reports
``on_train_epoch_end``.
Specifically, on each train epoch end, it
- collects all the logged metrics from ``trainer.callback_metrics``
- saves a checkpoint via ``trainer.save_checkpoint``
- reports to Ray Train via :func:`ray.train.report(metrics, checkpoint) <ray.train.report>`
.. literalinclude:: ../doc_code/checkpoints.py
:language: python
:start-after: __lightning_save_example_start__
:end-before: __lightning_save_example_end__
You can always get the saved checkpoint path from :attr:`result.checkpoint <ray.train.Result>` and
:attr:`result.best_checkpoints <ray.train.Result>`.
For more advanced usage (e.g. reporting at different frequency, reporting
customized checkpoint files), you can implement your own customized callback.
Here is a simple example that reports a checkpoint every 3 epochs:
.. literalinclude:: ../doc_code/checkpoints.py
:language: python
:start-after: __lightning_custom_save_example_start__
:end-before: __lightning_custom_save_example_end__
.. tab-item:: Hugging Face Transformers
Ray Train leverages Hugging Face Transformers Trainer's ``Callback`` interface
to report metrics and checkpoints.
**Option 1: Use Ray Train's default report callback**
We provide a simple callback implementation :class:`~ray.train.huggingface.transformers.RayTrainReportCallback` that
reports on checkpoint save. You can change the checkpointing frequency by ``save_strategy`` and ``save_steps``.
It collects the latest logged metrics and report them together with the latest saved checkpoint.
.. literalinclude:: ../doc_code/checkpoints.py
:language: python
:start-after: __transformers_save_example_start__
:end-before: __transformers_save_example_end__
Note that :class:`~ray.train.huggingface.transformers.RayTrainReportCallback`
binds the latest metrics and checkpoints together,
so users can properly configure ``logging_strategy``, ``save_strategy`` and ``evaluation_strategy``
to ensure the monitoring metric is logged at the same step as checkpoint saving.
For example, the evaluation metrics (``eval_loss`` in this case) are logged during
evaluation. If users want to keep the best 3 checkpoints according to ``eval_loss``, they
should align the saving and evaluation frequency. Below are two examples of valid configurations:
.. testcode::
:skipif: True
args = TrainingArguments(
...,
evaluation_strategy="epoch",
save_strategy="epoch",
)
args = TrainingArguments(
...,
evaluation_strategy="steps",
save_strategy="steps",
eval_steps=50,
save_steps=100,
)
# And more ...
**Option 2: Implement your customized report callback**
If you feel that Ray Train's default :class:`~ray.train.huggingface.transformers.RayTrainReportCallback`
is not sufficient for your use case, you can also implement a callback yourself!
Below is a example implementation that collects latest metrics
and reports on checkpoint save.
.. literalinclude:: ../doc_code/checkpoints.py
:language: python
:start-after: __transformers_custom_save_example_start__
:end-before: __transformers_custom_save_example_end__
You can customize when (``on_save``, ``on_epoch_end``, ``on_evaluate``) and
what (customized metrics and checkpoint files) to report by implementing your own
Transformers Trainer callback.
.. _train-distributed-checkpointing:
Saving checkpoints from multiple workers (distributed checkpointing)
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
In model parallel training strategies where each worker only has a shard of the full-model,
you can save and report checkpoint shards in parallel from each worker.
.. figure:: ../images/persistent_storage_checkpoint.png
Distributed checkpointing in Ray Train. Each worker uploads its own checkpoint shard
to persistent storage independently.
Distributed checkpointing is the best practice for saving checkpoints
when doing model-parallel training (e.g., DeepSpeed, FSDP, Megatron-LM).
There are two major benefits:
1. **It is faster, resulting in less idle time.** Faster checkpointing incentivizes more frequent checkpointing!
Each worker can upload its checkpoint shard in parallel,
maximizing the network bandwidth of the cluster. Instead of a single node
uploading the full model of size ``M``, the cluster distributes the load across
``N`` nodes, each uploading a shard of size ``M / N``.
2. **Distributed checkpointing avoids needing to gather the full model onto a single worker's CPU memory.**
This gather operation puts a large CPU memory requirement on the worker that performs checkpointing
and is a common source of OOM errors.
Here is an example of distributed checkpointing with PyTorch:
.. literalinclude:: ../doc_code/checkpoints.py
:language: python
:start-after: __distributed_checkpointing_start__
:end-before: __distributed_checkpointing_end__
.. note::
Checkpoint files with the same name will collide between workers.
You can get around this by adding a rank-specific suffix to checkpoint files.
Note that having filename collisions does not error, but it will result in the last
uploaded version being the one that is persisted. This is fine if the file
contents are the same across all workers.
Model shard saving utilities provided by frameworks such as DeepSpeed will create
rank-specific filenames already, so you usually do not need to worry about this.
.. _train-checkpoint-upload-modes:
Checkpoint upload modes
-----------------------
By default, when you call :func:`~ray.train.report`, Ray Train synchronously pushes
your checkpoint from ``checkpoint.path`` on local disk to ``checkpoint_dir_name`` on
your ``storage_path``. This is equivalent to calling :func:`~ray.train.report` with
:class:`~ray.train.CheckpointUploadMode` set to ``ray.train.CheckpointUploadMode.SYNC``.
.. literalinclude:: ../doc_code/checkpoints.py
:language: python
:start-after: __checkpoint_upload_mode_sync_start__
:end-before: __checkpoint_upload_mode_sync_end__
.. _train-checkpoint-upload-mode-async:
Asynchronous checkpoint uploading
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
You may want to upload your checkpoint asynchronously instead so that
the next training step can start in parallel. If so, you should use
``ray.train.CheckpointUploadMode.ASYNC``, which kicks off a new thread
to upload the checkpoint. This is helpful for larger
checkpoints that might take longer to upload, but might add unnecessary
complexity (see below) if you want to immediately upload only a small checkpoint.
Each ``report`` blocks until the previous ``report``\'s checkpoint
upload completes before starting a new checkpoint upload thread. Ray Train does this
to avoid accumulating too many upload threads and potentially running out of memory.
Because ``report`` returns without waiting for the checkpoint upload to complete,
you must ensure that the local checkpoint directory stays alive until the checkpoint
upload completes. This means you can't use a temporary directory that Ray Train may
delete before the upload finishes, for example from ``tempfile.TemporaryDirectory``.
``report`` also exposes the ``delete_local_checkpoint_after_upload`` parameter, which
defaults to ``True`` if ``checkpoint_upload_mode`` is ``ray.train.CheckpointUploadMode.ASYNC``.
.. literalinclude:: ../doc_code/checkpoints.py
:language: python
:start-after: __checkpoint_upload_mode_async_start__
:end-before: __checkpoint_upload_mode_async_end__
.. figure:: ../images/sync_vs_async_checkpointing.png
This figure illustrates the difference between synchronous and asynchronous
checkpoint uploading.
Custom checkpoint uploading
~~~~~~~~~~~~~~~~~~~~~~~~~~~
:func:`~ray.train.report` defaults to uploading from disk to the remote ``storage_path``
with the PyArrow filesystem copying utilities before reporting the checkpoint to Ray Train.
If you would rather upload the checkpoint manually or with a third-party library
such as `Torch Distributed Checkpointing <https://docs.pytorch.org/docs/stable/distributed.checkpoint.html>`_,
you have the following options:
.. tab-set::
.. tab-item:: Synchronous
If you want to upload the checkpoint synchronously, you can first upload the checkpoint
to the ``storage_path`` and then report a reference to the uploaded checkpoint with
``ray.train.CheckpointUploadMode.NO_UPLOAD``.
.. literalinclude:: ../doc_code/checkpoints.py
:language: python
:start-after: __checkpoint_upload_mode_no_upload_start__
:end-before: __checkpoint_upload_mode_no_upload_end__
.. tab-item:: Asynchronous
If you want to upload the checkpoint asynchronously, you can set ``checkpoint_upload_mode``
to ``ray.train.CheckpointUploadMode.ASYNC`` and pass a ``checkpoint_upload_fn`` to
``ray.train.report``. This function takes the ``Checkpoint`` and ``checkpoint_dir_name``
passed to ``ray.train.report`` and returns the persisted ``Checkpoint``.
.. literalinclude:: ../doc_code/checkpoints.py
:language: python
:start-after: __checkpoint_upload_fn_start__
:end-before: __checkpoint_upload_fn_end__
.. warning::
In your ``checkpoint_upload_fn``, you should not call ``ray.train.report``, which may
lead to unexpected behavior. You should also avoid collective operations, such as
:func:`~ray.train.report` or ``model.state_dict()``, which can cause deadlocks. Finally,
the upload function should only the return a checkpoint object once all checkpoint data
has been saved.
.. note::
Do not pass a ``checkpoint_upload_fn`` with ``checkpoint_upload_mode=ray.train.CheckpointUploadMode.NO_UPLOAD``
because Ray Train will simply ignore ``checkpoint_upload_fn``. You can pass a ``checkpoint_upload_fn`` with
``checkpoint_upload_mode=ray.train.CheckpointUploadMode.SYNC``, but this is equivalent to uploading the
checkpoint yourself and reporting the checkpoint with ``ray.train.CheckpointUploadMode.NO_UPLOAD``.
.. _train-dl-configure-checkpoints:
Configure checkpointing
-----------------------
Ray Train provides some configuration options for checkpointing via :class:`~ray.train.CheckpointConfig`.
The primary configuration is keeping only the top ``K`` checkpoints with respect to a metric.
Lower-performing checkpoints are deleted to save storage space. By default, all checkpoints are kept.
.. literalinclude:: ../doc_code/key_concepts.py
:language: python
:start-after: __checkpoint_config_start__
:end-before: __checkpoint_config_end__
.. note::
If you want to save the top ``num_to_keep`` checkpoints with respect to a metric via
:py:class:`~ray.train.CheckpointConfig`,
please ensure that the metric is always reported together with the checkpoints.
Using checkpoints during training
----------------------------------
During training, you may want to access checkpoints you've reported and their associated metrics
from training workers for a variety of reasons, such as
reporting the best checkpoint so far to an experiment tracker. You can do this by calling
:func:`~ray.train.get_all_reported_checkpoints` from within your training function. This function returns
a list of :class:`~ray.train.ReportedCheckpoint` objects that represent all the
:class:`~ray.train.Checkpoint`\s and their associated metrics that you've reported so far
and have been kept based on the :ref:`checkpoint configuration <train-dl-configure-checkpoints>`.
This function supports two consistency modes:
- ``CheckpointConsistencyMode.COMMITTED``: Block until the checkpoint from the latest ``ray.train.report``
has been uploaded to persistent storage and committed.
- ``CheckpointConsistencyMode.VALIDATED``: Block until the checkpoint from the latest ``ray.train.report``
has been uploaded to persistent storage, committed, and validated (see :ref:`train-validating-checkpoints`).
This is the default consistency mode and has the same behavior as ``CheckpointConsistencyMode.COMMITTED``
if your report did not kick off validation.
.. literalinclude:: ../doc_code/checkpoints.py
:language: python
:start-after: __get_all_reported_checkpoints_example_start__
:end-before: __get_all_reported_checkpoints_example_end__
Using checkpoints after training
--------------------------------
The latest saved checkpoint can be accessed with :attr:`Result.checkpoint <ray.train.Result>`.
The full list of persisted checkpoints can be accessed with :attr:`Result.best_checkpoints <ray.train.Result>`.
If :class:`CheckpointConfig(num_to_keep) <ray.train.CheckpointConfig>` is set, this list will contain the best ``num_to_keep`` checkpoints.
See :ref:`train-inspect-results` for a full guide on inspecting training results.
:meth:`Checkpoint.as_directory <ray.train.Checkpoint.as_directory>`
and :meth:`Checkpoint.to_directory <ray.train.Checkpoint.to_directory>`
are the two main APIs to interact with Train checkpoints:
.. literalinclude:: ../doc_code/checkpoints.py
:language: python
:start-after: __inspect_checkpoint_example_start__
:end-before: __inspect_checkpoint_example_end__
For Lightning and Transformers, if you are using the default `RayTrainReportCallback` for checkpoint saving in your training function,
you can retrieve the original checkpoint files as below:
.. tab-set::
.. tab-item:: PyTorch Lightning
.. literalinclude:: ../doc_code/checkpoints.py
:language: python
:start-after: __inspect_lightning_checkpoint_example_start__
:end-before: __inspect_lightning_checkpoint_example_end__
.. tab-item:: Transformers
.. literalinclude:: ../doc_code/checkpoints.py
:language: python
:start-after: __inspect_transformers_checkpoint_example_start__
:end-before: __inspect_transformers_checkpoint_example_end__
.. _train-dl-loading-checkpoints:
Restore training state from a checkpoint
----------------------------------------
In order to enable fault tolerance, you should modify your training loop to restore
training state from a :class:`~ray.train.Checkpoint`.
The :class:`Checkpoint <ray.train.Checkpoint>` to restore from can be accessed in the
training function with :func:`ray.train.get_checkpoint <ray.train.get_checkpoint>`.
The checkpoint returned by :func:`ray.train.get_checkpoint <ray.train.get_checkpoint>` is populated
as the latest reported checkpoint during :ref:`automatic failure recovery <train-fault-tolerance>`.
See :ref:`train-fault-tolerance` for more details on restoration and fault tolerance.
.. tab-set::
.. tab-item:: Native PyTorch
.. literalinclude:: ../doc_code/checkpoints.py
:language: python
:start-after: __pytorch_restore_start__
:end-before: __pytorch_restore_end__
.. tab-item:: PyTorch Lightning
.. literalinclude:: ../doc_code/checkpoints.py
:language: python
:start-after: __lightning_restore_example_start__
:end-before: __lightning_restore_example_end__
.. note::
In these examples, :meth:`Checkpoint.as_directory <ray.train.Checkpoint.as_directory>`
is used to view the checkpoint contents as a local directory.
*If the checkpoint points to a local directory*, this method just returns the
local directory path without making a copy.
*If the checkpoint points to a remote directory*, this method will download the
checkpoint to a local temporary directory and return the path to the temporary directory.
**If multiple processes on the same node call this method simultaneously,**
only a single process will perform the download, while the others
wait for the download to finish. Once the download finishes, all processes receive
the same local (temporary) directory to read from.
Once all processes have finished working with the checkpoint, the temporary directory
is cleaned up.
@@ -0,0 +1,755 @@
.. _data-ingest-torch:
Data Loading and Preprocessing
==============================
Ray Train integrates with :ref:`Ray Data <data>` to offer a performant and scalable streaming solution for loading and preprocessing large datasets.
Key advantages include:
- Streaming data loading and preprocessing, scalable to petabyte-scale data.
- Scaling out heavy data preprocessing to CPU nodes, to avoid bottlenecking GPU training.
- Automatic and fast failure recovery.
- Automatic on-the-fly data splitting across distributed training workers.
For more details about Ray Data, check out the :ref:`Ray Data documentation<data>`.
.. note::
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.
In this guide, we will cover how to incorporate Ray Data into your Ray Train script, and different ways to customize your data ingestion pipeline.
.. TODO: Replace this image with a better one.
.. figure:: ../images/train_ingest.png
:align: center
:width: 300px
Quickstart
----------
Install Ray Data and Ray Train:
.. code-block:: bash
pip install -U "ray[data,train]"
Data ingestion can be set up with four basic steps:
1. Create a Ray Dataset from your input data.
2. Apply preprocessing operations to your Ray Dataset.
3. Input the preprocessed Dataset into the Ray Train Trainer, which internally splits the dataset equally in a streaming way across the distributed training workers.
4. Consume the Ray Dataset in your training function.
.. tab-set::
.. 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, batch_size="auto")
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()
.. _train-datasets-load:
Loading data
~~~~~~~~~~~~
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
~~~~~~~~~~~~~~~~~~
Ray Data supports a wide range of preprocessing operations that you can use to transform data prior to training.
- For general preprocessing, see :ref:`Transforming Data <transforming_data>`.
- For tabular data, see :ref:`Preprocessing Structured Data <preprocessing_structured_data>`.
- For PyTorch tensors, see :ref:`Transformations with torch tensors <transform_pytorch>`.
- For optimizing expensive preprocessing operations, see :ref:`Caching the preprocessed dataset <dataset_cache_performance>`.
.. _train-datasets-input:
Inputting and splitting 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)
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
For use cases not covered by the default config class, you can also fully customize exactly how your input datasets are split. Define a custom :class:`DataConfig <ray.train.DataConfig>` class (DeveloperAPI). The :class:`DataConfig <ray.train.DataConfig>` class is responsible for that shared setup and splitting of data across nodes.
.. testcode::
# 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()
The subclass must be serializable, since Ray Train copies it from the driver script to the driving actor of the Trainer. Ray Train calls its :meth:`configure <ray.train.DataConfig.configure>` method on the main actor of the Trainer group to create the data iterators for each worker.
In general, you can use :class:`DataConfig <ray.train.DataConfig>` for any shared setup that has to occur ahead of time before the workers start iterating over data. The setup runs at the start of each Trainer run.
Random shuffling
----------------
Randomly shuffling data for each epoch can be important for model quality depending on what model you are training.
Ray Data provides multiple options for random shuffling, see :ref:`Shuffling Data <shuffling_data>` for more details.
Enabling reproducibility
------------------------
When developing or hyperparameter tuning models, reproducibility is important during data ingest so that data ingest does not affect model quality. Follow these three steps to enable reproducibility:
**Step 1:** Enable deterministic execution in Ray Datasets by setting the `preserve_order` flag in the :class:`DataContext <ray.data.context.DataContext>`.
.. testcode::
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"
)
**Step 2:** Set a seed for any shuffling operations:
* `seed` argument to :meth:`random_shuffle <ray.data.Dataset.random_shuffle>`
* `seed` argument to :meth:`randomize_block_order <ray.data.Dataset.randomize_block_order>`
* `local_shuffle_seed` argument to :meth:`iter_batches <ray.data.DataIterator.iter_batches>`
**Step 3:** Follow the best practices for enabling reproducibility for your training framework of choice. For example, see the `Pytorch reproducibility guide <https://pytorch.org/docs/stable/notes/randomness.html>`_.
.. _preprocessing_structured_data:
Preprocessing structured data
-----------------------------
.. note::
This section is for tabular/structured data. The recommended way for preprocessing unstructured data is to use
Ray Data operations such as `map_batches`. See the :ref:`Ray Data Working with Pytorch guide <working_with_pytorch>` for more details.
For tabular data, use Ray Data :ref:`preprocessors <preprocessor-ref>`, which implement common data preprocessing operations.
You can use this with Ray Train Trainers by applying them on the dataset before passing the dataset into a Trainer. For example:
.. testcode::
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))
This example persists the fitted preprocessor using the ``Trainer(metadata={...})`` constructor argument. This arg specifies a dict that is available from ``TrainContext.get_metadata()`` and ``checkpoint.get_metadata()`` for checkpoints that the Trainer saves. This design enables the recreation of the fitted preprocessor for inference.
Performance tips
----------------
Prefetching batches
~~~~~~~~~~~~~~~~~~~
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
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
The ``collate_fn`` parameter in :meth:`iter_batches <ray.data.DataIterator.iter_batches>` or :meth:`iter_torch_batches <ray.data.DataIterator.iter_torch_batches>` allows you to transform data before feeding it to the model. This operation happens locally in the training workers. Avoid adding a heavy transformation in this function as it may become the bottleneck. Instead, :ref:`apply the transformation with map or map_batches <transforming_data>` before passing the dataset to the Trainer. When your expensive transformation requires batch_size as input, such as text tokenization, you can :ref:`scale it out to Ray Data <scaling_collation_functions>` for better performance.
.. _dataset_cache_performance:
Caching the preprocessed dataset
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
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, batch_size="auto")
# 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, batch_size="auto")
# Pass train_ds to the Trainer
Adding CPU-only nodes to your cluster
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
If the GPU training is bottlenecked on expensive CPU preprocessing and the preprocessed Dataset is too large to fit in object store memory, then materializing the dataset doesn't work. In this case, Ray's native support for heterogeneous resources enables you to simply add more CPU-only nodes to your cluster, and Ray Data automatically scales out CPU-only preprocessing tasks to CPU-only nodes, making GPUs more saturated.
In general, adding CPU-only nodes can help in two ways:
* Adding more CPU cores helps further parallelize preprocessing. This approach is helpful when CPU compute time is the bottleneck.
* Increasing object store memory, which 1) allows Ray Data to buffer more data in between preprocessing and training stages, and 2) provides more memory to make it possible to :ref:`cache the preprocessed dataset <dataset_cache_performance>`. This approach is helpful when memory is the bottleneck.
Isolating Ray Data worker processes from training nodes
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
You may sometimes want to prevent Ray Data CPU tasks from running on training worker nodes
when training workers themselves run CPU or RAM-heavy operations
such as storing large local shuffle buffers or running expensive collate functions.
Launching more Ray Data processes would oversubscribe the training worker nodes.
Instead, the Ray Data tasks should run on a separate set of CPU nodes in your heterogeneous
cluster (for example, 4 GPU training nodes and 4 CPU-only nodes).
One workaround is to force full-node exclusion by reserving all CPUs per training worker via
``resources_per_worker={"CPU": node_cpus // num_gpus_per_node, "GPU": 1}`` in ``ScalingConfig``.
This method is fragile since it's tied to node shapes, and Ray Data also doesn't exclude other
resources such as object store memory properly, since the typical configuration is to only take up logical
CPUs and GPUs.
The recommended approach is to use :ref:`subclusters <data_concurrent_execution>` to pin the
training Dataset to CPU-only nodes. This correctly scopes the memory budget to only the nodes
where data tasks can actually run. It requires adding labels to your worker node configurations and setting the
``label_selector`` in two places:
.. code-block:: python
import ray
from ray.data import ExecutionOptions
from ray.train import DataConfig
from ray.train.torch import TorchTrainer
# (1) Pin construction-time tasks (schema inference, file listing).
ctx = ray.data.DataContext.get_current().copy()
ctx.execution_options.label_selector = {"ray-subcluster": "data"}
with ray.data.DataContext.current(ctx):
train_dataset = ray.data.read_parquet(...)
# (2) Pin per-worker ingest — Train replaces ds.context options
# wholesale, so the selector must be restated here.
trainer = TorchTrainer(
...,
datasets={"train": train_dataset},
dataset_config=DataConfig(
datasets_to_split=["train"],
execution_options={
"train": ExecutionOptions(
label_selector={"ray-subcluster": "data"}
),
},
),
)
.. tip::
Before resorting to isolating Ray Data tasks from training nodes, consider offloading that
heavy work from training workers to the data pipeline instead:
:ref:`scale out expensive collation <scaling_collation_functions>`
and use :ref:`map_batches-based shuffling <map_batches_shuffle>` in place of large local
shuffle buffers. These reduce CPU pressure on training workers and often eliminate the
need for node isolation entirely.
.. _balancing-data-production-consumption:
Balancing data production and consumption
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
Ideally, data production (dataset processing) and data consumption (training ingestion) happen at
the same rate. When production outpaces consumption, excess data is written to the
:ref:`object store <object-spilling-internals>`, which can spill to disk, which in turn decreases
Ray Data throughput and leads to out of disk errors. Ray Data's backpressure system automatically
balances production and consumption, but if you are still running into issues, you can try
tuning the following:
* **Use fewer CPUs for data production**: If you are using :func:`~ray.data.Dataset.map_batches`,
you can set the number of workers with ``compute`` and the number of CPUs per worker with ``num_cpus``.
* **Limit object store usage per dataset**: Set a per-dataset object store memory limit using
each dataset's execution options. Ray Data's backpressure system will slow down production
once the object store memory limit is reached.
.. code-block:: python
train_ds = ray.data.read_parquet("s3://bucket/train")
val_ds = ray.data.read_parquet("s3://bucket/val")
train_ds.context.execution_options.resource_limits = ray.data.ExecutionResources(
object_store_memory=50 * 1024**3,
)
val_ds.context.execution_options.resource_limits = ray.data.ExecutionResources(
object_store_memory=50 * 1024**3,
)
See :ref:`data_performance_tips` for more info on how to tune Ray Data.
More data ingest guides
-----------------------
- :ref:`Weighted Dataset Mixing <mixing_data>` — combine multiple datasets with target row ratios for training.
- :ref:`Scaling Collation Functions <scaling_collation_functions>` — scale out expensive collation functions to Ray Data.
@@ -0,0 +1,123 @@
.. _train-elastic-training:
Elastic training
================
Ray Train supports elastic training, enabling jobs to seamlessly adapt to changes in resource availability. This behavior ensures continuous execution despite hardware failures or node preemptions, avoiding idle or wasted time. As more nodes become available, the cluster dynamically scales up to speed up training with more worker processes.
To enable elastic training, use :attr:`~ray.train.ScalingConfig.num_workers` to specify ``(min_workers, max_workers)`` as a tuple instead of a fixed worker group size. You should also set :attr:`~ray.train.FailureConfig.max_failures` so that training can recover from worker failures instead of exiting immediately.
The following example shows how to configure elastic training with a range of 18 workers:
.. code-block:: python
from ray.train import RunConfig, FailureConfig
from ray.train.torch import TorchTrainer, ScalingConfig
def train_func():
# Your training code here
...
# Elastic training with 1-8 workers
scaling_config = ScalingConfig(num_workers=(1, 8), use_gpu=True)
# Allow retries so training survives worker failures
run_config = RunConfig(failure_config=FailureConfig(max_failures=3))
trainer = TorchTrainer(
train_func,
scaling_config=scaling_config,
run_config=run_config,
)
trainer.fit()
How it works
------------
Starting with available workers
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
Ray Train always requests ``max_workers`` number of workers. If it can't get all of them, it starts when ``min_workers`` is available so training can begin without waiting for the full set of resources.
When failures happen
~~~~~~~~~~~~~~~~~~~~
If any failures happen (for example, a worker crashes or a node is preempted), Ray Train restarts with fewer workers. It then attempts again to bring the worker group back up to ``max_workers``. Without a retry limit, the run would exit on the first such failure. To allow the run to retry when worker failures occur, configure :attr:`~ray.train.RunConfig.failure_config` with :attr:`~ray.train.FailureConfig.max_failures`:
.. code-block:: python
:emphasize-lines: 4
from ray.train import RunConfig, FailureConfig
# Retry up to 3 times on worker failures (e.g. preemption, node loss)
run_config = RunConfig(failure_config=FailureConfig(max_failures=3))
trainer = TorchTrainer(
train_func,
scaling_config=scaling_config,
run_config=run_config,
)
When more nodes become available
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
If the cluster gets more nodes eventually, Ray Train can resize the worker group and restart with the new workers added, so training can use the extra capacity. By default, the controller considers resizing every 60 seconds while the worker group is healthy. To change how often resize decisions are made, set :attr:`~ray.train.ScalingConfig.elastic_resize_monitor_interval_s` in your scaling config:
.. code-block:: python
# Consider resizing the worker group every 30 seconds (default is 60)
scaling_config = ScalingConfig(
num_workers=(1, 8),
use_gpu=True,
elastic_resize_monitor_interval_s=30.0,
)
Configure cluster autoscaling
-----------------------------
For elastic training to scale up when more resources become available, the cluster autoscaler must be configured to match your elastic training settings. Specifically, the cluster should be able to provision up to ``max_workers`` nodes and scale down to ``min_workers`` nodes.
.. tab-set::
.. tab-item:: KubeRay
Set the ``minReplicas`` and ``maxReplicas`` fields on your worker group to match the elastic training range. The following example configures a worker group that can scale between 1 and 8 nodes:
.. code-block:: yaml
:emphasize-lines: 3,4
workerGroupSpecs:
- groupName: gpu-workers
minReplicas: 1
maxReplicas: 8
replicas: 1
template:
spec:
containers:
- name: ray-worker
image: rayproject/ray:latest
.. note::
If the Kubernetes cluster itself doesn't have enough physical nodes, you also need to configure a Kubernetes-level autoscaler (such as the Cluster Autoscaler or Karpenter) so that new Kubernetes nodes are provisioned for the Ray worker pods. See :ref:`kuberay-autoscaling-config` for more details.
.. tab-item:: VMs
Set the ``min_workers`` and ``max_workers`` fields in your cluster config to match the elastic training range:
.. code-block:: yaml
:emphasize-lines: 5,6
max_workers: 8
available_node_types:
gpu_worker:
min_workers: 1
max_workers: 8
See :ref:`vms-autoscaling` for more details.
Limitations
-----------
Elastic training is supported for CPU and GPU backends only. It isn't supported yet for TPU training.
@@ -0,0 +1,399 @@
.. _train-experiment-tracking-native:
===================
Experiment Tracking
===================
Most experiment tracking libraries work out-of-the-box with Ray Train.
This guide provides instructions on how to set up the code so that your favorite experiment tracking libraries
can work for distributed training with Ray Train. The end of the guide has common errors to aid in debugging
the setup.
The following pseudo code demonstrates how to use the native experiment tracking library calls
inside of Ray Train:
.. testcode::
:skipif: True
from ray.train.torch import TorchTrainer
from ray.train import ScalingConfig
def train_func():
# Training code and native experiment tracking library calls go here.
scaling_config = ScalingConfig(num_workers=2, use_gpu=True)
trainer = TorchTrainer(train_func, scaling_config=scaling_config)
result = trainer.fit()
Ray Train lets you use native experiment tracking libraries by customizing the tracking
logic inside the :ref:`train_func<train-overview-training-function>` function.
In this way, you can port your experiment tracking logic to Ray Train with minimal changes.
Getting Started
===============
Let's start by looking at some code snippets.
The following examples uses Weights & Biases (W&B) and MLflow but it's adaptable to other frameworks.
.. tab-set::
.. tab-item:: W&B
.. testcode::
:skipif: True
import ray
from ray import train
import wandb
# Step 1
# This ensures that all ray worker processes have `WANDB_API_KEY` set.
ray.init(runtime_env={"env_vars": {"WANDB_API_KEY": "your_api_key"}})
def train_func():
# Step 1 and 2
if train.get_context().get_world_rank() == 0:
wandb.init(
name=...,
project=...,
# ...
)
# ...
loss = optimize()
metrics = {"loss": loss}
# Step 3
if train.get_context().get_world_rank() == 0:
# Only report the results from the rank 0 worker to W&B to avoid duplication.
wandb.log(metrics)
# ...
# Step 4
# Make sure that all loggings are uploaded to the W&B backend.
if train.get_context().get_world_rank() == 0:
wandb.finish()
.. tab-item:: MLflow
.. testcode::
:skipif: True
from ray import train
import mlflow
# Run the following on the head node:
# $ databricks configure --token
# mv ~/.databrickscfg YOUR_SHARED_STORAGE_PATH
# This function assumes `databricks_config_file` is specified in the Trainer's `train_loop_config`.
def train_func(config):
# Step 1 and 2
os.environ["DATABRICKS_CONFIG_FILE"] = config["databricks_config_file"]
mlflow.set_tracking_uri("databricks")
mlflow.set_experiment_id(...)
mlflow.start_run()
# ...
loss = optimize()
metrics = {"loss": loss}
# Step 3
if train.get_context().get_world_rank() == 0:
# Only report the results from the rank 0 worker to MLflow to avoid duplication.
mlflow.log_metrics(metrics)
.. tip::
A major difference between distributed and non-distributed training is that in distributed training,
multiple processes are running in parallel and under certain setups they have the same results. If all
of them report results to the tracking backend, you may get duplicated results. To address that,
Ray Train lets you apply logging logic to only the rank 0 worker with the following method:
:meth:`ray.train.get_context().get_world_rank() <ray.train.context.TrainContext.get_world_rank>`.
.. testcode::
:skipif: True
from ray import train
def train_func():
...
if train.get_context().get_world_rank() == 0:
# Add your logging logic only for rank0 worker.
...
The interaction with the experiment tracking backend within the :ref:`train_func<train-overview-training-function>`
has 4 logical steps:
#. Set up the connection to a tracking backend
#. Configure and launch a run
#. Log metrics
#. Finish the run
More details about each step follows.
Step 1: Connect to your tracking backend
----------------------------------------
First, decide which tracking backend to use: W&B, MLflow, TensorBoard, Comet, etc.
If applicable, make sure that you properly set up credentials on each training worker.
.. tab-set::
.. tab-item:: W&B
W&B offers both *online* and *offline* modes.
**Online**
For *online* mode, because you log to W&B's tracking service, ensure that you set the credentials
inside of :ref:`train_func<train-overview-training-function>`. See :ref:`Set up credentials<set-up-credentials>`
for more information.
.. testcode::
:skipif: True
# This is equivalent to `os.environ["WANDB_API_KEY"] = "your_api_key"`
wandb.login(key="your_api_key")
**Offline**
For *offline* mode, because you log towards a local file system,
point the offline directory to a shared storage path that all nodes can write to.
See :ref:`Set up a shared file system<set-up-shared-file-system>` for more information.
.. testcode::
:skipif: True
os.environ["WANDB_MODE"] = "offline"
wandb.init(dir="some_shared_storage_path/wandb")
.. tab-item:: MLflow
MLflow offers both *local* and *remote* (for example, to Databrick's MLflow service) modes.
**Local**
For *local* mode, because you log to a local file
system, point offline directory to a shared storage path. that all nodes can write
to. See :ref:`Set up a shared file system<set-up-shared-file-system>` for more information.
.. testcode::
:skipif: True
mlflow.set_tracking_uri(uri="file://some_shared_storage_path/mlruns")
mlflow.start_run()
**Remote, hosted by Databricks**
Ensure that all nodes have access to the Databricks config file.
See :ref:`Set up credentials<set-up-credentials>` for more information.
.. testcode::
:skipif: True
# The MLflow client looks for a Databricks config file
# at the location specified by `os.environ["DATABRICKS_CONFIG_FILE"]`.
os.environ["DATABRICKS_CONFIG_FILE"] = config["databricks_config_file"]
mlflow.set_tracking_uri("databricks")
mlflow.start_run()
.. _set-up-credentials:
Set up credentials
~~~~~~~~~~~~~~~~~~
Refer to each tracking library's API documentation on setting up credentials.
This step usually involves setting an environment variable or accessing a config file.
The easiest way to pass an environment variable credential to training workers is through
:ref:`runtime environments <runtime-environments>`, where you initialize with the following code:
.. testcode::
:skipif: True
import ray
# This makes sure that training workers have the same env var set
ray.init(runtime_env={"env_vars": {"SOME_API_KEY": "your_api_key"}})
For accessing the config file, ensure that the config file is accessible to all nodes.
One way to do this is by setting up a shared storage. Another way is to save a copy in each node.
.. _set-up-shared-file-system:
Set up a shared file system
~~~~~~~~~~~~~~~~~~~~~~~~~~~
Set up a network filesystem accessible to all nodes in the cluster.
For example, AWS EFS or Google Cloud Filestore.
Step 2: Configure and start the run
-----------------------------------
This step usually involves picking an identifier for the run and associating it with a project.
Refer to the tracking libraries' documentation for semantics.
.. To conveniently link back to Ray Train run, you may want to log the persistent storage path
.. of the run as a config.
..
.. testcode::
def train_func():
if ray.train.get_context().get_world_rank() == 0:
wandb.init(..., config={"ray_train_persistent_storage_path": "TODO: fill in when API stabilizes"})
.. tip::
When performing **fault-tolerant training** with auto-restoration, use a
consistent ID to configure all tracking runs that logically belong to the same training run.
Step 3: Log metrics
-------------------
You can customize how to log parameters, metrics, models, or media contents, within
:ref:`train_func<train-overview-training-function>`, just as in a non-distributed training script.
You can also use native integrations that a particular tracking framework has with
specific training frameworks. For example, ``mlflow.pytorch.autolog()``,
``lightning.pytorch.loggers.MLFlowLogger``, etc.
Step 4: Finish the run
----------------------
This step ensures that all logs are synced to the tracking service. Depending on the implementation of
various tracking libraries, sometimes logs are first cached locally and only synced to the tracking
service in an asynchronous fashion.
Finishing the run makes sure that all logs are synced by the time training workers exit.
.. tab-set::
.. tab-item:: W&B
.. testcode::
:skipif: True
# https://docs.wandb.ai/ref/python/finish
wandb.finish()
.. tab-item:: MLflow
.. testcode::
:skipif: True
# https://mlflow.org/docs/1.2.0/python_api/mlflow.html
mlflow.end_run()
.. tab-item:: Comet
.. testcode::
:skipif: True
# https://www.comet.com/docs/v2/api-and-sdk/python-sdk/reference/Experiment/#experimentend
Experiment.end()
Examples
========
The following are runnable examples for PyTorch and PyTorch Lightning.
PyTorch
-------
.. dropdown:: Log to W&B
.. literalinclude:: ../../../../python/ray/train/examples/experiment_tracking//torch_exp_tracking_wandb.py
:emphasize-lines: 16, 19-21, 59-60, 62-63
:language: python
:start-after: __start__
.. dropdown:: Log to file-based MLflow
.. literalinclude:: ../../../../python/ray/train/examples/experiment_tracking/torch_exp_tracking_mlflow.py
:emphasize-lines: 22-25, 58-59, 61-62, 68
:language: python
:start-after: __start__
:end-before: __end__
PyTorch Lightning
-----------------
You can use the native Logger integration in PyTorch Lightning with W&B, CometML, MLFlow,
and Tensorboard, while using Ray Train's TorchTrainer.
The following example walks you through the process. The code here is runnable.
.. dropdown:: W&B
.. literalinclude:: ../../../../python/ray/train/examples/experiment_tracking/lightning_exp_tracking_model_dl.py
:language: python
:start-after: __model_dl_start__
.. literalinclude:: ../../../../python/ray/train/examples/experiment_tracking/lightning_exp_tracking_wandb.py
:language: python
:start-after: __lightning_experiment_tracking_wandb_start__
.. dropdown:: MLflow
.. literalinclude:: ../../../../python/ray/train/examples/experiment_tracking/lightning_exp_tracking_model_dl.py
:language: python
:start-after: __model_dl_start__
.. literalinclude:: ../../../../python/ray/train/examples/experiment_tracking/lightning_exp_tracking_mlflow.py
:language: python
:start-after: __lightning_experiment_tracking_mlflow_start__
:end-before: __lightning_experiment_tracking_mlflow_end__
.. dropdown:: Comet
.. literalinclude:: ../../../../python/ray/train/examples/experiment_tracking/lightning_exp_tracking_model_dl.py
:language: python
:start-after: __model_dl_start__
.. literalinclude:: ../../../../python/ray/train/examples/experiment_tracking/lightning_exp_tracking_comet.py
:language: python
:start-after: __lightning_experiment_tracking_comet_start__
.. dropdown:: TensorBoard
.. literalinclude:: ../../../../python/ray/train/examples/experiment_tracking/lightning_exp_tracking_model_dl.py
:language: python
:start-after: __model_dl_start__
.. literalinclude:: ../../../../python/ray/train/examples/experiment_tracking/lightning_exp_tracking_tensorboard.py
:language: python
:start-after: __lightning_experiment_tracking_tensorboard_start__
:end-before: __lightning_experiment_tracking_tensorboard_end__
Common Errors
=============
Missing Credentials
-------------------
**I have already called `wandb login` cli, but am still getting**
.. code-block:: none
wandb: ERROR api_key not configured (no-tty). call wandb.login(key=[your_api_key]).
This is probably due to wandb credentials are not set up correctly
on worker nodes. Make sure that you run ``wandb.login``
or pass ``WANDB_API_KEY`` to each training function.
See :ref:`Set up credentials <set-up-credentials>` for more details.
Missing Configurations
----------------------
**I have already run `databricks configure`, but am still getting**
.. code-block:: none
databricks_cli.utils.InvalidConfigurationError: You haven't configured the CLI yet!
This is usually caused by running ``databricks configure`` which
generates ``~/.databrickscfg`` only on head node. Move this file to a shared
location or copy it to each node.
See :ref:`Set up credentials <set-up-credentials>` for more details.
@@ -0,0 +1,235 @@
.. _train-fault-tolerance:
Handling Failures and Node Preemption
=====================================
.. important::
This user guide shows how to configure fault tolerance for the revamped Ray Train V2
available starting from Ray 2.43 by enabling the environment variable ``RAY_TRAIN_V2_ENABLED=1``.
**This user guide assumes that the environment variable has been enabled.**
Please see :ref:`here <train-fault-tolerance-deprecation-info>` for information about the deprecation and migration.
Ray Train provides fault tolerance at three levels:
1. **Worker process fault tolerance** handles errors that happen to one or more Train worker processes while they are executing the user defined training function.
2. **Worker node fault tolerance** handles node failures that may occur during training.
3. **Job driver fault tolerance** handles the case where Ray Train driver process crashes, and training needs to be kicked off again, possibly from a new cluster.
This user guide covers how to configure and use these fault tolerance mechanisms.
.. _train-worker-fault-tolerance:
Worker Process and Node Fault Tolerance
---------------------------------------
**Worker process failures** are errors that occur within the user defined training function of a training worker,
such as GPU out-of-memory (OOM) errors, cloud storage access errors, or other runtime errors.
**Node failures** are errors that bring down the entire node, including node preemption, OOM, network partitions, or other hardware failures.
This section covers worker node failures. Recovery from head node failures is discussed in the :ref:`next section <train-job-driver-fault-tolerance>`.
Ray Train can be configured to automatically recover from worker process and worker node failures.
When a failure is detected, all the workers are shut down, new nodes are added if necessary, and a new set of workers is started.
The restarted training worker processes can resume training by loading the latest checkpoint.
In order to retain progress upon recovery, your training function
should implement logic for both :ref:`saving <train-dl-saving-checkpoints>`
*and* :ref:`loading checkpoints <train-dl-loading-checkpoints>`.
Otherwise, the training will just start from scratch.
Each recovery from a worker process or node failure is considered a retry. The
number of retries is configurable through the ``max_failures`` attribute of the
:class:`~ray.train.FailureConfig` argument set in the :class:`~ray.train.RunConfig`
passed to the ``Trainer``. By default, worker fault tolerance is disabled with ``max_failures=0``.
.. literalinclude:: ../doc_code/fault_tolerance.py
:language: python
:start-after: __failure_config_start__
:end-before: __failure_config_end__
Altogether, this is what an example Torch training script with worker fault tolerance looks like:
.. literalinclude:: ../doc_code/fault_tolerance.py
:language: python
:start-after: __worker_fault_tolerance_start__
:end-before: __worker_fault_tolerance_end__
Which checkpoint will be restored?
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
Ray Train will populate :func:`ray.train.get_checkpoint() <ray.train.get_checkpoint>` with the latest available
:ref:`checkpoint reported to Ray Train <train-checkpointing>`.
The :class:`~ray.train.Checkpoint` object returned by this method has the
:meth:`~ray.train.Checkpoint.as_directory` and :meth:`~ray.train.Checkpoint.to_directory` methods
to download the checkpoint from the :class:`RunConfig(storage_path) <ray.train.RunConfig>` to local disk.
.. note::
:meth:`~ray.train.Checkpoint.as_directory` and :meth:`~ray.train.Checkpoint.to_directory`
will only download the checkpoint once per node even if there are multiple workers on the node.
The workers share the same checkpoint directory on local disk.
Illustrated Example
~~~~~~~~~~~~~~~~~~~
Consider the following example of a cluster containing a CPU head node and 2 GPU worker nodes.
There are 4 GPU training workers running on the 2 worker nodes.
The :ref:`storage path has been configured <persistent-storage-guide>` to use cloud storage, which is where checkpoints are saved.
.. figure:: ../images/fault_tolerance/worker_failure_start.png
:align: left
Training has been running for some time, and the latest checkpoint has been saved to cloud storage.
.. figure:: ../images/fault_tolerance/worker_node_failure.png
:align: left
One of the worker GPU nodes fails due to a hardware fault. Ray Train detects this failure and shuts down all the workers.
Since the number of failures detected so far is less than the configured ``max_failures``, Ray Train will attempt to restart training,
rather than exiting and raising an error.
.. figure:: ../images/fault_tolerance/worker_node_replacement.png
:align: left
Ray Train has requested a new worker node to join the cluster and is waiting for it to come up.
.. figure:: ../images/fault_tolerance/worker_group_recovery.png
:align: left
The new worker node has joined the cluster.
Ray Train restarts all the worker processes and provides them with the latest checkpoint.
The workers download the checkpoint from storage and use it to resume training.
.. _train-restore-guide:
.. _train-job-driver-fault-tolerance:
Job Driver Fault Tolerance
--------------------------
Job driver fault tolerance is to handle cases where the Ray Train driver process is interrupted.
The Ray Train driver process is the process that calls ``trainer.fit()`` and is usually located on the head node of the cluster.
The driver process may be interrupted due to one of the following reasons:
- The run is manually interrupted by a user (e.g., Ctrl+C).
- The node where the driver process is running (head node) crashes (e.g., out of memory, out of disk).
- The entire cluster goes down (e.g., network error affecting all nodes).
In these cases, the Ray Train driver (which calls ``trainer.fit()``) needs to be launched again.
The relaunched Ray Train driver needs to find a minimal amount of run state in order to pick up where the previous run left off.
This state includes the latest reported checkpoints, which are located at the :ref:`storage path <persistent-storage-guide>`.
Ray Train fetches the latest checkpoint information from storage and passes it to the newly launched worker processes to resume training.
To find this run state, Ray Train relies on passing in the **same** :class:`RunConfig(storage_path, name) <ray.train.RunConfig>` pair as the previous run.
If the ``storage_path`` or ``name`` do not match, Ray Train will not be able to find the previous run state and will start a new run from scratch.
.. warning::
If ``name`` is reused unintentionally, Ray Train will fetch the previous run state, even if the user is trying to start a new run.
Therefore, always pass a unique run name when launching a new run. In other words, ``name`` should be a unique identifier for a training job.
.. note::
Job driver crashes and interrupts do not count toward the ``max_failures`` limit of :ref:`worker fault tolerance <train-worker-fault-tolerance>`.
Here's an example training script that highlights best practices for job driver fault tolerance:
.. literalinclude:: ../doc_code/fault_tolerance.py
:language: python
:start-after: __job_driver_fault_tolerance_start__
:end-before: __job_driver_fault_tolerance_end__
Then, the entrypoint script can be launched with the following command:
.. code-block:: bash
python entrypoint.py --storage_path s3://my_bucket/ --run_name unique_run_id=da823d5
If the job is interrupted, the same command can be used to resume training.
This example shows a ``da823d5`` id, which is determined by the one launching the job. The id can often be used for other purposes such as setting the ``wandb`` or ``mlflow`` run id.
Illustrated Example
~~~~~~~~~~~~~~~~~~~
Consider the following example of a cluster containing a CPU head node and 2 GPU worker nodes. There are 4 GPU training workers running on the 2 worker nodes. The storage path has been configured to use cloud storage, which is where checkpoints are saved.
.. figure:: ../images/fault_tolerance/cluster_failure_start.png
:align: left
Training has been running for some time, and the latest checkpoints and run state has been saved to storage.
.. figure:: ../images/fault_tolerance/head_node_failure.png
:align: left
The head node crashes for some reason (e.g., an out-of-memory error), and the Ray Train driver process is interrupted.
.. figure:: ../images/fault_tolerance/cluster_failure.png
:align: left
The entire cluster goes down due to the head node failure.
.. figure:: ../images/fault_tolerance/cluster_recovery.png
:align: left
A manual cluster restart or some job submission system brings up a new Ray cluster.
The Ray Train driver process runs on a new head node.
Ray Train fetches the run state information from storage at ``{storage_path}/{name}`` (e.g., ``s3://my_bucket/my_run_name``)
and passes the latest checkpoint to the newly launched worker processes to resume training.
.. _train-fault-tolerance-deprecation-info:
Fault Tolerance API Deprecations
--------------------------------
``<Framework>Trainer.restore`` API Deprecation
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
The ``<Framework>Trainer.restore`` and ``<Framework>Trainer.can_restore`` APIs are deprecated as of Ray 2.43 and will be removed in a future release.
Motivation
**********
This API change provides several benefits:
1. **Avoid saving user code to pickled files**: The old API saved user code to pickled files, which could lead to issues with deserialization, leading to unrecoverable runs.
2. **Improved configuration experience**: While some configurations were loaded from the pickled files, certain arguments were required to be re-specified, and another subset of arguments could even be optionally re-specified. This confused users about the set of configurations that are actually being used in the restored run.
Migration Steps
***************
To migrate from the old ``<Framework>Trainer.restore`` API to the new pattern:
1. Enable the environment variable ``RAY_TRAIN_V2_ENABLED=1``.
2. Replace ``<Framework>Trainer.restore`` with the regular ``<Framework>Trainer`` constructor, making sure to pass in the same ``storage_path`` and ``name`` as the previous run.
``<Framework>Trainer(restore_from_checkpoint)`` API Deprecation
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
The ``<Framework>Trainer(restore_from_checkpoint)`` API is deprecated as of Ray 2.43 and will be removed in a future release.
Motivation
**********
This API was a common source of confusion that provided minimal value. It was only used to set the initial value of ``ray.train.get_checkpoint()`` but did not load any other run state.
Migration Steps
***************
Simply pass in the initial checkpoint through the ``train_loop_config`` argument. See the migration guide linked below for a code example.
Additional Resources
~~~~~~~~~~~~~~~~~~~~
* `Train V2 Migration Guide <https://github.com/ray-project/ray/issues/49454>`_: Full migration guide for Train V2
* `Train V2 REP <https://github.com/ray-project/enhancements/blob/main/reps/2024-10-18-train-tune-api-revamp/2024-10-18-train-tune-api-revamp.md>`_: Technical details about the API change
* :ref:`train-fault-tolerance-deprecated-api`: Documentation for the old API
@@ -0,0 +1,209 @@
.. _train-tune:
Hyperparameter Tuning with Ray Tune
===================================
.. important::
This user guide shows how to integrate Ray Train and Ray Tune to tune over distributed hyperparameter runs
for the revamped Ray Train V2 available starting from Ray 2.43 by enabling the environment variable ``RAY_TRAIN_V2_ENABLED=1``.
**This user guide assumes that the environment variable has been enabled.**
Please see :ref:`here <train-tune-deprecation>` for information about the deprecation and migration.
Ray Train can be used together with Ray Tune to do hyperparameter sweeps of distributed training runs.
This is often useful when you want to do a small sweep over critical hyperparameters,
before launching a run with the best performing hyperparameters on all available cluster resources for a long duration.
Quickstart
----------
In the example below:
* :class:`~ray.tune.Tuner` launches the tuning job, which runs trials of ``train_driver_fn`` with different hyperparameter configurations.
* ``train_driver_fn``, which (1) takes in a hyperparameter configuration, (2) instantiates a ``TorchTrainer`` (or some other framework trainer), and (3) launches the distributed training job.
* :class:`~ray.train.ScalingConfig` defines the number of training workers and resources per worker for a single Ray Train run.
* ``train_fn_per_worker`` is the Python code that executes on each distributed training worker for a trial.
.. literalinclude:: ../doc_code/train_tune_interop.py
:language: python
:start-after: __quickstart_start__
:end-before: __quickstart_end__
What does Ray Tune provide?
---------------------------
Ray Tune provides utilities for:
* :ref:`Defining hyperparameter search spaces <tune-search-space-tutorial>` and :ref:`launching multiple trials concurrently <tune-parallel-experiments-guide>` on a Ray cluster
* :ref:`Using search algorithms <tune-search-alg>`
* :ref:`Early stopping runs based on metrics <tune-stopping-guide>`
This user guide only focuses on the integration layer between Ray Train and Ray Tune. For more details on how to use Ray Tune, refer to the :ref:`Ray Tune documentation <tune-main>`.
Configuring resources for multiple trials
-----------------------------------------
Ray Tune launches multiple trials which :ref:`run a user-defined function in a remote Ray actor <tune-function-api>`, where each trial gets a different sampled hyperparameter configuration.
When using Ray Tune by itself, trials do computation directly inside the Ray actor. For example, each trial could request 1 GPU and do some single-process model
training within the remote actor itself. When using Ray Train inside Ray Tune functions, the Tune trial is actually not doing extensive computation inside this actor
-- instead it just acts as a driver process to launch and monitor the Ray Train workers running elsewhere.
Ray Train requests its own resources via the :class:`~ray.train.ScalingConfig`.
See :ref:`train_scaling_config` for more details.
.. figure:: ../images/hyperparameter_optimization/train_without_tune.png
:align: center
A single Ray Train run to showcase how using Ray Tune in the next figure just adds a layer of hierarchy to this tree of processes.
.. figure:: ../images/hyperparameter_optimization/train_tune_interop.png
:align: center
Example of Ray Train runs being launched from within Ray Tune trials.
Limit the number of concurrent Ray Train runs
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
Ray Train runs can only start when resources for all workers can be acquired at once.
This means that multiple Tune trials spawning Train runs will be competing for the logical resources available in the Ray cluster.
If there is a limiting cluster resource such as GPUs, then it won't be possible to run training for all hyperparameter configurations concurrently.
Since the cluster only has enough resources for a handful of trials to run concurrently,
set :class:`tune.TuneConfig(max_concurrent_trials) <ray.tune.TuneConfig>` on the Tuner to limit the number of “in-flight” Train runs so that no trial is being starved of resources.
.. literalinclude:: ../doc_code/train_tune_interop.py
:language: python
:start-after: __max_concurrent_trials_start__
:end-before: __max_concurrent_trials_end__
As a concrete example, consider a fixed sized cluster with 128 CPUs and 8 GPUs.
* The ``Tuner(param_space)`` sweeps over 4 hyperparameter configurations with a grid search: ``param_space={“train_loop_config”: {“batch_size”: tune.grid_search([8, 16, 32, 64])}}``
* Each Ray Train run is configured to train with 4 GPU workers: ``ScalingConfig(num_workers=4, use_gpu=True)``. Since there are only 8 GPUs, only 2 Train runs can acquire their full set of resources at a time.
* However, since there are many CPUs available in the cluster, the 4 total Ray Tune trials (which default to requesting 1 CPU) can be launched immediately.
This results in 2 extra Ray Tune trial processes being launched, even though their inner Ray Train run just waits for resources until one of the other trials finishes.
This introduces some spammy log messages when Train waits for resources. There may also be an excessive number of Ray Tune trial processes if the total number of hyperparameter configurations is large.
* To fix this issue, set ``Tuner(tune_config=tune.TuneConfig(max_concurrent_trials=2))``. Now, only two Ray Tune trial processes will be running at a time.
This number can be calculated based on the limiting cluster resource and the amount of that resources required by each trial.
Advanced: Set Train driver resources
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
The default Train driver runs as a Ray Tune function with 1 CPU. Ray Tune will schedule these functions to run anywhere on the cluster that has free logical CPU resources.
**Recommendation:** If you are launching longer-running training jobs or using spot instances, these Tune functions which act as the Ray Train driver process should be run on “safe nodes” that are at lower risk of going down. For example, they should not be scheduled to run on preemptible spot instances and should not be colocated with training workers. This could be the head node or a dedicated CPU node in your cluster.
This is because the Ray Train driver process is responsible for handling fault tolerance of the worker processes, which are more likely to error. Nodes that are running Train workers can crash due to spot preemption or other errors that come up due to the user-defined model training code.
* If a Train worker node dies, the Ray Train driver process that is still alive on a different node can gracefully handle the error.
* On the other hand, if the driver process dies, then all Ray Train workers will ungracefully exit and some of the run state may not be committed fully.
One way to achieve this behavior is to set custom resources on certain node types and configure the Tune functions to request those resources.
.. literalinclude:: ../doc_code/train_tune_interop.py
:language: python
:start-after: __trainable_resources_start__
:end-before: __trainable_resources_end__
Reporting metrics and checkpoints
---------------------------------
Both Ray Train and Ray Tune provide utilities to help upload and track checkpoints via the :func:`ray.train.report <ray.train.report>` and :func:`ray.tune.report <ray.tune.report>` APIs.
See the :ref:`train-checkpointing` user guide for more details.
If the Ray Train workers report checkpoints, saving another Ray Tune checkpoint at the Train driver level is not needed because it does not hold any extra training state. The Ray Train driver process will already periodically snapshot its status to the configured storage_path, which is further described in the next section on fault tolerance.
In order to access the checkpoints from the Tuner output, you can append the checkpoint path as a metric. The provided :class:`~ray.tune.integration.ray_train.TuneReportCallback`
does this by propagating reported Ray Train results over to Ray Tune, where the checkpoint path is attached as a separate metric.
Advanced: Fault Tolerance
~~~~~~~~~~~~~~~~~~~~~~~~~
In the event that the Ray Tune trials running the Ray Train driver process crash, you can enable trial fault tolerance on the Ray Tune side via:
:class:`ray.tune.Tuner(run_config=ray.tune.RunConfig(failure_config)) <ray.tune.FailureConfig>`.
Fault tolerance on the Ray Train side is configured and handled separately. See the :ref:`train-fault-tolerance` user guide for more details.
.. literalinclude:: ../doc_code/train_tune_interop.py
:language: python
:start-after: __fault_tolerance_start__
:end-before: __fault_tolerance_end__
.. _train-with-tune-callbacks:
Advanced: Using Ray Tune callbacks
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
Ray Tune callbacks should be passed into the :class:`ray.tune.RunConfig(callbacks) <ray.tune.RunConfig>` at the Tuner level.
For Ray Train users that depend on behavior of built-in or custom Ray Tune callbacks, it's possible to use them by running Ray Train as a single trial Tune run
and passing in the callbacks to the Tuner.
If any callback functionality depends on reported metrics, make sure to pass the :class:`ray.tune.integration.ray_train.TuneReportCallback` to the trainer callbacks,
which propagates results to the Tuner.
.. testcode::
:skipif: True
import ray.tune
from ray.tune.integration.ray_train import TuneReportCallback
from ray.tune.logger import TBXLoggerCallback
def train_driver_fn(config: dict):
trainer = TorchTrainer(
...,
run_config=ray.train.RunConfig(..., callbacks=[TuneReportCallback()])
)
trainer.fit()
tuner = ray.tune.Tuner(
train_driver_fn,
run_config=ray.tune.RunConfig(callbacks=[TBXLoggerCallback()])
)
.. _train-tune-deprecation:
``Tuner(trainer)`` API Deprecation
----------------------------------
The ``Tuner(trainer)`` API which directly takes in a Ray Train trainer instance is deprecated as of Ray 2.43 and will be removed in a future release.
Motivation
~~~~~~~~~~
This API change provides several benefits:
1. **Better separation of concerns**: Decouples Ray Train and Ray Tune responsibilities
2. **Improved configuration experience**: Makes hyperparameter and run configuration more explicit and flexible
Migration Steps
~~~~~~~~~~~~~~~
To migrate from the old ``Tuner(trainer)`` API to the new pattern:
1. Enable the environment variable ``RAY_TRAIN_V2_ENABLED=1``.
2. Replace ``Tuner(trainer)`` with a function-based approach where Ray Train is launched inside a Tune trial.
3. Move your training logic into a driver function that Tune will call with different hyperparameters.
Additional Resources
~~~~~~~~~~~~~~~~~~~~
* `Train V2 REP <https://github.com/ray-project/enhancements/blob/main/reps/2024-10-18-train-tune-api-revamp/2024-10-18-train-tune-api-revamp.md>`_: Technical details about the API change
* `Train V2 Migration Guide <https://github.com/ray-project/ray/issues/49454>`_: Full migration guide for Train V2
* :ref:`train-tune-deprecated-api`: Documentation for the old API
+388
View File
@@ -0,0 +1,388 @@
.. _train-local-mode:
Local Mode
==========
.. important::
This user guide shows how to use local mode with Ray Train V2 only.
For information about migrating from Ray Train V1 to V2, see the Train V2 migration guide: https://github.com/ray-project/ray/issues/49454
What is local mode?
-------------------
Local mode in Ray Train runs your training function without launching Ray Train worker actors.
Instead of distributing your training code across multiple Ray actors, local mode executes your
training function directly in the current process. This provides a simplified debugging environment
where you can iterate quickly on your training logic.
Local mode supports two execution modes:
* **Single-process mode**: Runs your training function in a single process, ideal for rapid iteration and debugging.
* **Multi-process mode with torchrun**: Launches multiple processes for multi-GPU training, useful for debugging distributed training logic with familiar tools.
How to enable local mode
-------------------------
You can enable local mode by setting ``num_workers=0`` in your :class:`~ray.train.ScalingConfig`:
.. testcode::
:skipif: True
from ray.train import ScalingConfig
from ray.train.torch import TorchTrainer
def train_func(config):
# Your training logic
pass
trainer = TorchTrainer(
train_loop_per_worker=train_func,
scaling_config=ScalingConfig(num_workers=0),
)
result = trainer.fit()
Local mode provides the same ``ray.train`` APIs you use in distributed training, so your
training code runs without any other modifications. This makes it simple to verify your
training logic locally before scaling to distributed training.
When to use local mode
----------------------
Use single-process local mode to:
* **Develop and iterate quickly**: Test changes to your training function locally.
* **Write unit tests**: Verify your training logic works correctly in a simplified environment.
* **Debug training logic**: Use standard Python debugging tools to step through your training code and identify issues.
Use multi-process local mode with ``torchrun`` to:
* **Test multi-GPU logic**: Verify your distributed training code works correctly across multiple GPUs using familiar ``torchrun`` commands.
* **Migrate existing code**: Bring existing ``torchrun`` based training scripts into Ray Train while preserving your development workflow.
* **Debug distributed behavior**: Isolate issues in your distributed training logic using ``torchrun``'s process management.
.. note::
In local mode, Ray Train doesn't launch worker actors, but your training code can still
use other Ray features such as Ray Data (in single-process mode) or launch Ray actors if needed.
Single-process local mode
--------------------------
The following example shows how to use single-process local mode with PyTorch:
.. testcode::
:skipif: True
import torch
from torch import nn
import ray
from ray.train import ScalingConfig
from ray.train.torch import TorchTrainer
def train_func(config):
model = nn.Linear(10, 1)
optimizer = torch.optim.SGD(model.parameters(), lr=config["lr"])
for epoch in range(config["epochs"]):
# Training loop
loss = model(torch.randn(32, 10)).sum()
loss.backward()
optimizer.step()
# Report metrics
ray.train.report({"loss": loss.item()})
trainer = TorchTrainer(
train_loop_per_worker=train_func,
train_loop_config={"lr": 0.01, "epochs": 3},
scaling_config=ScalingConfig(num_workers=0),
)
result = trainer.fit()
print(f"Final loss: {result.metrics['loss']}")
.. note::
Local mode works with all Ray Train framework integrations, including PyTorch Lightning,
Hugging Face Transformers, LightGBM, XGBoost, TensorFlow, and others.
Testing with local mode
~~~~~~~~~~~~~~~~~~~~~~~
The following example shows how to write a unit test with local mode:
.. testcode::
:skipif: True
import pytest
import ray
from ray.train import ScalingConfig
from ray.train.torch import TorchTrainer
def test_training_runs():
def train_func(config):
# Report minimal training result
ray.train.report({"loss": 0.5})
trainer = TorchTrainer(
train_loop_per_worker=train_func,
scaling_config=ScalingConfig(num_workers=0),
)
result = trainer.fit()
assert result.error is None
assert result.metrics["loss"] == 0.5
Using local mode with Ray Data
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
Single-process local mode works seamlessly with Ray Data for data loading and preprocessing.
When you use Ray Data with local mode, Ray Data processes your data and provides it back to your
training function in the local process.
The following example shows how to use Ray Data with single-process local mode:
.. testcode::
:skipif: True
import ray
from ray.train import ScalingConfig
from ray.train.torch import TorchTrainer
def train_func(config):
# Get the dataset shard
train_dataset = ray.train.get_dataset_shard("train")
# Iterate over batches
for batch in train_dataset.iter_batches(batch_size=32):
# Training logic
pass
# Create a Ray Dataset
dataset = ray.data.read_csv("s3://bucket/data.csv")
trainer = TorchTrainer(
train_loop_per_worker=train_func,
scaling_config=ScalingConfig(num_workers=0),
datasets={"train": dataset},
)
result = trainer.fit()
.. warning::
Ray Data isn't supported when using ``torchrun`` for multi-process training in local mode.
For multi-process training, use standard PyTorch data loading mechanisms such as DataLoader
with DistributedSampler.
Multi-process local mode with ``torchrun``
-------------------------------------------
Local mode supports multi-GPU training through ``torchrun``, allowing you to develop and debug using ``torchrun``'s process management.
Single-node multi-GPU training
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
The following example shows how to use ``torchrun`` with local mode for multi-GPU training on a single node.
This approach is useful when migrating existing PyTorch training code or when you want to debug
distributed training logic using ``torchrun``'s familiar process management. The example uses standard
PyTorch ``DataLoader`` for data loading, making it easy to adapt your existing PyTorch training code.
First, create your training script (``train_script.py``):
.. testcode::
:skipif: True
import os
import tempfile
import torch
import torch.distributed as dist
from torch import nn
from torch.utils.data import DataLoader
from torchvision.datasets import FashionMNIST
from torchvision.transforms import ToTensor, Normalize, Compose
from filelock import FileLock
import ray
from ray.train import Checkpoint, ScalingConfig, get_context
from ray.train.torch import TorchTrainer
def train_func(config):
# Load dataset with file locking to avoid multiple downloads
transform = Compose([ToTensor(), Normalize((0.5,), (0.5,))])
data_dir = "./data"
# Only local rank 0 downloads the dataset
local_rank = get_context().get_local_rank()
if local_rank == 0:
with FileLock(os.path.join(data_dir, "fashionmnist.lock")):
train_dataset = FashionMNIST(
root=data_dir, train=True, download=True, transform=transform
)
# Wait for rank 0 to finish downloading
dist.barrier()
# Now all ranks can safely load the dataset
train_dataset = FashionMNIST(
root=data_dir, train=True, download=False, transform=transform
)
train_loader = DataLoader(
train_dataset, batch_size=config["batch_size"], shuffle=True
)
# Prepare dataloader for distributed training
train_loader = ray.train.torch.prepare_data_loader(train_loader)
# Prepare model for distributed training
model = nn.Sequential(
nn.Flatten(),
nn.Linear(28 * 28, 128),
nn.ReLU(),
nn.Linear(128, 10)
)
model = ray.train.torch.prepare_model(model)
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=config["lr"])
# Training loop
for epoch in range(config["epochs"]):
# Set epoch for distributed sampler
if ray.train.get_context().get_world_size() > 1:
train_loader.sampler.set_epoch(epoch)
epoch_loss = 0.0
for batch_idx, (images, labels) in enumerate(train_loader):
outputs = model(images)
loss = criterion(outputs, labels)
optimizer.zero_grad()
loss.backward()
optimizer.step()
epoch_loss += loss.item()
avg_loss = epoch_loss / len(train_loader)
# Report metrics and checkpoint
with tempfile.TemporaryDirectory() as temp_dir:
torch.save(model.state_dict(), os.path.join(temp_dir, "model.pt"))
ray.train.report(
{"loss": avg_loss, "epoch": epoch},
checkpoint=Checkpoint.from_directory(temp_dir)
)
# Configure trainer for local mode
trainer = TorchTrainer(
train_loop_per_worker=train_func,
train_loop_config={"lr": 0.001, "epochs": 10, "batch_size": 32},
scaling_config=ScalingConfig(num_workers=0, use_gpu=True),
)
result = trainer.fit()
Then, launch training with ``torchrun``:
.. code-block:: bash
# Train on 4 GPUs on a single node
torchrun --nproc-per-node=4 train_script.py
Ray Train automatically detects the ``torchrun`` environment variables and configures the distributed
training accordingly. You can access distributed training information through :func:`ray.train.get_context()`:
.. testcode::
:skipif: True
from ray.train import get_context
context = get_context()
print(f"World size: {context.get_world_size()}")
print(f"World rank: {context.get_world_rank()}")
print(f"Local rank: {context.get_local_rank()}")
.. warning::
Ray Data isn't supported when using ``torchrun`` for multi-process training in local mode.
For multi-process training, use standard PyTorch data loading mechanisms such as DataLoader with
DistributedSampler.
Multi-node multi-GPU training
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
You can also use ``torchrun`` to launch multi-node training with local mode. The following example shows
how to launch training across 2 nodes with 4 GPUs each:
On the master node (``192.168.1.1``):
.. code-block:: bash
RAY_TRAIN_V2_ENABLED=1 torchrun \
--nnodes=2 \
--nproc-per-node=4 \
--node_rank=0 \
--rdzv_backend=c10d \
--rdzv_endpoint=192.168.1.1:29500 \
--rdzv_id=job_id \
train_script.py
On the worker node:
.. code-block:: bash
RAY_TRAIN_V2_ENABLED=1 torchrun \
--nnodes=2 \
--nproc-per-node=4 \
--node_rank=1 \
--rdzv_backend=c10d \
--rdzv_endpoint=192.168.1.1:29500 \
--rdzv_id=job_id \
train_script.py
Transitioning from local mode to distributed training
-----------------------------------------------------
When you're ready to scale from local mode to distributed training, simply change ``num_workers``
to a value greater than 0:
.. code-block:: diff
trainer = TorchTrainer(
train_loop_per_worker=train_func,
train_loop_config=config,
- scaling_config=ScalingConfig(num_workers=0),
+ scaling_config=ScalingConfig(num_workers=4, use_gpu=True),
)
Your training function code remains the same, and Ray Train handles the distributed coordination automatically.
Limitations and API differences
--------------------------------
Local mode provides simplified implementations of Ray Train APIs to enable rapid debugging without distributed orchestration. However, this means some features behave differently or aren't available.
Features not available in local mode
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
The following Ray Train features aren't available in local mode:
* **Worker-level fault tolerance**: Ray Train's automatic fault tolerance features, such as worker restart on failure, aren't available. If you configured :class:`~ray.train.FailureConfig`, the settings don't apply in local mode.
* **Callbacks**: User-defined callbacks specified in :class:`~ray.train.RunConfig` aren't invoked in local mode.
* **Ray Data with multi-process training**: Ray Data isn't supported when using ``torchrun`` with local mode for multi-process training. Use standard PyTorch data loading mechanisms instead.
API behavior differences
~~~~~~~~~~~~~~~~~~~~~~~~
The following table summarizes how ``ray.train`` APIs behave differently in local mode:
.. list-table::
:header-rows: 1
:widths: 30 70
* - API
- Behavior in local mode
* - :func:`ray.train.report`
- Stores checkpoints in memory only (not persisted to storage). Ignores ``checkpoint_upload_mode``, ``checkpoint_upload_fn``, ``validation``, and ``delete_local_checkpoint_after_upload`` parameters. Logs metrics locally instead of through the reporting pipeline. Doesn't invoke a synchronization barrier across workers.
* - :func:`ray.train.get_checkpoint`
- Returns the last checkpoint from memory. Doesn't load checkpoints from persistent storage.
* - :func:`ray.train.get_all_reported_checkpoints`
- Always returns an empty list. Doesn't track checkpoint history.
* - :func:`ray.train.collective.barrier`
- No-op.
* - :func:`ray.train.collective.broadcast_from_rank_zero`
- Returns data as-is.
* - :meth:`ray.train.get_context().get_storage() <ray.train.TrainContext.get_storage>`
- Raises ``NotImplementedError``
@@ -0,0 +1,30 @@
.. _train-metrics:
Ray Train Metrics
-----------------
Ray Train exports Prometheus metrics including the Ray Train controller state, worker group start times, checkpointing times and more. You can use these metrics to monitor Ray Train runs.
The Ray dashboard displays these metrics in the Ray Train Grafana Dashboard. See :ref:`Ray Dashboard documentation<observability-getting-started>` for more information.
The Ray Train dashboard also displays a subset of Ray Core metrics that are useful for monitoring training but are not listed in the table below.
For more information about these metrics, see the :ref:`System Metrics documentation<system-metrics>`.
The following table lists the Prometheus metrics emitted by Ray Train:
.. list-table:: Train Metrics
:header-rows: 1
* - Prometheus Metric
- Labels
- Description
* - `ray_train_controller_state`
- `ray_train_run_name`, `ray_train_run_id`, `ray_train_controller_state`
- Current state of the Ray Train controller.
* - `ray_train_worker_group_start_total_time_s`
- `ray_train_run_name`, `ray_train_run_id`
- Total time taken to start the worker group.
* - `ray_train_worker_group_shutdown_total_time_s`
- `ray_train_run_name`, `ray_train_run_id`
- Total time taken to shut down the worker group.
* - `ray_train_report_total_blocked_time_s`
- `ray_train_run_name`, `ray_train_run_id`, `ray_train_worker_world_rank`, `ray_train_worker_actor_id`
- Cumulative time in seconds to report a checkpoint to storage.
@@ -0,0 +1,72 @@
.. _train-monitoring-and-logging:
Monitoring and Logging Metrics
==============================
Ray Train provides an API for attaching metrics to :ref:`checkpoints <train-checkpointing>` from the training function by calling :func:`ray.train.report(metrics, checkpoint) <ray.train.report>`.
The results will be collected from the distributed workers and passed to the Ray Train driver process for book-keeping.
The primary use cases for reporting are:
* metrics (accuracy, loss, etc.) at the end of each training epoch. See :ref:`train-dl-saving-checkpoints` for usage examples.
* validating checkpoints on a validation set with a user-defined validation function. See :ref:`train-validating-checkpoints` for usage examples.
Only the result reported by the rank 0 worker is attached to the checkpoint.
However, in order to ensure consistency, ``train.report()`` acts as a barrier and must be called on each worker.
To aggregate results from multiple workers, see :ref:`train-aggregating-results`.
.. _train-aggregating-results:
How to obtain and aggregate results from different workers?
-----------------------------------------------------------
In real applications, you may want to calculate optimization metrics besides accuracy and loss: recall, precision, Fbeta, etc.
You may also want to collect metrics from multiple workers. While Ray Train currently only reports metrics from the rank 0
worker, you can use third-party libraries or distributed primitives of your machine learning framework to report
metrics from multiple workers.
.. tab-set::
.. tab-item:: Native PyTorch
Ray Train natively supports `TorchMetrics <https://torchmetrics.readthedocs.io/en/latest/>`_, which provides a collection of machine learning metrics for distributed, scalable PyTorch models.
Here is an example of reporting both the aggregated R2 score and mean train and validation loss from all workers.
.. literalinclude:: ../doc_code/metric_logging.py
:language: python
:start-after: __torchmetrics_start__
:end-before: __torchmetrics_end__
.. _train-metric-only-reporting-deprecation:
(Deprecated) Reporting free-floating metrics
--------------------------------------------
Reporting metrics with ``ray.train.report(metrics, checkpoint=None)`` from every worker writes the metrics to a Ray Tune log file (``progress.csv``, ``result.json``)
and is accessible via the ``Result.metrics_dataframe`` on the :class:`~ray.train.Result` returned by ``trainer.fit()``.
As of Ray 2.43, this behavior is deprecated and will not be supported in Ray Train V2,
which is an overhaul of Ray Train's implementation and select APIs.
Ray Train V2 only keeps a slim set of experiment tracking features that are necessary for fault tolerance, so it does not support reporting free-floating metrics that are not attached to checkpoints.
The recommendation for metric tracking is to report metrics directly from the workers to experiment tracking tools such as MLFlow and WandB.
See :ref:`train-experiment-tracking-native` for examples.
In Ray Train V2, reporting only metrics from all workers is a no-op. However, it is still possible to access the results reported by all workers to implement custom metric-handling logic.
.. literalinclude:: ../doc_code/metric_logging.py
:language: python
:start-after: __report_callback_start__
:end-before: __report_callback_end__
To use Ray Tune :class:`Callbacks <ray.tune.Callback>` that depend on free-floating metrics reported by workers, :ref:`run Ray Train as a single Ray Tune trial. <train-with-tune-callbacks>`
See the following resources for more information:
* `Train V2 REP <https://github.com/ray-project/enhancements/blob/main/reps/2024-10-18-train-tune-api-revamp/2024-10-18-train-tune-api-revamp.md>`_: Technical details about the API changes in Train V2
* `Train V2 Migration Guide <https://github.com/ray-project/ray/issues/49454>`_: Full migration guide for Train V2
@@ -0,0 +1,594 @@
.. _persistent-storage-guide:
.. _train-log-dir:
Configuring Persistent Storage
==============================
A Ray Train run produces :ref:`checkpoints <train-checkpointing>` that can be saved to a persistent storage location.
.. figure:: ../images/persistent_storage_checkpoint.png
:align: center
:width: 600px
An example of multiple workers spread across multiple nodes uploading checkpoints to persistent storage.
**Ray Train expects all workers to be able to write files to the same persistent storage location.**
Therefore, Ray Train requires some form of external persistent storage such as
cloud object storage (for example, S3, GCS, or Azure Blob Storage) or a shared filesystem
(for example, AWS EFS, Google Cloud Filestore, Azure Files, or HDFS) for multi-node training.
Here are some capabilities that persistent storage enables:
- **Checkpointing and fault tolerance**: Saving checkpoints to a persistent storage location
allows you to resume training from the last checkpoint in case of a node failure.
See :ref:`train-checkpointing` for a detailed guide on how to set up checkpointing.
- **Post-experiment analysis**: A consolidated location storing data such as the best checkpoints and
hyperparameter configs after the Ray cluster has already been terminated.
- **Bridge training/fine-tuning with downstream serving and batch inference tasks**: You can easily access the models
and artifacts to share them with others or use them in downstream tasks.
Cloud object storage
--------------------
The Ray team recommends using cloud object storage such as S3, GCS, or Azure Blob Storage to persist Ray Train checkpoint files.
Use cloud object storage by specifying a storage container URI as the :class:`RunConfig(storage_path) <ray.train.RunConfig>`:
.. tab-set::
.. tab-item:: AWS S3
Specify a URI with the ``s3://`` scheme. Ray Train uses pyarrow's default
:class:`S3FileSystem <pyarrow.fs.S3FileSystem>` for upload and download.
.. testcode::
:skipif: True
from ray import train
from ray.train.torch import TorchTrainer
trainer = TorchTrainer(
...,
run_config=train.RunConfig(
storage_path="s3://bucket-name/sub-path/",
name="experiment_name",
)
)
.. tab-item:: Google Cloud Storage
Specify a URI with the ``gs://`` scheme. Ray Train uses pyarrow's default
:class:`GcsFileSystem <pyarrow.fs.GcsFileSystem>` for upload and download.
.. testcode::
:skipif: True
from ray import train
from ray.train.torch import TorchTrainer
trainer = TorchTrainer(
...,
run_config=train.RunConfig(
storage_path="gs://bucket-name/sub-path/",
name="experiment_name",
)
)
.. tab-item:: Azure Blob Storage
Ray Train uses ``pyarrow.fs`` for storage I/O, so wrap
``adlfs.AzureBlobFileSystem`` in a ``pyarrow.fs.PyFileSystem`` and pass
it as :class:`RunConfig(storage_filesystem) <ray.train.RunConfig>`.
Use the ``abfss://`` scheme (TLS-enforced) for the URI:
.. testcode::
:skipif: True
import adlfs
from pyarrow.fs import FSSpecHandler, PyFileSystem
from ray import train
from ray.train.torch import TorchTrainer
azure_fs = PyFileSystem(
FSSpecHandler(adlfs.AzureBlobFileSystem(account_name="account-name"))
)
trainer = TorchTrainer(
...,
run_config=train.RunConfig(
storage_filesystem=azure_fs,
storage_path="abfss://container@account.dfs.core.windows.net/sub-path/",
name="experiment_name",
)
)
See :ref:`custom-storage-filesystem` for more on ``storage_filesystem``.
Ensure that all nodes in the Ray cluster have access to the storage container, so outputs from workers can be uploaded to a shared location.
In the AWS S3 example above, all files are uploaded to shared storage at ``s3://bucket-name/sub-path/experiment_name`` for further processing.
Shared filesystem
-----------------
You can use shared filesystems such as AWS EFS, Google Cloud Filestore, Azure Files, HDFS, or NFS.
Either mount the filesystem so that it appears at a common path on every node in the Ray cluster, or specify a fully qualified URI.
In either case, ensure that networking rules and security permissions allow access from all nodes.
Specify the shared storage location as the :class:`RunConfig(storage_path) <ray.train.RunConfig>`:
.. tab-set::
.. tab-item:: Mounted filesystem
Mount the filesystem on every node in the cluster, then point
``storage_path`` at the mount. This works for AWS EFS, Google Cloud
Filestore, Azure Files, and NFS.
.. testcode::
:skipif: True
from ray import train
from ray.train.torch import TorchTrainer
trainer = TorchTrainer(
...,
run_config=train.RunConfig(
# Example for Azure Files mounted at /mnt/azure-fileshare on every node;
# AWS EFS, Google Cloud Filestore, and NFS work the same way.
storage_path="/mnt/cluster_storage",
name="experiment_name",
)
)
.. tab-item:: HDFS
Specify a fully qualified ``hdfs://`` URI.
.. testcode::
:skipif: True
from ray import train
from ray.train.torch import TorchTrainer
trainer = TorchTrainer(
...,
run_config=train.RunConfig(
storage_path=f"hdfs://{hostname}:{port}/subpath",
name="experiment_name",
)
)
In the mounted example above, all files are saved to ``/mnt/cluster_storage/experiment_name`` for further processing.
Local storage
-------------
Using local storage for a single-node cluster
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
If you're just running an experiment on a single node (e.g., on a laptop), Ray Train will use the
local filesystem as the storage location for checkpoints and other artifacts.
Results are saved to ``~/ray_results`` in a sub-directory with a unique auto-generated name by default,
unless you customize this with ``storage_path`` and ``name`` in :class:`~ray.train.RunConfig`.
.. testcode::
:skipif: True
from ray import train
from ray.train.torch import TorchTrainer
trainer = TorchTrainer(
...,
run_config=train.RunConfig(
storage_path="/tmp/custom/storage/path",
name="experiment_name",
)
)
In this example, all experiment results can found locally at ``/tmp/custom/storage/path/experiment_name`` for further processing.
.. _multinode-local-storage-warning:
Using local storage for a multi-node cluster
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
.. warning::
When running on multiple nodes, using the local filesystem of the head node as the persistent storage location is no longer supported.
If you save checkpoints with :meth:`ray.train.report(..., checkpoint=...) <ray.train.report>`
and run on a multi-node cluster, Ray Train will raise an error if NFS or cloud storage is not setup.
This is because Ray Train expects all workers to be able to write the checkpoint to
the same persistent storage location.
If your training loop does not save checkpoints, the reported metrics will still
be aggregated to the local storage path on the head node.
See `this issue <https://github.com/ray-project/ray/issues/37177>`_ for more information.
.. _custom-storage-filesystem:
Custom storage
--------------
If the cases above don't suit your needs, Ray Train can support custom filesystems and perform custom logic.
Ray Train standardizes on the ``pyarrow.fs.FileSystem`` interface to interact with storage
(`see the API reference here <https://arrow.apache.org/docs/python/generated/pyarrow.fs.FileSystem.html>`_).
By default, passing ``storage_path=s3://bucket-name/sub-path/`` will use pyarrow's
`default S3 filesystem implementation <https://arrow.apache.org/docs/python/generated/pyarrow.fs.S3FileSystem.html>`_
to upload files. (`See the other default implementations. <https://arrow.apache.org/docs/python/api/filesystems.html#filesystem-implementations>`_)
Implement custom storage upload and download logic by providing an implementation of
``pyarrow.fs.FileSystem`` to :class:`RunConfig(storage_filesystem) <ray.train.RunConfig>`.
.. warning::
When providing a custom filesystem, the associated ``storage_path`` is expected
to be a qualified filesystem path *without the protocol prefix*.
For example, if you provide a custom S3 filesystem for ``s3://bucket-name/sub-path/``,
then the ``storage_path`` should be ``bucket-name/sub-path/`` with the ``s3://`` stripped.
See the example below for example usage.
.. testcode::
:skipif: True
import pyarrow.fs
from ray import train
from ray.train.torch import TorchTrainer
fs = pyarrow.fs.S3FileSystem(
endpoint_override="http://localhost:9000",
access_key=...,
secret_key=...
)
trainer = TorchTrainer(
...,
run_config=train.RunConfig(
storage_filesystem=fs,
storage_path="bucket-name/sub-path",
name="unique-run-id",
)
)
``fsspec`` filesystems
~~~~~~~~~~~~~~~~~~~~~~~
`fsspec <https://filesystem-spec.readthedocs.io/en/latest/>`_ offers many filesystem implementations,
such as ``s3fs``, ``gcsfs``, etc.
You can use any of these implementations by wrapping the ``fsspec`` filesystem with a ``pyarrow.fs`` utility:
.. testcode::
:skipif: True
# Make sure to install: `pip install -U s3fs`
import s3fs
import pyarrow.fs
s3_fs = s3fs.S3FileSystem(
key='miniokey...',
secret='asecretkey...',
endpoint_url='https://...'
)
custom_fs = pyarrow.fs.PyFileSystem(pyarrow.fs.FSSpecHandler(s3_fs))
run_config = RunConfig(storage_path="minio_bucket", storage_filesystem=custom_fs)
.. seealso::
See the API references to the ``pyarrow.fs`` wrapper utilities:
* https://arrow.apache.org/docs/python/generated/pyarrow.fs.PyFileSystem.html
* https://arrow.apache.org/docs/python/generated/pyarrow.fs.FSSpecHandler.html
S3-compatible storage (Backblaze B2, MinIO, etc.)
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
For S3-compatible stores like `Backblaze B2 <https://www.backblaze.com/cloud-storage>`_
or `MinIO <https://min.io/>`_, follow the
:ref:`custom-filesystem examples above <custom-storage-filesystem>`, or pass
the endpoint as a query parameter in the ``storage_path`` URI:
.. testcode::
:skipif: True
from ray import train
from ray.train.torch import TorchTrainer
trainer = TorchTrainer(
...,
run_config=train.RunConfig(
# Backblaze B2 (substitute your bucket's region):
storage_path="s3://bucket-name/sub-path?endpoint_override=https://s3.us-west-001.backblazeb2.com",
# MinIO running locally:
# storage_path="s3://bucket-name/sub-path?endpoint_override=http://localhost:9000",
name="unique-run-id",
)
)
Alternatively, configure the endpoint and credentials through the environment
variables Arrow reads (see
`Arrow's S3 environment variables <https://arrow.apache.org/docs/cpp/env_vars.html>`_)
and use a plain ``storage_path="s3://bucket/path"``. For Backblaze B2, set
``AWS_ENDPOINT_URL_S3`` to your bucket's endpoint, and ``AWS_ACCESS_KEY_ID`` /
``AWS_SECRET_ACCESS_KEY`` to your B2 application key ID and key.
See `this end-to-end notebook <https://github.com/backblaze-b2-samples/notebooks/tree/main/ray-train-tune-checkpoints>`_ for a worked Backblaze B2 example.
Overview of Ray Train outputs
-----------------------------
So far, we covered how to configure the storage location for Ray Train outputs.
Let's walk through a concrete example to see what exactly these outputs are,
and how they're structured in storage.
.. seealso::
This example includes checkpointing, which is covered in detail in :ref:`train-checkpointing`.
.. testcode::
:skipif: True
import os
import tempfile
import ray.train
from ray.train import Checkpoint
from ray.train.torch import TorchTrainer
def train_fn(config):
for i in range(10):
# Training logic here
metrics = {"loss": ...}
with tempfile.TemporaryDirectory() as temp_checkpoint_dir:
torch.save(..., os.path.join(temp_checkpoint_dir, "checkpoint.pt"))
train.report(
metrics,
checkpoint=Checkpoint.from_directory(temp_checkpoint_dir)
)
trainer = TorchTrainer(
train_fn,
scaling_config=ray.train.ScalingConfig(num_workers=2),
run_config=ray.train.RunConfig(
storage_path="s3://bucket-name/sub-path/",
name="unique-run-id",
)
)
result: train.Result = trainer.fit()
last_checkpoint: Checkpoint = result.checkpoint
Here's a rundown of all files that will be persisted to storage:
.. code-block:: text
{RunConfig.storage_path} (ex: "s3://bucket-name/sub-path/")
└── {RunConfig.name} (ex: "unique-run-id") <- Train run output directory
├── *_snapshot.json <- Train run metadata files (DeveloperAPI)
├── checkpoint_epoch=0/ <- Checkpoints
├── checkpoint_epoch=1/
└── ...
The :class:`~ray.train.Result` and :class:`~ray.train.Checkpoint` objects returned by
``trainer.fit`` are the easiest way to access the data in these files:
.. testcode::
:skipif: True
result.filesystem, result.path
# S3FileSystem, "bucket-name/sub-path/unique-run-id"
result.checkpoint.filesystem, result.checkpoint.path
# S3FileSystem, "bucket-name/sub-path/unique-run-id/checkpoint_epoch=0"
See :ref:`train-inspect-results` for a full guide on interacting with training :class:`Results <ray.train.Result>`.
.. _train-storage-advanced:
Advanced configuration
----------------------
.. _train-working-directory:
Keep the original current working directory
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
Ray Train changes the current working directory of each worker to the same path.
By default, this path is a sub-directory of the Ray session directory (e.g., ``/tmp/ray/session_latest``),
which is also where other Ray logs and temporary files are dumped.
The location of the Ray session directory :ref:`can be customized <temp-dir-log-files>`.
To disable the default behavior of Ray Train changing the current working directory,
set the ``RAY_CHDIR_TO_TRIAL_DIR=0`` environment variable.
This is useful if you want your training workers to access relative paths from the
directory you launched the training script from.
.. tip::
When running in a distributed cluster, you will need to make sure that all workers
have a mirrored working directory to access the same relative paths.
One way to achieve this is setting the
:ref:`working directory in the Ray runtime environment <workflow-local-files>`.
.. testcode::
import os
import ray
import ray.train
from ray.train.torch import TorchTrainer
os.environ["RAY_CHDIR_TO_TRIAL_DIR"] = "0"
# Write some file in the current working directory
with open("./data.txt", "w") as f:
f.write("some data")
# Set the working directory in the Ray runtime environment
ray.init(runtime_env={"working_dir": "."})
def train_fn_per_worker(config):
# Check that each worker can access the working directory
# NOTE: The working directory is copied to each worker and is read only.
assert os.path.exists("./data.txt"), os.getcwd()
trainer = TorchTrainer(
train_fn_per_worker,
scaling_config=ray.train.ScalingConfig(num_workers=2),
run_config=ray.train.RunConfig(
# storage_path=...,
),
)
trainer.fit()
Deprecated
----------
The following sections describe behavior that is deprecated as of Ray 2.43 and will not be supported in Ray Train V2,
which is an overhaul of Ray Train's implementation and select APIs.
See the following resources for more information:
* `Train V2 REP <https://github.com/ray-project/enhancements/blob/main/reps/2024-10-18-train-tune-api-revamp/2024-10-18-train-tune-api-revamp.md>`_: Technical details about the API change
* `Train V2 Migration Guide <https://github.com/ray-project/ray/issues/49454>`_: Full migration guide for Train V2
(Deprecated) Persisting training artifacts
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
.. note::
This feature of persisting training worker artifacts is deprecated as of Ray 2.43.
The feature relied on Ray Tune's local working directory abstraction,
where the local files of each worker would be copied to storage.
Ray Train V2 decouples the two libraries, so this API, which already provided limited value, has been deprecated.
In the example above, we saved some artifacts within the training loop to the worker's
*current working directory*.
If you were training a stable diffusion model, you could save
some sample generated images every so often as a training artifact.
By default, Ray Train changes the current working directory of each worker to be inside the run's
:ref:`local staging directory <train-local-staging-dir>`.
This way, all distributed training workers share the same absolute path as the working directory.
See :ref:`below <train-working-directory>` for how to disable this default behavior,
which is useful if you want your training workers to keep their original working directories.
If :class:`RunConfig(SyncConfig(sync_artifacts=True)) <ray.train.SyncConfig>`, then
all artifacts saved in this directory will be persisted to storage.
The frequency of artifact syncing can be configured via :class:`SyncConfig <ray.train.SyncConfig>`.
Note that this behavior is off by default.
Here's an example of what the Train run output directory looks like, with the worker artifacts:
.. code-block:: text
s3://bucket-name/sub-path (RunConfig.storage_path)
└── experiment_name (RunConfig.name) <- The "experiment directory"
├── experiment_state-*.json
├── basic-variant-state-*.json
├── trainer.pkl
├── tuner.pkl
└── TorchTrainer_46367_00000_0_... <- The "trial directory"
├── events.out.tfevents... <- Tensorboard logs of reported metrics
├── result.json <- JSON log file of reported metrics
├── checkpoint_000000/ <- Checkpoints
├── checkpoint_000001/
├── ...
├── artifact-rank=0-iter=0.txt <- Worker artifacts
├── artifact-rank=1-iter=0.txt
└── ...
.. warning::
Artifacts saved by *every worker* will be synced to storage. If you have multiple workers
co-located on the same node, make sure that workers don't delete files within their
shared working directory.
A best practice is to only write artifacts from a single worker unless you
really need artifacts from multiple.
.. testcode::
:skipif: True
from ray import train
if train.get_context().get_world_rank() == 0:
# Only the global rank 0 worker saves artifacts.
...
if train.get_context().get_local_rank() == 0:
# Every local rank 0 worker saves artifacts.
...
.. _train-local-staging-dir:
(Deprecated) Setting the local staging directory
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
.. note::
This section describes behavior depending on Ray Tune implementation details that no longer applies to Ray Train V2.
.. warning::
Prior to 2.10, the ``RAY_AIR_LOCAL_CACHE_DIR`` environment variable and ``RunConfig(local_dir)``
were ways to configure the local staging directory to be outside of the home directory (``~/ray_results``).
**These configurations are no longer used to configure the local staging directory.
Please instead use** ``RunConfig(storage_path)`` **to configure where your
run's outputs go.**
Apart from files such as checkpoints written directly to the ``storage_path``,
Ray Train also writes some logfiles and metadata files to an intermediate
*local staging directory* before they get persisted (copied/uploaded) to the ``storage_path``.
The current working directory of each worker is set within this local staging directory.
By default, the local staging directory is a sub-directory of the Ray session
directory (e.g., ``/tmp/ray/session_latest``), which is also where other temporary Ray files are dumped.
Customize the location of the staging directory by :ref:`setting the location of the
temporary Ray session directory <temp-dir-log-files>`.
Here's an example of what the local staging directory looks like:
.. code-block:: text
/tmp/ray/session_latest/artifacts/<ray-train-job-timestamp>/
└── experiment_name
├── driver_artifacts <- These are all uploaded to storage periodically
│ ├── Experiment state snapshot files needed for resuming training
│ └── Metrics logfiles
└── working_dirs <- These are uploaded to storage if `SyncConfig(sync_artifacts=True)`
└── Current working directory of training workers, which contains worker artifacts
.. warning::
You should not need to look into the local staging directory.
The ``storage_path`` should be the only path that you need to interact with.
The structure of the local staging directory is subject to change
in future versions of Ray Train -- do not rely on these local staging files in your application.
@@ -0,0 +1,50 @@
.. _train-reproducibility:
Reproducibility
---------------
.. tab-set::
.. tab-item:: PyTorch
To limit sources of nondeterministic behavior, add
:func:`ray.train.torch.enable_reproducibility` to the top of your training
function.
.. code-block:: diff
def train_func():
+ train.torch.enable_reproducibility()
model = NeuralNetwork()
model = train.torch.prepare_model(model)
...
.. warning:: :func:`ray.train.torch.enable_reproducibility` can't guarantee
completely reproducible results across executions. To learn more, read
the `PyTorch notes on randomness <https://pytorch.org/docs/stable/notes/randomness.html>`_.
..
import ray
from ray import tune
def training_func(config):
dataloader = ray.train.get_dataset()\
.get_shard(torch.rank())\
.iter_torch_batches(batch_size=config["batch_size"])
for i in config["epochs"]:
ray.train.report(...) # use same intermediate reporting API
# Declare the specification for training.
trainer = Trainer(backend="torch", num_workers=12, use_gpu=True)
dataset = ray.dataset.window()
# Convert this to a trainable.
trainable = trainer.to_tune_trainable(training_func, dataset=dataset)
tuner = tune.Tuner(trainable,
param_space={"lr": tune.uniform(), "batch_size": tune.randint(1, 2, 3)},
tune_config=tune.TuneConfig(num_samples=12))
results = tuner.fit()
+155
View File
@@ -0,0 +1,155 @@
.. _train-inspect-results:
Inspecting Training Results
===========================
The return value of ``trainer.fit()`` is a :class:`~ray.train.Result` object.
The :class:`~ray.train.Result` object contains, among other information:
- The last reported checkpoint (to load the model) and its attached metrics
- Error messages, if any errors occurred
- Any data returned by the training function (on worker 0 only)
Viewing metrics
---------------
You can retrieve reported metrics that were attached to a checkpoint from the :class:`~ray.train.Result` object.
Common metrics include the training or validation loss, or prediction accuracies.
The metrics retrieved from the :class:`~ray.train.Result` object
correspond to those you passed to :func:`train.report <ray.train.report>`
as an argument :ref:`in your training function <train-monitoring-and-logging>`.
.. note::
Persisting free-floating metrics reported via ``ray.train.report(metrics, checkpoint=None)`` is deprecated.
This also means that retrieving these metrics from the :class:`~ray.train.Result` object is deprecated.
Only metrics attached to checkpoints are persisted. See :ref:`train-metric-only-reporting-deprecation` for more details.
Last reported metrics
~~~~~~~~~~~~~~~~~~~~~
Use :attr:`Result.metrics <ray.train.Result>` to retrieve the
metrics attached to the last reported checkpoint.
.. literalinclude:: ../doc_code/key_concepts.py
:language: python
:start-after: __result_metrics_start__
:end-before: __result_metrics_end__
Dataframe of all reported metrics
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
Use :attr:`Result.metrics_dataframe <ray.train.Result>` to retrieve
a pandas DataFrame of all metrics reported alongside checkpoints.
.. literalinclude:: ../doc_code/key_concepts.py
:language: python
:start-after: __result_dataframe_start__
:end-before: __result_dataframe_end__
Returned data from train function
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
Use :attr:`Result.return_value <ray.train.Result>` to retrieve any data
returned from worker 0's train function.
.. literalinclude:: ../doc_code/key_concepts.py
:language: python
:start-after: __result_return_value_start__
:end-before: __result_return_value_end__
Retrieving checkpoints
----------------------
You can retrieve checkpoints reported to Ray Train from the :class:`~ray.train.Result`
object.
:ref:`Checkpoints <train-checkpointing>` contain all the information that is needed
to restore the training state. This usually includes the trained model.
You can use checkpoints for common downstream tasks such as
:doc:`offline batch inference with Ray Data </data/data>` or
:doc:`online model serving with Ray Serve </serve/index>`.
The checkpoints retrieved from the :class:`~ray.train.Result` object
correspond to those you passed to :func:`train.report <ray.train.report>`
as an argument :ref:`in your training function <train-monitoring-and-logging>`.
Last saved checkpoint
~~~~~~~~~~~~~~~~~~~~~
Use :attr:`Result.checkpoint <ray.train.Result>` to retrieve the
last checkpoint.
.. literalinclude:: ../doc_code/key_concepts.py
:language: python
:start-after: __result_checkpoint_start__
:end-before: __result_checkpoint_end__
Other checkpoints
~~~~~~~~~~~~~~~~~
Sometimes you want to access an earlier checkpoint. For instance, if your loss increased
after more training due to overfitting, you may want to retrieve the checkpoint with
the lowest loss.
You can retrieve a list of all available checkpoints and their metrics with
:attr:`Result.best_checkpoints <ray.train.Result>`
.. literalinclude:: ../doc_code/key_concepts.py
:language: python
:start-after: __result_best_checkpoint_start__
:end-before: __result_best_checkpoint_end__
.. seealso::
See :ref:`train-checkpointing` for more information on checkpointing.
Accessing storage location
---------------------------
If you need to retrieve the results later, you can get the storage location
of the training run with :attr:`Result.path <ray.train.Result>`.
This path will correspond to the :ref:`storage_path <train-log-dir>` you configured
in the :class:`~ray.train.RunConfig`. It will be a
(nested) subdirectory within that path, usually
of the form `TrainerName_date-string/TrainerName_id_00000_0_...`.
The result also contains a :class:`pyarrow.fs.FileSystem` that can be used to
access the storage location, which is useful if the path is on cloud storage.
.. literalinclude:: ../doc_code/key_concepts.py
:language: python
:start-after: __result_path_start__
:end-before: __result_path_end__
You can restore a result with :meth:`Result.from_path <ray.train.Result.from_path>`:
.. literalinclude:: ../doc_code/key_concepts.py
:language: python
:start-after: __result_restore_start__
:end-before: __result_restore_end__
Catching Errors
---------------
If an error occurred during training,
:attr:`Result.error <ray.train.Result>` will be set and contain the exception
that was raised.
.. literalinclude:: ../doc_code/key_concepts.py
:language: python
:start-after: __result_error_start__
:end-before: __result_error_end__
Finding results on persistent storage
-------------------------------------
All training results including reported metrics and checkpoints
are stored on the configured :ref:`persistent storage <train-log-dir>`.
See :ref:`the persistent storage guide <train-log-dir>` to configure this location
for your training run.
@@ -0,0 +1,384 @@
.. _train_scaling_config:
Configuring Scale and Accelerators
==================================
Increasing the scale of a Ray Train training run is simple and can be done in a few lines of code.
The main interface for this is the :class:`~ray.train.ScalingConfig`,
which configures the number of workers and the resources they should use.
In this guide, a *worker* refers to a Ray Train distributed training worker,
which is a :ref:`Ray Actor <actor-key-concept>` that runs your training function.
Increasing the number of workers
--------------------------------
The main interface to control parallelism in your training code is to set the
number of workers. This can be done by passing the ``num_workers`` attribute to
the :class:`~ray.train.ScalingConfig`:
.. testcode::
from ray.train import ScalingConfig
scaling_config = ScalingConfig(
num_workers=8
)
Using accelerators
------------------
.. tab-set::
.. tab-item:: GPU
:sync: GPU
To use GPUs, pass ``use_gpu=True`` to the :class:`~ray.train.ScalingConfig`.
This requests one GPU per training worker. In the following example, training
runs on 8 GPUs (8 workers, each using one GPU).
.. testcode::
from ray.train import ScalingConfig
scaling_config = ScalingConfig(
num_workers=8,
use_gpu=True
)
.. tab-item:: TPU
:sync: TPU
To use TPUs, pass ``use_tpu=True`` to the :class:`~ray.train.ScalingConfig`.
You also need to specify ``topology`` and ``accelerator_type``.
Each ``num_workers`` maps to one TPU VM host. The total number of
workers must be a multiple of the number of hosts in a single slice.
For example, a ``v6e`` TPU slice with a ``4x4`` topology has 4 hosts,
so valid values include ``num_workers=4`` (one slice) or
``num_workers=8`` (two slices).
For details on how TPU topologies map to the number of hosts, see
`Plan TPUs in GKE <https://cloud.google.com/kubernetes-engine/docs/concepts/plan-tpus>`_.
.. testcode::
:skipif: True
from ray.train import ScalingConfig
# Single slice: 4 v6e VMs in a 4x4 topology
scaling_config = ScalingConfig(
num_workers=4,
use_tpu=True,
topology="4x4",
accelerator_type="TPU-V6E",
)
# Multi-slice: 2 v6e slices, 8 VMs total
scaling_config = ScalingConfig(
num_workers=8,
use_tpu=True,
topology="4x4",
accelerator_type="TPU-V6E",
)
Using accelerators in the training function
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
.. tab-set::
.. tab-item:: GPU
:sync: GPU
When ``use_gpu=True`` is set, Ray Train automatically sets up environment variables
in your training function so that the GPUs can be detected and used
(such as ``CUDA_VISIBLE_DEVICES``).
You can get the associated devices with :meth:`ray.train.torch.get_device`.
.. testcode::
import torch
from ray.train import ScalingConfig
from ray.train.torch import TorchTrainer, get_device
def train_func():
assert torch.cuda.is_available()
device = get_device()
assert device == torch.device("cuda:0")
trainer = TorchTrainer(
train_func,
scaling_config=ScalingConfig(
num_workers=1,
use_gpu=True
)
)
trainer.fit()
.. tab-item:: TPU
:sync: TPU
When ``use_tpu=True`` is set, Ray Train configures the distributed
environment for TPU execution on each worker. The specific initialization
depends on the trainer you use (such as :class:`~ray.train.v2.jax.JaxTrainer`).
The following example shows a basic TPU training setup with
:class:`~ray.train.v2.jax.JaxTrainer`:
.. testcode::
:skipif: True
import ray.train
from ray.train import ScalingConfig
from ray.train.v2.jax import JaxTrainer
def train_func():
import jax
devices = jax.devices()
ray.train.report({"num_devices": len(devices)})
trainer = JaxTrainer(
train_func,
scaling_config=ScalingConfig(
num_workers=4,
use_tpu=True,
topology="4x4",
accelerator_type="TPU-V6E",
)
)
trainer.fit()
Assigning multiple accelerators to a worker
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
.. tab-set::
.. tab-item:: GPU
:sync: GPU
Sometimes you might want to allocate multiple GPUs for a worker. For example,
you can specify ``resources_per_worker={"GPU": 2}`` in the ``ScalingConfig`` if you want to
assign 2 GPUs for each worker.
You can get a list of associated devices with :meth:`ray.train.torch.get_devices`.
.. testcode::
import torch
from ray.train import ScalingConfig
from ray.train.torch import TorchTrainer, get_device, get_devices
def train_func():
assert torch.cuda.is_available()
device = get_device()
devices = get_devices()
assert device == torch.device("cuda:0")
assert devices == [torch.device("cuda:0"), torch.device("cuda:1")]
trainer = TorchTrainer(
train_func,
scaling_config=ScalingConfig(
num_workers=1,
use_gpu=True,
resources_per_worker={"GPU": 2}
)
)
trainer.fit()
.. tab-item:: TPU
:sync: TPU
Each TPU VM host has multiple TPU chips. By default, when ``topology``
and ``accelerator_type`` are specified, Ray Train auto-detects the
correct ``resources_per_worker`` for the given TPU slice configuration.
To override the default, specify the number of chips explicitly in
``resources_per_worker``. Supported chip counts are 1, 2, 4, and 8.
For example, to use only 2 of the 4 chips on a ``ct6e-standard-4t``
host:
.. testcode::
:skipif: True
from ray.train import ScalingConfig
scaling_config = ScalingConfig(
num_workers=4,
use_tpu=True,
topology="4x4",
accelerator_type="TPU-V6E",
resources_per_worker={"TPU": 2},
)
Setting the accelerator type
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
Ray Train allows you to specify the accelerator type for each worker.
This is useful if you want to use a specific accelerator type for model training.
In a heterogeneous Ray cluster, this means that your training workers are forced to run on the specified accelerator type,
rather than on any arbitrary accelerator node. You can get a list of supported ``accelerator_type`` from
:ref:`the available accelerator types <accelerator_types>`.
.. tab-set::
.. tab-item:: GPU
:sync: GPU
The following example specifies ``accelerator_type="A100"`` to assign each worker
a NVIDIA A100 GPU.
.. tip::
Ensure that your cluster has instances with the specified accelerator type
or is able to autoscale to fulfill the request.
.. testcode::
ScalingConfig(
num_workers=1,
use_gpu=True,
accelerator_type="A100"
)
.. tab-item:: TPU
:sync: TPU
For TPUs, ``accelerator_type`` specifies the TPU generation.
See :ref:`the available accelerator types <accelerator_types>` for
the full list of supported values.
.. testcode::
:skipif: True
ScalingConfig(
num_workers=4,
use_tpu=True,
topology="2x2x4",
accelerator_type="TPU-V4",
)
(PyTorch) Setting the communication backend
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
PyTorch Distributed supports multiple `backends <https://pytorch.org/docs/stable/distributed.html#backends>`__
for communicating tensors across workers. By default Ray Train uses NCCL when ``use_gpu=True`` and Gloo otherwise.
If you explicitly want to override this setting, you can configure a :class:`~ray.train.torch.TorchConfig`
and pass it into the :class:`~ray.train.torch.TorchTrainer`.
.. testcode::
:hide:
num_training_workers = 1
.. testcode::
from ray.train.torch import TorchConfig, TorchTrainer
trainer = TorchTrainer(
train_func,
scaling_config=ScalingConfig(
num_workers=num_training_workers,
use_gpu=True, # Defaults to NCCL
),
torch_config=TorchConfig(backend="gloo"),
)
(NCCL) Setting the communication network interface
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
When using NCCL for distributed training, you can configure the network interface cards
that are used for communicating between GPUs by setting the
`NCCL_SOCKET_IFNAME <https://docs.nvidia.com/deeplearning/nccl/user-guide/docs/env.html#nccl-socket-ifname>`__
environment variable.
To ensure that the environment variable is set for all training workers, you can pass it
in a :ref:`Ray runtime environment <runtime-environments>`:
.. testcode::
:skipif: True
import ray
runtime_env = {"env_vars": {"NCCL_SOCKET_IFNAME": "ens5"}}
ray.init(runtime_env=runtime_env)
trainer = TorchTrainer(...)
Setting the resources per worker
--------------------------------
If you want to allocate more than one CPU or accelerator per training worker, or if you
defined :ref:`custom cluster resources <cluster-resources>`, set
the ``resources_per_worker`` attribute:
.. testcode::
from ray.train import ScalingConfig
scaling_config = ScalingConfig(
num_workers=8,
resources_per_worker={
"CPU": 4,
"GPU": 2,
},
use_gpu=True,
)
.. note::
If you specify GPUs in ``resources_per_worker``, you also need to set
``use_gpu=True``.
You can also instruct Ray Train to use fractional GPUs. In that case, multiple workers
are assigned the same CUDA device.
.. testcode::
from ray.train import ScalingConfig
scaling_config = ScalingConfig(
num_workers=8,
resources_per_worker={
"CPU": 4,
"GPU": 0.5,
},
use_gpu=True,
)
(Deprecated) Trainer resources
------------------------------
.. important::
This API is deprecated. See `this migration guide <https://github.com/ray-project/ray/issues/49454>`_ for more details.
So far we've configured resources for each training worker. Technically, each
training worker is a :ref:`Ray Actor <actor-guide>`. Ray Train also schedules
an actor for the trainer object when you call ``trainer.fit()``.
This object often only manages lightweight communication between the training workers.
By default, a trainer uses 1 CPU. If you have a cluster with 8 CPUs and want
to start 4 training workers at 2 CPUs each, this won't work, as the total number
of required CPUs is 9 (4 * 2 + 1). In that case, you can specify the trainer
resources to use 0 CPUs:
.. testcode::
from ray.train import ScalingConfig
scaling_config = ScalingConfig(
num_workers=4,
resources_per_worker={
"CPU": 2,
},
trainer_resources={
"CPU": 0,
}
)