[megatron] feat: use yaml to manage mbridge args

[Kite0011]

What does this PR do?

Use dict to manage mbridge args, support modifying mbridge input args, and avoid mbridge version compatibility issues.

Checklist Before Starting

• Search for similar PRs. Paste at least one query link here: ...
• Format the PR title as [{modules}] {type}: {description} (This will be checked by the CI)
  • {modules} include fsdp, megatron, sglang, vllm, rollout, trainer, ci, training_utils, recipe, hardware, deployment, ray, worker, single_controller, misc, perf, model, algo, env, tool, ckpt, doc, data, cfg, reward
  • If this PR involves multiple modules, separate them with , like [megatron, fsdp, doc]
  • {type} is in feat, fix, refactor, chore, test
  • If this PR breaks any API (CLI arguments, config, function signature, etc.), add [BREAKING] to the beginning of the title.
  • Example: [BREAKING][fsdp, megatron] feat: dynamic batching

Test

For changes that can not be tested by CI (e.g., algorithm implementation, new model support), validate by experiment(s) and show results like training curve plots, evaluation results, etc.

API and Usage Example

Demonstrate how the API changes if any, and provide usage example(s) if possible.

# Add code snippet or script demonstrating how to use this

Design & Code Changes

Demonstrate the high-level design if this PR is complex, and list the specific changes.

Checklist Before Submitting

Important

Please check all the following items before requesting a review, otherwise the reviewer might deprioritize this PR for review.

• Read the Contribute Guide.
• Apply pre-commit checks: pre-commit install && pre-commit run --all-files --show-diff-on-failure --color=always
• Add / Update the documentation.
• Add unit or end-to-end test(s) to the CI workflow to cover all the code. If not feasible, explain why: ...
• Once your PR is ready for CI, send a message in the ci-request channel in the verl Slack workspace. (If not accessible, please try the Feishu group (飞书群).)

[gemini-code-assist]

Code Review

This pull request refactors how parameters are passed to mbridge.save_weights by using a dictionary from the configuration. While this increases flexibility, the current implementation has some critical issues. It inadvertently changes the default behavior by dropping the distributed_filesystem=True parameter and changing memory_efficient from True to False. Additionally, it introduces a risk of a TypeError at runtime if mbridge_config is set to null in the configuration. My review includes suggestions to restore the original default behavior, fix the potential crash, and improve the configuration's robustness.

[Kite0011]

/gemini review

[gemini-code-assist]

Code Review

This pull request refactors how arguments are passed to mbridge.save_weights by introducing a mbridge_config dictionary. This is a good improvement for flexibility and maintainability. The implementation is clean and correctly handles the new configuration. However, I've noticed that in some of the generated configuration files, the value of memory_efficient is set to false, whereas it was previously hardcoded to true. This could lead to increased memory usage during checkpointing. My review includes comments highlighting these specific locations. Please verify if this change is intentional.

[Kite0011]

/gemini review

[gemini-code-assist]

Code Review

This pull request refactors the management of mbridge arguments by moving them from hardcoded values into a configurable dictionary in YAML files. This is a good improvement for flexibility. The implementation is mostly correct, but I've found a critical issue where the code does not handle the case where checkpoint_config is None, which could lead to a crash when saving checkpoints. I've provided suggestions to make the code more robust against this potential AttributeError.

[Kite0011]

/gemini review

[gemini-code-assist]

Code Review

This pull request refactors the management of mbridge arguments by moving them from hardcoded values into YAML configuration files. This is a good improvement for flexibility and maintainability. The changes across the configuration files and the CheckpointConfig dataclass are well-implemented. I've found one minor issue in megatron_checkpoint_manager.py regarding inconsistent and potentially unsafe access to the new configuration property, for which I've left a specific comment.

[Kite0011]

/gemini review

[gemini-code-assist]

Code Review

This pull request refactors the management of mbridge arguments by moving them into YAML configurations, which is a great step towards more flexible and maintainable code. My main feedback is about the handling of default values for mbridge_config. By setting the default to an empty dictionary in CheckpointConfig, the default behavior of save_weights is altered unless mbridge_config is explicitly defined in every relevant YAML file. This has led to code duplication across several configuration files and makes the system fragile. I've suggested setting the default values directly within the CheckpointConfig dataclass. This will ensure the previous behavior is maintained by default, eliminate the need for duplicated configuration blocks in the YAML files, and make the system more robust for future extensions.

[Kite0011]

/gemini review

[gemini-code-assist]

Code Review

This pull request refactors how mbridge arguments are managed by moving them into a mbridge_config dictionary within the checkpoint configuration. This is a good change that improves configurability and helps with version compatibility. I've found a critical issue in the implementation where the code does not handle the case where checkpoint_config might be None, which would lead to a crash. I've provided suggestions to fix this potential AttributeError.

[Kite0011]

cc @ISEEKYAN

[ISEEKYAN]

please don't add these yaml configs.

[ISEEKYAN]

recommend to modify like this:

import inspect
extended_args = {}
if "distributed_filesystem" in inspect.signature(self.bridge.save_weights):
    extended_args["distributed_filesystem"] = True
    extended_args["memory_efficient"] = True
self.bridge.save_weights(self.model, hf_ckpt_path, **extended_args)

[ISEEKYAN]

same as above

[Kite0011]

In the argument passing section, some modifications were made to better support backward compatibility with older versions and adapt to newer versions of mbridge.
I used inspect to retrieve the "distributed_filesystem" and "memory_efficient" args for backward compatibility with older versions, while also adding a new field to facilitate compatibility with future versions or internal mbridge versions.
@ISEEKYAN

[Kite0011]

We can use +checkpoint.mbridge_config.memory_efficient to add/set mbridge.save_weights func.