Hi team, Was wondering if I could kindly ask for some assistance. I'm working on a demo for LLM Fine Tuning using PyTorch Lightning that leverages Metaflow's

@pytorch_parallel

decorator. I'm still rather new to distributed training, so learning the ropes. I've tried

fsdp

to do sharded training, but it was having some trouble with the model parameters. I was getting this error -

valueerror: optimizer got an empty parameter list

. When I printed out the params though, they're definitely there. I also tried

ddp

as the strategy, which didn't give any explicit error, but it just stalled and didn't have any progress output in the stdout logs... I have a reproducible example here - https://github.com/rileyhun/llm_finetuning_metaflow/blob/main/gpt-j-8bit-flow.py. Any pointers or guidance would be greatly appreciated.

hey! few quick questions. I see you are using PTL 1.8.6. And you are using

@pytorch_parallel

so you have multiple Jobs spawned with gang scheduling to run the distributed training (over multiple nodes (with multiple GPUS per node?). 1. Did other tasks start in the image you shared. AWS batch can take up some time to start all the nodes you need and then place the containers on the nodes. Can you verify if other tasks started coz it seems the "control task" is waiting on the child tasks to start and stuck in the barrier. "control task" is the main that that will get the MASTER_IP and to which other tasks will communicate with. 2. I think the DDP support for PT Lightning has changed since we built the

@pytorch_parallel

. This short version is that now it calls the parent script under the hood on its own and it also has introduced abstractions to handle how a ClusterEnviornment performs DDP or SFDP. you can try using LightningEnvironment if that helps? `ClusterEnvironment`s can be passed to PT lightning Trainer as a plugins.

Hi @little-apartment-49355, Thanks for your response. I am using 16 nodes, each with just 1 GPU (

g5.4xlarge

), and they all started up and were running. I can look into

ClusterEnvironment

One thing I am curious about, given the

valueerror: optimizer got an empty parameter list

and DDP issue is the

Adam8bit

parts. I wonder if setting this arg when HuggingFace loads model will help:

self.model = GPTJForCausalLM.from_pretrained(..., load_in_8bit=True)

I see in the DDP version of the flow output it is using 16bit AMP, so maybe the optimizer is getting confused and not reading the optimizer state correctly since the instantiated Adam is 8bit. FYI I haven't had a chance to reproduce yet, so this is just a guess.

Thanks Eddie! Good catch. I'll add that in and see if that helps.

Hey Riley what are you network and VPC settings. It seems the jobs are unable to communicate with each other.

Hi @hallowed-glass-14538, I think the network settings should be OK. I tried a simple pytorch distributed program using the MNIST data and that seemed to work.

Can you change the backend from nccl to

gloo

and try it out ? I noticed a comment in our code that says that nccl didn't play nice with AWS batch.

Ok, will do - I sort of gave up and switched models now (I just need this for a DEMO 😅), so I'm working on pre-training a BERT model right now instead and *cross fingers...*I think it might be working now. I think I had to set these environment variables -

os.environ["MASTER_PORT"] = "12345"
os.environ["MASTER_ADDR"] = "localhost"
os.environ["WORLD_SIZE"] = str(current.parallel.num_nodes)
os.environ["NODE_RANK"] = str(current.parallel.node_index)

I also removed the cluster environment/lightning environment.

Wait so if there are 16 machines then the master should have a different IP than localhost

Oh very nice! Thanks for the warning Valay. Let me plug that in. 🙏

and

master_port

can be set in

@pytorch_parallel

. So something like :

@pytorch_parallel(master_port=9001)

Okay got it. Will fix that too.

what is the value of

devices

and

num_nodes

in your

Trainer

?

devices

is 1 and

num_nodes

is 4

can you try setting

LOCAL_RANK=0

once

Yup.

os.environ["MASTER_ADDR"] = str(current.parallel.main_ip)
os.environ["WORLD_SIZE"] = str(current.parallel.num_nodes)
os.environ["NODE_RANK"] = str(current.parallel.node_index)
os.environ["LOCAL_RANK"] = 0

can you print

os.environ

in your task processes get get all the actual values of

MASTER_ADDR

,

GLOBAL_RANK

etc. I wanted to ensure that we are not modifying something from the MF side at runtime. can you also share your code by any chance ? Wanted to see all configurations once.

Here's my reproducible code example: https://github.com/rileyhun/llm_finetuning_metaflow/blob/main/fsdp/fsdp_flow.py I'm currently trying it with a Lightning Environment to see if that helps

can you change here to:

env.global_rank = lambda: int(current.parallel.node_index)

Yup - I think it has to be a

str

though otherwise it errors out.

sure make it

str

.

Ah ok. Will try it now

can you print out your PT related env vars here ?

Printing them out on the next flow:

os.environ["MASTER_ADDR"] = str(current.parallel.main_ip)
env = LightningEnvironment()
env.world_size = lambda: int(current.parallel.num_nodes)
env.global_rank = lambda: int(current.parallel.num_nodes)
env.node_rank = lambda: int(current.parallel.node_index)
env.local_rank = lambda: int(os.environ.get("LOCAL_RANK", 0))
        
print(os.environ)
print(env)

coz it seems that its calling

[10.14.52.10]:5558

so may even have to pluck and replace the main_port property

ah ok. Got it.

and just return the port you are opening

environ variables:

'MF_PARALLEL_MAIN_IP': '10.14.51.206', 'MF_PARALLEL_NUM_NODES': '4', 'MF_PARALLEL_NODE_INDEX': '0', 'MASTER_PORT': '9001', 'MASTER_ADDR': '10.14.51.206', 'NODE_RANK': '0', 'WORLD_SIZE': '4', 'NUM_NODES': '4', 'PL_TORCH_DISTRIBUTED_BACKEND': 'gloo', 'PYTORCH_NVML_BASED_CUDA_CHECK': '1', 'CRC32C_SW_MODE': 'auto'

I think this error indicates

torch.distributed

sees worldsize as 1. Does this specific error go away if you uncomment the line:

os.environ["WORLD_SIZE"] = str(current.parallel.num_nodes)

Let me try it @crooked-jordan-29960. Thanks so much for all the help by the way guys. Really appreciate it. 🙏

of course, we appreciate you putting these feature through their paces!

To clarify, I'm using a

LightningEnvironment

to set up the cluster config: commenting out

world_size

os.environ["MASTER_ADDR"] = str(current.parallel.main_ip)
os.environ["MASTER_PORT"] = "9001"
        
env = LightningEnvironment()
# env.world_size = lambda: int(current.parallel.num_nodes)
env.global_rank = lambda: int(current.parallel.node_index)
env.node_rank = lambda: int(current.parallel.node_index)
env.local_rank = lambda: int(os.environ.get("LOCAL_RANK", 0))

fsdp_native = DDPFullyShardedNativeStrategy(
            cpu_offload=CPUOffload(offload_params=True),
            cluster_environment=env,
            process_group_backend="gloo"
        )

ahh I see, and actually I also now see that var is already in your list above, so something else must be off in the chain of events that launches the underlying torch.distributed job from the lightning Trainer/Environment 🤔

I'm afraid commenting it out doesn't work -

RuntimeError: [enforce fail at /opt/conda/conda-bld/pytorch_1656352645774/work/third_party/gloo/gloo/context.cc:27] rank < size. 8 vs 1

can you use this as your

ClusterEnvironment

:

import logging
import os
from metaflow import current

from lightning.fabric.plugins.environments.cluster_environment import ClusterEnvironment

log = logging.getLogger(__name__)


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)

    def set_world_size(self, size: int) -> None:
        log.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)

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

    def local_rank(self) -> int:
        return 0

    def node_rank(self) -> int:
        return self.global_rank()

