Add hetero grid args and MoE process groups for MIMO example - #5375

Merged

yashaswikarnati merged 6 commits into

NVIDIA:mainfrom

yashaswikarnati:x2-args-topology

Jun 24, 2026

Commits on Jun 16, 2026

mimo example: add hetero grid/topology args and MoE process groups
yashaswikarnati
committed
225f4ea

Commits on Jun 19, 2026

mimo example: address review comments on hetero grid args

yashaswikarnati
and
claude
committed
c3f466d
mimo example: derive train_iters from --global-batch-size directly

yashaswikarnati
and
claude
committed
82e07c7
mimo example: take encoder module name as a build_module_grid_specs param

yashaswikarnati
and
claude
committed
e1af6e3
mimo example: drop train_samples->train_iters from grid validation

yashaswikarnati
and
claude
committed
e3a0b7c

Commits on Jun 23, 2026

mimo example: narrow hetero grid args to supported surface
yashaswikarnati
committed
4f0a1f9