support optional megatron-bridge backend#9725
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces support for the megatron-bridge backend as an alternative to mcore-bridge in Megatron-SWIFT, including validation constraints, weight loading/saving adapters, and JIT warmup adjustments. Feedback on these changes highlights two issues: first, the MoE check for num_moe_experts in _get_megatron_bridge_model should also check the provider configuration to avoid incorrectly overriding parallel sizes; second, additional_saved_files should be retrieved from processor.model_meta instead of args to ensure custom files are properly copied during checkpoint saving.
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.
|
/gemini review |
There was a problem hiding this comment.
Code Review
This pull request introduces support for the megatron-bridge backend as an alternative to mcore-bridge in Megatron-SWIFT. It adds a new bridge_backend argument, implements the MegatronBridgeBackend adapter class to interface with NVIDIA's AutoBridge, and updates model initialization and validation logic accordingly. Feedback on the changes highlights two issues: first, _check_mcore_bridge() is called unconditionally, which incorrectly enforces mcore-bridge version constraints even when using the megatron-bridge backend; second, the torch module is missing an import in swift/megatron/model/utils.py despite being referenced in type annotations.
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.
|
test with sft tp2pp2 |
| bias.requires_grad, input_tensor.requires_grad = bias_grad, input_grad | ||
| for _ in range(5): | ||
| if config.swiglu: | ||
| if getattr(config, 'swiglu', getattr(config, 'gated_linear_unit', True)): |
There was a problem hiding this comment.
| os.environ.pop('VLLM_USE_MODELSCOPE', None) | ||
| logging_level = logging.root.level | ||
| _patch_unified_memory() | ||
| _patch_mcore_bridge() |
There was a problem hiding this comment.
移到MegatronArguments._check_bridge_backend里了
3 participants
No description provided.