How to gather all validation_step_outputs at validation_epoch_end and run in rank_zero properly without deadlock?

[allanchan339]

Hi,

I am trying to gather all the output and label pairs in the validation epoch end to run a simple validation process.
First, all validation data will be separated into different devices (controlled by DDP).
The validation step is simple as follows:

def validation_step(self, batch, batch_idx):
        # use valid_metric
        feat, label = batch
        output = self.model(feat)
        del batch, feat
        torch.cuda.empty_cache()


        for i, metric in enumerate(self.metrics):
            metric.update(output[:,i], label[:,i])

        return {"label": label, "output": output}

To calculate an accurate metrics, I need to gather all outputs to a single device and log on rank 0 only as follows:

 def validation_epoch_end(self, validation_step_outputs):
        def _valid_epoch_end(validation_step_outputs):
            enable_Flag = False
            # AP refers to average_precision_score in torchmetrics
            all_labels = list(map(itemgetter('label'), validation_step_outputs))
            all_labels = torch.cat(all_labels).cpu().detach().numpy()
            all_outputs = list(map(itemgetter('output'), validation_step_outputs))
            all_outputs = torch.cat(all_outputs).cpu().detach().numpy()

            AP = []
            for i in range(1, 17+1):
                AP.append(np.nan_to_num(average_precision_score(all_labels
                                                                [:, i], all_outputs[:,  i])))
            mAP = np.mean(AP)
            tmp = EVENT_DICTIONARY_V2
            tmp = tmp.copy()
            for k, v in tmp.items():
                #dict["kick-off"] = AP[0]
                tmp[k] = AP[v]
            
            self.log('Valid/mAP', mAP, logger=True, prog_bar=True,
                     rank_zero_only=True if self.args.strategy != 'dp' and enable_Flag else False)

            label_cls = (list(EVENT_DICTIONARY_V2.keys()))
            zip_iterator = zip(label_cls, AP)
            AP_dictionary = dict(zip_iterator)
            self.log('Valid/AP', AP_dictionary, logger=True, prog_bar=False,
                     rank_zero_only=True if self.args.strategy != 'dp' and enable_Flag else False)

        _valid_epoch_end(validation_step_outputs)

The program works fine. However, the metric result is always different from DP. Therefore, I need to gather all data first.
However, the program doesn't even pass the validation step as the deadlock. Any good idea to modify this program so that I can gather

[rohitgr7]

with DP, you need to configure the metric computation differently
https://pytorch-lightning.readthedocs.io/en/stable/common/lightning_module.html#validating-with-dataparallel

with DDP to gather all outputs you can use self.all_gather under the condition of rank_zero

if self.trainer.is_global_zero:
    all_val_outs = self.all_gather(...)

[ecolss]

Same issue here, is_global_zero and all_gather make the code hang forever.

[furkanbiten]

Same issue here, have a metric that needs to be computed across entire dataset. Branching with if self.trainer.is_global_zero leads to deadlock in the other ranks. In my case EarlyStopping on that process complains that it's tracking metric is not available. This also begs the questions why early stopping is triggered at all in each rank - shouldn't it be a global check?

Same happened for me in ModelCheckpoint.

[vgthengane]

anybody has solved this yet?

[Vishwas-Venkatachalapathy]

I am running into the same issue. Let me know if there is a fix for it

[akihironitta]

Just to clarify, the collective all_gather needs to be called in all processes but not only in the main process, otherwise it leads to hanging.

[allanchan339]

The solution is as follows:
Instead of using

if self.trainer.is_global_zero:
    all_val_outs = self.all_gather(...)

The function above will hang as card 0 is trying to communicate other cards

You should write code in this way

all_val_out = self.all_gather(...)

if self.trainer.is_global_zero:
    # merge output and process

self.trainer.strategy.barrier() #to let other cards to wait

[kenanseyidov]

Can you please provide documentation which explains how data moves from card to card? Right now I struggle to understand what all_gather, is_global_zero and barrier() achieve.

[akihironitta]

For collective operations, such as all gather/reduce, these blog posts helped me understand how it works.

• https://towardsdatascience.com/visual-intuition-on-ring-allreduce-for-distributed-deep-learning-d1f34b4911da
• https://tech.preferred.jp/en/blog/technologies-behind-distributed-deep-learning-allreduce

is_global_zero is just a bool value that represents whether the process is running on the 0th device of all devices or not:

• https://lightning.ai/docs/pytorch/2.0.1/common/trainer.html#is-global-zero

barrier() makes the running process wait until other processes get to the same code line. Some references I find helpful are:

• https://pytorch.org/docs/2.0/distributed.html#torch.distributed.barrier
• https://discuss.pytorch.org/t/torch-distributed-barrier-used-in-multi-node-distributed-data-parallel-training/89711
• Using torch.distributed.barrier() makes the whole code hang pytorch/pytorch#54059