[megatron] feat: use yaml to manage mbridge args (verl-project#4584)

### 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.

```python
# 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](https://github.com/volcengine/verl/blob/main/CONTRIBUTING.md).
- [ ] Apply [pre-commit
checks](https://github.com/volcengine/verl/blob/main/CONTRIBUTING.md#code-linting-and-formatting):
`pre-commit install && pre-commit run --all-files --show-diff-on-failure
--color=always`
- [ ] Add / Update [the
documentation](https://github.com/volcengine/verl/tree/main/docs).
- [ ] Add unit or end-to-end test(s) to [the CI
workflow](https://github.com/volcengine/verl/tree/main/.github/workflows)
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](https://verl-project.slack.com/archives/C091TCESWB1) in [the
`verl` Slack
workspace](https://join.slack.com/t/verl-project/shared_invite/zt-3855yhg8g-CTkqXu~hKojPCmo7k_yXTQ).
(If not accessible, please try [the Feishu group
(飞书群)](https://applink.larkoffice.com/client/chat/chatter/add_by_link?link_token=772jd4f1-cd91-441e-a820-498c6614126a).)

---------

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
Co-authored-by: Yan Bai <[email protected]>

Author: 3 people

verl/trainer/config/_generated_ppo_megatron_trainer.yaml

@@ -107,6 +107,7 @@ actor_rollout_ref:
       - extra
       load_contents: ${.save_contents}
       async_save: false
+      mbridge_config: {}
     use_fused_kernels: ${oc.select:actor_rollout_ref.model.use_fused_kernels,false}
     profiler:
       _target_: verl.utils.profiler.ProfilerConfig
@@ -551,6 +552,7 @@ critic:
     - extra
     load_contents: ${.save_contents}
     async_save: false
+    mbridge_config: {}
   profiler:
     _target_: verl.utils.profiler.ProfilerConfig
     tool: ${oc.select:global_profiler.tool,null}

verl/trainer/config/_generated_ppo_trainer.yaml

@@ -87,6 +87,7 @@ actor_rollout_ref:
       - extra
       load_contents: ${.save_contents}
       async_save: false
+      mbridge_config: {}
     use_fused_kernels: ${oc.select:actor_rollout_ref.model.use_fused_kernels,false}
     profiler:
       _target_: verl.utils.profiler.ProfilerConfig
@@ -469,6 +470,7 @@ critic:
     - extra
     load_contents: ${.save_contents}
     async_save: false
+    mbridge_config: {}
   profiler:
     _target_: verl.utils.profiler.ProfilerConfig
     tool: ${oc.select:global_profiler.tool,null}

verl/trainer/config/actor/actor.yaml

@@ -132,6 +132,13 @@ checkpoint:
 
   # Whether to save checkpoints asynchronously. Only effective for Megatron as of now.
   async_save: False
+  
+  # Mbridge config extension.
+  # when vanilla_mbridge=True, and your filesystem is a distributed filesystem,(which means you write a file in node A
+  # and you can read the file in node B immediately)
+  # set `mbridge_config.distributed_filesystem=True` and `mbridge_config.memory_efficient=True` to 
+  # speed up the checkpoint saving by 10x speed.
+  mbridge_config: {}
 
 # optimizer configs
 optim:

verl/trainer/config/config.py

@@ -36,6 +36,7 @@ class CheckpointConfig(BaseConfig):
     save_contents: list[str] = field(default_factory=lambda: ["model", "optimizer", "extra"])
     load_contents: list[str] = field(default_factory=lambda: ["model", "optimizer", "extra"])
     async_save: bool = False
+    mbridge_config: dict[str, Any] = field(default_factory=dict)
 
 
 @dataclass

verl/trainer/config/critic/critic.yaml

@@ -98,6 +98,9 @@ checkpoint:
   # Whether to save checkpoints asynchronously. Only effective for Megatron as of now.
   async_save: False
 
+  # Mbridge config extension.
+  mbridge_config: {}
+
 # profile the critic model in `update_critic`
 profiler:
 

verl/trainer/config/sft_trainer_engine.yaml

@@ -55,6 +55,8 @@ checkpoint:
 
   # For more flexibility, you can specify the contents to load from the checkpoint.
   load_contents: ${checkpoint.save_contents}
+  # Mbridge config extension.
+  mbridge_config: {}
 
 trainer:
   default_local_dir: checkpoints/${trainer.project_name}/${trainer.experiment_name}

verl/utils/checkpoint/megatron_checkpoint_manager.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import inspect
 import json
 import logging
 import os
@@ -578,9 +579,14 @@ def save_checkpoint(self, local_path: str, hdfs_path: str = None, global_step: i
                 log_with_rank(f"Saving HF model checkpoint to {local_path} with bridge", rank=self.rank, logger=logger)
                 hf_ckpt_path = get_hf_model_checkpoint_path(local_path)
                 if self.vanilla_bridge:
-                    self.bridge.save_weights(
-                        self.model, hf_ckpt_path, distributed_filesystem=True, memory_efficient=True
-                    )
+                    extended_args = {}
+                    mbridge_config = getattr(self.checkpoint_config, "mbridge_config", None) or {}
+                    for sig in inspect.signature(self.bridge.save_weights).parameters:
+                        if sig == "weights_path" or sig == "models":
+                            continue
+                        if sig in mbridge_config:
+                            extended_args[sig] = mbridge_config[sig]
+                    self.bridge.save_weights(self.model, hf_ckpt_path, **extended_args)
                 else:
                     self.bridge.save_hf_weights(self.model, hf_ckpt_path)
 
@@ -651,9 +657,14 @@ def save_checkpoint(self, local_path: str, hdfs_path: str = None, global_step: i
             if self.bridge is not None:
                 hf_model_ckpt_path = get_hf_model_checkpoint_path(local_path)
                 if self.vanilla_bridge:
-                    self.bridge.save_weights(
-                        self.model, hf_model_ckpt_path, distributed_filesystem=True, memory_efficient=True
-                    )
+                    extended_args = {}
+                    mbridge_config = getattr(self.checkpoint_config, "mbridge_config", None) or {}
+                    for sig in inspect.signature(self.bridge.save_weights).parameters:
+                        if sig == "weights_path" or sig == "models":
+                            continue
+                        if sig in mbridge_config:
+                            extended_args[sig] = mbridge_config[sig]
+                    self.bridge.save_weights(self.model, hf_model_ckpt_path, **extended_args)
                 else:
                     self.bridge.save_hf_weights(self.model, hf_model_ckpt_path)
             else: