Add MimoOptimizer for heterogeneous parallelism#4018

Closed

yashaswikarnati wants to merge 35 commits intoNVIDIA:mainfrom

yashaswikarnati:yash/mimo-optimizer

Commits on Jan 27, 2026

Commits on Jan 29, 2026

Merge branch 'main' into yash/1f1b_changes
dimapihtar
authored

Commits on Feb 4, 2026

Merge branch 'main' into yash/mimo_non_colocated
dimapihtar
authored

Commits on Mar 19, 2026

Merge upstream main into yash/mimo_non_colocated

yashaswikarnati
and
claude
committed
Simplify MIMO model and consolidate submodule logic

yashaswikarnati
and
claude
committed
Simplify and deduplicate MIMO model tests

yashaswikarnati
and
claude
committed
Make RankRole always-present and simplify encoder dispatch

yashaswikarnati
and
claude
committed
Replace configurable language_module_key with fixed LANGUAGE_MODULE_KEY

yashaswikarnati
and
claude
committed
Add MimoOptimizer for heterogeneous parallelism
yashaswikarnati
committed
Simplify _get_pg_collection_for_optimizer to only fetch required groups
yashaswikarnati
committed
Fix MimoOptimizer to fail fast on missing modules

yashaswikarnati
and
claude
committed
[MIMO] Configure distributed optimizer process group in optimizer setup
yashaswikarnati
committed
Improve MIMO pipeline abstraction and type safety

yashaswikarnati
and
claude
committed
Rename PipelineMode to ModuleLayout

yashaswikarnati
and
claude
committed

Commits on Mar 20, 2026

Fix non-colocated forward pass and rewrite MIMO 1F1B schedule tests

yashaswikarnati
and
claude
committed

Commits on Mar 21, 2026

Merge remote-tracking branch 'origin/main' into yash/mimo_non_colocated
yashaswikarnati
committed

Commits on Mar 22, 2026

Commits on Mar 23, 2026

Commits on Mar 24, 2026

Fix grad norm computation for MIMO DistributedOptimizer

yashaswikarnati
and
claude
committed