Ray Train overview#

To use Ray Train effectively, understand the following four main concepts:

  • Training function: A Python function that contains your model training logic.

  • Worker: A process that runs the training function.

  • Scaling configuration: A configuration of the number of workers and the compute resources, such as CPUs or GPUs.

  • Trainer: A Python class that ties together the training function, workers, and scaling configuration to run a distributed training job.

../_images/overview.png

Training function#

The training function is a Python function that you define, and it contains the end-to-end model training loop. When you launch a distributed training job, each worker runs this training function.

Ray Train documentation uses the following two conventions:

  • train_func is a function that you define, and it contains the training code.

  • You pass train_func to the trainer’s train_loop_per_worker parameter.

def train_func():
    """User-defined training function that runs on each distributed worker process.

    This function typically contains logic for loading the model,
    loading the dataset, training the model, saving checkpoints,
    and logging metrics.
    """
    ...

Worker#

Ray Train distributes model training compute to individual worker processes across the cluster. Each worker is a process that runs train_func. The number of workers determines the parallelism of the training job, and you set it in the ScalingConfig.

Scaling configuration#

The ScalingConfig defines the scale of the training job. Specify the following two basic parameters for worker parallelism and compute resources:

  • num_workers: The number of workers to launch for a distributed training job.

  • use_gpu: Whether each worker uses a GPU.

from ray.train import ScalingConfig

# Single worker with a CPU
scaling_config = ScalingConfig(num_workers=1, use_gpu=False)

# Single worker with a GPU
scaling_config = ScalingConfig(num_workers=1, use_gpu=True)

# Multiple workers, each with a GPU
scaling_config = ScalingConfig(num_workers=4, use_gpu=True)

Trainer#

The trainer ties the previous three concepts together to launch distributed training jobs. Ray Train provides trainer classes for different frameworks. When you call the fit() method, the trainer does the following to run the training job:

  1. Launches workers according to the scaling_config.

  2. Sets up the framework’s distributed environment on all workers.

  3. Runs train_func on all workers.

from ray.train.torch import TorchTrainer

trainer = TorchTrainer(train_func, scaling_config=scaling_config)
trainer.fit()