Skip to content

feat(megatron): inherit sweep MBS/GBS/precision from train_params and pass Primus TP/PP - #450

Merged
sukesh-amd merged 3 commits into
mainfrom
skalla/primus_update
Sep 18, 2026
Merged

sukesh-amd merged 3 commits into
mainfrom
skalla/primus_update

Conversation

@sukesh-amd

Copy link
Copy Markdown
Collaborator

Summary

  • Sweep cells take MBS / GBS / PRECISION from the combination key (and matching body overlays). When sweep is omitted, the implicit default cell requires train_params.micro_batch_size, global_batch_size, and precision instead of falling back to 2 / 128 / BF16.
  • Declared combinations keep optional body overlays (training_iterations, tensor_parallelism, pipeline_parallelism); Primus now passes --tensor_model_parallel_size and --pipeline_model_parallel_size.
  • Empty sweep.runs with non-empty combinations fails at load. Packaged Megatron JSON and docs match this behavior.

Test plan

  • Load a packaged config with a sweep: combo key values win; body overlays apply
  • Omit sweep with complete train_params MBS/GBS/precision: one default cell
  • Omit those train_params fields with no sweep: load error
  • Empty sweep.runs with declared combinations: load error
  • Primus run includes TP/PP CLI flags from train_params / combo overlay

… pass Primus TP/PP

Require those fields on train_params, fill combo gaps from the cell key, and emit Primus tensor/pipeline parallel CLI flags.
…omitted

Drop the hardcoded 2/128/BF16 fallback from docs and configs, and reject empty sweep.runs at load.
@sukesh-amd
sukesh-amd requested a review from solaiys September 17, 2026 19:56

@solaiys solaiys left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

LGTM.

Same patch as #449 (release/cvs-0.2.0): identical file list and diff, targeting main. No blocking issues.

@sukesh-amd
sukesh-amd merged commit 65687b9 into main Sep 18, 2026
2 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants