Skip to content

[Feature] Add --mg-fl-prefer argument for Megatron-LM-FL vendor selection#1183

Merged
lxd-cumt merged 1 commit intoflagos-ai:mainfrom
lihongyang1990:feature/mg-fl-prefer
Apr 3, 2026
Merged

[Feature] Add --mg-fl-prefer argument for Megatron-LM-FL vendor selection#1183
lxd-cumt merged 1 commit intoflagos-ai:mainfrom
lihongyang1990:feature/mg-fl-prefer

Conversation

@lihongyang1990
Copy link
Copy Markdown
Contributor

Summary

  • Add --mg-fl-prefer CLI argument to select the preferred vendor backend (cuda/musa/txda) for Megatron-LM-FL plugin override dispatch
  • Sync the argument value to the MG_FL_PREFER environment variable at training startup, following the same pattern as --te-fl-prefer / TE_FL_PREFER
  • Rename the argument group from "flagscale transformer engine fl" to "flagscale fl" to reflect the broader scope

Changes

flagscale/train/megatron/training/arguments_fs.py

  • Add --mg-fl-prefer argument with choices ['cuda', 'musa', 'txda'], default empty string
  • Rename argument group title to "flagscale fl"

flagscale/train/megatron/training/training.py

  • Sync args.mg_fl_prefer to os.environ['MG_FL_PREFER'] in pretrain(), alongside the existing TE_FL_PREFER sync logic

Usage

CLI:

python train.py --mg-fl-prefer musa --te-fl-prefer reference

YAML config:

model:
  mg_fl_prefer: musa
  te_fl_prefer: reference

Test plan

  • Verify --mg-fl-prefer musa sets MG_FL_PREFER=musa in the environment
  • Verify omitting --mg-fl-prefer does not set MG_FL_PREFER
  • Verify invalid values are rejected by argparse choices validation
  • Verify compatibility with existing --te-fl-prefer argument

…tion

Add a new command-line argument --mg-fl-prefer to select the preferred
vendor backend (cuda/musa/txda) for Megatron-LM-FL plugin override
dispatch. The value is synced to the MG_FL_PREFER environment variable
at training startup, enabling YAML-based configuration of vendor
selection alongside the existing --te-fl-prefer mechanism.
Copy link
Copy Markdown
Collaborator

@lxd-cumt lxd-cumt left a comment

Choose a reason for hiding this comment

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

LGTM

@lxd-cumt lxd-cumt merged commit 142d3ea into flagos-ai:main Apr 3, 2026
25 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.

3 participants