Hello team, I'm currently running a PyTorch Light...
# ask-metaflow
a
Hello team, I'm currently running a PyTorch Lightning distributed training job using the
@pytorch_parallel
decorator. One thing I noticed is that Metaflow seems to be suppressing the logs from the training session so I don't actually know what the training progress is. Any advice? I have a reproducible example, if needed, as well. Thanks in advance!
v
hmm, task-level logs should be available. Can you share the example?
a
Yup here's the reproducible example: https://github.com/rileyhun/llm_finetuning_metaflow/tree/main/pytorch-deepspeed And the below was the last output I got from the logs.
👍 1
v
thanks, will take a look
a
Appreciate it Ville!
u
Looks like the training actually happens here:
Copy code
with tqdm_logging_redirect():
            model.train(
                train_df=self.train_df, 
                eval_df=self.val_df, 
                args=arguments
            )
Do you know if
tqdm_logging_redirect
is writing logs to stdout? Or is it redirecting the logs to some log file?
👍 1
a
It's generally supposed to redirect to stdout. I have used that in the past to view the logs because I know Metaflow has some trouble w/ tqdm.
u
Can you try this? I ran the flow mentioned in the solution and it worked fine for me.
a
Cool - thanks Shri. Let me try it. Thanks!
Oh wait - I think that's what I'm currently doing. It works for training on a single GPU, but not for distributed training it seems.
c
Hey Riley! Have you tried setting the sync_dist_group arg in the logging calls in your lightning module?
a
Hi Eddie! I have not, but I'll definitely look into that.
c
Not sure that is the ticket here, but I suspect there is some modification in the Lightning Module definition that can make this work. If not, my next approach would be to try the non-tqdm progress bars in lightning, and if that doesn't work to use some other way to do the logging to stdout inside the callbacks in training_step and validation_step.
a
Great pointers, thanks Eddie! I'll give the
sync_dist_group
arg a try first. Does this look correct?
Copy code
def training_step(self, batch, batch_size):
        """ training step """
        input_ids = batch["source_text_input_ids"]
        attention_mask = batch["source_text_attention_mask"]
        labels = batch["labels"]
        labels_attention_mask = batch["labels_attention_mask"]

        loss, outputs = self(
            input_ids=input_ids,
            attention_mask=attention_mask,
            decoder_attention_mask=labels_attention_mask,
            labels=labels,
        )
        
        group = dist.init_process_group("nccl", rank=self.global_rank, world_size=self.world_size)

        self.log(
            "train_loss", loss, prog_bar=True, logger=True, on_epoch=True, on_step=True, sync_dist_group=group
        )
        return loss
c
yep I think so
thankyou 1
a
Another thing I noticed unrelated to the logging as well. If you specify > 6 nodes, it starts to hang due to the following:
Copy code
INFO:torch.distributed.distributed_c10d:Waiting in store based barrier to initialize process group for rank: 2, key: store_based_barrier_key:1 (world_size=8, worker_count=6, timeout=0:30:00)
I have only ever been able to get up to 6 nodes, otherwise I'll end up waiting idefinitely for all the nodes to initialize
🤔 1
No luck with
sync_dist
unfortunately
c
@acoustic-van-30942 is it possible to share the data splits you are using in the
/data
folder? I'd like to reproduce the issue before lobbing more suggestions at you and want to minimize moving parts.
a
Absolutely I'll upload the data shortly.
Hi @crooked-jordan-29960 - I ended up just uploading the data to Github to make things easier - https://github.com/rileyhun/llm_finetuning_metaflow/tree/main/pytorch-deepspeed
thankyou 1
For benchmarking purposes, I used a smaller model
t5-small
just to make sure everything is ok. The t5-small is only about 60M parameters. On 6 nodes (each node being a
g5.4xlarge
), it took about 30 minutes using DDP. Conversely, it took 60 minutes using DeepSpeed. I'm not an expert, but I guess CPUOffload makes things slightly slower, although the docs also indicate that
DeepSpeedCPUAdam
is a faster optimizer? In summary there are three issues: 1. The logs (progress bar especially) are being suppressed by either PyTorch Lightning Distributed Multi-Node or Metaflow, so it's impossible to track training progress 2. For DeepSpeed Strategy, choosing number of nodes greater than 6 results in a long wait for store based barrier to initialize process group... 3. DeepSpeed with CPU Offload takes 20 minutes longer than DDP on the same hardware, for the same model, and batch size.
Okay I managed to resolve [2] I think by adding this environment variable
"NCCL_IB_DISABLE": "1"
and then removing
"NCCL_P2P_DISABLE": "1"
Hi Eddie, Just wondering if you had any luck getting the progress bar to display?
c
Hey Riley, I haven't had any luck yet with the progress bar + distrib training but plan to do some more debugging later today to see if I can understand the issue better. Have you tried any other approaches since last week?
a
Thanks Eddie! I tried a different progress bar other than tqdm but no luck with that
👍 1