Hi all, Could I please get some help on how to do multi-node multi-GPU distributed training on Metaflow? I tried adapting this PyTorch Lightning

ClusterEnvironment

for multi-node multi-GPU training? I'm stuck on deriving the

LOCAL_RANK

, which isn't provided as an environment variable using

@pytorch_parallel

. Thanks in advance!

Could you add in the following env variables

TORCH_DISTRIBUTED_DEBUG=DETAIL

&

TORCH_SHOW_CPP_STACKTRACES=1

to get more logs?

I'll give it a try, thanks. I've only ever tried multi-node single GPU distributed training. But realistically, most of our training jobs will need multi-node multi-GPU training.

Same here, I am looking to try out some multi-node training soon. Hopefully you manage to solve it for the rest of us using Metaflow.

are you able to use

NODE_RANK

which is set by metaflow?

Thanks Savin ! Actually haven't tried with DDP. I'll give that a shot next But I think @hallowed-glass-14538 and I discovered that

FSDP

and

Deepspeed

strategies within PTL mess up the environment variables and thus require a custom PTL

ClusterEnvironment

. That set-up I'm a little unsure about w/ regards to multi-node multi-GPU training. @hallowed-glass-14538 was able to figure out the right

ClusterEnvironment

for multi-node single GPU, but we will need to do multi-node multi-gpu in the near future for training LLMs.

Hey Riley!, I can have a look into this today. I think if we figure out the environment we maybe able to crack multi-node distributed training! 😄

Thanks so much Valay! In the meantime, I'll give DDP a shot just to test it out.

can you try this :

import os 
from lightning.pytorch.plugins.environments import ClusterEnvironment
from metaflow import current
class MetaflowEnvironment(ClusterEnvironment):
    """
    Quick Dirty MF environment in PTL for Single GPU Multi-node training.
    """

    @property
    def creates_processes_externally(self) -> bool:
        return True

    @property
    def main_address(self) -> str:
        return current.parallel.main_ip

    @property
    def main_port(self) -> int:
        return 9001 # Fix me

    @staticmethod
    def detect() -> bool:
        return True

    def world_size(self) -> int:
        return int(current.parallel.num_nodes) * int(N_GPU)

    def set_world_size(self, size: int) -> None:
        logging.debug("MetaflowEnvironment.set_world_size was called, but setting world size is not allowed. Ignored.")

    def global_rank(self) -> int:
        return int(current.parallel.node_index) * int(N_GPU) + int(os.environ.get("LOCAL_RANK", 0))

    def set_global_rank(self, rank: int) -> None:
        logging.debug("MetaflowEnvironment.set_global_rank was called, but setting global rank is not allowed. Ignored.")

    def local_rank(self) -> int:
        return os.environ.get("LOCAL_RANK", 0)

    def node_rank(self) -> int:
        return int(current.parallel.node_index)

I forgot where LOCAL_RANK was set in the code. do you know?

I'll give that a try, thanks. Where LOCAL_RANK is set in the FSDP codebase?

yes

I'll try running it first and if it errors out due to LOCAL_RANK, I'll provide the logs which hopefully should point to where it's set in the code

Are you still using your old source code ?

I am using this slightly different one - https://github.com/rileyhun/llm_finetuning_metaflow/tree/main/pytorch-deepspeed

can you change the line here to set

num_nodes = current.parallel.num_nodes

. You can even just use the

@parallel

decorator since

@pytorch_parallel

maybe setting worldsize which might be incorrect.

@pytorch_parallel

is only syntactic sugar over @parallel; All the environment information is anyways being set in the

ClusterEnvironment

Got it! Thanks so much. Will do. Will keep you posted in about an hour :)

Let me readup a bit about this today. I think there is a way here. We just need to peel through how PTL is starting the processses for the code you are running and setting those rank values.

Ok - sounds good. Thank you!

Hey Riley, a few questions on the code base after spending a little understanding what is happening under the hood : • are you using deepspeed or ddp. based on the code I deepspeed. ? • The screen shot is showing 3 workers and this is showing 4. Don;t know why this is the way it is • Can you also share full logs for control process and for 1 worker process. Apparently based on some reading I did, the error may have appeared much earlier in the logs The one thing I dont know currently is the following; For Multi-process multi-node DDP, do we have to launch separate subprocess per GPU or does lightning does that for us ? Based on the docs, it seems it maybe doing it for us since the docs show that we just need to instantiate the trainer on each node and apparently it figure out the GPU based multi-process creation on its own

Hello @hallowed-glass-14538, Sorry for the confusion - What happened was we were originally using DeepSpeed to test (we are wanting to fine-tune LLMs), but we gradually stripped away complexity just to test multi-node multi-GPU training. So in a sense, we are doing both. I'll create the

ddp

repo so it's not so confusing. But essentially, we are testing the following strategies w/ Metaflow:

ddp

,

fsp

,

deepspeed

with PTL

Finally I now understanding what is going on under the hood. Here is a quick rundown after peeling through PTL abstractions. What was the issue When we run a metaflow task and tell PTL that we are using multi-node DDP, PTL will automatically create subprocesses using the caller script. For metaflow's case it is the metaflow's flow file. PTL calling this file, would fail in mysterious ways since metaflow has additional CLI abstractions over the flow file and what ever arguments the file is being called with will not work as-is. So all the processes that need to start for DDP to work will not start and we will be stuck at that barrier since the master process started and the other child processes didint start on each node. This code block sets the subprocess launcher and here is the entry point class doing the automatic subprocess launching magic. Whats our quickest workaround We will have to call the model running code via subprocess and have that logic in a separate file. See a small example here. You can even do it without

torchrun

; I think this can unblock you. We can find better ways of handling this once we can get something working. This is so that PTL can launch additional subprocesses from the same script. Let me know if this helps

Ohhh I see. Interesting! Thanks for uncovering the details @hallowed-glass-14538. Really appreciate it! I'll poke around with

subprocess

and run the training as a

script

. Do you know why we were able to get multi-node single GPU working w/ DDP/FSDP without the need for

subprocess

?

My assumption is that because you were only having one GPU per-node, we did not face this issue. Is my assumption, correct?

Yup that's correct. So only multiple GPUs per node face this problem?

Yeah, and that is because it has to start the sub processes with local rank environment, variable set.

Oh I see. That is rather interesting! Well thank you so much for this startling discovery. I cannot wait to test your hypothesis. So all I need to do is move the training code into a script file and use

argparse

Yes ! That’s should be fine ! What’s your PTL version ?

I think

1.7.7

Hey let me have a look soon. We are so close 🏁

Thanks so much @hallowed-glass-14538!

Worldsize should be 16 given 4*4

Yup absolutely. I'm free tomorrow after 10:30 AM PST or anytime Friday. Will set

TORCH_DISTRIBUTED_DEBUG=INFO

as an environment variable now and send logs shortly.

Curious why it says worker count = 3

I think because there are 4 nodes - 3 of which are the workers and 1 is the control

I need to check but I was assuming worker count is number of procs per node (I think)

Based on the single-GPU multi-node training POC that we did, that wasn't the case. It was always (number of nodes - 1) for workers

Oh so worker count == num nodes in that scenario ?

No worker count == num nodes - 1 in that scenario.

Do you see similar logs for any other ran outside

8

?

Yes they are all the same logs for rank 4 and 12

i am assuming even 13,14,15

Okay will try that now

let me know how it goes excited

Will do - having trouble spinning up the GPU instances now. Likely because of resource availability (sigh)

can you share full logs ?

error.txt

can you share logs of the other tasks too ?

This message contains interactive elements.

From my understanding the problem is coming from this LOC; Where it's re-executing that for each worker and those workers crash because LOCAL_RANK is greater than 0 for that worker. And this happened because of our switch of creates_processes_externally == false. How can we fix this : 1. Set the creates_processes_externally == True Few options on what to do post setting creates_processes_externally == True 1. Manually launching process (No recommended) 2. Using torchrun or `python -m torch.distributed.launch`; You will just have to update the subprocess script like this :

from metaflow import current
subprocess.run(
    [
        "torchrun",
        f"--nproc_per_node={str(self.n_gpu)}",
        f"--nnodes={str(self.num_nodes)}",
        f"--rdzv-id={current.run_id}",
        "--rdzv-backend=c10d",
        f"--rdzv-endpoint={current.parallel.main_ip}:29400",
        "ddp_trainer.py",
        "--output-dir", self.output_dir,
        "--source-max-token-length", self.source_max_token_length,
        "--target-max-token-length", self.target_max_token_length,
        "--batch-size", self.batch_size,
        "--max-epochs", self.max_epochs,
        "--learning-rate", self.learning_rate,
        "--weight-decay", self.weight_decay,
        "--adam-epsilon", self.adam_epsilon,
        "--warmup-steps", self.warmup_steps,
        "--gradient-accumulation-steps", self.gradient_accumulation_steps,
        "--n-gpu", self.n_gpu,
        "--num-nodes=%d" % self.num_nodes,
        "--early-stopping-patience-epochs", self.early_stopping_patience_epochs,
        "--precision", self.precision,
        "--logger", self.logger,
        "--dataloader-num-workers", self.dataloader_num_workers,
        "--opt-level", self.opt_level,
        "--max-grad-norm", self.max_grad_norm,
        "--seed", self.seed,
    ]
    + (["--early-stop-callback"] if self.early_stop_callback else [])
    + (["--save-only-last-epoch"] if self.save_only_last_epoch else [])
    + (["--fp-16"] if self.fp_16 else [])
    + (["--use-gpu"] if self.use_gpu else []),
    check=True,
)

Oh right good catch. I'll give option 2 a shot right now. 🙂

Oops. Wrong args by mistake : Here's the right one

