[megatron] fix: Qwen3.5 LoRA support

- move `.base_layer` normalization out of the Megatron sender and into the vLLM receiver
- make vLLM weight-name normalization robust to packed modules and fused MoE logical aliases
- fix bucketed LoRA IPC updates so multi-bucket adapters are applied only after the final bucket arrives
- avoid incorrect engine offload behavior when param offload is disabled
- improve MTP compatibility for nested HF text configs and newer Megatron-LM APIs
- add targeted regression coverage for the new weight-sync and worker behaviors

The previous weight sync path relied on sender-side name rewriting, including hard-coded `.base_layer` handling. That made the Megatron/vLLM boundary brittle for newer models, especially packed
projections and fused MoE modules.

In addition, the async LoRA update path assumed each adapter arrived in a single IPC bucket, which is not guaranteed by the bucketed transport. That could produce incomplete `add_lora` requests when
LoRA tensors were split across multiple buckets.

There were also a few compatibility issues around:
- Qwen-style nested `text_config` MTP fields
- newer Megatron-LM `process_mtp_loss` API
- actor engine offload behavior when automatic reload is not enabled

Instead of mutating exported Megatron parameter names on the sender side, the vLLM colocate worker now resolves incoming names against the live vLLM parameter namespace.

This includes:
- generic add/remove resolution for `.base_layer` leaf params (`weight` / `bias`)
- stripping Bridge-inserted `.base_layer` from non-leaf fused-MoE logical aliases
- packed-owner lookup for aliases such as `q_proj -> qkv_proj`

This keeps the sender simpler and lets the receiver normalize names based on the actual loaded vLLM model structure.

`BucketedWeightReceiver` now forwards `is_last` to the callback.

The async rollout weight update path uses that signal to:
- keep standard base-weight loading per bucket
- accumulate LoRA tensors across buckets
- call `add_lora` only once, after the final bucket is received

This aligns VERL's bucketed transport semantics with vLLM's expectation that one `add_lora` request contains a complete adapter tensor dict.

The old hard-coded stacked-parameter suffix list was removed.

`megatron_peft_utils` now exposes generic helpers to:
- add `.base_layer`
- remove `.base_layer`
- resolve the correct name by probing the target namespace

The Megatron-to-HF module mapping was also expanded for GDN modules, including:
- `in_proj -> [in_proj_qkv, in_proj_z, in_proj_b, in_proj_a]`
- `out_proj -> [out_proj]`

MTP checks and disabling logic now work with nested HF text configs as well as configs that use `mtp_num_hidden_layers` instead of only `num_nextn_predict_layers`.

The Megatron MTP patch also prefers the newer upstream `process_mtp_loss` helper when available, while preserving a fallback path for older Megatron-LM versions.

- `ActorRolloutRefWorker.update_weights` now only offloads the actor engine back to CPU when param offload is actually enabled.
- `vision_config` is now preserved in Megatron checkpoint manager backup state.

Signed-off-by: Hollow Man <hollowman@opensuse.org>

Author: HollowMan6

tests/utils/test_bucketed_weight_transfer.py

@@ -86,7 +86,12 @@ def _receiver_fn(zmq_handle, use_shm, result_queue):
         use_shm=use_shm,
     )
     received = []
-    receiver.receive_weights(on_bucket_received=lambda w: received.extend([(name, t.clone()) for name, t in w]))
+
+    def on_bucket_received(weights, *, is_last):
+        del is_last
+        received.extend([(name, t.clone()) for name, t in weights])
+
+    receiver.receive_weights(on_bucket_received=on_bucket_received)
     # Only send lightweight metadata + checksum back through the queue
     summaries = [(name, t.dtype, tuple(t.shape), t.float().sum().item()) for name, t in received]
     result_queue.put(summaries)

tests/utils/test_megatron_peft_utils.py

@@ -0,0 +1,54 @@
+# Copyright 2026 Bytedance Ltd. and/or its affiliates
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from verl.utils.megatron_peft_utils import convert_megatron_to_hf_target_modules, resolve_base_layer_name
+
+
+def test_convert_megatron_to_hf_target_modules_expands_gdn_in_proj():
+    converted = convert_megatron_to_hf_target_modules(["in_proj", "out_proj"])
+
+    assert converted == [
+        "in_proj_qkv",
+        "in_proj_z",
+        "in_proj_b",
+        "in_proj_a",
+        "out_proj",
+    ]
+
+
+def test_resolve_base_layer_name_adds_suffix_when_target_requires_it():
+    resolved_name = resolve_base_layer_name(
+        "model.layers.0.self_attn.q_proj.weight",
+        exists=lambda candidate: candidate == "model.layers.0.self_attn.q_proj.base_layer.weight",
+    )
+
+    assert resolved_name == "model.layers.0.self_attn.q_proj.base_layer.weight"
+
+
+def test_resolve_base_layer_name_removes_suffix_when_target_does_not_use_it():
+    resolved_name = resolve_base_layer_name(
+        "model.visual.merger.linear_fc1.base_layer.weight",
+        exists=lambda candidate: candidate == "model.visual.merger.linear_fc1.weight",
+    )
+
+    assert resolved_name == "model.visual.merger.linear_fc1.weight"
+
+
+def test_resolve_base_layer_name_keeps_existing_name():
+    resolved_name = resolve_base_layer_name(
+        "model.visual.merger.linear_fc1.weight",
+        exists=lambda candidate: candidate == "model.visual.merger.linear_fc1.weight",
+    )
+
+    assert resolved_name == "model.visual.merger.linear_fc1.weight"

