DeepspeedTorchDistributor¶
-
class
pyspark.ml.deepspeed.deepspeed_distributor.DeepspeedTorchDistributor(numGpus: int = 1, nnodes: int = 1, localMode: bool = True, useGpu: bool = True, deepspeedConfig: Union[str, Dict[str, Any], None] = None)[source]¶ Methods
run(train_object, *args, **kwargs)Runs distributed training.
Methods Documentation
-
run(train_object: Union[Callable, str], *args: Any, **kwargs: Any) → Optional[Any][source]¶ Runs distributed training.
- Parameters
- train_objectcallable object or str
Either a PyTorch function, PyTorch Lightning function, or the path to a python file that launches distributed training.
- args :
If train_object is a python function and not a path to a python file, args need to be the input parameters to that function. It would look like
>>> model = distributor.run(train, 1e-3, 64)
where train is a function and 1e-3 and 64 are regular numeric inputs to the function.
If train_object is a python file, then args would be the command-line arguments for that python file which are all in the form of strings. An example would be
>>> distributor.run("/path/to/train.py", "--learning-rate=1e-3", "--batch-size=64")
where since the input is a path, all of the parameters are strings that can be handled by argparse in that python file.
- kwargs :
If train_object is a python function and not a path to a python file, kwargs need to be the key-word input parameters to that function. It would look like
>>> model = distributor.run(train, tol=1e-3, max_iter=64)
where train is a function of 2 arguments tol and max_iter.
If train_object is a python file, then you should not set kwargs arguments.
- Returns
- Returns the output of train_object called with args inside spark rank 0 task if the
train_object is a Callable with an expected output. Returns None if train_object is a file.
-