feat(megatron): expose optimizer_cuda_graph for the optimizer

[Nas01010101]

What

Megatron-SWIFT's MegatronArguments exposes several optimizer-level toggles (optimizer_cpu_offload, optimizer_offload_fraction, use_precision_aware_optimizer, …) but is missing optimizer_cuda_graph, which Megatron-Core supports:

• Megatron CLI: --optimizer-cuda-graph (megatron/training/arguments.py, action='store_true')
• Megatron config: OptimizerConfig.optimizer_cuda_graph: bool = False (megatron/core/optimizer/optimizer_config.py) — "If true, enables CUDA graph for optimizer step."

optimizer_cuda_graph CUDA-graphs the optimizer step, reducing per-step CPU launch overhead. Without it exposed, recipes that want the CUDA-graphed optimizer step cannot enable it through Megatron-SWIFT.

Change

• Add optimizer_cuda_graph: bool = False to MegatronArguments, alongside the other optimizer_* toggles (matching Megatron's OptimizerConfig field name and default).
• Document it in the zh/en command-line-parameters docs.

No forwarding code is needed: swift/megatron/trainers/base.py builds the Megatron OptimizerConfig from MegatronArguments via a field-name-keyed comprehension over dataclasses.fields(config_cls) ({f.name: getattr(args, f.name) for f in dataclasses.fields(config_cls) if hasattr(args, f.name) ...}), so exposing a field whose name matches OptimizerConfig.optimizer_cuda_graph is sufficient for it to be forwarded — exactly like the existing optimizer toggles.

Notes

Typed as bool (default False) to mirror Megatron's own CLI (store_true) and OptimizerConfig field default. Verified against current main: the field is absent in Megatron-SWIFT (grep -r optimizer_cuda_graph → 0 refs) while present in Megatron-Core, and the ms-swift attribute name matches Megatron's OptimizerConfig field name exactly (so the forwarding comprehension picks it up with no additional code).

[gemini-code-assist]

Code Review

This pull request introduces a new configuration parameter, optimizer_cuda_graph, which allows users to enable CUDA graphs for optimizer steps. The parameter is added to the MegatronArguments class with a default value of False, and its documentation has been updated in both the English and Chinese command-line parameter guides. There are no review comments to address.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.