From 617f4eafe008fd918861a9c8ca1002e851d8887b Mon Sep 17 00:00:00 2001 From: Saud Kamran Date: Mon, 13 Jul 2026 10:30:54 +0500 Subject: [PATCH 1/2] fix(loss): use autograd-aware all_reduce in BarlowTwinsLoss The raw torch.distributed.all_reduce overwrote the cross-correlation matrix in place and was not registered in the autograd graph. Under DDP gradient averaging this scaled the backbone gradient down by 1/world_size relative to single-GPU Barlow Twins on the global batch. Swap it for the autograd-aware torch.distributed.nn.all_reduce, whose backward is itself an all_reduce, so the gradient is scaled back up by world_size. The forward value is unchanged. Fixes #1977. Same fix as SIGReg in #1923. --- lightly/loss/barlow_twins_loss.py | 7 ++- tests/loss/test_barlow_twins_loss.py | 74 ++++++++++++++++++++++++++++ 2 files changed, 79 insertions(+), 2 deletions(-) diff --git a/lightly/loss/barlow_twins_loss.py b/lightly/loss/barlow_twins_loss.py index cca09fcd9..e7aff40e5 100644 --- a/lightly/loss/barlow_twins_loss.py +++ b/lightly/loss/barlow_twins_loss.py @@ -4,6 +4,7 @@ import torch.distributed as dist import torch.nn.functional as F from torch import Tensor +from torch.distributed import nn as dist_nn class BarlowTwinsLoss(torch.nn.Module): @@ -76,8 +77,10 @@ def forward(self, z_a: torch.Tensor, z_b: torch.Tensor) -> torch.Tensor: if self.gather_distributed and dist.is_initialized(): world_size = dist.get_world_size() if world_size > 1: - c = c / world_size - dist.all_reduce(c) + # Use the autograd-aware torch.distributed.nn.all_reduce so the + # backward pass scales gradients by world_size, matching + # single-GPU behaviour. ref #1977 + c = dist_nn.all_reduce(c / world_size) invariance_loss = torch.diagonal(c).add_(-1).pow_(2).sum() redundancy_reduction_loss = _off_diagonal(c).pow_(2).sum() diff --git a/tests/loss/test_barlow_twins_loss.py b/tests/loss/test_barlow_twins_loss.py index 99c3d503f..5d6b4848d 100644 --- a/tests/loss/test_barlow_twins_loss.py +++ b/tests/loss/test_barlow_twins_loss.py @@ -4,6 +4,7 @@ from pytest_mock import MockerFixture from torch import Tensor from torch import distributed as dist +from torch.distributed import nn as dist_nn from torch.nn import Module from lightly.loss.barlow_twins_loss import BarlowTwinsLoss @@ -92,3 +93,76 @@ def test__loss_is_affine_invariant(self) -> None: # Loss should be invariant to affine transformations. assert torch.allclose(loss(x, x), loss(x, 2 * x + 4)) + + def test__gather_distributed_forward_matches_non_distributed( + self, mocker: MockerFixture + ) -> None: + """Distributed forward equals non-distributed forward on the global batch. + + Simulates world_size=2 with identical data on both ranks. The global + batch is two copies of the local batch, so the cross-correlation matrix + after the all_reduce equals the non-distributed matrix on the global batch. + """ + world_size = 2 + torch.manual_seed(0) + z_a_local = torch.randn(16, 64) + z_b_local = torch.randn(16, 64) + z_a_global = torch.cat([z_a_local, z_a_local], dim=0) + z_b_global = torch.cat([z_b_local, z_b_local], dim=0) + + loss_truth = BarlowTwinsLoss(gather_distributed=False)(z_a_global, z_b_global) + + mocker.patch.object(dist, "is_initialized", return_value=True) + mocker.patch.object(dist, "get_world_size", return_value=world_size) + mocker.patch.object( + dist_nn, + "all_reduce", + side_effect=lambda tensor, *args, **kwargs: tensor * world_size, + ) + loss_dist = BarlowTwinsLoss(gather_distributed=True)(z_a_local, z_b_local) + + assert torch.allclose(loss_dist, loss_truth, atol=1e-5) + + def test__gather_distributed_gradient_matches_non_distributed( + self, mocker: MockerFixture + ) -> None: + """Gradient under gather_distributed=True equals gradient on the global batch. + + Verifies the fix for the autograd bug where raw dist.all_reduce left the + backward pass unaware of the cross-rank reduction, producing gradients + scaled down by 1/world_size. + """ + world_size = 2 + torch.manual_seed(0) + weight = torch.randn(64, 64) + z_a_local = torch.randn(16, 64) + z_b_local = torch.randn(16, 64) + z_a_global = torch.cat([z_a_local, z_a_local], dim=0) + z_b_global = torch.cat([z_b_local, z_b_local], dim=0) + + def grad_through_weight( + loss_fn: BarlowTwinsLoss, + z_a: torch.Tensor, + z_b: torch.Tensor, + ) -> torch.Tensor: + w = weight.clone().requires_grad_(True) + loss_fn(z_a @ w, z_b @ w).backward() + assert w.grad is not None + return w.grad.clone() + + truth = grad_through_weight( + BarlowTwinsLoss(gather_distributed=False), z_a_global, z_b_global + ) + + mocker.patch.object(dist, "is_initialized", return_value=True) + mocker.patch.object(dist, "get_world_size", return_value=world_size) + mocker.patch.object( + dist_nn, + "all_reduce", + side_effect=lambda tensor, *args, **kwargs: tensor * world_size, + ) + dist_grad = grad_through_weight( + BarlowTwinsLoss(gather_distributed=True), z_a_local, z_b_local + ) + + assert torch.allclose(dist_grad, truth, atol=1e-5) From 4bf0050eb74dee5212faece0ed26e4e397798fb0 Mon Sep 17 00:00:00 2001 From: Saud Kamran Date: Tue, 14 Jul 2026 19:15:35 +0500 Subject: [PATCH 2/2] test(loss): remove mock-based gather_distributed tests Mocking all_reduce only tested the mock, not the real collective, and would break on any change to the reduction. Real multi-rank coverage is tracked in #1982. --- tests/loss/test_barlow_twins_loss.py | 74 ---------------------------- 1 file changed, 74 deletions(-) diff --git a/tests/loss/test_barlow_twins_loss.py b/tests/loss/test_barlow_twins_loss.py index 5d6b4848d..99c3d503f 100644 --- a/tests/loss/test_barlow_twins_loss.py +++ b/tests/loss/test_barlow_twins_loss.py @@ -4,7 +4,6 @@ from pytest_mock import MockerFixture from torch import Tensor from torch import distributed as dist -from torch.distributed import nn as dist_nn from torch.nn import Module from lightly.loss.barlow_twins_loss import BarlowTwinsLoss @@ -93,76 +92,3 @@ def test__loss_is_affine_invariant(self) -> None: # Loss should be invariant to affine transformations. assert torch.allclose(loss(x, x), loss(x, 2 * x + 4)) - - def test__gather_distributed_forward_matches_non_distributed( - self, mocker: MockerFixture - ) -> None: - """Distributed forward equals non-distributed forward on the global batch. - - Simulates world_size=2 with identical data on both ranks. The global - batch is two copies of the local batch, so the cross-correlation matrix - after the all_reduce equals the non-distributed matrix on the global batch. - """ - world_size = 2 - torch.manual_seed(0) - z_a_local = torch.randn(16, 64) - z_b_local = torch.randn(16, 64) - z_a_global = torch.cat([z_a_local, z_a_local], dim=0) - z_b_global = torch.cat([z_b_local, z_b_local], dim=0) - - loss_truth = BarlowTwinsLoss(gather_distributed=False)(z_a_global, z_b_global) - - mocker.patch.object(dist, "is_initialized", return_value=True) - mocker.patch.object(dist, "get_world_size", return_value=world_size) - mocker.patch.object( - dist_nn, - "all_reduce", - side_effect=lambda tensor, *args, **kwargs: tensor * world_size, - ) - loss_dist = BarlowTwinsLoss(gather_distributed=True)(z_a_local, z_b_local) - - assert torch.allclose(loss_dist, loss_truth, atol=1e-5) - - def test__gather_distributed_gradient_matches_non_distributed( - self, mocker: MockerFixture - ) -> None: - """Gradient under gather_distributed=True equals gradient on the global batch. - - Verifies the fix for the autograd bug where raw dist.all_reduce left the - backward pass unaware of the cross-rank reduction, producing gradients - scaled down by 1/world_size. - """ - world_size = 2 - torch.manual_seed(0) - weight = torch.randn(64, 64) - z_a_local = torch.randn(16, 64) - z_b_local = torch.randn(16, 64) - z_a_global = torch.cat([z_a_local, z_a_local], dim=0) - z_b_global = torch.cat([z_b_local, z_b_local], dim=0) - - def grad_through_weight( - loss_fn: BarlowTwinsLoss, - z_a: torch.Tensor, - z_b: torch.Tensor, - ) -> torch.Tensor: - w = weight.clone().requires_grad_(True) - loss_fn(z_a @ w, z_b @ w).backward() - assert w.grad is not None - return w.grad.clone() - - truth = grad_through_weight( - BarlowTwinsLoss(gather_distributed=False), z_a_global, z_b_global - ) - - mocker.patch.object(dist, "is_initialized", return_value=True) - mocker.patch.object(dist, "get_world_size", return_value=world_size) - mocker.patch.object( - dist_nn, - "all_reduce", - side_effect=lambda tensor, *args, **kwargs: tensor * world_size, - ) - dist_grad = grad_through_weight( - BarlowTwinsLoss(gather_distributed=True), z_a_local, z_b_local - ) - - assert torch.allclose(dist_grad, truth, atol=1e-5)