Question: sync multi-node distributed processes

I'm using BertBing example to train BERT on a multi-node setup, I've also added a custom validation step that runs after X steps and not at the end of the epoch (as in the example). I'm using the dataset provider (by nvidia) in order to run in parallel on all GPUs.
The loss is synced to rank = 0 but once the training loop ends the whole process freezes.

What is the proper way of syncing the tensors in distributed mode in Deepspeed?
I've tried dist.reduce, dist.all_reduce, adding dist.barrier. Gets stuck every time.
Thanks!

for b_i, batch_index in enumerate(tqdm(dataset_iterator, smoothing=1)):
    batch = validation_dataset_provider.get_batch(batch_index)
    batch = tuple(t.to(args.device) for t in batch)  # Move to GPU
    loss = model.network(batch)
    dist.reduce(loss, 0)
    curr_eval_loss += (loss / dist.get_world_size()).mean().item()

Start with the BertBing example and the custom validation loop shown in the issue, then trace how the dataset provider and distributed ranks enter and leave validation. Reproduce the multi-node freeze while comparing the attempted reduce, all_reduce, and barrier calls. Done means the validation loop completes on every rank and the loss is synchronized without hanging.