DDP and BackboneFinetuning: model weights get out of sync when unfreezing layers for training

Bug description

When model training using DDP and pl.callbacks.BackboneFinetuning, it seems that model weights start to get out of sync across the processes after the backbone is unfrozen. Prior to unfreezing, model weights stay in sync across processes as expected.

I discovered this issue when trying to adopt DDP. I saw that on rank 0 process, validation loss trended downward while training, while on rank > 1 processes validation loss increased steadily. This led to the suspicion that model weights were different across nodes, which was confirmed by printing out the hash of model weights on the different processes on each epoch.

What version are you seeing the problem on?

v2.4

How to reproduce the bug

The below example is programmed to check that model weights are in sync after every epoch. It fails the assertion after epoch 3 (unfreeze_backbone_at_epoch).

import hashlib
import pytorch_lightning as pl
import torch
from torch import nn
import torch.distributed as dist
from torch.utils.data import DataLoader, Dataset


# 1. Define a simple dataset
class RandomDataset(Dataset):
    def __init__(self, size, length):
        self.len = length
        self.data = torch.randn(length, size)

    def __getitem__(self, index):
        return self.data[index]

    def __len__(self):
        return self.len


# 2. Define a LightningModule
class SimpleModel(pl.LightningModule):
    def __init__(self):
        super().__init__()
        self.backbone = nn.Linear(32, 16)
        self.layer = nn.Linear(16, 2)

    def forward(self, x):
        x = torch.relu(self.backbone(x))
        x = self.layer(x)
        return x

    def training_step(self, batch, batch_idx):
        x = batch
        y_hat = self(x)
        loss = torch.nn.functional.mse_loss(y_hat, torch.ones_like(y_hat))
        return loss

    def configure_optimizers(self):
        return torch.optim.SGD(
            filter(lambda p: p.requires_grad, self.parameters()), lr=0.1
        )

    def on_train_epoch_end(self):
        # Compute hash of model weights and checks if they're equal across processes.
        hasher = hashlib.sha256()
        for param in self.parameters():
            hasher.update(param.data.cpu().numpy().tobytes())
        param_hash = hasher.hexdigest()
        all_param_hashes = [None] * dist.get_world_size()
        dist.all_gather_object(all_param_hashes, param_hash)
        if self.trainer.is_global_zero:
            assert len(set(all_param_hashes)) == 1, "Model weights not in sync :("
            print("Model weights in sync!")


# 3. Create data loaders
pl.seed_everything(0)
train_loader = DataLoader(RandomDataset(32, 64), batch_size=2)

# 4. Initialize the model and trainer
model = SimpleModel()
trainer = pl.Trainer(
    accelerator="cpu",
    strategy="ddp",
    devices=2,
    callbacks=[
        pl.callbacks.BackboneFinetuning(unfreeze_backbone_at_epoch=3, verbose=True)
    ],
)

# 5. Train the model
trainer.fit(model, train_loader)

Output:

Epoch 0: 100%|████████████████████████████| 16/16 [00:00<00:00, 163.83it/s, loss=0.295, v_num=22]
Model weights in sync!
Epoch 1: 100%|███████████████████████████| 16/16 [00:00<00:00, 269.84it/s, loss=0.0667, v_num=22]
Model weights in sync!
Epoch 2: 100%|███████████████████████████| 16/16 [00:00<00:00, 257.68it/s, loss=0.0361, v_num=22]
Model weights in sync!
Current lr: 0.1, Backbone lr: 0.01
Current lr: 0.1, Backbone lr: 0.01
Epoch 3: 100%|███████████████████████████| 16/16 [00:00<00:00, 244.17it/s, loss=0.0243, v_num=22]Current lr: 0.1, Backbone lr: 0.02
[rank0]: Traceback (most recent call last):
...
[rank0]:   File "/home/ksikka/lightning-pose/example2.py", line 58, in _assert_model_weights_in_sync
[rank0]:     assert len(set(all_param_hashes)) == 1, "Model weights not in sync :("
[rank0]: AssertionError: Model weights not in sync :(

Error messages and logs

No warning or error. Validation loss with sync_dist=True increases after unfreezing, while with sync_dist=False, it decreases although at a lower rate than single process.

Environment

I originally noticed the issue in a multi-GPU linux environment in lightning studio, but I reproduced with the example code above on the following environment.

More info

No response

cc @justusschock @lantiga

Start by running the provided two-device DDP reproduction with BackboneFinetuning and inspect the unfreezing transition at epoch 3. Trace the DDP strategy and BackboneFinetuning callback behavior, then add a regression check showing that model weights remain synchronized after unfreezing and run the relevant test suite.