Add Sinkhorn functions by mkhona-nvidia · Pull Request #298 · NVIDIA-NeMo/Emerging-Optimizers

mkhona-nvidia · 2026-09-10T19:05:13Z

No description provided.

Signed-off-by: mkhona <mkhona@nvidia.com>

copy-pr-bot · 2026-09-10T19:05:17Z

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

greptile-apps · 2026-09-10T19:08:41Z

Confidence Score: 5/5

The PR appears safe to merge; no actionable correctness, security, or repository-rule violations remain.

Summary

Masks inactive or near-zero rows before normalization.
Alternates row and column L2 normalization in an FP32 workspace.
Validates matrix shape, dtype, and numerical parameters.
Exports and documents the helper and adds reference, numerical, validation, and contract tests.

Diagram

%%{init: {'theme': 'neutral'}}%%
flowchart LR
    U[Signed 2D update] --> V[Validate shape, dtype, and parameters]
    V --> F[Copy to FP32 workspace]
    F --> M[Mask near-zero rows]
    M --> R[Normalize rows]
    R --> C[Alternate column and row normalization]
    C --> S[Scale by sqrt of column count]
    S --> O[Return new tensor in input dtype and device]

_{Reviews (10) · Last reviewed commit: "Clarify Sinkhorn normalization step coun..."}

Signed-off-by: mkhona <mkhona@nvidia.com>

skyw · 2026-09-11T16:52:07Z

