[AHMETHAKANBEZIR1]

Description

SIGReg reduces its global sample count in the projection dtype. This rounds valid counts in float16/bfloat16 (e.g. two float16 batches of 2049 become 4096 rather than 4098), changing the characteristic-function normalization and statistic scale. Larger float16 counts overflow to infinity and int(...item()) raises OverflowError.

Construct the count as torch.int64 from the integer batch size. This keeps metadata exact without changing embedding precision, sampled directions, differentiable reductions or the SIGReg formula.

Tests

• Added and ran regression tests.
• 9 added cases: unchanged master has 4 failures and 5 passing controls; the fix passes all 9.
• SIGReg/LeJEPA modules: 32 passed.
• Loss suite: 920 passed, 98 skipped, 466 existing warnings.
• An explicit repeated global batch provides the numerical reference for two simulated identical ranks. Cases cover bfloat16 rounding, float16 rounding/overflow, float32/float64 and exactly representable low-precision controls.
• Locked Python 3.12.14 / torch 2.8.0 environment: lock check, global Ruff lint/format, full mypy (238 files), changed-file and commit hooks passed.

The new tests simulate the collective operations, following the existing SIGReg tests; they are not real distributed integration tests. An additional local two-process CPU/Gloo attempt failed during init_process_group, before model execution, with makeDeviceForHostname(): unsupported gloo device on the Windows wheel, including the official UV transport selection. No backend code or test gate was changed. CUDA/NCCL and real DDP correctness are not claimed.

Documentation

No public API or docstring changes. The existing all-files codespell false positive on the untouched ord="fro" FroSSL line is recorded in #2080; no spelling rule was disabled.

AI disclosure: code, tests and description prepared autonomously with OpenAI Codex; no independent human review was performed.

SIGReg.forward now reduces the global sample count as torch.int64, preventing rounding and overflow in low-precision projection dtypes. A parameterized regression test checks nine dtypes and batch-size cases by comparing simulated distributed loss with loss on the repeated global batch.

Minimal CI compatibility follow-up (16 October)

The minimal Linux/Python 3.8 job failed three bfloat16 cases before checking the regression: its older PyTorch lacks the bfloat16 torch.eye kernel used by the deterministic test-vector fixture. Construct the same exact identity vectors in float32, then convert to the requested dtype. The production fix, low-precision cases, numerical reference, tolerances and assertions are unchanged.

Locally emulating this specific unavailable kernel reproduced the same three failures with six passing controls; the corrected fixture passes all nine under that emulation. This is not a local run of the entire minimal Linux environment. The normal locked environment passes all 32 SIGReg/LeJEPA tests. Global Ruff lint/format, full configured pre-push mypy and all commit/push hooks passed. The previous loss-suite result applies to the unchanged production fix; no new full loss-suite or real DDP/GPU run is claimed. Refreshed remote CI remains to be checked.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

• Configuration used: Organization UI
• Review profile: CHILL
• Plan: Advanced
• Run ID: 42c2a4be-5b44-44b2-8af2-660188a83297

📥 Commits

Reviewing files that changed from the base of the PR and between b8630f0 and 56c4ce4.

📒 Files selected for processing: 1

• tests/loss/test_sigreg_sample_count.py

Included review availability: This review used your included allowance. Your plan provides up to 8 included reviews per hour; 7 remain after this review.

📝 Walkthrough

SIGReg now creates the distributed sample-count tensor with torch.int64. A parameterized test checks the reduced count and compares distributed loss with a reference loss across batch sizes and floating-point dtypes.

Changes: SIGReg distributed sample count

Layer / File(s)	Summary
Sample count and loss validation lightly/loss/lejepa_loss.py, tests/loss/test_sigreg_sample_count.py	SIGReg creates the distributed sample-count tensor with torch.int64. Parameterized tests check the reduced sample count and compare distributed loss with a reference loss across the tested batch sizes and dtypes.

Priority: ⬇️ Low

Estimated code review effort: 2 (Simple) | ~10 minutes

Change: Bug fix

Merge Risk: ⚪ Minimal · up to 56c4c

No actionable merge-blocking risk was identified.

Pre-merge checks: 4 · 1

Failed checks: 1 warning

Check name	Status	Explanation	Resolution
Docstring Coverage		Docstring coverage is 33.33% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 6 functions across 2 files.	Write docstrings for the functions missing them to satisfy the coverage threshold.

Passed checks: 4 passed

Check name	Status	Explanation
Linked Issues check		Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check		Check skipped because no linked issues were found for this pull request.
Description Check		Check skipped - CodeRabbit’s high-level summary is enabled.
Title check		The title clearly and concisely describes the main change: keeping distributed SIGReg sample counts in an integral dtype.

• Fix all pre-merge checks with AI

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create a new PR

• Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[mrpositron]

/review