Logging synchronization and confusion matrix log

[scy-helbling]

I run a distributed training on 2 GPUs and want to use auto-synchronized torchmetrics for logging as well as a confusion matrix.
Single-script minimal example:

import torch
import lightning as L
import time
from torchmetrics import ConfusionMatrix
from torchmetrics.classification import BinaryAccuracy


class ToyExample(L.LightningModule):
    def __init__(self, model):
        super().__init__()
        self.model = model
        # sync_on_compute=True is the default
        self.train_confmat = ConfusionMatrix(num_classes=2, task="binary", sync_on_compute=True)
        self.train_acc = BinaryAccuracy()

    def training_step(self, batch):
        out = self.model(batch)
        loss = out.sum()
        time.sleep(0.05)
        self.train_confmat.update(out[:, 0] > 0.5, out[:, 0] > 0.5)

        # Create different accuracies in the two different processes
        if self.trainer.global_rank == 0:
            self.train_acc.update(out[:, 0] * 0, out[:, 0] * 0 + 1)  # All wrong
        else:
            self.train_acc.update(out[:, 0] * 0, out[:, 0] * 0)  # All correct

        self.log("train/acc", self.train_acc, on_step=False, on_epoch=True)  # PROBLEM A
        return loss

    def validation_step(self, batch):
        loss = self.model(batch).sum()
        time.sleep(0.05)
        return loss

    def configure_optimizers(self):
        return torch.optim.Adam(self.model.parameters())

    def on_train_epoch_end(self):
        fig, ax = self.train_confmat.plot()  # PROBLEM B
        self.logger.experiment.add_figure("train/Confusion Matrix", fig, self.current_epoch)
        self.train_confmat.reset()


if __name__ == "__main__":
    model = torch.nn.Linear(32, 2)
    pl_module = ToyExample(model)
    train_dataloader = torch.utils.data.DataLoader(torch.randn(100, 32), batch_size=10)
    val_dataloader = torch.utils.data.DataLoader(torch.randn(50, 32), batch_size=10)
    trainer = L.Trainer(max_epochs=10, devices=2)
    trainer.fit(pl_module, train_dataloader, val_dataloader)

When one of the two problem lines (A or B) is commented out, the script runs fine. Howerver, when I use both options (problem lines A and B both uncommented), training halts at the end of the first epoch, presumeably due to a race condition because there are two different synchronization systems at work: one for the TorchMetric epoch-aggregation and one for the confusion matrix' .compute() (which I guess is called internally when I use .plot)

Versions and environment:

• Python 3.10
• lightning==2.6.0
• torch==2.6.0
• linux (EC2 instance with 4x NVIDIA T4)