Skip to content

Normalized Riemannian optimizer#36

Merged
skyw merged 29 commits into
NVIDIA-NeMo:mainfrom
mkhona-nvidia:mkhona/normalized_opt
Oct 7, 2025
Merged

Normalized Riemannian optimizer#36
skyw merged 29 commits into
NVIDIA-NeMo:mainfrom
mkhona-nvidia:mkhona/normalized_opt

Conversation

@mkhona-nvidia

@mkhona-nvidia mkhona-nvidia commented Oct 2, 2025

Copy link
Copy Markdown
Contributor

Added an optimizer for normalized weights, i.e. weights whose columns or rows sum to 1. This is referred to as the "Oblique manifold" and this optimizer performs Riemannian descent on the oblique manifold (see An Introduction to Optimization on Smooth Manifolds by Nicolas Boumal at https://www.nicolasboumal.net/book/ for details)

@copy-pr-bot

copy-pr-bot Bot commented Oct 2, 2025

Copy link
Copy Markdown

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.

@mkhona-nvidia mkhona-nvidia changed the title Mkhona/normalized opt Normalized Riemannian optimizer Oct 3, 2025
@mkhona-nvidia mkhona-nvidia requested a review from skyw October 3, 2025 17:00
skyw
skyw requested changes Oct 3, 2025
View reviewed changes

@skyw skyw left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Haven't finished reviewing all the tests. Will do another path when requested changes are done.

Comment thread emerging_optimizers/riemannian_optimizers/normalized_optimizer.py
Comment thread emerging_optimizers/riemannian_optimizers/normalized_optimizer.py Outdated
Comment thread emerging_optimizers/riemannian_optimizers/normalized_optimizer.py Outdated
Comment thread emerging_optimizers/riemannian_optimizers/normalized_optimizer.py Outdated
Comment thread emerging_optimizers/riemannian_optimizers/normalized_optimizer.py
Comment thread tests/test_normalized_optimizer.py
Comment thread tests/test_normalized_optimizer.py Outdated
Comment thread tests/test_normalized_optimizer.py Outdated
Comment thread tests/test_normalized_optimizer.py Outdated
Comment thread tests/test_normalized_optimizer.py Outdated
@mkhona-nvidia mkhona-nvidia force-pushed the mkhona/normalized_opt branch from 2f8543f to dd2194d Compare October 3, 2025 22:38
skyw
skyw requested changes Oct 3, 2025
View reviewed changes

@skyw skyw left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Couple of left overs need further fix, parameter for example.

I didn't check convergence test very carefully. Also L1 tests haven't been set up in CI, post a local run result in the chat before final approval.

Comment thread tests/test_normalized_optimizer.py
Comment thread tests/test_normalized_optimizer.py Outdated
Comment thread tests/test_normalized_optimizer.py Outdated
Comment thread tests/test_normalized_optimizer.py Outdated
Comment thread tests/normalized_optimizer_convergence_test.py Outdated
Comment thread tests/normalized_optimizer_convergence_test.py Outdated
@mkhona-nvidia

mkhona-nvidia commented Oct 4, 2025

Copy link
Copy Markdown
Contributor Author

Couple of left overs need further fix, parameter for example.

I didn't check convergence test very carefully. Also L1 tests haven't been set up in CI, post a local run result in the chat before final approval.

Here are results of convergence test:

Running tests under Python 3.12.8: /Users/mkhona/miniconda3/envs/pytorch_env/bin/python
[ RUN ] NormalizedOptimizerConvergenceTest.test_oblique_adam_convergence
[ OK ] NormalizedOptimizerConvergenceTest.test_oblique_adam_convergence
[ RUN ] NormalizedOptimizerConvergenceTest.test_oblique_sgd_convergence
[ OK ] NormalizedOptimizerConvergenceTest.test_oblique_sgd_convergence
[ RUN ] NormalizedOptimizerConvergenceTest.test_optimizer_modes_convergence_adam_col
Final accuracy: 89.2
[ OK ] NormalizedOptimizerConvergenceTest.test_optimizer_modes_convergence_adam_col
[ RUN ] NormalizedOptimizerConvergenceTest.test_optimizer_modes_convergence_adam_row
Final accuracy: 80.4
[ OK ] NormalizedOptimizerConvergenceTest.test_optimizer_modes_convergence_adam_row
[ RUN ] NormalizedOptimizerConvergenceTest.test_optimizer_modes_convergence_sgd_col
Final accuracy: 74.8
[ OK ] NormalizedOptimizerConvergenceTest.test_optimizer_modes_convergence_sgd_col
[ RUN ] NormalizedOptimizerConvergenceTest.test_optimizer_modes_convergence_sgd_row
Final accuracy: 100.0
[ OK ] NormalizedOptimizerConvergenceTest.test_optimizer_modes_convergence_sgd_row

@mkhona-nvidia mkhona-nvidia force-pushed the mkhona/normalized_opt branch from b9a92ce to 25c7835 Compare October 4, 2025 00:04
@mkhona-nvidia

mkhona-nvidia commented Oct 6, 2025

Copy link
Copy Markdown
Contributor Author

/ok to test 25c7835

@mkhona-nvidia mkhona-nvidia force-pushed the mkhona/normalized_opt branch from 25c7835 to a9f16b0 Compare October 6, 2025 20:59
@mkhona-nvidia

mkhona-nvidia commented Oct 6, 2025

Copy link
Copy Markdown
Contributor Author

/ok to test f87162d

@mkhona-nvidia
added normalized optimizers and fixed docstrings and formatting
0a68eb3 
Signed-off-by: mikail <mkhona@nvidia.com>
@mkhona-nvidia
added tests for normalized optimizer
81d2a64 
Signed-off-by: mikail <mkhona@nvidia.com>
@mkhona-nvidia
separated out riemannian grad as a private function
b32a71e 
Signed-off-by: mikail <mkhona@nvidia.com>
@mkhona-nvidia
added a simpek convergence test
b13a403 
Signed-off-by: mikail <mkhona@nvidia.com>
@mkhona-nvidia
added test to L1
0170de8 
Signed-off-by: mikail <mkhona@nvidia.com>
@mkhona-nvidia
added tests to L0
134e08b 
Signed-off-by: mikail <mkhona@nvidia.com>
@mkhona-nvidia
changed docstring
3ae3147 
Signed-off-by: mikail <mkhona@nvidia.com>
@mkhona-nvidia
changed to one torch.add for memory pressure
f130aa6 
Signed-off-by: mikail <mkhona@nvidia.com>
@mkhona-nvidia
added missing type hints
6859351 
Signed-off-by: mikail <mkhona@nvidia.com>
@mkhona-nvidia
added missing types for optimizer args
017cc39 
Signed-off-by: mikail <mkhona@nvidia.com>
@mkhona-nvidia
added some more missing type hints
1adfebb 
Signed-off-by: mikail <mkhona@nvidia.com>
@mkhona-nvidia
cleaned up test cases
ed90ff6 
Signed-off-by: mikail <mkhona@nvidia.com>
@mkhona-nvidia
fixed testing as per PR
2b4cb4e 
Signed-off-by: mikail <mkhona@nvidia.com>
@mkhona-nvidia
fixed inplace init
5bc9a49 
Signed-off-by: mikail <mkhona@nvidia.com>
@mkhona-nvidia
fixed formatting
07c52bd 
Signed-off-by: mikail <mkhona@nvidia.com>
@mkhona-nvidia
added missing device movement
344689f 
Signed-off-by: mikail <mkhona@nvidia.com>
@mkhona-nvidia
moved model to device appropriately
6eab498 
Signed-off-by: mikail <mkhona@nvidia.com>
@mkhona-nvidia mkhona-nvidia force-pushed the mkhona/normalized_opt branch from c4db61a to 6eab498 Compare October 7, 2025 18:38
@mkhona-nvidia

mkhona-nvidia commented Oct 7, 2025

Copy link
Copy Markdown
Contributor Author

/ok to test 6eab498

@mkhona-nvidia
added edm2 paper in docstring
5d7dcf2 
Signed-off-by: mikail <mkhona@nvidia.com>
skyw
skyw requested changes Oct 7, 2025
View reviewed changes

@skyw skyw left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Rerun test after addressing last a few comment, plz.

Comment thread tests/test_normalized_optimizer.py Outdated
# Set seed for CUDA if available
if torch.cuda.is_available():
torch.cuda.manual_seed_all(1234)
self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

@skyw skyw Oct 7, 2025

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Don't auto select device, especially in test. If the test is accidently assigned to a machine without GPU, it can end up never being tested on GPU. Same for other test cases.

@mkhona-nvidia
added explicit device flag to control device
b8a9bdc 
Signed-off-by: mikail <mkhona@nvidia.com>
@mkhona-nvidia
added flags to test
f97ef77 
Signed-off-by: mikail <mkhona@nvidia.com>
@mkhona-nvidia

mkhona-nvidia commented Oct 7, 2025

Copy link
Copy Markdown
Contributor Author

/ok to test f97ef77

@skyw skyw merged commit fb1add8 into NVIDIA-NeMo:main Oct 7, 2025
12 checks passed
mkhona-nvidia added a commit to mkhona-nvidia/Emerging-Optimizers that referenced this pull request Oct 7, 2025
@mkhona-nvidia
Add normalized Riemannian optimizer (NVIDIA-NeMo#36)
907d4f5 
* added normalized optimizers and fixed docstrings and formatting

Signed-off-by: mikail <mkhona@nvidia.com>
mkhona-nvidia added a commit to mkhona-nvidia/Emerging-Optimizers that referenced this pull request Oct 7, 2025
@mkhona-nvidia
Add normalized Riemannian optimizer (NVIDIA-NeMo#36)
de8f3f8 
* added normalized optimizers and fixed docstrings and formatting

Signed-off-by: mikail <mkhona@nvidia.com>
@mkhona-nvidia mkhona-nvidia deleted the mkhona/normalized_opt branch October 15, 2025 02:58
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@skyw skyw skyw approved these changes

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

2 participants

@mkhona-nvidia @skyw