all_gather with ddp is possibly corrupting data?
🐛 Bug
I'm working on a segmentation model using huge images from the current Kaggle HUBMAP competition. These images are typically something like (30000, 40000). For training, the approach is to take a tile that is large enough to capture the features of interest (e.g. 2048x2048) then resize that tile to something smaller, say (512x512),
For validation, things get a little crazy, since I want to calculate the dice score on the full image, not just tiles, since many tiles won't have a mask present:

In validation_step I infer on the 512x512 tiles to get y_pred, and also capture other details like the image id (label encoded) and the coords from which that tile came from
In validation_epoch_end I gather all the y_preds, resize to the original tile size (2048x2048) then reconstruct the full-size mask (30000, 40000) in a numpy array using the coords captured above
I then calculate the dice score against the ground truth mask

Using a single GPU I'm able to reconstruct a mask fine, so I know that the tiling method is working:

However when I switch to ddp, things get weird (images made from global rank 0 & 1):

I have no idea where that noise is coming from. It seems like if we take the good parts from both images, I could get to the result from the top image, but shouldn't all_gather take care of that?
Please reproduce using the BoringModel

To Reproduce
Use following BoringModel and post here. Can we use BoringModel for DDP bugs?
My validation code:
def validation_step(self, batch, batch_idx):
    img = batch["image"].float()
    mask = batch["mask"].float()
    out = self.forward(img)
    loss = self.loss_fn(out, mask)

    metric = dice_score(torch.squeeze(out), torch.squeeze(mask))

    return {
        "val_loss": loss,
        "metric": metric,
        "y_pred": out.sigmoid(),
        "coords": batch["coords"],
        "img_num": batch["img_num"],
    }

def validation_epoch_end(self, outputs):

    outputs = self.all_gather(outputs)

    loss_val = torch.stack([x["val_loss"] for x in outputs]).mean()
    metric = torch.stack([x["metric"] for x in outputs]).mean()

    coords = torch.cat([x["coords"] for x in outputs], dim=1)
    img_num = torch.cat([x["img_num"] for x in outputs], dim=1)
    y_pred = torch.cat([x["y_pred"] for x in outputs], dim=1)

    # Concat the world_size dim after all_gather
    coords = torch.cat([x for x in coords]).cpu().numpy()
    img_num = torch.cat([x for x in img_num]).cpu().numpy()
    y_pred = torch.cat([x for x in y_pred]).cpu().float()
    img_name = [num2img[x] for x in img_num]

    tile_size = self.hparams.img_size
    pred_masks = {}

    # Make an empty numpy array for each image at full resolution
    for img in list(set(img_name)):
        pred_masks[img] = np.zeros(img_dims[img], dtype=np.uint8)

    # Resize the tiles back to original size (512 > 2048), apply threshold and put
    # predicted tile into the full res numpy array
    for tile_pred, coord, name in zip(y_pred, coords, img_name):
        x1, x2, y1, y2 = coord
        y_pred_mask = (
            torch.nn.functional.interpolate(tile_pred[(None,)], size=tile_size)
            > 0.5
        )
        pred_masks[name][x1:x2, y1:y2] = y_pred_mask.squeeze().numpy()


    # Compare against ground truth mask
    scores = []
    for img, prediction in pred_masks.items():
        target = rle_decode(self.df.loc[img, "encoding"], img_dims[img])
        scores.append(dice_score_numpy(prediction, target))

        # Log images
        prediction = cv2.resize(prediction, (512, 512))
        self.log("mask", wandb.Image(prediction, caption=img))
        plt.imshow(prediction)
        plt.axis("off")
        plt.savefig(
            f"png/{img}_mask_qc_{self.global_rank}_ddp.png", bbox_inches="tight"
        )
        plt.close()

    dice = torch.tensor(np.mean(scores)).type_as(metric)

    self.log("metric_micro", metric)

    self.log_dict(
        {"loss/valid": loss_val, "metric": dice}, prog_bar=True, sync_dist=False,
    )
Expected behavior
Mask should look the same as the single GPU version
Environment
Note: Bugs with code are solved faster ! Colab Notebook should be made public !


IDE: Please, use our python bug_report_model.py template.


Colab Notebook: Please copy and paste the output from our environment collection script (or fill out the checklist below manually).


You can get the script and run it with:
wget https://raw.githubusercontent.com/PyTorchLightning/pytorch-lightning/master/tests/collect_env_details.py
# For security purposes, please check the contents of collect_env_details.py before running it.
python collect_env_details.py


CUDA:
- GPU:
- GeForce RTX 2080 Ti
- GeForce RTX 2080 Ti
- available:         True
- version:           10.2
Packages:
- numpy:             1.20.1
- pyTorch_debug:     False
- pyTorch_version:   1.8.0
- pytorch-lightning: 1.2.4
- tqdm:              4.58.0
System:
- OS:                Linux
- architecture:
- 64bit
- ELF
- processor:         x86_64
- python:            3.8.8
- version:           #74-Ubuntu SMP Wed Jan 27 22:54:38 UTC 2021

Additional context