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

• [x] I am familiar with the Contributing Guidelines.
• [x] New or existing tests cover these changes.
• [x] The documentation is up to date with these changes.
• [x] 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[bot]]

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