Add optimizer-provider and pg_collection hooks to setup_model_and_optimizer - #5332
Closed
yashaswikarnati wants to merge 1 commit into
Closed
Add optimizer-provider and pg_collection hooks to setup_model_and_optimizer#5332yashaswikarnati wants to merge 1 commit into
yashaswikarnati wants to merge 1 commit into
Conversation
yashaswikarnati
commented
Jun 14, 2026
| model_provider_func, | ||
| model_type, | ||
| checkpointing_context=None, | ||
| optimizer_provider_func=None, |
Contributor
Author
There was a problem hiding this comment.
instead of having a provider function, can we instead change get_megatron_optimizer to handle mimo training ?
Contributor
Author
There was a problem hiding this comment.
Done — dropped the provider hook. get_megatron_optimizer now dispatches a single MimoModel chunk to get_mimo_optimizer directly (lazy import to avoid the circular dep, since get_mimo_optimizer already calls get_megatron_optimizer per module). setup_model_and_optimizer keeps only the pg_collection→get_model forwarding. (b3913cc)
Add optional optimizer_provider_func and pg_collection parameters to setup_model_and_optimizer so a MIMO entry can build a per-module MimoOptimizer and pass per-module process groups while reusing the stock pretrain()/setup path end-to-end. Both parameters default to None and preserve byte-identical behavior for existing callers: when optimizer_provider_func is None it falls back to get_megatron_optimizer, and pg_collection=None is forwarded to get_model as before. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: ykarnati <ykarnati@nvidia.com>
yashaswikarnati
force-pushed
the
ykarnati/upstream-trainloop-tl-opt
branch
from
June 14, 2026 00:41
88f2d35 to
b3913cc
Compare
Contributor
Author
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Adds two optional hooks to setup_model_and_optimizer:
Both default to the current behavior => byte-for-byte unchanged for existing callers. Enables reusing stock pretrain()/setup_model_and_optimizer for heterogeneous MIMO.
Part of the NMFW-516 MIMO-on-stock-trainloop series.
cog: per-rank verified green on cw-dfw 8-GPU.