tests/utils/test_vllm_weight_name_normalization_on_cpu.py

@@ -0,0 +1,135 @@
+# Copyright 2026 Bytedance Ltd. and/or its affiliates
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from types import SimpleNamespace
+
+import torch
+
+from verl.workers.rollout.vllm_rollout.utils import VLLM_LORA_INT_ID, vLLMColocateWorkerExtension
+
+
+class _FakeMapper:
+    def __init__(self, mapping: dict[str, str]):
+        self.mapping = mapping
+
+    def apply_list(self, names: list[str]) -> list[str]:
+        return [self.mapping.get(name, name) for name in names]
+
+
+class _FakeModel:
+    def __init__(self):
+        self.hf_to_vllm_mapper = _FakeMapper(
+            {
+                "model.language_model.layers.0.self_attn.qkv_proj.base_layer.weight": (
+                    "language_model.model.layers.0.self_attn.qkv_proj.base_layer.weight"
+                ),
+            }
+        )
+        self.packed_modules_mapping = {
+            "qkv_proj": ["q_proj", "k_proj", "v_proj"],
+        }
+
+    def named_parameters(self, remove_duplicate: bool = False):
+        del remove_duplicate
+        yield "language_model.model.layers.0.mlp.experts.base_layer.w13_weight", torch.empty(0)
+        yield "language_model.model.layers.0.mlp.experts.base_layer.w2_weight", torch.empty(0)
+        yield "language_model.model.layers.0.self_attn.qkv_proj.base_layer.weight", torch.empty(0)
+
+
+def _make_worker(model):
+    worker = object.__new__(vLLMColocateWorkerExtension)
+    worker.model_runner = SimpleNamespace(model=model)
+    return worker
+
+
+def test_normalize_base_sync_weight_names_preserves_expert_logical_aliases():
+    worker = _make_worker(_FakeModel())
+    tensor = torch.empty(0)
+
+    normalized_weights = worker._normalize_base_sync_weight_names(
+        [
+            ("model.language_model.layers.0.mlp.experts.gate_up_proj", tensor),
+            ("model.language_model.layers.0.mlp.experts.down_proj", tensor),
+            ("model.language_model.layers.0.self_attn.q_proj.weight", tensor),
+        ]
+    )
+
+    assert [name for name, _ in normalized_weights] == [
+        "model.language_model.layers.0.mlp.experts.gate_up_proj",
+        "model.language_model.layers.0.mlp.experts.down_proj",
+        "model.language_model.layers.0.self_attn.q_proj.base_layer.weight",
+    ]
+
+
+def test_normalize_base_sync_weight_names_handles_bridge_inserted_base_layer_on_fused_experts():
+    worker = _make_worker(_FakeModel())
+    tensor = torch.empty(0)
+
+    normalized_weights = worker._normalize_base_sync_weight_names(
+        [
+            ("model.language_model.layers.0.mlp.experts.base_layer.gate_up_proj", tensor),
+            ("model.language_model.layers.0.mlp.experts.base_layer.down_proj", tensor),
+        ]
+    )
+
+    assert [name for name, _ in normalized_weights] == [
+        "model.language_model.layers.0.mlp.experts.gate_up_proj",
+        "model.language_model.layers.0.mlp.experts.down_proj",
+    ]
+
+
+def test_update_weights_from_ipc_accumulates_lora_tensors_across_buckets(monkeypatch):
+    import verl.workers.rollout.vllm_rollout.bucketed_weight_transfer as bucketed_weight_transfer
+
+    class _FakeBucketReceiver:
+        def __init__(self, zmq_handle, device, use_shm):
+            del zmq_handle, device, use_shm
+
+        def receive_weights(self, on_bucket_received):
+            on_bucket_received(
+                [("layers.0.self_attn.q_proj.lora_A.weight", torch.ones(1))],
+                is_last=False,
+            )
+            on_bucket_received(
+                [("layers.0.self_attn.q_proj.lora_B.weight", torch.zeros(1))],
+                is_last=True,
+            )
+
+    monkeypatch.setattr(bucketed_weight_transfer, "BucketedWeightReceiver", _FakeBucketReceiver)
+
+    worker = _make_worker(_FakeModel())
+    worker.model_runner.vllm_config = SimpleNamespace()
+    worker.device = torch.device("cpu")
+    worker.local_rank = 0
+    worker._is_qat_model = False
+    worker._get_zmq_handle = lambda: "ipc:///tmp/test-bucketed-lora.sock"
+
+    removed_loras = []
+    added_requests = []
+    worker.remove_lora = removed_loras.append
+
+    def _add_lora(lora_request):
+        added_requests.append(lora_request)
+        return True
+
+    worker.add_lora = _add_lora
+
+    worker.update_weights_from_ipc(peft_config={"r": 1}, base_sync_done=True)
+
+    assert removed_loras == [VLLM_LORA_INT_ID]
+    assert len(added_requests) == 1
+    assert set(added_requests[0].lora_tensors) == {
+        "layers.0.self_attn.q_proj.lora_A.weight",
+        "layers.0.self_attn.q_proj.lora_B.weight",
+    }

