Skip to content

Refactor the optimizer override function so that users can swap in their own #2010

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom
Merged
jstjohn merged 3 commits into from
Jan 23, 2026
Merged

Refactor the optimizer override function so that users can swap in their own #2010

jstjohn merged 3 commits into from
Jan 23, 2026

Conversation

@jstjohn
Copy link
Contributor

@jstjohn jstjohn commented Jan 21, 2026 •
edited
Loading

What does this PR do ?

Makes the private optimizer.py:_build_config_overrides into a builder object that a user can swap in their config. This allows new models (eg evo2) to have their own custom override function without adding booleans/modifying the core function in mbridge optimizer.py.

Changelog

GitHub Actions CI

See the CI sectionin the Contributing doc for how to trigger the CI. A Nvidia developer will need to approve and trigger the CI for external contributors.

Before your PR is "Ready for review"

Pre checks:

  • Make sure you read and followed Contributor guidelines
  • Did you write any new necessary tests?
  • Did you add or update any necessary documentation?
  • Does the PR affect components that are optional to install? (Ex: Numba, Pynini, Apex etc)
    • Reviewer: Does the PR have correct import guards for all optional libraries?

If you haven't finished some of the above items you can still open "Draft" PR.

Additional Information

  • Related to # (issue)

FDecaYed
FDecaYed reviewed Jan 23, 2026
View reviewed changes
Copy link
Contributor

@FDecaYed FDecaYed left a comment

Choose a reason for hiding this comment

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

Agree on making the config override not hard coded and each model need to change something create and pass their own. we probably also want to handle qwen the same way instead of hardcode if block in default

@jstjohn jstjohn merged commit a2c11f8 into main Jan 23, 2026
49 checks passed
@jstjohn jstjohn deleted the jstjohn/opt_config_override_provider branch January 23, 2026 17:21
aroshanghias-nvd pushed a commit to aroshanghias-nvd/Megatron-Bridge that referenced this pull request Jan 29, 2026
@jstjohn @aroshanghias-nvd
Refactor the optimizer override function so that users can swap in th…
7c1869d 
…eir own (NVIDIA-NeMo#2010)

Signed-off-by: John St. John <jstjohn@nvidia.com>
Signed-off-by: Ali Roshan Ghias <aroshanghias@nvidia.com>
aroshanghias-nvd pushed a commit to aroshanghias-nvd/Megatron-Bridge that referenced this pull request Jan 29, 2026
@jstjohn @aroshanghias-nvd
Refactor the optimizer override function so that users can swap in th…
3bf5fd2 
…eir own (NVIDIA-NeMo#2010)

Signed-off-by: John St. John <jstjohn@nvidia.com>
Signed-off-by: Ali Roshan Ghias <aroshanghias@nvidia.com>
github-merge-queue bot pushed a commit to NVIDIA/bionemo-framework that referenced this pull request Feb 3, 2026
@jstjohn
Refactor the way that we do weight decay skipping for hyena to follow…
5778321 
… ToT mbridge. (#1429)

### Description

* No change to user level API but under the hood make use of the new
mbridge API for definining custom weight decay skips. Depends on
NVIDIA-NeMo/Megatron-Bridge#2010
* Update to tokenizer to support the new mbridge API for tokenizer init
that no longer requires a path for path object for based inputs. Path
objects no longer work with megatron using this path, so switching to
strings in the recipe.
* Remove unused nemo2 code/files that were left over in the refactor.

### Type of changes

<!-- Mark the relevant option with an [x] -->

- [ ] Bug fix (non-breaking change which fixes an issue)
- [ ] New feature (non-breaking change which adds functionality)
- [x] Refactor
- [ ] Documentation update
- [ ] Other (please describe):

### CI Pipeline Configuration

Configure CI behavior by applying the relevant labels. By default, only
basic unit tests are run.

-
[ciflow:skip](https://github.com/NVIDIA/bionemo-framework/blob/main/docs/docs/main/contributing/contributing.md#ciflow:skip)
- Skip all CI tests for this PR
-
[ciflow:notebooks](https://github.com/NVIDIA/bionemo-framework/blob/main/docs/docs/main/contributing/contributing.md#ciflow:notebooks)
- Run Jupyter notebooks execution tests for bionemo2
-
[ciflow:slow](https://github.com/NVIDIA/bionemo-framework/blob/main/docs/docs/main/contributing/contributing.md#ciflow:slow)
- Run slow single GPU integration tests marked as @pytest.mark.slow for
bionemo2
-
[ciflow:all](https://github.com/NVIDIA/bionemo-framework/blob/main/docs/docs/main/contributing/contributing.md#ciflow:all)
- Run all tests (unit tests, slow tests, and notebooks) for bionemo2.
This label can be used to enforce running tests for all bionemo2.
-
[ciflow:all-recipes](https://github.com/NVIDIA/bionemo-framework/blob/main/docs/docs/main/contributing/contributing.md#ciflow:all-recipes)
- Run tests for all recipes (under bionemo-recipes). This label can be
used to enforce running tests for all recipes.

Unit tests marked as `@pytest.mark.multi_gpu` or
`@pytest.mark.distributed` are not run in the PR pipeline.

For more details, see [CONTRIBUTING](CONTRIBUTING.md)

> [!NOTE]
> By default, only basic unit tests are run. Add appropriate labels to
enable an additional test coverage.

#### Authorizing CI Runs

We use
[copy-pr-bot](https://docs.gha-runners.nvidia.com/apps/copy-pr-bot/#automation)
to manage authorization of CI
runs on NVIDIA's compute resources.

- If a pull request is opened by a trusted user and contains only
trusted changes, the pull request's code will
automatically be copied to a pull-request/ prefixed branch in the source
repository (e.g. pull-request/123)
- If a pull request is opened by an untrusted user or contains untrusted
changes, an NVIDIA org member must leave an
`/ok to test` comment on the pull request to trigger CI. This will need
to be done for each new commit.

#### Triggering Code Rabbit AI Review

To trigger a code review from code rabbit, comment on a pull request
with one of these commands:

- @coderabbitai review - Triggers a standard review
- @coderabbitai full review - Triggers a comprehensive review

See https://docs.coderabbit.ai/reference/review-commands for a full list
of commands.

### Pre-submit Checklist

<!--- Ensure all items are completed before submitting -->

- [ ] I have tested these changes locally
- [ ] I have updated the documentation accordingly
- [ ] I have added/updated tests as needed
- [ ] All existing tests pass successfully


<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->

## Summary by CodeRabbit

## Release Notes

* **New Features**
* Added `no_weight_decay_embeddings` configuration parameter for Evo2
training recipes to control embedding weight decay behavior.

* **Chores**
  * Updated Megatron-related dependency versions.

* **Tests**
  * Improved test fixture scoping for better test isolation.

<sub>✏️ Tip: You can customize this high-level summary in your review
settings.</sub>

<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Signed-off-by: John St. John <jstjohn@nvidia.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@FDecaYed FDecaYed FDecaYed left review comments

@ko3n1g ko3n1g ko3n1g 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.

3 participants

@jstjohn @ko3n1g @FDecaYed