[SGD] Add PTL Docs (#12440)

Co-authored-by: Richard Liaw <rliaw@berkeley.edu>

doc/source/conf.py

@@ -44,6 +44,13 @@ MOCK_MODULES = [
     "mxnet",
     "mxnet.model",
     "psutil",
+    "pytorch_lightning.core.step_result",
+    "pytorch_lightning.overrides.data_parallel",
+    "pytorch_lightning.utilities.model_utils",
+    "pytorch_lightning.trainer.model_hooks",
+    "pytorch_lightning.trainer.optimizers",
+    "pytorch_lightning.utilities.exceptions",
+    "pytorch_lightning.utilities.memory",
     "ray._raylet",
     "ray.core.generated",
     "ray.core.generated.common_pb2",
@@ -76,6 +83,7 @@ MOCK_MODULES = [
     "wandb",
     "zoopt",
 ]
+
 import scipy.stats
 import scipy.linalg
 
@@ -87,6 +95,30 @@ sys.modules["tensorflow"].VERSION = "9.9.9"
 sys.modules["tensorflow.keras.callbacks"] = ChildClassMock()
 sys.modules["pytorch_lightning"] = ChildClassMock()
 
+
+class SimpleClass(object):
+    pass
+
+
+class SimpleClass2(object):
+    pass
+
+
+# ray.util.sgd.torch.lightning_operator.LightningOperator extends
+# TrainingOperator, pytorch_lightning.TrainerOptimizersMixin,
+# and pytorch_lightning.TrainerModelHooksMixin.
+# But, we are mocking all pytorch_lightning modules, causing the ptl base
+# classes to have a different metaclass than TrainingOperator.
+# To fix this, we replace the base classes with dummy classes that extend
+# object.
+# We have to create 2 dummy classes, one for TrainerOptimizersMixin and one
+# for TrainerModelHooksMixin so that we don't extend from the same base
+# class twice.
+setattr(sys.modules["pytorch_lightning.trainer.optimizers"],
+        "TrainerOptimizersMixin", SimpleClass)
+setattr(sys.modules["pytorch_lightning.trainer.model_hooks"],
+        "TrainerModelHooksMixin", SimpleClass2)
+
 # If extensions (or modules to document with autodoc) are in another directory,
 # add these directories to sys.path here. If the directory is relative to the
 # documentation root, use os.path.abspath to make it absolute, like shown here.

doc/source/index.rst

@@ -274,6 +274,7 @@ Papers
    raysgd/raysgd_pytorch.rst
    raysgd/raysgd_tensorflow.rst
    raysgd/raysgd_dataset.rst
+   raysgd/raysgd_ptl.rst
    raysgd/raysgd_tune.rst
    raysgd/raysgd_ref.rst
 

doc/source/raysgd/raysgd_ptl.rst

@@ -0,0 +1,117 @@
+Pytorch Lightning with RaySGD
+==============================
+.. image:: /images/sgd_ptl.png
+  :align: center
+  :scale: 50 %
+
+
+RaySGD includes an integration with Pytorch Lightning's `LightningModule <https://pytorch-lightning.readthedocs.io/en/latest/lightning_module.html>`_.
+Easily take your existing ``LightningModule``, and use it with Ray SGD's ``TorchTrainer`` to take advantage of all of Ray SGD's distributed training features with minimal code changes.
+
+.. tip:: This LightningModule integration is currently under active development. If you encounter any bugs, please raise an issue on `Github <https://github.com/ray-project/ray/issues>`_!
+
+.. note:: Not all Pytorch Lightning features are supported. A full list of unsupported model hooks is listed down :ref:`below <ptl-unsupported-features>`. Please post any feature requests on `Github <https://github.com/ray-project/ray/issues>`_ and we will get to it shortly!
+
+.. contents::
+    :local:
+
+Quick Start
+-----------
+Step 1: Define your ``LightningModule`` just like how you would with Pytorch Lightning.
+
+.. code-block:: python
+
+    from pytorch_lightning.core.lightning import LightningModule
+
+    class MyLightningModule(LightningModule):
+        ...
+
+Step 2: Use the ``TrainingOperator.from_ptl`` method to convert the ``LightningModule`` to a Ray SGD compatible ``LightningOperator``.
+
+.. code-block:: python
+
+    from ray.util.sgd.torch import TrainingOperator
+
+    MyLightningOperator = TrainingOperator.from_ptl(MyLightningModule)
+
+Step 3: Use the Operator with Ray SGD's ``TorchTrainer``, just like how you would normally. See :ref:`torch-guide` for a more full guide on ``TorchTrainer``.
+
+.. code-block:: python
+
+    import ray
+    from ray.util.sgd.torch import TorchTrainer
+
+    ray.init()
+    trainer = TorchTrainer(training_operator_cls=MyLightningOperator, num_workers=4, use_gpu=True)
+    train_stats = trainer.train()
+
+And that's it! For a more comprehensive guide, see the MNIST tutorial :ref:`below <ptl-mnist>`.
+
+.. _ptl-mnist:
+
+MNIST Tutorial
+--------------
+In this walkthrough we will go through how to train an MNIST classifier with Pytorch Lightning's ``LightningModule`` and Ray SGD.
+
+We will follow `this tutorial from the PyTorch Lightning documentation
+<https://pytorch-lightning.readthedocs.io/en/latest/introduction_guide.html>`_ for specifying our MNIST LightningModule.
+
+Setup / Imports
+~~~~~~~~~~~~~~~
+Let's start with some basic imports:
+
+.. literalinclude:: /../../python/ray/util/sgd/torch/examples/pytorch-lightning/mnist-ptl.py
+   :language: python
+   :start-after: __import_begin__
+   :end-before: __import_end__
+
+Most of these imports are needed for building our Pytorch model and training components.
+Only a few additional imports are needed for Ray and Pytorch Lightning.
+
+MNIST LightningModule
+~~~~~~~~~~~~~~~~~~~~~
+We now define our Pytorch Lightning ``LightningModule``:
+
+.. literalinclude:: /../../python/ray/util/sgd/torch/examples/pytorch-lightning/mnist-ptl.py
+   :language: python
+   :start-after: __ptl_begin__
+   :end-before: __ptl_end__
+
+This is the same code that would normally be used in Pytorch Lightning, and is taken directly from `this PTL guide <https://pytorch-lightning.readthedocs.io/en/latest/introduction_guide.html>`_.
+The only difference here is that the ``__init__`` method can optionally take in a ``config`` argument,
+as a way to pass in hyperparameters to your model, optimizer, or schedulers. The ``config`` will be passed in directly from
+the TorchTrainer. Or if using Ray SGD in conjunction with Tune (:ref:`raysgd-tune`), it will come directly from the config in your
+``tune.run`` call.
+
+Training with Ray SGD
+~~~~~~~~~~~~~~~~~~~~~
+We now can define our training function using our LitMNIST module and Ray SGD.
+
+.. literalinclude:: /../../python/ray/util/sgd/torch/examples/pytorch-lightning/mnist-ptl.py
+   :language: python
+   :start-after: __train_begin__
+   :end-before: __train_end__
+
+With just a single ``from_ptl`` call, we can convert our LightningModule to a ``TrainingOperator`` class that's compatible
+with Ray SGD. Now we can take full advantage of all of Ray SGD's distributed trainign features without having to rewrite our existing
+LightningModule.
+
+The last thing to do is initialize Ray, and run our training function!
+
+.. code-block:: python
+
+    # Use ray.init(address="auto") if running on a Ray cluster.
+    ray.init()
+    train_mnist(num_workers=32, use_gpu=True, num_epochs=5)
+
+.. _ptl-unsupported-features:
+
+Unsupported Features
+--------------------
+This integration is currently under active development, so not all Pytorch Lightning features are supported.
+Please post any feature requests on `Github
+<https://github.com/ray-project/ray/issues>`_ and we will get to it shortly!
+
+A list of unsupported model hooks (as of v1.0.0) is as follows:
+``test_dataloader``, ``on_test_batch_start``, ``on_test_epoch_start``, ``on_test_batch_end``, ``on_test_epoch_start``,
+``get_progress_bar_dict``, ``on_fit_end``, ``on_pretrain_routine_end``, ``manual_backward``, ``tbtt_split_batch``.

doc/source/raysgd/raysgd_pytorch.rst

@@ -1,3 +1,5 @@
+.. _torch-guide:
+
 Distributed PyTorch
 ===================
 
@@ -467,7 +469,6 @@ After connecting, you can scale up the number of workers seamlessly across multi
     trainer.train()
     model = trainer.get_model()
 
-
 Advanced: Fault Tolerance
 -------------------------
 

doc/source/raysgd/raysgd_ref.rst

@@ -1,10 +1,13 @@
-RaySGD API Documentation
-========================
+RaySGD API Reference
+====================
+
+PyTorch
+-------
 
 .. _ref-torch-trainer:
 
 TorchTrainer
-------------
+~~~~~~~~~~~~
 
 .. autoclass:: ray.util.sgd.torch.TorchTrainer
     :members:
@@ -12,40 +15,66 @@ TorchTrainer
 .. _ref-torch-operator:
 
 PyTorch TrainingOperator
-------------------------
+~~~~~~~~~~~~~~~~~~~~~~~~
 
 .. autoclass:: ray.util.sgd.torch.TrainingOperator
     :members:
 
+.. _ref-creator-operator:
+
+CreatorOperator
+~~~~~~~~~~~~~~~~
+
+.. autoclass:: ray.util.sgd.torch.training_operator.CreatorOperator
+    :members:
+    :exclude-members: setup
+
+.. _ref-lightning-operator:
+
+Pytorch Lightning LightningOperator
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autoclass:: ray.util.sgd.torch.lightning_operator.LightningOperator
+    :members:
+    :exclude-members: setup, train_epoch, train_batch, validate, validate_batch, state_dict, load_state_dict
+
 .. _BaseTorchTrainable-doc:
 
 BaseTorchTrainable
-------------------
+~~~~~~~~~~~~~~~~~~
 
 .. autoclass:: ray.util.sgd.torch.BaseTorchTrainable
     :members:
     :private-members:
 
+Tensorflow
+----------
+
 TFTrainer
----------
+~~~~~~~~~
 
 .. autoclass:: ray.util.sgd.tf.TFTrainer
     :members:
 
     .. automethod:: __init__
 
+RaySGD Dataset
+---------------
+
 Dataset
--------
+~~~~~~~
 
 .. autoclass:: ray.util.sgd.data.Dataset
     :members:
 
     .. automethod:: __init__
 
+RaySGD Utils
+-------------
 .. _ref-utils:
 
 Utils
------
+~~~~~
 
 .. autoclass:: ray.util.sgd.utils.AverageMeter
     :members:

doc/source/raysgd/raysgd_tune.rst

@@ -1,3 +1,5 @@
+.. _raysgd-tune:
+
 RaySGD Hyperparameter Tuning
 ============================