from metaflow import current
subprocess.run(
    [
        "torchrun",
        f"--nproc_per_node={str(self.n_gpu)}",
        f"--nnodes={str(self.num_nodes)}",
        f"--rdzv_id=metaflow_{current.run_id}",
        "--rdzv_backend=c10d",
        f"--rdzv_endpoint={current.parallel.main_ip}:29400",
        "ddp_trainer.py",
        "--output-dir", self.output_dir,
        "--source-max-token-length", self.source_max_token_length,
        "--target-max-token-length", self.target_max_token_length,
        "--batch-size", self.batch_size,
        "--max-epochs", self.max_epochs,
        "--learning-rate", self.learning_rate,
        "--weight-decay", self.weight_decay,
        "--adam-epsilon", self.adam_epsilon,
        "--warmup-steps", self.warmup_steps,
        "--gradient-accumulation-steps", self.gradient_accumulation_steps,
        "--n-gpu", self.n_gpu,
        "--num-nodes=%d" % self.num_nodes,
        "--early-stopping-patience-epochs", self.early_stopping_patience_epochs,
        "--precision", self.precision,
        "--logger", self.logger,
        "--dataloader-num-workers", self.dataloader_num_workers,
        "--opt-level", self.opt_level,
        "--max-grad-norm", self.max_grad_norm,
        "--seed", self.seed,
    ]
    + (["--early-stop-callback"] if self.early_stop_callback else [])
    + (["--save-only-last-epoch"] if self.save_only_last_epoch else [])
    + (["--fp-16"] if self.fp_16 else [])
    + (["--use-gpu"] if self.use_gpu else []),
    check=True,
)

This message contains interactive elements.

Oh whoops sorry

Ah fair. I thought it was

nproc-per-node

; But apparently it isn't.

This message contains interactive elements.

'--nproc_per_node=-1'

in the calledProcessError

Oh yes it is. I'll set it to 4.

Totally! Once you have it working end-to-end, let’s work on abstracting the complexity away from the users.

@hallowed-glass-14538 - Whoohoo! I got Deepspeed working as well w/ PTL using Metaflow's

@parallel

decorator. It's using 4

g5.12xlarge

instances. Thanks so much for your help I added a reproducible example here if it helps anyone else: https://github.com/rileyhun/llm_finetuning_metaflow/tree/main/pytorch-deepspeed @crooked-jordan-29960 By sheer coincidence, I also was able to get the progress bar to show up as well. I'm not really sure how or why though...All I did was print the loss (see here) and it allowed me to see the progress. If I remove the print statements though, the progress bar is suppressed. Any insight on this?

Amazing ! Few ideas on optimizing what is currently present; 1. Lightning has this interesting pattern where it allows you to expose the full training / validation / test loops via CLI a. You can use LightningCLI to create the CLI for the model you wish to train. b. You can export all arguments of your model to a yaml based config file with LightningCLI c. This pattern has tons of advantages: i. You can convert the model related args to a config file so that it avoids to CLI argument bloat. ii. You can only expose CLI parameters from metaflow which are really needed for parameterizing the infrastructure part of distributed training and use the config file to configure which lies in scope of the training loop. What I mean by this is that metaflow related

Parameters

would only be things like

num_gpus

,

loggers

,

seed

,

use_gpu

, `dataloader_num_workers`etc. The

config_file

can be on additional parameter that holds everything needed by PTL. iii. You can reduce the deps you install like

click

; iv. The internal teams can follow the pattern of using LightningCLI to create Models but the "training pipeline" is semi-general purpose to plug-in different models which expose a lightning CLI. v. Potentially this can be a way of abstracting out torchrun from users since users develop isolated modules that expose CLIs and config files; While the training pipeline takes care of passing things into the training code and running it at scale. (this will surely have it's own gotchas as I have not yet played much with PTL's cli) 2. When it comes abstracting out

torchrun

, one way to do this is to have a ray like approach of creating a strategy but one thing I am still not clear about is how we will launch sub-processes since torch distributed expects us to call the same script multiple times with different env variables. If the training loop is in the "main-process" (ie. the process launched by metaflow on remote) then we in a bind (of sorts) since Metaflow's cli won't play nice with calling itself. 3. The key distinction between ray and MF here is that: a. ray will run inside PTL's context meaning ray will launch jobs from within PTL's codebase; The compute is configured once outside PTL's context but the compute is accessed within PTL's context. b. MF is currently running outside PTL's codebase and also launches jobs outside PTL's codebase. MF will also dynamically configure/launch compute (batch-array specs based on resources in decorator); There is much we can explore here. Can we have a call sometime next week to discuss all the hurdles you faced while getting this running? This will help us understand better the abstractions we should provide to build this.

Thanks for providing this summary. This is really useful information. Absolutely! What is your availability like next week?

can Monday / tuesday / wednesday work ?

Yup - I'm free any time after 11AM PST on Monday or any time after 9AM PST on Wednesday

Did this ever turn into a more widely available wrapper / process with metaflow?

yep - https://outerbounds.com/blog/distributed-training-with-metaflow/