Add: torch compile support for DINOv31#839
Open
IgorSusmelj wants to merge 8 commits into
Open
Conversation
|
Codex usage limits have been reached for code reviews. Please check with the admins of this repo to increase the limits by adding credits. |
|
|
DINOv31 continues a DINOv2-pretrained backbone with the full DINOv2 objective (DINO + iBOT + KoLeo) plus an auxiliary Patch Kernel Alignment (PaKA / CKA) loss that aligns the relational structure of student and teacher dense patch tokens. It is a thin subclass of DINOv2 and does not modify it; the PaKA term is added on top and the DINOv2 path is left byte-identical. - DINOv31 method: PaKA projection heads (reusing dinov2 `_build_mlp`), ROI- aligned cross-view CKA on clean-teacher globals vs high-overlap student locals, EMA teacher head update, lenient checkpoint load (tolerates missing `*_paka_head*` keys), teacher-temperature warmup cap for short post-trains, and an `on_fit_start` override that suppresses DINOv2's 125k-step warning. - DINOv31 transform: view layout [global0, global1, dino_local.., clean_global0/1, paka_local..] with geometry-tracked crops (clean globals reuse the global crop geometry; paka locals are true small locals nested inside alternating globals). - Shared geometry plumbing (opt-in via `record_geometry`): TrackedRandomResized Crop / TrackedHorizontalFlip / TrackedVerticalFlip, ViewTransform records an 8-element geometry tensor, optional `geometries` on the transform output, dataset item and batch, and ImageDataset propagates it only when every view records one. - Lightly SSL source pinned to the PatchKernelAlignmentLoss merge commit (lightly-ai/lightly#1953) until the loss is released. Co-Authored-By: Claude <noreply@anthropic.com>
DINOv31 PaKA trains at 518px (pos_embed grid 37x37) but inits from a 224px DINOv2 checkpoint (grid 16x16). The strict load_state_dict raised `size mismatch for _model.pos_embed [1,257,384] vs [1,1370,384]`, blocking the high-resolution PaKA recipe. Add `_interpolate_pos_embed(state_dict, target_state_dict)`: a shape-driven, in-place bicubic resize of any `*pos_embed` whose square patch grid differs from the target, preserving the leading cls token. Call it before each of the four strict loads in `train_helpers.load_state_dict` (wrapped model, model, embedding model, method). It composes with `DINOv31.load_state_dict`, which filters only `*paka_head*` missing keys. This is a general load-time capability (DINOv3 uses RoPE -> no-op), needed first by dinov31's 518px training. Assumes a single leading cls token (correct for dinov2/vits14 with no register tokens). Co-Authored-By: Claude <noreply@anthropic.com>
- dinov31 docs: drop implementation-detail phrasing; init from model="dinov2/vits14" (pretrained Meta weights) instead of a .ckpt - ViewTransform: move record_geometry off the generic ViewTransformArgs onto a ViewTransform.__init__ kwarg set by the method transform, so each method keeps its whole configuration in one place - DINOv2 ViT: register a load_state_dict pre-hook for pos_embed interpolation on DinoVisionTransformer (moved out of the generic trainHelpers loader), mirroring the patch_embed channel hook - uv.lock: regenerate via `make lock` to restore the exclude-newer flag Co-Authored-By: Claude <noreply@anthropic.com>
IgorSusmelj
force-pushed
the
pr2/adamw8bit
branch
from
July 10, 2026 12:08
a675e33 to
d5704f7
Compare
IgorSusmelj
force-pushed
the
pr2/adamw8bit
branch
from
July 10, 2026 13:11
d5704f7 to
aeeff30
Compare
Adds an opt-in torch.compile path for pretraining, used by the DINOv31
large-batch recipe. Disabled by default; enable with
pretrain(..., torch_compile_args={"disable": False, "mode": "max-autotune"}).
- compile_method_backbones compiles only wrapped_model.forward_features on the
method's embedding models (student/teacher/momentum, deduped); the Lightning
training_step is never compiled (filenames/logging are graph-hostile).
- Methods are patched on the class, not the instance, so models stay picklable
for DDP/Lightning; sentinels prevent double patching.
- interpolate_pos_encoding is excluded from the compiled graph (bicubic resize
is graph-hostile); documented assumption: input resolution matches the
checkpoint's pos_embed grid.
- try_compile gains a one-shot runtime fallback: if the compiled callable
fails at runtime, it logs once and permanently falls back to eager.
- torch._dynamo.config.optimize_ddp=False so pre-compiled backbones survive
DDP wrapping (DDPOptimizer fails on DINOv2 stochastic-depth graphs).
Co-Authored-By: Claude <noreply@anthropic.com>
IgorSusmelj
force-pushed
the
pr3/torch-compile
branch
from
July 10, 2026 14:06
92d36ce to
bf40617
Compare
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.
What has changed and why?
switches back to normal (eager) execution and logs a notice, so a compile problem never breaks your training run.
How has it been tested?
Unit tests + manual testing
Did you update CHANGELOG.md?
Did you update the documentation?