Ray Train API

This page covers framework specific integrations with Ray Train and Ray Train Developer APIs.

For core Ray AIR APIs, take a look at the AIR package reference.

Ray Train Base Classes (Developer APIs)

Trainer Base Classes

BaseTrainer(*args, **kwargs)	Defines interface for distributed training on Ray.
DataParallelTrainer(*args, **kwargs)	A Trainer for data parallel training.
DataConfig([datasets_to_split, ...])	Class responsible for configuring Train dataset preprocessing.
GBDTTrainer(*args, **kwargs)	Abstract class for scaling gradient-boosting decision tree (GBDT) frameworks.

BaseTrainer API

fit()	Runs training.
setup()	Called during fit() to perform initial setup on the Trainer.
preprocess_datasets()	Called during fit() to preprocess dataset attributes with preprocessor.
training_loop()	Loop called by fit() to run training and report results to Tune.
as_trainable()	Converts self to a tune.Trainable class.

Train Backend Base Classes

Backend(*args, **kwargs)	Singleton for distributed communication backend.
BackendConfig()	Parent class for configurations of training backend.

Ray Train Config

DataConfig([datasets_to_split, ...])

Class responsible for configuring Train dataset preprocessing.

Ray Train Loop

TrainContext()	Context for Ray training executions.
get_context()	Get or create a singleton training context.
get_dataset_shard([dataset_name])	Returns the ray.data.DataIterator shard for this worker.
report(metrics, *[, checkpoint])	Report metrics and optionally save a checkpoint.

Ray Train Checkpoints

Checkpoint([local_path, data_dict, uri])

Ray AIR Checkpoint.

Ray Train Context

get_context()	Get or create a singleton training context.
TrainContext()	Context for Ray training executions.

Ray Train Integrations

PyTorch

TorchTrainer(*args, **kwargs)	A Trainer for data parallel PyTorch training.
TorchConfig([backend, init_method, timeout_s])	Configuration for torch process group setup.
TorchCheckpoint([local_path, data_dict, uri])	A Checkpoint with Torch-specific functionality.

PyTorch Training Loop Utilities

prepare_model(model[, move_to_device, ...])	Prepares the model for distributed execution.
prepare_optimizer(optimizer)	Wraps optimizer to support automatic mixed precision.
prepare_data_loader(data_loader[, ...])	Prepares DataLoader for distributed execution.
get_device()	Gets the correct torch device configured for this process.
accelerate([amp])	Enables training optimizations.
backward(tensor)	Computes the gradient of the specified tensor w.r.t.
enable_reproducibility([seed])	Limits sources of nondeterministic behavior.

PyTorch Lightning

LightningTrainer(*args, **kwargs)	A Trainer for data parallel PyTorch Lightning training.
LightningConfigBuilder()	Configuration Class to pass into LightningTrainer.
LightningCheckpoint(*args, **kwargs)	A Checkpoint with Lightning-specific functionality.
LightningPredictor(model[, preprocessor, ...])	A predictor for PyTorch Lightning modules.

Tensorflow/Keras

TensorflowTrainer(*args, **kwargs)	A Trainer for data parallel Tensorflow training.
TensorflowConfig()	PublicAPI (beta): This API is in beta and may change before becoming stable.
TensorflowCheckpoint(*args, **kwargs)	A Checkpoint with TensorFlow-specific functionality.

Tensorflow/Keras Training Loop Utilities

prepare_dataset_shard(tf_dataset_shard)

A utility function that overrides default config for Tensorflow Dataset.

ReportCheckpointCallback([checkpoint_on, ...])

Keras callback for Ray AIR reporting and checkpointing.

Horovod

HorovodTrainer(*args, **kwargs)	A Trainer for data parallel Horovod training.
HorovodConfig([nics, verbose, key, ...])	Configurations for Horovod setup.

XGBoost

XGBoostTrainer(*args, **kwargs)	A Trainer for data parallel XGBoost training.
XGBoostCheckpoint([local_path, data_dict, uri])	A Checkpoint with XGBoost-specific functionality.

LightGBM

LightGBMTrainer(*args, **kwargs)	A Trainer for data parallel LightGBM training.
LightGBMCheckpoint([local_path, data_dict, uri])	A Checkpoint with LightGBM-specific functionality.

Hugging Face

Transformers

TransformersTrainer(*args, **kwargs)	A Trainer for data parallel HuggingFace Transformers on PyTorch training.
TransformersCheckpoint([local_path, ...])	A Checkpoint with HuggingFace-specific functionality.

Accelerate

AccelerateTrainer(*args, **kwargs)

A Trainer for data parallel HuggingFace Accelerate training with PyTorch.

Scikit-Learn

SklearnTrainer(*args, **kwargs)	A Trainer for scikit-learn estimator training.
SklearnCheckpoint([local_path, data_dict, uri])	A Checkpoint with sklearn-specific functionality.

Mosaic

MosaicTrainer(*args, **kwargs)

A Trainer for data parallel Mosaic Composers on PyTorch training.

Ray Train Experiment Restoration

train.trainer.BaseTrainer.restore(path[, ...])

Restores a Train experiment from a previously interrupted/failed run.

Note

All trainer classes have a restore method that takes in a path pointing to the directory of the experiment to be restored. restore also exposes a subset of construtor arguments that can be re-specified. See Restoration API for Built-in Trainers below for details on restore arguments for different AIR trainer integrations.

Restoration API for Built-in Trainers

train.data_parallel_trainer.DataParallelTrainer.restore(path)

Restores a DataParallelTrainer from a previously interrupted/failed run.

train.huggingface.TransformersTrainer.restore(path)

Restores a TransformersTrainer from a previously interrupted/failed run.

Note

TorchTrainer.restore, TensorflowTrainer.restore, and HorovodTrainer.restore can take in the same parameters as their parent class’s DataParallelTrainer.restore.

Unless otherwise specified, other trainers will accept the same parameters as BaseTrainer.restore.

See also

See Restore a Ray Train Experiment for more details on when and how trainer restore should be used.