#3614 [Core feature] Support torch elastic training/torchrun Issue created by fg91 Motivation: Why do you think this is important? Flyte supports pytorch distributed training using the kf-pytorch plugin. This plugin, currently, does not support pytorch elastic training (`torchrun`). Elastic training has multiple benefits over vanilla torch-distributed training. For instance: • > Worker failures are handled gracefully by restarting all workers (source). • > Number of nodes is allowed to change between minimum and maximum sizes (elasticity) (source). • With

torchrun

one can perform distributed training on a single machine with multiple GPUs by starting a local process group in sub-processes. (With the kf-pytorch plugin one currently cannot perform distributed training on a single machine with multiple GPUs.) In addition, many open source projects or libraries building on-top of pytorch now simply assume that one uses

torchrun

. For instance: • Finetuning Stanford's new Llama model requires `torchrun`. • Pytorch ignite assumes that the environment variable `LOCAL_RANK` (which is set by `torchrun`) is set when initializing distributed training even though this is not required for native pytorch distributed training. * * * The kubeflow training operator which is used by Flyte to perform pytorch distributed training in Kubernetes clusters now supports elastic training. Flyte should make use of this to give users the benefits of elastic training. Goal: What should the final outcome look like, ideally? Currently, distributed training is configured as follows in `flytekit`:

from flytekitplugins.kfpytorch import PyTorch

@task(
    task_config=PyTorch(
        num_workers=2,
    ),
    ....
)
def train(...):
    ...

Users of

torchrun

typically start their training e.g. like this:

torchrun
    ....
    --nproc-per-node=4
    train.py

Flyte users should be able to configure elastic training e.g. like this:

from flytekitplugins.kfpytorch import Elastic

@task(
    task_config=Elastic(
        replicas=4,
        nproc_per_node=4,
        ...
    ),
    ...
)
def train(...):
    ...

Behind the scenes a task with such an

Elastic

task config should use

torch.distributed.launcher.api.elastic_launch

to start the task function which is also used by `torchrun`. Additionally, the values in the

Elastic

task config should also be used to configure the `ElasticPolicy` of the kubeflow PytorchJob used to perform multi-node distributed training in a Kubernetes cluster. Describe alternatives you've considered Propose: Link/Inline OR Additional context I implemented the required changes: • `flytekit`: flyteorg/flytekit#1603 • `flyteidl`: flyteorg/flyteidl#394 • `flyteplugins`: flyteorg/flyteplugins#343 • Docs: flyteorg/flytesnacks#987 Are you sure this issue hasn't been raised already? ☑︎ Yes Have you read the Code of Conduct? ☑︎ Yes flyteorg/flyte