tests/workers/test_engine_workers_update_weights.py

@@ -0,0 +1,97 @@
+# Copyright 2026 Bytedance Ltd. and/or its affiliates
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import asyncio
+from types import SimpleNamespace
+
+import torch
+from omegaconf import OmegaConf
+
+from verl.workers.engine_workers import ActorRolloutRefWorker
+
+
+class _DummyEngine:
+    def __init__(self, *, is_param_offload_enabled: bool):
+        self.is_param_offload_enabled = is_param_offload_enabled
+        self.get_per_tensor_param_calls = []
+        self.to_calls = []
+
+    def get_per_tensor_param(self, **kwargs):
+        self.get_per_tensor_param_calls.append(kwargs)
+
+        def _weights():
+            yield ("model.embed_tokens.weight", torch.tensor([1.0]))
+
+        return _weights(), None
+
+    def to(self, device: str, model: bool = True, optimizer: bool = True, grad: bool = True):
+        self.to_calls.append((device, model, optimizer, grad))
+
+
+class _DummyRollout:
+    def __init__(self):
+        self.sleep_level = 0
+        self.update_calls = []
+        self.resume_calls = []
+
+    async def resume(self, tags):
+        self.resume_calls.append(tags)
+
+    async def update_weights(self, weights, **kwargs):
+        self.update_calls.append({"weights": list(weights), **kwargs})
+
+
+def _build_worker(*, is_param_offload_enabled: bool):
+    worker = object.__new__(ActorRolloutRefWorker)
+    worker.config = OmegaConf.create(
+        {
+            "rollout": {
+                "checkpoint_engine": {"backend": "naive"},
+                "free_cache_engine": False,
+            }
+        }
+    )
+    worker.actor = SimpleNamespace(engine=_DummyEngine(is_param_offload_enabled=is_param_offload_enabled))
+    worker.rollout = _DummyRollout()
+    worker.base_sync_done = True
+    worker.layered_summon = False
+    worker.peft_merge = False
+    return worker
+
+
+def test_update_weights_does_not_offload_actor_when_param_offload_disabled(monkeypatch):
+    monkeypatch.setattr("verl.workers.engine_workers.set_expandable_segments", lambda *_: None)
+    monkeypatch.setattr("verl.workers.engine_workers.log_gpu_memory_usage", lambda *args, **kwargs: None)
+    monkeypatch.setattr("verl.workers.engine_workers.aggressive_empty_cache", lambda *args, **kwargs: None)
+
+    worker = _build_worker(is_param_offload_enabled=False)
+
+    asyncio.run(ActorRolloutRefWorker.update_weights(worker))
+
+    assert worker.actor.engine.to_calls == []
+    assert worker.actor.engine.get_per_tensor_param_calls == [{"layered_summon": False, "base_sync_done": True}]
+    assert len(worker.rollout.update_calls) == 1
+    assert worker.rollout.update_calls[0]["base_sync_done"] is True
+
+
+def test_update_weights_offloads_actor_when_param_offload_enabled(monkeypatch):
+    monkeypatch.setattr("verl.workers.engine_workers.set_expandable_segments", lambda *_: None)
+    monkeypatch.setattr("verl.workers.engine_workers.log_gpu_memory_usage", lambda *args, **kwargs: None)
+    monkeypatch.setattr("verl.workers.engine_workers.aggressive_empty_cache", lambda *args, **kwargs: None)
+
+    worker = _build_worker(is_param_offload_enabled=True)
+
+    asyncio.run(ActorRolloutRefWorker.update_weights(worker))
+
+    assert worker.actor.engine.to_calls == [("cpu", True, False, False)]

verl/utils/checkpoint/megatron_checkpoint_manager.py

@@ -645,6 +645,7 @@ def save_checkpoint(self, local_path: str, hdfs_path: str = None, global_step: i
                     "grad_sync_func",
                     "param_sync_func",
                     "generation_config",
+                    "vision_config",
                     "_pg_collection",
                 ]
                 backup = {}