chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,533 @@
|
||||
.. _train-pytorch-lightning:
|
||||
|
||||
Get Started with Distributed Training using PyTorch Lightning
|
||||
=============================================================
|
||||
|
||||
This tutorial walks through the process of converting an existing PyTorch Lightning script to use Ray Train.
|
||||
|
||||
Learn how to:
|
||||
|
||||
1. Configure the Lightning Trainer so that it runs distributed with Ray and on the correct CPU or GPU device.
|
||||
2. Configure :ref:`training function <train-overview-training-function>` to report metrics and save checkpoints.
|
||||
3. Configure :ref:`scaling <train-overview-scaling-config>` and CPU or GPU resource requirements for a training job.
|
||||
4. Launch a distributed training job with a :class:`~ray.train.torch.TorchTrainer`.
|
||||
|
||||
Quickstart
|
||||
----------
|
||||
|
||||
For reference, the final code is as follows:
|
||||
|
||||
.. testcode::
|
||||
:skipif: True
|
||||
|
||||
from ray.train.torch import TorchTrainer
|
||||
from ray.train import ScalingConfig
|
||||
|
||||
def train_func():
|
||||
# Your PyTorch Lightning training code here.
|
||||
|
||||
scaling_config = ScalingConfig(num_workers=2, use_gpu=True)
|
||||
trainer = TorchTrainer(train_func, scaling_config=scaling_config)
|
||||
result = trainer.fit()
|
||||
|
||||
1. `train_func` is the Python code that executes on each distributed training worker.
|
||||
2. :class:`~ray.train.ScalingConfig` defines the number of distributed training workers and whether to use GPUs.
|
||||
3. :class:`~ray.train.torch.TorchTrainer` launches the distributed training job.
|
||||
|
||||
Compare a PyTorch Lightning training script with and without Ray Train.
|
||||
|
||||
.. tab-set::
|
||||
|
||||
.. tab-item:: PyTorch Lightning + Ray Train
|
||||
|
||||
.. code-block:: python
|
||||
:emphasize-lines: 11-12, 38, 52-57, 59, 63, 66-73
|
||||
|
||||
import os
|
||||
import tempfile
|
||||
|
||||
import torch
|
||||
from torch.utils.data import DataLoader
|
||||
from torchvision.models import resnet18
|
||||
from torchvision.datasets import FashionMNIST
|
||||
from torchvision.transforms import ToTensor, Normalize, Compose
|
||||
import lightning.pytorch as pl
|
||||
|
||||
import ray.train.lightning
|
||||
from ray.train.torch import TorchTrainer
|
||||
|
||||
# Model, Loss, Optimizer
|
||||
class ImageClassifier(pl.LightningModule):
|
||||
def __init__(self):
|
||||
super(ImageClassifier, self).__init__()
|
||||
self.model = resnet18(num_classes=10)
|
||||
self.model.conv1 = torch.nn.Conv2d(
|
||||
1, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False
|
||||
)
|
||||
self.criterion = torch.nn.CrossEntropyLoss()
|
||||
|
||||
def forward(self, x):
|
||||
return self.model(x)
|
||||
|
||||
def training_step(self, batch, batch_idx):
|
||||
x, y = batch
|
||||
outputs = self.forward(x)
|
||||
loss = self.criterion(outputs, y)
|
||||
self.log("loss", loss, on_step=True, prog_bar=True)
|
||||
return loss
|
||||
|
||||
def configure_optimizers(self):
|
||||
return torch.optim.Adam(self.model.parameters(), lr=0.001)
|
||||
|
||||
|
||||
def train_func():
|
||||
# Data
|
||||
transform = Compose([ToTensor(), Normalize((0.28604,), (0.32025,))])
|
||||
data_dir = os.path.join(tempfile.gettempdir(), "data")
|
||||
train_data = FashionMNIST(root=data_dir, train=True, download=True, transform=transform)
|
||||
train_dataloader = DataLoader(train_data, batch_size=128, shuffle=True)
|
||||
|
||||
# Training
|
||||
model = ImageClassifier()
|
||||
# [1] Configure PyTorch Lightning Trainer.
|
||||
trainer = pl.Trainer(
|
||||
max_epochs=10,
|
||||
devices="auto",
|
||||
accelerator="auto",
|
||||
strategy=ray.train.lightning.RayDDPStrategy(),
|
||||
plugins=[ray.train.lightning.RayLightningEnvironment()],
|
||||
callbacks=[ray.train.lightning.RayTrainReportCallback()],
|
||||
# [1a] Optionally, disable the default checkpointing behavior
|
||||
# in favor of the `RayTrainReportCallback` above.
|
||||
enable_checkpointing=False,
|
||||
)
|
||||
trainer = ray.train.lightning.prepare_trainer(trainer)
|
||||
trainer.fit(model, train_dataloaders=train_dataloader)
|
||||
|
||||
# [2] Configure scaling and resource requirements.
|
||||
scaling_config = ray.train.ScalingConfig(num_workers=2, use_gpu=True)
|
||||
|
||||
# [3] Launch distributed training job.
|
||||
trainer = TorchTrainer(
|
||||
train_func,
|
||||
scaling_config=scaling_config,
|
||||
# [3a] If running in a multi-node cluster, this is where you
|
||||
# should configure the run's persistent storage that is accessible
|
||||
# across all worker nodes.
|
||||
# run_config=ray.train.RunConfig(storage_path="s3://..."),
|
||||
)
|
||||
result: ray.train.Result = trainer.fit()
|
||||
|
||||
# [4] Load the trained model.
|
||||
with result.checkpoint.as_directory() as checkpoint_dir:
|
||||
model = ImageClassifier.load_from_checkpoint(
|
||||
os.path.join(
|
||||
checkpoint_dir,
|
||||
ray.train.lightning.RayTrainReportCallback.CHECKPOINT_NAME,
|
||||
),
|
||||
)
|
||||
|
||||
.. tab-item:: PyTorch Lightning
|
||||
|
||||
.. This snippet isn't tested because it doesn't use any Ray code.
|
||||
|
||||
.. testcode::
|
||||
:skipif: True
|
||||
|
||||
import torch
|
||||
from torchvision.models import resnet18
|
||||
from torchvision.datasets import FashionMNIST
|
||||
from torchvision.transforms import ToTensor, Normalize, Compose
|
||||
from torch.utils.data import DataLoader
|
||||
import lightning.pytorch as pl
|
||||
|
||||
# Model, Loss, Optimizer
|
||||
class ImageClassifier(pl.LightningModule):
|
||||
def __init__(self):
|
||||
super(ImageClassifier, self).__init__()
|
||||
self.model = resnet18(num_classes=10)
|
||||
self.model.conv1 = torch.nn.Conv2d(
|
||||
1, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False
|
||||
)
|
||||
self.criterion = torch.nn.CrossEntropyLoss()
|
||||
|
||||
def forward(self, x):
|
||||
return self.model(x)
|
||||
|
||||
def training_step(self, batch, batch_idx):
|
||||
x, y = batch
|
||||
outputs = self.forward(x)
|
||||
loss = self.criterion(outputs, y)
|
||||
self.log("loss", loss, on_step=True, prog_bar=True)
|
||||
return loss
|
||||
|
||||
def configure_optimizers(self):
|
||||
return torch.optim.Adam(self.model.parameters(), lr=0.001)
|
||||
|
||||
# Data
|
||||
transform = Compose([ToTensor(), Normalize((0.28604,), (0.32025,))])
|
||||
train_data = FashionMNIST(root='./data', train=True, download=True, transform=transform)
|
||||
train_dataloader = DataLoader(train_data, batch_size=128, shuffle=True)
|
||||
|
||||
# Training
|
||||
model = ImageClassifier()
|
||||
trainer = pl.Trainer(max_epochs=10)
|
||||
trainer.fit(model, train_dataloaders=train_dataloader)
|
||||
|
||||
|
||||
Set up a training function
|
||||
--------------------------
|
||||
|
||||
.. include:: ./common/torch-configure-train_func.rst
|
||||
|
||||
Ray Train sets up your distributed process group on each worker. You only need to
|
||||
make a few changes to your Lightning Trainer definition.
|
||||
|
||||
.. code-block:: diff
|
||||
|
||||
import lightning.pytorch as pl
|
||||
-from pl.strategies import DDPStrategy
|
||||
-from pl.plugins.environments import LightningEnvironment
|
||||
+import ray.train.lightning
|
||||
|
||||
def train_func():
|
||||
...
|
||||
model = MyLightningModule(...)
|
||||
datamodule = MyLightningDataModule(...)
|
||||
|
||||
trainer = pl.Trainer(
|
||||
- devices=[0, 1, 2, 3],
|
||||
- strategy=DDPStrategy(),
|
||||
- plugins=[LightningEnvironment()],
|
||||
+ devices="auto",
|
||||
+ accelerator="auto",
|
||||
+ strategy=ray.train.lightning.RayDDPStrategy(),
|
||||
+ plugins=[ray.train.lightning.RayLightningEnvironment()]
|
||||
)
|
||||
+ trainer = ray.train.lightning.prepare_trainer(trainer)
|
||||
|
||||
trainer.fit(model, datamodule=datamodule)
|
||||
|
||||
The following sections discuss each change.
|
||||
|
||||
Configure the distributed strategy
|
||||
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
|
||||
|
||||
Ray Train offers several sub-classed distributed strategies for Lightning.
|
||||
These strategies retain the same argument list as their base strategy classes.
|
||||
Internally, they configure the root device and the distributed
|
||||
sampler arguments.
|
||||
|
||||
- :class:`~ray.train.lightning.RayDDPStrategy`
|
||||
- :class:`~ray.train.lightning.RayFSDPStrategy`
|
||||
- :class:`~ray.train.lightning.RayDeepSpeedStrategy`
|
||||
|
||||
|
||||
.. code-block:: diff
|
||||
|
||||
import lightning.pytorch as pl
|
||||
-from pl.strategies import DDPStrategy
|
||||
+import ray.train.lightning
|
||||
|
||||
def train_func():
|
||||
...
|
||||
trainer = pl.Trainer(
|
||||
...
|
||||
- strategy=DDPStrategy(),
|
||||
+ strategy=ray.train.lightning.RayDDPStrategy(),
|
||||
...
|
||||
)
|
||||
...
|
||||
|
||||
Configure the Ray cluster environment plugin
|
||||
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
|
||||
|
||||
Ray Train also provides a :class:`~ray.train.lightning.RayLightningEnvironment` class
|
||||
as a specification for the Ray Cluster. This utility class configures the worker's
|
||||
local, global, and node rank and world size.
|
||||
|
||||
|
||||
.. code-block:: diff
|
||||
|
||||
import lightning.pytorch as pl
|
||||
-from pl.plugins.environments import LightningEnvironment
|
||||
+import ray.train.lightning
|
||||
|
||||
def train_func():
|
||||
...
|
||||
trainer = pl.Trainer(
|
||||
...
|
||||
- plugins=[LightningEnvironment()],
|
||||
+ plugins=[ray.train.lightning.RayLightningEnvironment()],
|
||||
...
|
||||
)
|
||||
...
|
||||
|
||||
|
||||
Configure parallel devices
|
||||
^^^^^^^^^^^^^^^^^^^^^^^^^^
|
||||
|
||||
In addition, Ray TorchTrainer has already configured the correct
|
||||
``CUDA_VISIBLE_DEVICES`` for you. One should always use all available
|
||||
GPUs by setting ``devices="auto"`` and ``acelerator="auto"``.
|
||||
|
||||
|
||||
.. code-block:: diff
|
||||
|
||||
import lightning.pytorch as pl
|
||||
|
||||
def train_func():
|
||||
...
|
||||
trainer = pl.Trainer(
|
||||
...
|
||||
- devices=[0,1,2,3],
|
||||
+ devices="auto",
|
||||
+ accelerator="auto",
|
||||
...
|
||||
)
|
||||
...
|
||||
|
||||
|
||||
|
||||
Report checkpoints and metrics
|
||||
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
|
||||
|
||||
To persist your checkpoints and monitor training progress, add a
|
||||
:class:`ray.train.lightning.RayTrainReportCallback` utility callback to your Trainer.
|
||||
|
||||
|
||||
.. code-block:: diff
|
||||
|
||||
import lightning.pytorch as pl
|
||||
from ray.train.lightning import RayTrainReportCallback
|
||||
|
||||
def train_func():
|
||||
...
|
||||
trainer = pl.Trainer(
|
||||
...
|
||||
- callbacks=[...],
|
||||
+ callbacks=[..., RayTrainReportCallback()],
|
||||
)
|
||||
...
|
||||
|
||||
|
||||
Reporting metrics and checkpoints to Ray Train enables you to support :ref:`fault-tolerant training <train-fault-tolerance>` and :ref:`hyperparameter optimization <train-tune>`.
|
||||
|
||||
You can also configure :ref:`asynchronous checkpointing <train-checkpoint-upload-mode-async>` and
|
||||
:ref:`asynchronous validation <train-validating-checkpoints>` through the callback.
|
||||
The ``checkpoint_upload_mode`` flag offloads checkpoint uploading to a Ray Train managed background thread
|
||||
instead of blocking the Lightning training loop. The ``validation`` flag launches an asynchronous
|
||||
Ray task to validate the checkpoint instead of running ``validation_step`` synchronously in the training workers.
|
||||
Note that this is incompatible with Lightning's `AsyncCheckpointIO <https://lightning.ai/docs/pytorch/stable/api/lightning.pytorch.plugins.io.AsyncCheckpointIO.html>`_ plugin
|
||||
because Ray Train needs to control the upload thread in order to wait for it to finish before committing the checkpoint.
|
||||
|
||||
.. code-block:: diff
|
||||
|
||||
import lightning.pytorch as pl
|
||||
from ray.train.lightning import RayTrainReportCallback
|
||||
+from ray.train import CheckpointUploadMode
|
||||
+from ray.train import ValidationConfig, ValidationTaskConfig
|
||||
|
||||
def train_func():
|
||||
...
|
||||
trainer = pl.Trainer(
|
||||
...
|
||||
- callbacks=[..., RayTrainReportCallback()],
|
||||
+ callbacks=[..., RayTrainReportCallback(
|
||||
+ checkpoint_upload_mode=CheckpointUploadMode.ASYNC,
|
||||
+ validation=ValidationTaskConfig(fn_kwargs={}),
|
||||
+ )],
|
||||
)
|
||||
...
|
||||
|
||||
+def validation_fn(checkpoint):
|
||||
+ # Load and validate the checkpoint, return metrics.
|
||||
+ return {"val_score": ...}
|
||||
+
|
||||
+trainer = TorchTrainer(
|
||||
+ train_func,
|
||||
+ validation_config=ValidationConfig(fn=validation_fn),
|
||||
+ ...
|
||||
+)
|
||||
|
||||
Note that the :class:`ray.train.lightning.RayTrainReportCallback` class only provides a simple implementation, and can be :ref:`further customized <train-dl-saving-checkpoints>`.
|
||||
|
||||
Prepare your Lightning Trainer
|
||||
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
|
||||
|
||||
Finally, pass your Lightning Trainer into
|
||||
:meth:`~ray.train.lightning.prepare_trainer` to validate
|
||||
your configurations.
|
||||
|
||||
|
||||
.. code-block:: diff
|
||||
|
||||
import lightning.pytorch as pl
|
||||
import ray.train.lightning
|
||||
|
||||
def train_func():
|
||||
...
|
||||
trainer = pl.Trainer(...)
|
||||
+ trainer = ray.train.lightning.prepare_trainer(trainer)
|
||||
...
|
||||
|
||||
|
||||
.. include:: ./common/torch-configure-run.rst
|
||||
|
||||
|
||||
Next steps
|
||||
----------
|
||||
|
||||
After you have converted your PyTorch Lightning training script to use Ray Train:
|
||||
|
||||
* See :ref:`User Guides <train-user-guides>` to learn more about how to perform specific tasks.
|
||||
* Browse the :doc:`Examples <examples>` for end-to-end examples of how to use Ray Train.
|
||||
* Consult the :ref:`API Reference <train-api>` for more details on the classes and methods from this tutorial.
|
||||
|
||||
Version Compatibility
|
||||
---------------------
|
||||
|
||||
Ray Train is tested with `pytorch_lightning` versions `1.6.5` and `2.1.2`. For full compatibility, use ``pytorch_lightning>=1.6.5`` .
|
||||
Earlier versions aren't prohibited but may result in unexpected issues. If you run into any compatibility issues, consider upgrading your PyTorch Lightning version or
|
||||
`file an issue <https://github.com/ray-project/ray/issues>`_.
|
||||
|
||||
.. note::
|
||||
|
||||
If you are using Lightning 2.x, please use the import path `lightning.pytorch.xxx` instead of `pytorch_lightning.xxx`.
|
||||
|
||||
.. _lightning-trainer-migration-guide:
|
||||
|
||||
LightningTrainer Migration Guide
|
||||
--------------------------------
|
||||
|
||||
Ray 2.4 introduced the `LightningTrainer`, and exposed a
|
||||
`LightningConfigBuilder` to define configurations for `pl.LightningModule`
|
||||
and `pl.Trainer`.
|
||||
|
||||
It then instantiates the model and trainer objects and runs a pre-defined
|
||||
training function in a black box.
|
||||
|
||||
This version of the LightningTrainer API was constraining and limited
|
||||
your ability to manage the training functionality.
|
||||
|
||||
Ray 2.7 introduced the newly unified :class:`~ray.train.torch.TorchTrainer` API, which offers
|
||||
enhanced transparency, flexibility, and simplicity. This API is more aligned
|
||||
with standard PyTorch Lightning scripts, ensuring users have better
|
||||
control over their native Lightning code.
|
||||
|
||||
|
||||
.. tab-set::
|
||||
|
||||
.. tab-item:: (Deprecating) LightningTrainer
|
||||
|
||||
.. This snippet isn't tested because it raises a hard deprecation warning.
|
||||
|
||||
.. testcode::
|
||||
:skipif: True
|
||||
|
||||
from ray.train.lightning import LightningConfigBuilder, LightningTrainer
|
||||
|
||||
config_builder = LightningConfigBuilder()
|
||||
# [1] Collect model configs
|
||||
config_builder.module(cls=MyLightningModule, lr=1e-3, feature_dim=128)
|
||||
|
||||
# [2] Collect checkpointing configs
|
||||
config_builder.checkpointing(monitor="val_accuracy", mode="max", save_top_k=3)
|
||||
|
||||
# [3] Collect pl.Trainer configs
|
||||
config_builder.trainer(
|
||||
max_epochs=10,
|
||||
accelerator="gpu",
|
||||
log_every_n_steps=100,
|
||||
)
|
||||
|
||||
# [4] Build datasets on the head node
|
||||
datamodule = MyLightningDataModule(batch_size=32)
|
||||
config_builder.fit_params(datamodule=datamodule)
|
||||
|
||||
# [5] Execute the internal training function in a black box
|
||||
ray_trainer = LightningTrainer(
|
||||
lightning_config=config_builder.build(),
|
||||
scaling_config=ScalingConfig(num_workers=4, use_gpu=True),
|
||||
run_config=RunConfig(
|
||||
checkpoint_config=CheckpointConfig(
|
||||
num_to_keep=3,
|
||||
checkpoint_score_attribute="val_accuracy",
|
||||
checkpoint_score_order="max",
|
||||
),
|
||||
)
|
||||
)
|
||||
result = ray_trainer.fit()
|
||||
|
||||
# [6] Load the trained model from an opaque Lightning-specific checkpoint.
|
||||
lightning_checkpoint = result.checkpoint
|
||||
model = lightning_checkpoint.get_model(MyLightningModule)
|
||||
|
||||
|
||||
|
||||
.. tab-item:: (New API) TorchTrainer
|
||||
|
||||
.. This snippet isn't tested because it runs with 4 GPUs, and CI is only run with 1.
|
||||
|
||||
.. testcode::
|
||||
:skipif: True
|
||||
|
||||
import os
|
||||
|
||||
import lightning.pytorch as pl
|
||||
|
||||
import ray.train
|
||||
from ray.train.torch import TorchTrainer
|
||||
from ray.train.lightning import (
|
||||
RayDDPStrategy,
|
||||
RayLightningEnvironment,
|
||||
RayTrainReportCallback,
|
||||
prepare_trainer
|
||||
)
|
||||
|
||||
def train_func():
|
||||
# [1] Create a Lightning model
|
||||
model = MyLightningModule(lr=1e-3, feature_dim=128)
|
||||
|
||||
# [2] Report Checkpoint with callback
|
||||
ckpt_report_callback = RayTrainReportCallback()
|
||||
|
||||
# [3] Create a Lighting Trainer
|
||||
trainer = pl.Trainer(
|
||||
max_epochs=10,
|
||||
log_every_n_steps=100,
|
||||
# New configurations below
|
||||
devices="auto",
|
||||
accelerator="auto",
|
||||
strategy=RayDDPStrategy(),
|
||||
plugins=[RayLightningEnvironment()],
|
||||
callbacks=[ckpt_report_callback],
|
||||
)
|
||||
|
||||
# Validate your Lightning trainer configuration
|
||||
trainer = prepare_trainer(trainer)
|
||||
|
||||
# [4] Build your datasets on each worker
|
||||
datamodule = MyLightningDataModule(batch_size=32)
|
||||
trainer.fit(model, datamodule=datamodule)
|
||||
|
||||
# [5] Explicitly define and run the training function
|
||||
ray_trainer = TorchTrainer(
|
||||
train_func,
|
||||
scaling_config=ray.train.ScalingConfig(num_workers=4, use_gpu=True),
|
||||
run_config=ray.train.RunConfig(
|
||||
checkpoint_config=ray.train.CheckpointConfig(
|
||||
num_to_keep=3,
|
||||
checkpoint_score_attribute="val_accuracy",
|
||||
checkpoint_score_order="max",
|
||||
),
|
||||
)
|
||||
)
|
||||
result = ray_trainer.fit()
|
||||
|
||||
# [6] Load the trained model from a simplified checkpoint interface.
|
||||
checkpoint: ray.train.Checkpoint = result.checkpoint
|
||||
with checkpoint.as_directory() as checkpoint_dir:
|
||||
print("Checkpoint contents:", os.listdir(checkpoint_dir))
|
||||
checkpoint_path = os.path.join(checkpoint_dir, "checkpoint.ckpt")
|
||||
model = MyLightningModule.load_from_checkpoint(checkpoint_path)
|
||||
Reference in New Issue
Block a user