feat(megatron-run): support command prefix

[saivishal1999]

Summary

Adds command_prefix support to MegatronRun so workloads can prepend a wrapper command before the Megatron training entrypoint.

This is intended to replace patch-based command wrapping used by some recipes, for example:

[cmd_args]
command_prefix = "/cloudai_install/bindpcie --"

With this field, the generated command becomes:

/cloudai_install/bindpcie -- python /megatron/pretrain_hybrid.py ...

The field is excluded from Megatron CLI arguments, so it controls CloudAI command generation only and is not forwarded to the training script.

Test / Verification Plan

Config validation:

• Added command_prefix = "/cloudai_install/bindpcie --" locally to representative MegatronRun test TOMLs.
• Ran verify-configs with this CloudAI branch:
  • 120B MegatronRun configs: 4 configuration files, all passed
  • 130B MegatronRun configs: 12 configuration files, all passed

Command-generation validation:

• Compared generated dry-run output against the previous patch-based command wrapping flow.

• Dry-ran representative MegatronRun scenarios:

  • 120B, 256-GPU scenario
  • 130B, 256-GPU scenario

• Result:

  • Generated test_cmd matched the patch-based flow after normalizing local output/install paths.
  • Generated srun_cmd matched the patch-based flow after normalizing local output/install paths.
  • Rendered command includes:
    nsys profile ... /cloudai_install/bindpcie -- python /megatron/pretrain_hybrid.py ...
  • Confirmed --command-prefix is not forwarded to Megatron CLI args.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: ASSERTIVE

Plan: Enterprise

Run ID: 2b9c70d2-392f-4951-a955-a87be6d2942f

📥 Commits

Reviewing files that changed from the base of the PR and between aae5c58 and a96bb7d.

📒 Files selected for processing (3)

• src/cloudai/workloads/megatron_run/megatron_run.py
• src/cloudai/workloads/megatron_run/slurm_command_gen_strategy.py
• tests/workloads/megatron_run/test_command_gen_strategy_slurm.py

Included review availability: Your plan provides up to 12 included reviews per hour; 11 remain after this review.

---

📝 Walkthrough

Walkthrough

Changes

Megatron command prefix

Layer / File(s)	Summary
Prefix contract and argument serialization src/cloudai/workloads/megatron_run/megatron_run.py	MegatronRunCmdArgs adds an optional command_prefix. Serialization excludes this field from Megatron CLI arguments.
Slurm prefix generation and validation src/cloudai/workloads/megatron_run/slurm_command_gen_strategy.py, tests/workloads/megatron_run/test_command_gen_strategy_slurm.py	Slurm command generation prepends the configured prefix. Tests cover default commands and the bindpcie prefix.

Estimated code review effort: 2 (Simple) | ~10 minutes

Merge Risk: ⚪ Minimal · up to a96bb

The optional command prefix is applied to generated Megatron Slurm commands without becoming a Megatron CLI option, while the default behavior remains unchanged. The change is mergeable with no concrete blocking risk identified.

🚥 Pre-merge checks | ✅ 4

✅ Passed checks (4 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly and concisely describes the main change: adding command prefix support to Megatron Run.
Description check	✅ Passed	The description directly explains command_prefix behavior, its purpose, generated command format, CLI argument handling, and validation results.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

_{Comment @coderabbitai help to get the list of available commands.}