Got it, first single node multi GPU 👍 I can implement this on Friday or Saturday and then ping you for review. Or do you need it before that?

Hmm that should be ok. What I want to do is train alpaca on Flyte and have that as a demo

Can you give me permissions to open a PR in flytekit please?

Ohh you don’t have perms

Ok, branch is ready to push 🙂

Ok you should have it

Perfect

This is going to be awesome

Yup, I pushed some Commits too

Saw them 👍

I will try to get alpaca working on it too

You mean write task in notebook and then just run task from there?

Yup

I’m not much of a notebook user ^^ But I guess for many data scientists this is a killer feature

Ya that’s my hope

I have a question about how to select the plugin for the task type. I have this:

class MultiNodePytorchElasticFunctionTask(PythonFunctionTask[Elastic]):
    _ELASTIC_TASK_TYPE = "torch-elastic"

    def __init__(self, task_config: Elastic, task_function: Callable, **kwargs):
        super(MultiNodePytorchElasticFunctionTask, self).__init__(
            task_type=self._ELASTIC_TASK_TYPE,
            **kwargs,

    def get_custom(...): ...

I also added this to helm values:

enabled_plugins:
    tasks:
      task-plugins:
        enabled-plugins:
          - ...
          - pytorch
        default-for-task-types:
          - ...
          pytorch: pytorch
          torch-elastic: pytorch

Propeller says:

{
  "json": {
    "exec_id": "f1c678faf0fd74fad828",
    "node": "n0",
    "ns": "flytesnacks-development",
    "res_ver": "23058",
    "routine": "worker-2",
    "tasktype": "torch-elastic",
    "wf": "flytesnacks:development:<http://wf.wf|wf.wf>"
  },
  "level": "warning",
  "msg": "No plugin found for Handler-type [torch-elastic], defaulting to [container]",
  "ts": "2023-04-08T22:26:30Z"
}

Do I need to configure this somewhere else as well? The existing pytorch plugin in flyteplugins just needs an additional if else whether to configure an ElasticPolicy.

AFK

I was thinking exactly the same. I feel like this should go into the existing pytorch plugin, not new ones, since also for the kubeflow training operator vanilla torch distributed training and torch elastic training only differs by the elastic config in the pytorchjob manifest. Same k8s kind though. This stays the same for backwards compatibility of course

pip install flytekitplugins-kfpytorch

from flytekitplugins.kfpytorch import Pytorch

@task(
    task_config=Pytorch(...)
)

But people could to

pip install flytekitplugins-kfpytorch[elastic]

(for the torch dependency) and then:

from flytekitplugins.kfpytorch import ElasticPytorch

@task(
    task_config=ElasticPytorch(nnodes=1) # single pod, no operator
)

@task(
    task_config=ElasticPytorch(nnodes=2) # pytorch operator
)

And in flyteplugins all the pytorch code can be reused as well, just an if whether we need to set elastic config in pytorchjob.

I can give your perms

I saw you simplified the models, didn’t know this was possible, nice 👍 Update from my side: I opened a draft PR in idl and in plugins. Built a propeller image, creating a distributed pytorchjob with elastic config works. What is not working reliably yet is the rendevouz when initiating the process group. We definitely need something similar to what I added here:

rdzv_endpoint=os.environ.get("PET_RDZV_ENDPOINT", f"localhost:0"),

Here,

localhost:0

means torchrun picks a free port (see docs). I’m currently working on making this minimal elastic example from the kubeflow training operator repo work:

import os
import logging

from flytekit import task, workflow
from flytekitplugins.kfpytorch import PyTorch, Elastic

logging.basicConfig(level=<http://logging.INFO|logging.INFO>)  # To see torchrun trying to establish the rendevouz


@task(
    task_config=Elastic(
        nnodes=2,
        nproc_per_node=2,
        start_method="fork",
    )
    #task_config=PyTorch(num_workers=2)
)
def train() -> str:
    import io
    import os
    import pprint
    import sys
    import time

    import torch.distributed as dist

    env_dict = {
        k: os.environ[k]
        for k in (
            "LOCAL_RANK",
            "RANK",
            "GROUP_RANK",
            "WORLD_SIZE",
            "MASTER_ADDR",
            "MASTER_PORT",
            "TORCHELASTIC_RESTART_COUNT",
            "TORCHELASTIC_MAX_RESTARTS",
        )
    }

    with io.StringIO() as buff:
        print("======================================================", file=buff)
        print(
            f"Environment variables set by the agent on PID {os.getpid()}:", file=buff
        )
        pprint.pprint(env_dict, stream=buff)
        print("======================================================", file=buff)
        print(buff.getvalue())
        sys.stdout.flush()

    dist.init_process_group(backend="gloo")
    dist.barrier()
    rank = dist.get_rank()
    print(
        (
            f"On PID {os.getpid()}, after init process group, "
            f"rank={dist.get_rank()}, world_size = {dist.get_world_size()}\n"
        )
    )
    dist.destroy_process_group()
    return f"foo-{rank}"

@workflow
def wf():
    train()


if __name__ == "__main__":
    print(f"Parent {os.getpid()}")
    print(wf())

Rendevouz sometimes fails, sometimes works, currently debugging why. Just as fyi where I’m at…