[GitHub] [spark] HyukjinKwon commented on a diff in pull request #42087: [SPARK-44264] Added Example to Deepspeed Distributor

["HyukjinKwon (via GitHub)" <gi...@apache.org>]

HyukjinKwon commented on code in PR #42087:
URL: https://github.com/apache/spark/pull/42087#discussion_r1274230031


##########
python/pyspark/ml/deepspeed/deepspeed_distributor.py:
##########
@@ -61,6 +61,30 @@ def __init__(
             The configuration file to be used for launching the deepspeed application.
             If it's a dictionary containing the parameters, then we will create the file.
             If None, deepspeed will fall back to default parameters.
+
+        Examples
+        --------
+        Run Deepspeed training function on a single node
+
+        >>> def train(learning_rate):
+                import deepspeed
+                # rest of training function
+                return model
+        >>> distributor = DeepspeedTorchDistributor(num_gpus=4,
+                                                    nnodes=1,
+                                                    use_gpu=True,
+                                                    local_mode=True,
+                                                    deepspeed_config="path/to/config.json")
+        >>> output = distributor.run(train, 0.01)
+
+        Run Deepspeed training function on multiple nodes
+
+        >>> distributor = DeepspeedTorchDistributor(num_gpus=4,
+                                                    nnodes=3,
+                                                    use_gpu=True,
+                                                    local_mode=False,
+                                                    deepspeed_config="path/to/config.json")

Review Comment:
   ```suggestion
           >>> def train(learning_rate):
           ...     import deepspeed
           ...     # rest of training function
           ...     return model
           >>> distributor = DeepspeedTorchDistributor(
           ...     num_gpus=4,
           ...     nnodes=1,
           ...     use_gpu=True,
           ...     local_mode=True,
           ...     deepspeed_config="path/to/config.json")
           >>> output = distributor.run(train, 0.01)
   
           Run Deepspeed training function on multiple nodes
   
           >>> distributor = DeepspeedTorchDistributor(
           ...     num_gpus=4,
           ...     nnodes=3,
           ...     use_gpu=True,
           ...     local_mode=False,
           ...     deepspeed_config="path/to/config.json")
   ```



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org