(prototype) Add lightly.optim with LARS, the schedulers and param_groups

[gabrielfruet]

Prototype. Do not merge. This is a design exploration for the LightlySSL 2.0
refactor, opened to be read and argued with, not to land on master. The whole
stack goes together or not at all, and the shape is still open.

4 of 7 in a stack. Base: #2035.

Optimiser pieces are scattered across lightly/utils/, which is where the grab-bag starts. This gives them a subpackage:

lightly/utils/lars.py       -> lightly/optim/lars.py
lightly/utils/scheduler.py  -> lightly/optim/schedulers.py

Both old paths stay as re-export shims with no deprecation warning, because 1.x is still shipping and roughly 20 benchmark and docs files import them. Nothing has to change today.

New: param_groups(*modules, weight_decay=...) returns the two groups an optimiser takes, with normalization parameters and biases in the group that is not decayed. Every SSL reference implementation does that split and the published numbers were produced with it, but each benchmark file spells it out again by hand. Both groups carry an explicit weight_decay, so the optimiser's own default never silently applies. It is built on the existing get_weight_decay_parameters, which stays where it is.

Testing: pytest tests/optim, which asserts the exact parameter names in each group rather than the counts, plus pytest tests -k "lars or scheduler" for the moved modules.

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: d3d0edf1f7

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

[chatgpt-codex-connector]

Deduplicate shared parameters before returning groups

When a model ties a Parameter between submodules, or when overlapping modules are passed (for example, param_groups(model, model.head, ...)), get_weight_decay_parameters returns the same object once per traversal. The resulting optimizer groups therefore contain duplicates: the bundled LARS loops over every entry and updates that parameter multiple times per step, while a parameter classified into different groups can make PyTorch reject the optimizer entirely. Deduplicate parameters by identity before returning the groups.

Useful? React with 👍 / 👎.