[megatron] feat: Integrate Megatron-Bridge and support LoRA/PEFT

[HollowMan6]

What does this PR do?

Add concise overview of what this PR aims to achieve or accomplish. Reference related GitHub issues and PRs that help with the review.

This PR aims to add LoRA/PEFT support by integrating Megatron-Bridge into Verl while maintaining compatibility with mbridge. As a result. LoRA/PEFT support can be added to Verl with megatron backend.

Resolves #3857
Resolves #3402
Resolves #3279

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

To be added later

API and Usage Example

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

To be added later

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

To be added later

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 (飞书群).)

_{✨ Presented to you with Mind Lab - A Lab for Experiential Intelligence.}

[gemini-code-assist]

Code Review

This pull request introduces significant new functionality by integrating Megatron-Bridge to support LoRA and other PEFT methods. The changes are extensive, touching configuration, model creation, and checkpointing logic. The addition of PEFT utilities for handling adapter weights is a good step towards efficient fine-tuning. Overall, the implementation is comprehensive, but I've identified a critical bug in the checkpoint loading logic and a significant maintainability issue due to code duplication that should be addressed.

[HollowMan6]

There are 2 remaining issues that might need some help. I'm not sure if this is caused by some bugs in Megatron-Bridge or in this PR, please let me know if you find anything, and I'll also continue debugging:

• When PEFT is enabled, we will get the following error when we call get_megatron_optimizer in actor:

User exception:
zer/distrib_optimizer.py", line 584, in __init__
 self._build_optimizer_group_ranges(self.optimizer.param_groups, self.gbuf_ranges)
 File "megatron/core/optimizer/distrib_optimizer.py", line 290, in _build_optimizer_group_ranges
 group_index = world_param_group_map[param]
 ~~~~~~~~~~~~~~~~~~~~~^^^^^^^
KeyError: Parameter containing:
tensor([1., 1., 1., ..., 1., 1., 1.], device='cuda:0', dtype=torch.bfloat16)

