fix(loss): use autograd-aware all_reduce in BarlowTwinsLoss#1980
Open
saud5150 wants to merge 3 commits into
Open
fix(loss): use autograd-aware all_reduce in BarlowTwinsLoss#1980saud5150 wants to merge 3 commits into
saud5150 wants to merge 3 commits into
Conversation
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 lightly-ai#1977. Same fix as SIGReg in lightly-ai#1923.
gabrielfruet
requested changes
Jul 13, 2026
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 lightly-ai#1982.
Contributor
Author
|
Kept the manual division — agreed, gloo not supporting |
gabrielfruet
approved these changes
Jul 15, 2026
gabrielfruet
left a comment
gabrielfruet
left a comment
Contributor
There was a problem hiding this comment.
LGTM!! Thank you for the nice work
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
closes #1977
Description
BarlowTwinsLoss(gather_distributed=True)reduced the cross-correlation matrixwith the raw
torch.distributed.all_reduce, which isn't autograd-aware. Theforward value was correct, but the backward pass scaled the gradient down by
1/world_sizeunder DDP — silently. Swapped it fortorch.distributed.nn.all_reduce, whose backward is itself an all_reduce, so thegradient is correct and the forward value is unchanged. Same fix #1923 landed for
SIGReg; Barlow Twins was the last loss with the pattern.
Tests
Added two tests (world_size=2, two identical ranks): forward matches the
non-distributed loss on the global batch, and the gradient matches single-GPU
Barlow Twins on the global batch. The gradient test fails on master and passes
with the fix.
make static-checksand the loss suite pass.Documentation
.rstfiles).Internal implementation only — no API, docstring, or
.rstchanges.Implications / comments / further issues
None. Public API, forward loss value, and single-GPU behavior unchanged; only the
multi-GPU gradient is corrected.