Repository navigation
Conversation
|
/review |
liopeer
left a comment
There was a problem hiding this comment.
It looks correct, I would mainly try to re-use the existing SwAV code though. I will have again a detailed look again later.
DINOv2 and DINOv3 replace the mean centering of the teacher output with the Sinkhorn-Knopp centering from SwAV. DINOLoss now accepts center_mode="sinkhorn_knopp", which normalizes the teacher logits so that every prototype receives the same total weight across the batch instead of subtracting a running center. Following the reference implementation, the normalization is applied jointly over all teacher views and no center is tracked. The center buffer stays registered so that the state dict does not depend on the center mode. The number of samples is reduced across processes rather than derived from the world size, so that the shared helper also works for the iBOT loss where the number of masked tokens differs between processes.
The DINOv3 paper applies Sinkhorn-Knopp centering to both the DINO and the iBOT objective, and the reference training code asserts that centering is set to sinkhorn_knopp. IBOTPatchLoss now accepts the same center_mode as DINOLoss. Teacher normalization and the center update move into helpers so that IBOTPlusPlusPatchLoss picks up the new mode without duplicating the branch.
DINOv3 applies the KoLeo regularizer to small batches of 16 samples. The released configs reach that by running with a per-GPU batch size of 16, which is not a setting lightly users are likely to train with. KoLeoLoss now splits the batch into consecutive groups of group_size and searches nearest neighbors within each group, so the paper setting is reproducible at any batch size. topk and gather_distributed are added for parity with the reference implementation. Defaults are unchanged: without group_size the whole batch forms one group and the loss is identical to before.
swav_loss.sinkhorn and the Sinkhorn-Knopp centering added for DINOv3 ran the same algorithm. The implementation now lives in center.py and swav_loss.sinkhorn delegates to it, keeping its signature and output dtype. Verified identical to the previous implementation over 3600 random input, epsilon and iteration combinations single-process, and over two gloo ranks with gather_distributed both True and False. The shared version calculates the exponential in float32, which fixes a NaN in the SwaV path for half-precision input, and all-reduces the number of samples instead of deriving it from the world size, which the iBOT loss needs because the number of masked tokens differs between processes. The docstring also no longer claims that DINOv2 uses Sinkhorn-Knopp centering. DINOv2 offers it as an option but keeps mean centering in its released configs, where it reports no difference on ImageNet-1k. DINOv3 uses it for both objectives.
KoLeoLoss now rejects gather_distributed when torch.distributed is unavailable, matching the ten other losses that do this, and reports a topk that is too large for the group size instead of failing inside torch.topk. The previous check skipped groups of size one, where a topk above one reached torch.topk with k larger than the dimension. DINOLoss and IBOTPatchLoss reject a negative sinkhorn_iterations, matching MSNLoss, instead of silently running no iterations. The Sinkhorn-Knopp probabilities are cast back to the dtype of the teacher output, which was already the case but is now covered by a half-precision test and a comment, because torch.einsum does not promote dtypes and the loss fails without it.
The gathered KoLeoLoss was only covered by a mocked world_size=2 test, which bypasses GatherLayer and therefore does not exercise its backward. The shared gloo pool from lightly-ai#1982 makes a real two-rank test possible. The forward test asserts that every rank sees the loss of the non-distributed run on the concatenated global batch. The gradient test asserts the local gradient is NUM_PROCESSES times the non-distributed one, because every rank computes the loss of the whole global batch and GatherLayer all_reduces the identical gradients before slicing out the local one. DDP cancels the factor when it averages parameter gradients across ranks.
8ce532a to
5336c4b
Compare
Replaces the CENTER_MODE_SINKHORN_KNOPP/VALID_CENTER_MODES module-level constants with a Literal["mean", "sinkhorn_knopp"] type on center_mode in DINOLoss and IBOTPatchLoss, and an explicit if/elif/else with a final ValueError instead of a membership check against the shared constant list. CENTER_MODE_TO_FUNCTION is kept since Center itself still relies on it. Literal is imported behind TYPE_CHECKING because it requires Python 3.8+ and this package still supports 3.7; the guard is exercised by deleting typing.Literal and reloading both modules, which succeeds because the annotation is never evaluated at runtime. sinkhorn_iterations must now be at least 1 rather than merely non-negative, since 0 iterations does not produce a valid probability distribution for DINOv3 centering. The bound only moves in DINOLoss and IBOTPatchLoss: the shared sinkhorn_knopp() function in center.py is left unconstrained, because SwaVLoss also calls it and has long-tested, pre-existing support for sinkhorn_iterations=0 as a deliberate configuration.
Two failures reproduced on the Python 3.7 CI job, both invisible locally because this machine resolves a newer torch than the one pinned there (torch==1.13.1). _nearest_neighbor_indices returned an expression assembled from tensor arithmetic directly; under torch 1.13.1's weaker operator stubs mypy widens that to Any, tripping no-any-return. Assigning to a Tensor-annotated variable before returning fixes it, matching the same gotcha hit on the VISReg PR (lightly-ai#1968). test_sinkhorn_knopp__half_precision ran the full DINOLoss forward pass in half precision, which calls F.log_softmax on the student output. That has no CPU kernel for Half in torch 1.13.1. The dtype-cast behavior the test guards against lives entirely in _teacher_probabilities, which never touches log_softmax, so the test now calls that method directly instead of going through forward(). Confirmed this still catches the regression it was written for: reverting the .to(teacher_out.dtype) cast makes the new assertion fail.
Resolves the import conflict in ibot_loss.py, where master added the cast import (lightly-ai#2016) next to this branch's TYPE_CHECKING import. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Python 3.8 is now the minimum supported version (lightly-ai#2028) and mypy only runs against the current torch stubs, so two workarounds from the Python 3.7 CI leg are no longer needed: - Literal is imported from typing directly instead of behind TYPE_CHECKING in DINOLoss and IBOTPatchLoss. - _nearest_neighbor_indices returns its expression directly again; the Tensor-annotated intermediate only existed for the old torch stubs. The half-precision DINOLoss test still calls _teacher_probabilities directly, because the oldest supported torch (1.10) has no half-precision CPU kernel for log_softmax. Its docstring now says so instead of pointing at the removed CI leg. sinkhorn_knopp's temperature uses a PEP 604 union with postponed annotations, following the repo's typing convention. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
The comment said GatherLayer averages the gradients. It sums them: its backward all-reduces the gradients of all processes without dividing, so the local gradient is world_size times the non-distributed one, as the multi-rank gradient test asserts. DDP's averaging of parameter gradients cancels the factor. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
The distributed KoLeoLoss tests used local batches that were exactly one group, so no group spanned two ranks and the order in which the gathered batches are concatenated was never checked: the tests still passed with the gathered batches reversed. They now use local batches of 6 with groups of 4, which makes one group span both ranks; with the reversed concatenation the pool tests fail. The pool workers now read their rank from the process group. The pool hands every task to whichever worker is free, so the rank passed with a task does not always match the worker's process group rank (5 in 200 calls in a local run). That was harmless while every group stayed within one rank, but with cross-rank groups it would make the tests flaky. The gradient worker returns its rank so that results are matched by rank rather than task order. sinkhorn_knopp's distributed path, which SwaV now shares, had no multi-process test. A pool test checks that it matches the non-distributed result on the concatenated batch, with equal and with unequal batch sizes per rank, as happens with the masked tokens of the iBOT loss. Deriving the number of samples from the local batch size and the world size instead of reducing it fails the unequal case. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
|
@liopeer the Aug 12 feedback is in ( |
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info
📝 Walkthrough
Merge Risk: ⚪ Minimal · up to
|
Part of #1881 — phase 1 of 3, so this does not close the issue.
Description
The three loss-level pieces of DINOv3's first training phase, scoped in #1881 (comment). Paper and reference evidence for each is in that thread; this is what changed.
center_modemirrors DINOv2/v3's owntrain.centeringkey. Sinkhorn-Knopp already existed in the repo asswav_loss.sinkhorn, so the shared implementation now lives inlightly/models/modules/center.pybesideCenterandswav_loss.sinkhorndelegates to it, keeping its signature and output dtype. In Sinkhorn-Knopp mode the center buffer stays registered but unused, so checkpoints load across modes, andIBOTPlusPlusPatchLossinherits the mode through the extracted_teacher_probabilities().KoLeoLosssearches nearest neighbors within consecutive groups ofgroup_size;topkandgather_distributedcomplete the referenceKoLeoLossDistributedsurface in ~15 lines.Defaults are unchanged and no existing test was adapted.
Sharing the implementation changes two things for SwaV, both improvements:
sinkhornpreviously returned NaN in fp16 for|out| > ~0.55at the defaultepsilon=0.05, which is reachable under AMP with cosine-similarity logits. The codes are cast back, so the returned dtype is unchanged.sum_Qall-reduce, so the number of collectives is unchanged.Two things not in the issue thread:
exp(logit / 0.04)overflows fp32 at|logit| >= ~3.6. DINO heads L2-normalize before the weight-normed prototype layer, so logits stay in [-1, 1] and this cannot trigger. The reference has the identical exposure, and stabilizing it under DDP needs a globally reduced max, i.e. an extra collective per step. I matched the reference.DINOLossalready matches DINOv3's loss composition. Upstream averages the global-crop and local-crop terms separately and recombines them withdino_global_scale/dino_local_scale, which are proportional to the term counts; for the pretrain defaultreweight_dino_local_loss: falsethat is algebraically the uniform mean over view pairsDINOLossalready computes. Checked numerically on identical inputs with 2 global and 6 local views: 5.9371023 vs 5.9371023 for mean centering, 5.6693587 vs 5.6693592 for Sinkhorn-Knopp.Commits are independently reviewable; the KoLeo commit shares no file with the others, and each review round is kept in its own commits on top so the incremental diff stays readable. Since the last round master is merged in, and three commits on top drop the Python 3.7 workarounds now that 3.8 is the minimum (#2028), correct a comment that said
GatherLayeraverages gradients (it sums them), and strengthen the distributed tests.Tests
30 new tests.
sinkhorn_knoppis verified against a copy of the DINOv3 reference across 9 temperature x iteration combinations, following theOriginalDINOLosspattern already intest_dino_loss.py, plus row-sum, prototype-marginal and detachment tests. The loss-level tests cover wiring rather than repeating the algorithm: Sinkhorn-Knopp applied jointly over all teacher views as upstream does, the center staying zero while remaining in the state dict,IBOTPlusPlusPatchLossinheriting the mode, and a half-precision test of the teacher probabilities that fails if the cast back to the input dtype is dropped, sincetorch.einsumdoes not promote dtypes. It calls_teacher_probabilitiesdirectly because torch 1.10, the minimum, has no half-precision CPUlog_softmax.KoLeoLossgets 14 tests, including two real two-rank tests on the shared gloo pool from #1982, markedDDP. Local batches of 6 with groups of 4 make one group span both ranks, so the tests also pin the rank order of the gathered batch; they fail if it is reversed. Workers read their rank from the process group, since the pool does not pin tasks to ranks. The forward test asserts every rank sees the loss of the non-distributed run on the concatenated global batch; the gradient test asserts the local gradient isNUM_PROCESSEStimes the non-distributed one, since every rank computes the loss of the whole global batch andGatherLayerall-reduces the identical gradients before slicing out the local one — the factor DDP cancels when it averages parameter gradients. A correct gathered forward can still have a wrong backward, as #1977 showed. Themocker-simulatedworld_size=2test is kept alongside them because it isolates grouping and neighbor indexing fromGatherLayer.The distributed path of
sinkhorn_knoppalso runs on the pool, against the non-distributed result on the concatenated batch, with equal and with unequal batch sizes per rank; the unequal case fails if the number of samples is derived from the world size.Manually verified that the shared helper is a no-op for SwaV: bit-identical to the previous
swav_loss.sinkhornover 3600 random (input, epsilon, iterations) combinations single-process, and over two gloo ranks withgather_distributedboth True and False.Also ran a toy convergence check, since the tests above are all unit-level. A student/teacher pair with a
DINOProjectionHeadoverfits 300 steps on a fixed batch of 256 synthetic views with 64 prototypes. Entropy is that of the mean teacher assignment as a fraction of its maximum, so 1.0 is perfectly balanced prototype usage and 0.0 is collapse:Both modes train in both losses. Sinkhorn-Knopp holds prototype usage essentially uniform where mean centering drifts to 0.80, which is the property it is chosen for; its slightly higher final loss is the expected cost of enforcing balance per batch rather than nudging a running center. This is a synthetic overfit check, not a benchmark, and says nothing about downstream quality.
KoLeoLossdecreases in the same setup with and without grouping (0.076 to -0.232 ungrouped, -0.125 to -0.412 withgroup_size=8), and grouping leaves the global nearest-neighbor cosine similarity higher (0.56 vs 0.21) because only within-group neighbors are pushed apart.Run on Python 3.9 with the lockfile's dependencies. The touched test files also pass at the floors of the minimal CI job, Python 3.8 with torch 1.10.0 and torchvision 0.11.1 (298 passed, 30 skipped for CUDA and the pool; the 6 pool tests pass with the pool on). The 3 failures are macOS-specific and unrelated to this PR: two PyAV
spawnpickling tests and the DCL Lightning-DDP numerics test, all known to fail on this machine on a cleanupstream/master.Documentation
.rstfiles).Every new parameter is documented on both the class
Attributesand the__init__Argsof the affected loss.lightly.loss.htmlis the only rendered page that changes, attached as a PDF printout.Improvements put into another issue:
DINOLosscannot express the scheduled local-crop reweighting (reweight_dino_local_loss: true) used by the DINOv3 gram-anchoring config. Not needed for phase 1; I will raise it with the refinement-phase work.DINOLoss(center_mode="mean")fails on half-precision input, because the float32centerbuffer promotes the teacher output while the student output stays fp16 andtorch.einsumdoes not promote. Pre-existing onmasterand untouched here; happy to open an issue.Issues covering the breaking change:
Summary
sinkhorn_knoppcentering toDINOLoss,IBOTPatchLoss, andIBOTPlusPlusPatchLoss.center.pyand keep SwaV compatibility.KoLeoLosswith grouped search,topk, and distributed gathering.Validation