• (I think I might made some mistakes in this PR for this one, as several CI failed with OOM, and it shouldn't be related to Megatron-Bridge) Without PEFT but with Megatron-Bridge, we will have "Tensor exists on more than one PP rank" error when vLLM updates weight, and this happens for all the params, not only embed_tokens (so this is not caused by tie_word_embedding)

File "verl/workers/megatron_workers.py", line 686, in rollout_mode
  for name, weight in self.bridge.export_hf_weights(self.actor.actor_module, cpu=True, show_progress=False):
File "megatron/bridge/models/conversion/model_bridge.py", line 623, in stream_weights_megatron_to_hf
  converted_weights_dict = task.mapping.megatron_to_hf(task.param_weight, task.megatron_module)
                           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "megatron/bridge/models/conversion/param_mapping.py", line 1237, in megatron_to_hf
  result = self._mapping.megatron_to_hf(megatron_weights, megatron_module)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "megatron/bridge/models/conversion/param_mapping.py", line 995, in megatron_to_hf
  megatron_weights = self.broadcast_from_pp_rank(megatron_weights, cache_key=str(self.hf_param))
                     ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "megatron/bridge/models/conversion/param_mapping.py", line 336, in broadcast_from_pp_rank
  raise ValueError(f"Tensor exists on more than one PP rank. Found on ranks {src_rank} and {rank}.")
ValueError: Tensor exists on more than one PP rank. Found on ranks 0 and 1.

cc @ISEEKYAN

[ISEEKYAN]

@HollowMan6 good job! @conver334 has also been developing this feature lately, I will check this PR with her tomorrow and leave feedbacks.

[HollowMan6]

/gemini review

[gemini-code-assist]

Code Review

This pull request integrates Megatron-Bridge to add LoRA/PEFT support, which is a significant and complex feature. The changes touch upon model creation, configuration, and checkpointing to accommodate the new backend and PEFT methods. While the overall direction is good, I've identified a few critical and high-severity issues. There's a critical bug where base model weights are not loaded when using Megatron-Bridge, which would lead to using an uninitialized model. Another critical issue is a placeholder implementation for make_value_model, which will cause problems for critic/value model training. Additionally, there are several instances of significant code duplication for configuration parsing that should be refactored into utility functions to improve maintainability. Addressing these points will be crucial for the stability and correctness of this new functionality.

[HollowMan6]

Thank you @ISEEKYAN ! I have solved 2 issues mentioned in #4063 (comment) by myself. The first one is that we should set PEFT first before we wrap with DDP, the second one is that I forgot to pass those pp, tp, ep configs to the Megatron-Bridge provider.

I think now this PR is ready to get some systematic review. Aside from #4063 (comment) (which I'm not sure if there's any equivalent option in Megatron-Bridge/how to modify), and the below issue related to export HF weights with PEFT adapters, I have tested with the following scripts with saving/loading checkpoints (with lora.rank=16 or 0 (full params)), and they all look good:

ray job submit -- python3 -m verl.trainer.main_ppo --config-path=config \
    --config-name='ppo_megatron_trainer.yaml'\
    algorithm.adv_estimator=grpo \
    trainer.val_before_train=False \
    data.train_files=$HOME/data/gsm8k/train.parquet \
    data.val_files=$HOME/data/gsm8k/test.parquet \
    data.train_batch_size=16 \
    data.max_prompt_length=512 \
    data.max_response_length=1024 \
    data.filter_overlong_prompts=True \
    data.truncation='error' \
    data.shuffle=False \
    actor_rollout_ref.model.path="Qwen/Qwen2.5-7B-Instruct" \
    actor_rollout_ref.model.lora.rank=16 \
    actor_rollout_ref.model.lora.alpha=32 \
    actor_rollout_ref.model.use_fused_kernels=True \
    actor_rollout_ref.actor.optim.lr=1e-6 \
    actor_rollout_ref.actor.ppo_mini_batch_size=16 \
    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=2 \
    actor_rollout_ref.actor.megatron.use_mbridge=True \
    actor_rollout_ref.actor.megatron.pipeline_model_parallel_size=1 \
    actor_rollout_ref.actor.megatron.tensor_model_parallel_size=2 \
    actor_rollout_ref.actor.use_kl_loss=True \
    actor_rollout_ref.actor.kl_loss_coef=0.001 \
    actor_rollout_ref.actor.kl_loss_type=low_var_kl \
    actor_rollout_ref.actor.entropy_coeff=0 \
    actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=4 \
    actor_rollout_ref.rollout.tensor_model_parallel_size=2 \
    actor_rollout_ref.rollout.name=vllm \
    actor_rollout_ref.rollout.gpu_memory_utilization=0.8 \
    actor_rollout_ref.rollout.n=4 \
    actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=4 \
    actor_rollout_ref.ref.megatron.pipeline_model_parallel_size=1 \
    actor_rollout_ref.ref.megatron.tensor_model_parallel_size=2 \
    algorithm.use_kl_in_reward=False \
    trainer.logger='["console"]' \
    trainer.project_name='verl_grpo_example_gsm8k_math' \
    trainer.experiment_name='qwen2_7b_megatron' \
    trainer.n_gpus_per_node=8 \
    trainer.nnodes=1 \
    trainer.save_freq=-1 \
    trainer.test_freq=5 \
    trainer.total_epochs=15 $@

except that I face the following error when I set pp=2, though it looks like to be model specific.

  File "verl/single_controller/ray/base.py", line 700, in func
    return getattr(self.worker_dict[key], name)(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "verl/single_controller/base/decorator.py", line 442, in inner
    return func(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^
  File "verl/utils/transferqueue_utils.py", line 199, in dummy_inner
    return func(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^
  File "verl/utils/profiler/performance.py", line 105, in f
    return self.log(decorated_function, *args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "verl/utils/profiler/performance.py", line 118, in log
    output = func(*args, **kwargs)
             ^^^^^^^^^^^^^^^^^^^^^
  File "verl/utils/profiler/profile.py", line 256, in wrapper
    return func(self_instance, *args, **kwargs_inner)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "verl/workers/megatron_workers.py", line 870, in compute_log_prob
    output, entropys = self.actor.compute_log_prob(data=data, calculate_entropy=True)
                       ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "verl/utils/profiler/performance.py", line 105, in f
    return self.log(decorated_function, *args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "verl/utils/profiler/performance.py", line 118, in log
    output = func(*args, **kwargs)
             ^^^^^^^^^^^^^^^^^^^^^
  File "verl/workers/actor/megatron_actor.py", line 217, in compute_log_prob
    output = self.forward_backward_batch(
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "verl/workers/actor/megatron_actor.py", line 617, in forward_backward_batch
    losses_reduced = forward_backward_func(
                     ^^^^^^^^^^^^^^^^^^^^^^
  File "megatron/core/pipeline_parallel/schedules.py", line 2178, in forward_backward_pipelining_without_interleaving
    output_tensor, num_tokens = forward_step(
                                ^^^^^^^^^^^^^
  File "megatron/core/pipeline_parallel/schedules.py", line 402, in forward_step
    output_tensor, loss_func = forward_step_func(data_iterator, model)
                               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "verl/workers/actor/megatron_actor.py", line 590, in forward_step
    output = forward_fn(
             ^^^^^^^^^^^
  File "verl/models/mcore/model_forward.py", line 75, in model_forward
    output_orig = model(**input_args)
                  ^^^^^^^^^^^^^^^^^^^
  File "torch/nn/modules/module.py", line 1773, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "torch/nn/modules/module.py", line 1784, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "megatron/core/distributed/data_parallel_base.py", line 22, in forward
    return self.module(*inputs, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "torch/nn/modules/module.py", line 1773, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "torch/nn/modules/module.py", line 1784, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "megatron/core/transformer/module.py", line 435, in forward
    outputs = self.module(*inputs, **kwargs)
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "torch/nn/modules/module.py", line 1773, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "torch/nn/modules/module.py", line 1784, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "megatron/core/models/gpt/gpt_model.py", line 453, in forward
    hidden_states = self.decoder(
                    ^^^^^^^^^^^^^
  File "torch/nn/modules/module.py", line 1773, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "torch/nn/modules/module.py", line 1784, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "megatron/core/transformer/transformer_block.py", line 649, in forward
    hidden_states, context = layer(
                             ^^^^^^
  File "megatron/core/transformer/transformer_layer.py", line 865, in __call__
    return super().__call__(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "megatron/core/transformer/module.py", line 311, in __call__
    return super().__call__(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "torch/nn/modules/module.py", line 1773, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "torch/nn/modules/module.py", line 1784, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "megatron/core/transformer/transformer_layer.py", line 444, in forward
    hidden_states, context = self._forward_attention(*args, **kwargs)
                             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "megatron/core/transformer/transformer_layer.py", line 509, in _forward_attention
    attention_output_with_bias = self.self_attention(
                                 ^^^^^^^^^^^^^^^^^^^^
  File "torch/nn/modules/module.py", line 1773, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "torch/nn/modules/module.py", line 1784, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "megatron/core/transformer/attention.py", line 828, in forward
    query = apply_rotary_pos_emb(
            ^^^^^^^^^^^^^^^^^^^^^
  File "megatron/core/models/common/embeddings/rope_utils.py", line 238, in apply_rotary_pos_emb
    return _apply_rotary_pos_emb_thd(
           ^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "megatron/core/models/common/embeddings/rope_utils.py", line 181, in _apply_rotary_pos_emb_thd
    for x in torch.split(t, seqlens)
             ^^^^^^^^^^^^^^^^^^^^^^^
  File "torch/functional.py", line 222, in split
    return tensor.split(split_size_or_sections, dim)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "torch/_tensor.py", line 1052, in split
    return torch._VF.split_with_sizes(self, split_size, dim)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
RuntimeError: split_with_sizes expects split_sizes to sum exactly to 6144 (input tensor's size at dimension 0), but got split_sizes=[1098, 1108, 1134, 1168]

cc: @conver334

[ccclyu]

thanks for the great contribution and one general comment: can we have lora related CI/CD test?

[HollowMan6]

can we have lora related CI/CD test?

Yes, sure, but I think the piority is that we get weight updating/loading working first: #4063 (comment) after that, we can add some test cases here.

[ISEEKYAN]

Thank you @ISEEKYAN ! I have solved 2 issues mentioned in #4063 (comment) by myself. The first one is that we should set PEFT first before we wrap with DDP, the second one is that I forgot to pass those pp, tp, ep configs to the Megatron-Bridge provider.

I think now this PR is ready to get some systematic review. Aside from #4063 (comment) (which I'm not sure if there's any equivalent option in Megatron-Bridge/how to modify), and the below issue related to export HF weights with PEFT adapters, I have tested with the following scripts with saving/loading checkpoints (with lora.rank=16 or 0 (full params)), and they all look good:

except that I face the following error when I set pp=2, though it looks like to be model specific.

cc: @conver334

Good job!
@conver334 will help with testing your great work and make sure the integration's correctness.

For more formal and advanced usage of megatron-bridge, I think nemo-rl is a very good reference.

[gemini-code-assist]

Code Review

This pull request integrates Megatron-Bridge to support LoRA/PEFT, which is a significant feature addition. The changes are extensive, touching configuration, model initialization, and checkpointing logic across multiple components. The introduction of a vanilla_mbridge flag to switch between bridge implementations and the use of a provider object for Megatron-Bridge are well-structured. The PEFT support, including LoRA configuration parsing, model transformation, and adapter-only checkpointing, appears to be correctly implemented.

My main feedback concerns code duplication. I've identified two instances of repeated logic across several files: one for parsing LoRA configurations and another for stripping LoRA-specific keys from model configurations. Consolidating these into shared utility functions would greatly improve the code's maintainability and reduce the risk of future inconsistencies. I've left specific comments pointing to these areas.

[gemini-code-assist]

Code Review

This pull request introduces a significant feature by integrating Megatron-Bridge to support LoRA/PEFT, which is a substantial enhancement for the project. The changes are extensive, touching configuration, model initialization, and checkpointing logic across many files. The implementation correctly distinguishes between the original mbridge and the new Megatron-Bridge and introduces a flexible PEFT configuration system. The core logic for applying PEFT transformations before DDP wrapping via Megatron-Bridge's provider hooks appears to be correctly implemented. However, I've identified a critical issue that would prevent LoRA from being enabled, and a high-severity performance issue related to checkpointing with PEFT. Addressing these points will help ensure the new functionality works as expected and is efficient.

[HollowMan6]

Thank you @conver334 and @ISEEKYAN again for reviewing! I've successfully integrated LoRA Merge and related mapping in Megatron-Bridge: NVIDIA-NeMo/Megatron-Bridge#1310, I've also tested this with TP, PP, and EP, so that now load_hf_weights and export_hf_weights works correctly for LoRA enabled models as well by merging weights. I recommend to test this PR with the main branch mcore due to NVIDIA-NeMo/Megatron-Bridge#1289 (comment) After NVIDIA-NeMo/Megatron-Bridge#1310 is merged and released, we can consider adding some CI test cases for megatron lora.

I've also tested this PR, both w/ and w/o lora using Megatron-Bridge, aside from the pp issue mentioned in #4063 (review), which I think is probably related to Megatron-Bridge as it works fine when I use vanilla mbridge, I find that when I combine tp with ep for Qwen3-30B-A3B (please check the script I use in this PR examples/grpo_trainer/run_qwen3moe-30b_megatron_lora.sh), it will stuck in the forward pass (see the stack trace when stuck happens below), which I'm not sure what's going on here. Since I'm not familiar with Megatron-Bridge, I'll probably need your help to resolve these issues. Please let me know if you have any insights about these 2 issues, thanks in advance!

    forward (core/tensor_parallel/layers.py:480)
    decorate_fwd (torch/amp/autocast_mode.py:517)
    apply (torch/autograd/function.py:576)
    linear_with_grad_accumulation_and_async_allreduce (core/tensor_parallel/layers.py:737)
    _forward_impl (core/tensor_parallel/layers.py:944)
    forward (core/tensor_parallel/layers.py:1014)
    _call_impl (torch/nn/modules/module.py:1784)
    _wrapped_call_impl (torch/nn/modules/module.py:1773)
    forward (bridge/peft/utils.py:581)
    _call_impl (torch/nn/modules/module.py:1784)
    _wrapped_call_impl (torch/nn/modules/module.py:1773)
    forward (bridge/peft/lora_layers.py:56)
    _call_impl (torch/nn/modules/module.py:1784)
    _wrapped_call_impl (torch/nn/modules/module.py:1773)
    forward (core/transformer/moe/experts.py:882)
    _call_impl (torch/nn/modules/module.py:1784)
    _wrapped_call_impl (torch/nn/modules/module.py:1773)
    routed_experts_compute (core/transformer/moe/moe_layer.py:231)
    custom_forward (core/transformer/moe/moe_layer.py:276)
    forward (core/transformer/moe/moe_layer.py:292)
    _call_impl (torch/nn/modules/module.py:1784)
    _wrapped_call_impl (torch/nn/modules/module.py:1773)
    _forward_mlp (core/transformer/transformer_layer.py:618)
    forward (core/transformer/transformer_layer.py:435)
    _call_impl (torch/nn/modules/module.py:1784)
    _wrapped_call_impl (torch/nn/modules/module.py:1773)
    __call__ (core/transformer/module.py:305)
    __call__ (core/transformer/transformer_layer.py:852)
    custom_forward (core/transformer/transformer_block.py:426)
    forward (core/tensor_parallel/random.py:426)
    apply (torch/autograd/function.py:576)
    checkpoint (core/tensor_parallel/random.py:480)
    checkpoint_handler (core/transformer/transformer_block.py:456)
    _checkpointed_forward (core/transformer/transformer_block.py:472)
    forward (core/transformer/transformer_block.py:669)
    _call_impl (torch/nn/modules/module.py:1784)
    _wrapped_call_impl (torch/nn/modules/module.py:1773)
    __call__ (core/transformer/module.py:305)
    __call__ (core/transformer/transformer_block.py:553)
    _fused_GPTModel_forward (verl/models/mcore/model_forward_fused.py:175)
    _call_impl (torch/nn/modules/module.py:1784)
    _wrapped_call_impl (torch/nn/modules/module.py:1773)
    forward (core/transformer/module.py:429)
    _call_impl (torch/nn/modules/module.py:1784)
    _wrapped_call_impl (torch/nn/modules/module.py:1773)
    forward (core/distributed/data_parallel_base.py:22)
    _call_impl (torch/nn/modules/module.py:1784)
    _wrapped_call_impl (torch/nn/modules/module.py:1773)
    fused_forward_model (verl/models/mcore/model_forward_fused.py:117)
    forward_step (verl/workers/actor/megatron_actor.py:561)
    forward_step (core/pipeline_parallel/schedules.py:402)
    forward_backward_no_pipelining (core/pipeline_parallel/schedules.py:615)
    forward_backward_batch (verl/workers/actor/megatron_actor.py:633)
    compute_log_prob (verl/workers/actor/megatron_actor.py:217)
    log (verl/utils/profiler/performance.py:118)
    f (verl/utils/profiler/performance.py:105)
    compute_log_prob (verl/workers/megatron_workers.py:801)
    wrapper (verl/utils/profiler/profile.py:256)
    log (verl/utils/profiler/performance.py:118)
    f (verl/utils/profiler/performance.py:105)
    dummy_inner (verl/utils/transferqueue_utils.py:199)
    inner (verl/single_controller/base/decorator.py:442)
    func (verl/single_controller/ray/base.py:700)
    _resume_span (ray/util/tracing/tracing_helper.py:461)
    actor_method_executor (ray/_private/function_manager.py:693)
    main_loop (ray/_private/worker.py:1042)
    <module> (ray/_private/workers/default_worker.py:322)

[ryxli]

Wondering if you have encountered this issue in your testing:

When loading a model with Megatron-Bridge using TP > 1 and PP > 1 across multiple nodes, it looks like each TP rank within a PP group ends up loading the full PP shard (all layers for that PP stage), instead of being further sharded by TP.

Encountering an issue like this example:

• Model: ~10B parameters, 8 layers
• TP8, PP2
• 2 nodes with 8 h100 gpus

This should end up with 2 PP groups, each responsible for 5 layers:

• PP stage 0: layers 1–4
• PP stage 1: layers 5-8

Each PP stage should own ~10B / 2 = 5B parameters.
This ends up behaving as expected when I inspect the loaded actor module.

For TP, I encounter the following -
Expected Behavior:

• Within each PP group, the 5B parameters should be sharded across TP=8 ranks, so that each TP rank only holds roughly 5B / 8 parameters for that PP stage.

Observed Behavior:

• every TP rank within a PP group appears to load the full 5B parameters for that PP stage. In other words, the PP sharding works (each PP group gets the correct subset of layers), but TP sharding does not appear to be applied on top of that for the loaded weights, and each TP rank in the PP group ends up with the full PP shard.

This results in OOM very quickly as the DDP wrapper initializes the param buffers for the entire PP stage on every TP rank

[HollowMan6]

Wondering if you have encountered this issue in your testing:

@ryxli
Are you using the latest commit of this PR? I used to have an issue on my side for my initial draft (see issue 2: #4063 (comment)) when I export to hf weight using Megatron Bridge, but it turned out to be that I forgot to pass the pp size value to the provider during initialization.

However, actually I haven't fully tested TP+PP as I encountered another shape mismatch issue as mentioned in: #4063 (review) so I'm not sure about your situation if it's not caused by using an older commit of this PR, and your issue might be related to the shape mismatch issue as well. Please let me know if you have any further insights, thanks!

[HollowMan6]

Hi @rtz1998，我建议你安装Megatron-Bridge的时候用NVIDIA-NeMo/Megatron-Bridge@a489bed 这个commit，具体可以参考https://github.com/volcengine/verl/pull/4533/changes#diff-ede636eb585931bd42b3b527bff08e1a3fab254eaf02fdd0fa33734e6dcafa0aR125-R126

脚本我用的examples/grpo_trainer/run_qwen2-7b_math_megatron_lora.sh测试的，你也可以看一下哪里不一样

[rtz1998]

Hi @rtz1998，我建议你安装Megatron-Bridge的时候用NVIDIA-NeMo/Megatron-Bridge@a489bed 这个commit，具体可以参考https://github.com/volcengine/verl/pull/4533/changes#diff-ede636eb585931bd42b3b527bff08e1a3fab254eaf02fdd0fa33734e6dcafa0aR125-R126

脚本我用的examples/grpo_trainer/run_qwen2-7b_math_megatron_lora.sh测试的，你也可以看一下哪里不一样

依然有问题，重新安装Megatron-Bridge和Megatron-LM后，脚本也一致。更奇怪的是Moe模型是正常的

安装命令如下：
pip3 uninstall megatron-bridge megatron-core -y

安装指定 commit 的 Megatron-Bridge

python3 -m pip install --no-deps --no-build-isolation
git+https://github.com/NVIDIA-NeMo/Megatron-Bridge.git@a489bed

安装指定 commit 的 Megatron-LM

python3 -m pip install --no-deps --no-build-isolation
git+https://github.com/NVIDIA/Megatron-LM.git@2d398b4

[HollowMan6]

@rtz1998 I just double checked locally and found that for dense, it's OK when setting pipeline_model_parallel_size > 1, and grad_norm will remain 0 if we set pipeline_model_parallel_size = 1. I will double check the logic in NVIDIA-NeMo/Megatron-Bridge#1603 and see what might cause this issue, meanwhile, you can set pipeline_model_parallel_size > 1 as a workaround.

I just found that it will also work if you disable activation checkpointing by removing these lines in the training script:

    +actor_rollout_ref.actor.megatron.override_transformer_config.recompute_method=uniform
    +actor_rollout_ref.actor.megatron.override_transformer_config.recompute_granularity=full
    +actor_rollout_ref.actor.megatron.override_transformer_config.recompute_num_layers=1

[rtz1998]

@rtz1998 I just double checked locally and found that for dense, it's OK when setting pipeline_model_parallel_size > 1, and grad_norm will remain 0 if we set pipeline_model_parallel_size = 1. I will double check the logic in NVIDIA-NeMo/Megatron-Bridge#1603 and see what might cause this issue, meanwhile, you can set pipeline_model_parallel_size > 1 as a workaround.

I just found that it will also work if you disable activation checkpointing by removing these lines in the training script:

    +actor_rollout_ref.actor.megatron.override_transformer_config.recompute_method=uniform
    +actor_rollout_ref.actor.megatron.override_transformer_config.recompute_granularity=full
    +actor_rollout_ref.actor.megatron.override_transformer_config.recompute_num_layers=1

ty, when setting pipeline_model_parallel_size to 2, grad norm returns to normal ！

[HollowMan6]

Setting activation checkpointing with PP>1 will still affect convergence a bit (refer to NVIDIA-NeMo/Megatron-Bridge#1750 (comment)), I just opened #4566 to fix LoRA with activation checkpointing under megatron, please feel free to try it out and let me know if you face any issues @rtz1998

[HollowMan6]

Now PEFT recompute fix has been merged into Megatron-Bridge main branch, you can install via python3 -m pip install --no-deps --no-build-isolation git+https://github.com/NVIDIA-NeMo/Megatron-Bridge.git@953aabf