From a4524046d96652bfeb6aa57a893ac0f53f576388 Mon Sep 17 00:00:00 2001 From: Sezgin Er Date: Tue, 15 Sep 2026 11:43:31 +0200 Subject: [PATCH] Use the actual global token count in distributed Sinkhorn The Sinkhorn-Knopp denominator assumed every rank holds the same number of rows (local_rows * world_size). iBOT block masking gives ranks different masked-token counts, so the prototype targets were biased. Sum the real local row counts across ranks instead; equal-count inputs (e.g. CLS logits) are bit-identical. A rank may now hold zero rows; a globally empty batch raises. Adds single-process and two-rank gloo tests: equal-count parity with the historical implementation, unequal-count parity with a plain Sinkhorn over all tokens, and empty-rank handling. --- dino-pretraining/mr_dino/objective.py | 28 +++++++-- dino-pretraining/tests/test_mr_dino.py | 83 ++++++++++++++++++++++++++ 2 files changed, 106 insertions(+), 5 deletions(-) diff --git a/dino-pretraining/mr_dino/objective.py b/dino-pretraining/mr_dino/objective.py index 0d6e82c..917cd30 100644 --- a/dino-pretraining/mr_dino/objective.py +++ b/dino-pretraining/mr_dino/objective.py @@ -50,16 +50,34 @@ def init_weights(self) -> None: @torch.no_grad() def distributed_sinkhorn(logits: Tensor, temperature: float, iterations: int = 3) -> Tensor: - """Official DINOv3 Sinkhorn-Knopp assignment with WORLD as the group.""" + """Official DINOv3 Sinkhorn-Knopp assignment with WORLD as the group. + + The global batch size is the SUM of the actual local row counts across all + ranks, not ``local_rows * world_size``. iBOT block masking gives every rank + a different number of masked tokens, so the equal-count shortcut biased the + prototype targets; with equal counts (e.g. CLS logits) both are identical. + A rank may hold zero rows; a globally empty batch raises on every rank. + """ + if logits.ndim != 2 or logits.shape[1] < 1: + raise ValueError("logits must have shape [local_tokens, prototypes]") + distributed = dist.is_initialized() + global_count = torch.tensor(logits.shape[0], dtype=torch.int64, device=logits.device) + if distributed: + dist.all_reduce(global_count, op=dist.ReduceOp.SUM) + global_batch = int(global_count.item()) + if global_batch == 0: + raise ValueError("Sinkhorn requires at least one token globally") # A global additive constant cancels during Sinkhorn normalization. Remove # it before exponentiation so a larger masked-token batch cannot overflow # merely because it is more likely to contain an extreme teacher logit. - max_logit = logits.detach().float().max() - if dist.is_initialized(): + max_logit = ( + logits.detach().float().max() + if logits.shape[0] + else torch.full((), -torch.inf, dtype=torch.float32, device=logits.device) + ) + if distributed: dist.all_reduce(max_logit, op=dist.ReduceOp.MAX) q = torch.exp((logits.float() - max_logit) / temperature).t() - world = dist.get_world_size() if dist.is_initialized() else 1 - global_batch = q.shape[1] * world prototypes = q.shape[0] total = q.sum() if dist.is_initialized(): diff --git a/dino-pretraining/tests/test_mr_dino.py b/dino-pretraining/tests/test_mr_dino.py index c44c341..0d91cf8 100644 --- a/dino-pretraining/tests/test_mr_dino.py +++ b/dino-pretraining/tests/test_mr_dino.py @@ -198,3 +198,86 @@ def test_synthetic_forward_backward_and_checkpoint(tmp_path): learner.load_state_dict(saved["model"]) optimizer.load_state_dict(saved["optimizer"]) assert saved["step"] == 1 + + + +# --------------------------------------------------------------------------- Sinkhorn count +def _reference_sinkhorn_single_process(logits, temperature, iterations=3): + """Plain (non-distributed) Sinkhorn over ALL tokens: the definition the distributed version must match.""" + q = torch.exp((logits.float() - logits.float().max()) / temperature).t() + global_batch, prototypes = q.shape[1], q.shape[0] + q /= q.sum().clamp_min(1e-12) + for _ in range(iterations): + q /= q.sum(dim=1, keepdim=True).clamp_min(1e-12) + q /= prototypes + q /= q.sum(dim=0, keepdim=True).clamp_min(1e-12) + q /= global_batch + return (q * global_batch).t() + + +def _legacy_sinkhorn(logits, temperature, world, iterations=3): + """The historical implementation (local rows x world size denominator) for equal-count parity.""" + import torch.distributed as dist + max_logit = logits.detach().float().max() + if dist.is_initialized(): + dist.all_reduce(max_logit, op=dist.ReduceOp.MAX) + q = torch.exp((logits.float() - max_logit) / temperature).t() + global_batch = q.shape[1] * world + prototypes = q.shape[0] + total = q.sum() + if dist.is_initialized(): + dist.all_reduce(total) + q /= total.clamp_min(1e-12) + for _ in range(iterations): + rows = q.sum(dim=1, keepdim=True) + if dist.is_initialized(): + dist.all_reduce(rows) + q /= rows.clamp_min(1e-12) + q /= prototypes + q /= q.sum(dim=0, keepdim=True).clamp_min(1e-12) + q /= global_batch + return (q * global_batch).t() + + +def test_sinkhorn_single_process_matches_legacy_and_rejects_empty(): + from mr_dino.objective import distributed_sinkhorn + torch.manual_seed(0) + logits = torch.randn(23, 7) + torch.testing.assert_close(distributed_sinkhorn(logits, 0.07), _legacy_sinkhorn(logits, 0.07, 1), atol=0, rtol=0) + with pytest.raises(ValueError): + distributed_sinkhorn(torch.zeros(0, 7), 0.07) + + +def _sinkhorn_worker(rank, rendezvous): + import torch.distributed as dist + from datetime import timedelta + from mr_dino.objective import distributed_sinkhorn + torch.set_num_threads(1) + dist.init_process_group("gloo", init_method="file://" + rendezvous, rank=rank, world_size=2, + timeout=timedelta(seconds=60)) + try: + torch.manual_seed(11) + logits = torch.randn(30, 9) * 3 + # Equal local counts: corrected == legacy bit for bit. + local = logits[rank * 15:(rank + 1) * 15] + torch.testing.assert_close(distributed_sinkhorn(local, 0.07), _legacy_sinkhorn(local, 0.07, 2), atol=0, rtol=0) + # Unequal local counts (masked tokens): corrected == single-process over all tokens; legacy differs. + for split in (7, 0, 30): + local = logits[:split] if rank == 0 else logits[split:] + got = distributed_sinkhorn(local, 0.07) + expected = _reference_sinkhorn_single_process(logits, 0.07) + expected_local = expected[:split] if rank == 0 else expected[split:] + assert got.shape == local.shape and got.dtype == torch.float32 + torch.testing.assert_close(got, expected_local, atol=1e-6, rtol=1e-5) + if split == 7: + legacy = _legacy_sinkhorn(local, 0.07, 2) + assert not torch.allclose(legacy, expected_local, atol=1e-6, rtol=1e-5) + # Globally empty raises on both ranks. + with pytest.raises(ValueError): + distributed_sinkhorn(torch.zeros(0, 9), 0.07) + finally: + dist.destroy_process_group() + + +def test_sinkhorn_uses_actual_global_token_count_across_ranks(tmp_path): + torch.multiprocessing.spawn(_sinkhorn_worker, args=(str(tmp_path / "rendezvous"),), nprocs=2, join=True)