Ray Train API#

PyTorch Ecosystem#

TorchTrainer

A Trainer for data parallel PyTorch training.

TorchConfig

Configuration for torch process group setup.

PyTorch#

get_device

Gets the correct torch device configured for this process.

prepare_model

Prepares the model for distributed execution.

prepare_data_loader

Prepares DataLoader for distributed execution.

enable_reproducibility

Limits sources of nondeterministic behavior.

PyTorch Lightning#

prepare_trainer

Prepare the PyTorch Lightning Trainer for distributed execution.

RayLightningEnvironment

Setup Lightning DDP training environment for Ray cluster.

RayDDPStrategy

Subclass of DDPStrategy to ensure compatibility with Ray orchestration.

RayFSDPStrategy

Subclass of FSDPStrategy to ensure compatibility with Ray orchestration.

RayDeepSpeedStrategy

Subclass of DeepSpeedStrategy to ensure compatibility with Ray orchestration.

RayTrainReportCallback

A simple callback that reports checkpoints to Ray on train epoch end.

Hugging Face Transformers#

prepare_trainer

Prepare your HuggingFace Transformer Trainer for Ray Train.

RayTrainReportCallback

A simple callback to report checkpoints and metrics to Ray Tarin.

More Frameworks#

Tensorflow/Keras#

TensorflowTrainer

A Trainer for data parallel Tensorflow training.

TensorflowConfig

PublicAPI (beta): This API is in beta and may change before becoming stable.

prepare_dataset_shard

A utility function that overrides default config for Tensorflow Dataset.

ReportCheckpointCallback

Keras callback for Ray Train reporting and checkpointing.

Horovod#

HorovodTrainer

A Trainer for data parallel Horovod training.

HorovodConfig

Configurations for Horovod setup.

XGBoost#

XGBoostTrainer

A Trainer for data parallel XGBoost training.

LightGBM#

LightGBMTrainer

A Trainer for data parallel LightGBM training.

Ray Train Configuration#

CheckpointConfig

Configurable parameters for defining the checkpointing strategy.

DataConfig

Class responsible for configuring Train dataset preprocessing.

FailureConfig

Configuration related to failure handling of each training/tuning run.

RunConfig

Runtime configuration for training and tuning runs.

ScalingConfig

Configuration for scaling training.

SyncConfig

Configuration object for Train/Tune file syncing to RunConfig(storage_path).

Ray Train Utilities#

Classes

Checkpoint

A reference to data persisted as a directory in local or remote storage.

TrainContext

Context for Ray training executions.

Functions

get_checkpoint

Access the session's last checkpoint to resume from if applicable.

get_context

Get or create a singleton training context.

get_dataset_shard

Returns the ray.data.DataIterator shard for this worker.

report

Report metrics and optionally save a checkpoint.

Ray Train Output#

Result

The final result of a ML training run or a Tune trial.

Ray Train Developer APIs#

Trainer Base Classes#

BaseTrainer

Defines interface for distributed training on Ray.

DataParallelTrainer

A Trainer for data parallel training.

GBDTTrainer

Abstract class for scaling gradient-boosting decision tree (GBDT) frameworks.

Train Backend Base Classes#

Backend

Singleton for distributed communication backend.

BackendConfig

Parent class for configurations of training backend.