+    # Apply the first row step separately, then express each remaining iteration as a column/row pair.
+    balanced_update.div_(row_norms.add_(eps))
+    for _ in range(num_steps // 2):
+        # Normalize columns along the row/token dimension.


remove not very meaningful comment, and potentially wrong if used in say some image model. code is clean enough.

skyw · 2026-09-11T16:53:08Z

+
+    # Apply the first row step separately, then express each remaining iteration as a column/row pair.
+    balanced_update.div_(row_norms.add_(eps))
+    for _ in range(num_steps // 2):


Q: why num_steps is set up this way? positive integer // 2?

I can see each loop normalize twice, but positive integer // 2 has no value in the code. num_xxx_steps would be a better name if num_steps to loop count may cause confusion.

skyw · 2026-09-11T16:55:10Z

+            torch.cuda.manual_seed_all(FLAGS.seed)
+
+
+def _sinkhorn_balance_reference(


AI copying code as reference? that's a bad reference in case both are wrong.
Implement reference in a different way, could be a very naive serial CPU version.

skyw · 2026-09-11T16:56:11Z

+
+
+class SinkhornBalanceTest(parameterized.TestCase):
+    def test_defaults_match_deepseek_v41(self) -> None:


Use inspect to check default value would be a better implementation for the purpose of this test.

skyw · 2026-09-11T16:56:37Z

+        ((17, 5), torch.bfloat16, 11, 1.0),
+        ((6, 6), torch.float32, 7, 0.0),
+    )
+    def test_matches_step_by_step_reference(self, shape, dtype, num_steps, zero_row_threshold) -> None:


if not exact match, test name should say "close_to"

skyw · 2026-09-11T16:57:35Z

+            rtol=0.05,
+        )
+
+    def test_single_step_converts_unit_row_l2_norm_to_unit_rms(self) -> None:


Too many tests don't really add more protections. Use your best judgement to remove and keep the essential ones.

skyw · 2026-09-11T16:58:47Z

+
+        self.assertFalse(result.requires_grad)
+
+    @parameterized.parameters(torch.float64, torch.int64, torch.complex64)


Nit: Combine input validation to the same test, otherwise it gets too noisy.

Signed-off-by: mkhona <mkhona@nvidia.com>

skyw · 2026-09-11T17:23:20Z

+    Args:
+        update: Dense signed update with shape ``(num_rows, num_columns)``.
+        eps: Numerical stability term added to each row or column norm.
+        num_steps: Total positive odd number of individual axis-normalization steps.


Nit: if we keep this, would be better to explain how steps are counted and why positive odd number here.

Signed-off-by: mkhona <mkhona@nvidia.com>

mkhona-nvidia · 2026-09-11T17:30:44Z

/ok to test aa26a46

Add Sinkhorn embedding optimizer

047023f

Signed-off-by: mkhona <mkhona@nvidia.com>

mkhona-nvidia requested a review from skyw September 10, 2026 19:05

greptile-apps Bot reviewed Sep 10, 2026

View reviewed changes

Comment thread emerging_optimizers/embedding_optimizers/__init__.py Outdated

mkhona-nvidia added 6 commits September 10, 2026 18:14

Restrict Sinkhorn to supported dtypes

85f2aed

Signed-off-by: mkhona <mkhona@nvidia.com>

Clarify Sinkhorn normalization sequence

cf0cc37

Signed-off-by: mkhona <mkhona@nvidia.com>

Align Sinkhorn loop with mapper convention

b821362

Signed-off-by: mkhona <mkhona@nvidia.com>

Document Sinkhorn balancing math

334cdf0

Signed-off-by: mkhona <mkhona@nvidia.com>

Refocus Sinkhorn change on balancing function

a3ac1f4

Signed-off-by: mkhona <mkhona@nvidia.com>

Move Sinkhorn balance to orthogonalized utilities

14fe009

Signed-off-by: mkhona <mkhona@nvidia.com>

skyw requested changes Sep 11, 2026

View reviewed changes

mkhona-nvidia added 2 commits September 11, 2026 10:03

Remove model-specific Sinkhorn comments

aac845d

Signed-off-by: mkhona <mkhona@nvidia.com>

Simplify Sinkhorn balance tests

d95016c

Signed-off-by: mkhona <mkhona@nvidia.com>

skyw previously approved these changes Sep 11, 2026

View reviewed changes

Clarify Sinkhorn normalization step count

aa26a46

Signed-off-by: mkhona <mkhona@nvidia.com>

mkhona-nvidia dismissed skyw’s stale review via aa26a46 September 11, 2026 17:29

copy-pr-bot Bot deployed to test September 11, 2026 17:31 Active

copy-pr-bot Bot deployed to public September 11, 2026 17:31 Active

skyw approved these changes Sep 11, 2026

View reviewed changes

copy-pr-bot Bot deployed to public September 11, 2026 17:37 Active

copy-pr-bot Bot deployed to public September 11, 2026 17:38 Active

skyw enabled auto-merge (squash) September 11, 2026 17:38

skyw changed the title ~~Add Sinkhorn embedding optimizer~~ Add Sinkhorn functions Sep 11, 2026

skyw merged commit 36f7033 into NVIDIA-NeMo:main Sep 11, 2026
31 checks passed

Navigation Menu

Repository navigation

Add Sinkhorn functions - #298

Add Sinkhorn functions#298
skyw merged 10 commits into
NVIDIA-NeMo:mainfrom
mkhona-nvidia:mkhona/add-sinkhorn-embedding-optimizer

mkhona-nvidia commented Sep 10, 2026

Uh oh!

copy-pr-bot Bot commented Sep 10, 2026

Uh oh!

greptile-apps Bot commented Sep 10, 2026 •

edited

Loading

Uh oh!

Uh oh!

skyw Sep 11, 2026

Uh oh!

skyw Sep 11, 2026

Uh oh!

skyw Sep 11, 2026

Uh oh!

skyw Sep 11, 2026

Uh oh!

skyw Sep 11, 2026

Uh oh!

skyw Sep 11, 2026

Uh oh!

skyw Sep 11, 2026

Uh oh!

skyw Sep 11, 2026

Uh oh!

mkhona-nvidia commented Sep 11, 2026

Uh oh!

Uh oh!

This branch was successfully deployed

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

		torch.cuda.manual_seed_all(FLAGS.seed)


		def _sinkhorn_balance_reference(



		class SinkhornBalanceTest(parameterized.TestCase):
		def test_defaults_match_deepseek_v41(self) -> None:


		self.assertFalse(result.requires_grad)

		@parameterized.parameters(torch.float64, torch.int64, torch.complex64)

Conversation

mkhona-nvidia commented Sep 10, 2026

Uh oh!

copy-pr-bot Bot commented Sep 10, 2026

Uh oh!

greptile-apps Bot commented Sep 10, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Confidence Score: 5/5

Summary

Diagram

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

mkhona-nvidia commented Sep 11, 2026

Uh oh!

Uh oh!

This branch was successfully deployed

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

greptile-apps Bot commented Sep 10, 2026 •

edited

Loading