[tune/sgd] Document func_trainable and add checkpoint context (#9739)

Co-authored-by: krfricke <krfricke@users.noreply.github.com>
Co-authored-by: Amog Kamsetty <amogkam@users.noreply.github.com>

doc/source/raysgd/raysgd_pytorch.rst

@@ -18,6 +18,8 @@ For end to end examples leveraging RaySGD TorchTrainer, jump to :ref:`raysgd-tor
 Setting up training
 -------------------
 
+.. tip:: If you want to leverage multi-node data parallel training with PyTorch while using RayTune *without* restructuring your code, check out the :ref:`Tune PyTorch user guide <tune-pytorch-cifar>` and Tune's  :ref:`distributed pytorch integrations <tune-ddp-doc>`.
+
 The ``TorchTrainer`` can be constructed with functions that wrap components of the training script. Specifically, it requires constructors for the Model, Data, Optimizer, Loss, and ``lr_scheduler`` to create replicated copies across different devices and machines.
 
 Under the hood, ``TorchTrainer`` will create *replicas* of your model (controlled by ``num_workers``), each of which is managed by a Ray actor. One of the replicas will be on the main process, which can simplify the debugging and logging experience.

doc/source/tune/_tutorials/tune-distributed.rst

@@ -213,16 +213,7 @@ In GCP, you can use the following configuration modification:
         scheduling:
           - preemptible: true
 
-Spot instances may be removed suddenly while trials are still running. Often times this may be difficult to deal with when using other distributed hyperparameter optimization frameworks. Tune allows users to mitigate the effects of this by preserving the progress of your model training through checkpointing.
-
-The easiest way to do this is to subclass the pre-defined ``Trainable`` class and implement ``save_checkpoint``, and ``load_checkpoint`` abstract methods, as seen in the example below:
-
-.. literalinclude:: /../../python/ray/tune/examples/mnist_pytorch_trainable.py
-   :language: python
-   :start-after: __trainable_example_begin__
-   :end-before: __trainable_example_end__
-
-This can then be used similarly to the Function API as before:
+Spot instances may be removed suddenly while trials are still running. Often times this may be difficult to deal with when using other distributed hyperparameter optimization frameworks. Tune allows users to mitigate the effects of this by preserving the progress of your model training through :ref:`checkpointing <tune-function-checkpointing>`.
 
 .. literalinclude:: /../../python/ray/tune/tests/tutorial.py
     :language: python

doc/source/tune/_tutorials/tune-pytorch-cifar.rst

@@ -25,6 +25,8 @@ need to
 3. add checkpointing (optional),
 4. and define the search space for the model tuning
 
+Optionally, you can seamlessly leverage :ref:`DistributedDataParallel training <tune-torch-ddp>` for each individual Pytorch model within Tune.
+
 .. note::
 
     To run this example, you will need to install the following:
@@ -74,15 +76,16 @@ The train function
 Now it gets interesting, because we introduce some changes to the example `from the PyTorch
 documentation <https://pytorch.org/tutorials/beginner/blitz/cifar10_tutorial.html>`_.
 
-We wrap the training script in a function ``train_cifar(config, checkpoint=None)``. As you
+We wrap the training script in a function ``train_cifar(config, checkpoint_dir=None)``. As you
 can guess, the ``config`` parameter will receive the hyperparameters we would like to
-train with. The ``checkpoint`` parameter is used to restore checkpoints.
+train with. The ``checkpoint_dir`` parameter is used to restore checkpoints.
 
 .. code-block:: python
 
     net = Net(config["l1"], config["l2"])
 
-    if checkpoint:
+    if checkpoint_dir:
+        checkpoint = os.path.join(checkpoint_dir, "checkpoint")
         net.load_state_dict(torch.load(checkpoint))
 
 The learning rate of the optimizer is made configurable, too:
@@ -97,6 +100,7 @@ with which we iterate through the training and test sets are configurable as wel
 
 Adding (multi) GPU support with DataParallel
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
 Image classification benefits largely from GPUs. Luckily, we can continue to use
 PyTorch's abstractions in Ray Tune. Thus, we can wrap our model in ``nn.DataParallel``
 to support data parallel training on multiple GPUs:
@@ -132,10 +136,9 @@ The most interesting part is the communication with Tune:
 
 .. code-block:: python
 
-    checkpoint_dir = tune.make_checkpoint_dir(epoch)
-    path = os.path.join(checkpoint_dir, "checkpoint")
-    torch.save((net.state_dict(), optimizer.state_dict()), path)
-    tune.save_checkpoint(path)
+    with tune.checkpoint_dir(epoch) as checkpoint_dir:
+        path = os.path.join(checkpoint_dir, "checkpoint")
+        torch.save((net.state_dict(), optimizer.state_dict()), path)
 
     tune.report(loss=(val_loss / val_steps), accuracy=correct / total)
 
@@ -273,6 +276,75 @@ be confirmed on the test set.
 
 So that's it! You can now tune the parameters of your PyTorch models.
 
+.. _tune-torch-ddp:
+
+Advanced: Distributed training with DistributedDataParallel
+-----------------------------------------------------------
+
+Some models require multiple nodes to train in a short amount of time. Ray Tune allows you to easily do distributed data parallel training in addition to distributed hyperparameter tuning.
+
+You can wrap your model in ``torch.nn.parallel.DistributedDataParallel`` to support distributed data parallel training:
+
+.. code-block:: python
+
+    from ray.util.sgd.torch import is_distributed_trainable
+    from torch.nn.parallel import DistributedDataParallel
+
+    def train_cifar(config, checkpoint_dir=None, data_dir=None):
+        net = Net(config["l1"], config["l2"])
+
+        device = "cpu"
+
+        #### Using distributed data parallel training
+        if is_distributed_trainable():
+            net = DistributedDataParallel(net)
+
+        if torch.cuda.is_available():
+            device = "cuda"
+
+        net.to(device)
+
+
+If using checkpointing, be sure to use a :ref:`special checkpoint context manager <tune-ddp-doc>`, ``distributed_checkpoint_dir`` that avoids redundant checkpointing across multiple processes:
+
+.. code-block:: python
+
+    from ray.util.sgd.torch import distributed_checkpoint_dir
+
+    #### Using distributed data parallel training
+    # Inside `def train_cifar(...)`,
+    # replace tune.checkpoint_dir() with the following
+    # Avoids redundant checkpointing on different processes.
+    with distributed_checkpoint_dir(step=epoch) as checkpoint_dir:
+        path = os.path.join(checkpoint_dir, "checkpoint")
+        torch.save((net.state_dict(), optimizer.state_dict()), path)
+
+
+Finally, we need to tell Ray Tune to start multiple distributed processes at once by using ``ray.tune.integration.torch.DistributedTrainableCreator`` (:ref:`docs <tune-ddp-doc>`). This is essentially equivalent to running ``torch.distributed.launch`` for each hyperparameter trial:
+
+.. code-block:: python
+
+    # You'll probably want to be running on a distributed Ray cluster.
+    # ray.init(address="auto")
+
+    from ray.util.sgd.integration.torch import DistributedTrainableCreator
+
+    distributed_train_cifar = DistributedTrainableCreator(
+      partial(train_cifar, data_dir=data_dir),
+      use_gpu=True,
+      num_workers=2,  # number of parallel workers to use
+      num_cpus_per_worker=8
+    )
+    tune.run(
+      distributed_train_cifar,
+      resources_per_trial=None,
+      config=config,
+      num_samples=num_samples,
+      ...
+    )
+
+See an :doc:`end-to-end example here </tune/examples/ddp_mnist_torch>`.
+
 If you consider switching to PyTorch Lightning to get rid of some of your boilerplate
 training code, please know that we also have a walkthrough on :doc:`how to use Tune with
-PyTorch Lightning models <tune-pytorch-lightning>`.
+PyTorch Lightning models <tune-pytorch-lightning>`.

doc/source/tune/api_docs/trainable.rst

@@ -17,6 +17,7 @@ For the sake of example, let's maximize this objective function:
 Function API
 ------------
 
+
 Here is a simple example of using the function API. You can report intermediate metrics by simply calling ``tune.report`` within the provided function.
 
 .. code-block:: python
@@ -40,38 +41,41 @@ Here is a simple example of using the function API. You can report intermediate
 
 Tune will run this function on a separate thread in a Ray actor process.
 
+.. tip:: If you want to leverage multi-node data parallel training with PyTorch while using parallel hyperparameter tuning, check out our :ref:PyTorch user guide and Tune's :ref:distributed pytorch integrations.
+
+.. _tune-function-checkpointing:
 
 Function API Checkpointing
 ~~~~~~~~~~~~~~~~~~~~~~~~~~
 
-Many Tune features rely on checkpointing, including the usage of certain Trial Schedulers and fault tolerance. To use Tune's checkpointing features, you must expose a ``checkpoint`` argument in the function signature, and call ``tune.make_checkpoint_dir`` and ``tune.save_checkpoint``:
+Many Tune features rely on checkpointing, including the usage of certain Trial Schedulers and fault tolerance. To use Tune's checkpointing features, you must expose a ``checkpoint_dir`` argument in the function signature, and call ``tune.checkpoint_dir`` :
 
 .. code-block:: python
 
         import time
         from ray import tune
 
-        def train_func(config, checkpoint=None):
+        def train_func(config, checkpoint_dir=None):
             start = 0
-            if checkpoint:
-                with open(checkpoint) as f:
+            if checkpoint_dir:
+                with open(os.path.join(checkpoint_dir, "checkpoint")) as f:
                     state = json.loads(f.read())
                     start = state["step"] + 1
 
             for iter in range(start, 100):
                 time.sleep(1)
 
-                #
-                checkpoint_dir = tune.make_checkpoint_dir(step=step)
-                path = os.path.join(checkpoint_dir, "checkpoint")
-                with open(path, "w") as f:
-                    f.write(json.dumps({"step": start}))
-                tune.save_checkpoint(path)
+                with tune.checkpoint_dir(step=step):
+                    path = os.path.join(checkpoint_dir, "checkpoint")
+                    with open(path, "w") as f:
+                        f.write(json.dumps({"step": start}))
 
                 tune.report(hello="world", ray="tune")
 
         tune.run(train_func)
 
+.. note:: ``checkpoint_freq`` and ``checkpoint_at_end`` will not work with Function API checkpointing.
+
 In this example, checkpoints will be saved by training iteration to ``local_dir/exp_name/trial_name/checkpoint_<step>``. You can restore a single trial checkpoint by using ``tune.run(restore=<checkpoint_dir>)``:
 
 .. code-block:: python
@@ -263,9 +267,7 @@ tune.report / tune.checkpoint (Function API)
 
 .. autofunction:: ray.tune.report
 
-.. autofunction:: ray.tune.make_checkpoint_dir
-
-.. autofunction:: ray.tune.save_checkpoint
+.. autofunction:: ray.tune.checkpoint_dir
 
 .. autofunction:: ray.tune.get_trial_dir
 
@@ -282,6 +284,21 @@ tune.Trainable (Class API)
     :private-members:
     :members:
 
+
+.. _tune-ddp-doc:
+
+Distributed Torch
+-----------------
+
+Ray also offers lightweight integrations to distribute your model training on Ray Tune.
+
+
+.. autofunction:: ray.tune.integration.torch.DistributedTrainableCreator
+
+.. autofunction:: ray.tune.integration.torch.distributed_checkpoint_dir
+
+.. autofunction:: ray.tune.integration.torch.is_distributed_trainable
+
 tune.DurableTrainable
 ---------------------
 

doc/source/tune/examples/ddp_mnist_torch.rst

@@ -0,0 +1,6 @@
+:orphan:
+
+ddp_mnist_torch
+~~~~~~~~~~~~~~~
+
+.. literalinclude:: /../../python/ray/tune/examples/ddp_mnist_torch.py

doc/source/tune/examples/index.rst

@@ -41,6 +41,7 @@ PyTorch Examples
 
 - :doc:`/tune/examples/mnist_pytorch`: Converts the PyTorch MNIST example to use Tune with the function-based API. Also shows how to easily convert something relying on argparse to use Tune.
 - :doc:`/tune/examples/mnist_pytorch_trainable`: Converts the PyTorch MNIST example to use Tune with Trainable API. Also uses the HyperBandScheduler and checkpoints the model at the end.
+- :doc:`/tune/examples/ddp_mnist_torch`: An example showing how to use DistributedDataParallel with Ray Tune. This enables both distributed training and distributed hyperparameter tuning.
 
 
 XGBoost Example

doc/source/tune/user-guide.rst

@@ -151,17 +151,18 @@ When running a hyperparameter search, Tune can automatically and periodically sa
 
 Checkpointing assumes that the model state will be saved to disk on whichever node the Trainable is running on.
 
-To use Tune's checkpointing features, you must expose a ``checkpoint`` argument in the function signature, and call ``tune.make_checkpoint_dir`` and ``tune.save_checkpoint``:
+To use Tune's checkpointing features, you must expose a ``checkpoint_dir`` argument in the function signature, and call ``tune.checkpoint_dir``:
 
 .. code-block:: python
 
+        import os
         import time
         from ray import tune
 
-        def train_func(config, checkpoint=None):
+        def train_func(config, checkpoint_dir=None):
             start = 0
-            if checkpoint:
-                with open(checkpoint) as f:
+            if checkpoint_dir:
+                with open(os.path.join(checkpoint_dir, "checkpoint")) as f:
                     state = json.loads(f.read())
                     start = state["step"] + 1
 
@@ -169,11 +170,10 @@ To use Tune's checkpointing features, you must expose a ``checkpoint`` argument
                 time.sleep(1)
 
                 # Obtain a checkpoint directory
-                checkpoint_dir = tune.make_checkpoint_dir(step=step)
-                path = os.path.join(checkpoint_dir, "checkpoint")
-                with open(path, "w") as f:
-                    f.write(json.dumps({"step": start}))
-                tune.save_checkpoint(path)
+                with tune.checkpoint_dir(step=step) as checkpoint_dir:
+                    path = os.path.join(checkpoint_dir, "checkpoint")
+                    with open(path, "w") as f:
+                        f.write(json.dumps({"step": start}))
 
                 tune.report(hello="world", ray="tune")