Yes - will do. Thanks @hallowed-glass-14538 !

Let me have a look in the morning. Has to be something really tiny that we need to change to fix all of this 😄

Okay thanks very much! thankyou

can you also put the stacktrace of the error for debugging

Control logs

Ok, based on the all the data I have some good news. 1. The bottleneck is not MF and connectivity, but the dataloader and how the sampler works. 2. The error comes from the distributed sampler and how PTL wraps the distributed sampler. 3. You are not using a distributed sampler and my hunch is that you will have to feed a DistributedSampler from top level to the DataLoaders (based on the FSDP tutorial)

Quick question about this - doesn't PyTorch Lightning wrap the

DataLoader

in a distributed sampler already?

Yes it does wrap it but somehow right now its passing the global rank to the sampler instead of the local rank. I don't know the exact code path its taking to pass the rank to the sampler but seems that ever strategy has a way to set the sampler's kwargs. Since DDP worked without even modifying anything else i assume its passing the correct rank to the sampler. I think you can try to instantiate the sampler yourself and pass it to the dataloader for FSDP or try going down the rabbithole of figuring out what kwargs are set via FSDP .

Got it - good to know. Thanks @hallowed-glass-14538. Will definitely investigate that. Another question - is there anything that can be done to use

nccl

as the process group backend? I don't think

gloo

is recommended for communication between GPUs, correct me if I'm wrong?

yes yes. nccl is way better but when we were developing this we noticed aws batch was not playing very nice with nccl. Currently @parallel only supports batch so we haven't researched or explored on paths to support nccl.

Update from my side -

fsdp

wasn't working unless process group is using

nccl

backend. But the

DistributedSampler

worked I believe because it got me out of that

rank

issue I was experiencing. This is the new error I was getting:

RuntimeError: no support for _allgather_base in Gloo process group

I did some research and currently looking into whether disabling IOMMU will help get NCCL working with AWS Batch.

Well done riley on NCCL support !

Oh interesting! Thanks so much. I will give that a try and report back. I could upgrade PTL but upgrading PyTorch beyond

1.12

in my

@conda_base

seems to result in the training session not being able to find cuda installed. So I just kept it at

1.12

.

Related reading : https://discuss.pytorch.org/t/number-of-parameters-in-model-summary-is-0-for-my-custom-model/97602/6

Ahh maybe. I'm downgrading to PL 1.7.7 to see if that helps - https://github.com/Lightning-AI/lightning/discussions/16402 Not a huge deal either way. Still want to figure out why CPU Offload results in significant hanging...

And by “not playing nice” I mean it’s not able to find the attributes seeked by summary function

Gotcha. Makes total sense.