Adds CombinedOptimizer

[peterdsharpe]

PhysicsNeMo Pull Request

Adds a new CombinedOptimizer utility, which is useful for the increasingly-popular "architecture-aware optimizers", such as Muon.

This PR targets the v2.0 refactor branch, so this should only be merged after #1235 .

Description

Checklist

• I am familiar with the Contributing Guidelines.
• New or existing tests cover these changes.
• The documentation is up to date with these changes.
• The CHANGELOG.md is up to date with these changes.
• An issue is linked to this pull request.

Dependencies

Review Process

All PRs are reviewed by the PhysicsNeMo team before merging.

Depending on which files are changed, GitHub may automatically assign a maintainer for review.

We are also testing AI-based code review tools (e.g., Greptile), which may add automated comments with a confidence score.
This score reflects the AI’s assessment of merge readiness and is not a qualitative judgment of your work, nor is
it an indication that the PR will be accepted / rejected.

AI-generated feedback should be reviewed critically for usefulness.
You are not required to respond to every AI comment, but they are intended to help both authors and reviewers.
Please react to Greptile comments with 👍 or 👎 to provide feedback on their accuracy.

[greptile-apps]

Greptile Overview

Greptile Summary

Adds CombinedOptimizer utility that wraps multiple PyTorch optimizers into a unified interface for architecture-aware optimization strategies. The implementation provides proper state management, serialization, and learning rate scheduler compatibility.

Critical Issue Found:

• The documentation claims closures are evaluated once (line 53-55), but the implementation passes the closure directly to each optimizer, causing multiple evaluations. This leads to redundant forward/backward passes and conflicts with optimizers like LBFGS that may call the closure multiple times internally.
• The test suite validates this incorrect behavior (test_combined_optimizer.py:158-174), which needs alignment with the intended design.

Recommendations:

• Decide on the desired closure behavior (single evaluation vs. per-optimizer evaluation)
• Update either the implementation or documentation to match
• Adjust the corresponding test case to validate the correct behavior

Important Files Changed

File Analysis

Filename	Score	Overview
physicsnemo/optim/combined_optimizer.py	2/5	New CombinedOptimizer wrapper for multiple optimizers. Critical bug: docstring claims closure evaluated once but implementation evaluates it per optimizer
test/optim/test_combined_optimizer.py	3/5	Comprehensive test coverage for CombinedOptimizer. Tests validate incorrect behavior (multiple closure calls) that contradicts documentation

[greptile-apps]

_{3 files reviewed, 2 comments}

_{Edit Code Review Agent Settings | Greptile}

[coreyjadams]

Overall, looks good. I appreciate extensive testing.

This adds a new package, physicsnemo.optim, which I think is overdue. As part of the refactor let's make sure to get it in the docs too :).

I left some comments but do have another question. What happens when users instantiate and use this and the parameter groups are not disjoint? Will it cause an error? Silent bugs? Should we include a check that each parameter group is completely disjoint?

[peterdsharpe]

/blossom-ci

[peterdsharpe]

/blossom-ci

[peterdsharpe]

/blossom-ci

[megnvidia]

looks good to me. Doc approval given.