[megatron] Convert fused-MoE expert LoRA to vLLM layout on adapter sync

skyrl/backends/skyrl_train/distributed/megatron/megatron_utils.py

@@ -21,7 +21,7 @@
 # limitations under the License.
 
 import gc
-from typing import Any, List, Optional, Union
+from typing import Any, Dict, List, Optional, Union
 
 import torch
 import torch.nn as nn
@@ -162,6 +162,38 @@ def freeze_moe_router(model_or_models: Union[nn.Module, List[nn.Module]]):
     return model_or_models
 
 
+def _convert_moe_experts_lora_to_vllm(
+    adapter_state: Dict[str, "torch.Tensor"],
+) -> Dict[str, "torch.Tensor"]:
+    """Rewrite fused-MoE expert LoRA tensors into the layout vLLM expects.
+
+    Megatron-Bridge exports fused experts as 3D tensors keyed
+    ``...mlp.experts.gate_up_proj`` (w13) / ``...mlp.experts.down_proj`` (w2),
+    with ``lora_A=(E, rank, in)`` and ``lora_B=(E, out, rank)``. vLLM's 3D-MoE
+    loader (``FusedMoE3DWithLoRA`` / ``_stack_moe_lora_weights``) instead expects
+    the flat PEFT layout keyed ``...experts.base_layer`` (w13) / ``...experts``
+    (w2), with ``lora_A=(rank*E, in)`` and ``lora_B=(out, rank*E)``. This is the
+    exact inverse of vLLM's per-expert reshape. Non-expert tensors pass through.
+    """
+    converted: Dict[str, "torch.Tensor"] = {}
+    for key, tensor in adapter_state.items():
+        is_gate_up = ".mlp.experts.gate_up_proj." in key
+        is_down = ".mlp.experts.down_proj." in key
+        if (is_gate_up or is_down) and tensor.ndim == 3:
+            if key.endswith(".lora_A.weight"):
+                # (E, rank, in) -> (rank*E [expert-major], in)
+                tensor = tensor.reshape(-1, tensor.shape[-1]).contiguous()
+            elif key.endswith(".lora_B.weight"):
+                # (E, out, rank) -> (out, rank*E [expert-minor])
+                tensor = tensor.permute(1, 2, 0).contiguous().reshape(tensor.shape[1], -1)
+            if is_gate_up:
+                key = key.replace(".mlp.experts.gate_up_proj.", ".mlp.experts.base_layer.")
+            else:
+                key = key.replace(".mlp.experts.down_proj.", ".mlp.experts.")
+        converted[key] = tensor
+    return converted
+
+
 @torch.no_grad()
 def offload_megatron_grads_to_cpu(models):
     for model_chunk in models:

skyrl/backends/skyrl_train/workers/megatron/megatron_worker.py

@@ -24,6 +24,7 @@
     MegatronStrategy,
 )
 from skyrl.backends.skyrl_train.distributed.megatron.megatron_utils import (
+    _convert_moe_experts_lora_to_vllm,
     broadcast_object_across_pp_ranks,
     freeze_moe_router,
     get_model_config,
@@ -1033,12 +1034,10 @@ def forward_backward(
                 }
             )
 
-        # Count microbatches that carry real (non-padding) samples. Token-based batching
-        # appends fully-padding microbatches so every DP rank runs the same number of
-        # forward passes; those contribute 0 to KL/entropy and to mean metrics but would
-        # otherwise inflate the denominators. A real microbatch can still have an all-zero
-        # loss_mask (for example, DAPO overlong filtering), so use the iterator's padding
-        # count rather than inferring from loss_mask.
+        # Count real (non-padding) microbatches. Token-based batching appends padding
+        # microbatches so every DP rank runs the same number of forward passes; they must
+        # not inflate the KL/entropy denominators. Use the iterator's padding count rather
+        # than loss_mask, since a real microbatch can be all-zero (e.g. DAPO overlong filtering).
         num_padding_microbatches = (
             getattr(microbatch_iterator, "num_padding_microbatches", 0) if microbatch_iterator is not None else 0
         )
@@ -1242,7 +1241,15 @@ async def _save_lora_adapters_and_sync(
         if torch.distributed.get_rank() == 0:
             os.makedirs(lora_sync_path, exist_ok=True)
 
-            target_modules = infer_target_modules_from_adapter_weights(adapter_state.keys())
+            # Rewrite fused-MoE expert LoRA into vLLM's flat PEFT layout so
+            # merge_lora=False on-policy sync is accepted (otherwise
+            # load_lora_adapter rejects `experts.down_proj`). See
+            # _convert_moe_experts_lora_to_vllm for the layout details.
+            adapter_state = _convert_moe_experts_lora_to_vllm(adapter_state)
+
+            target_modules = sorted(
+                set(infer_target_modules_from_adapter_weights(adapter_state.keys())) - {"base_layer"}
+            )
             base_model_name_or_path = str(
                 getattr(self.bridge.hf_pretrained, "model_name_or_path", "")
                 or getattr(self.bridge.hf_pretrained, "name_or_path", "")

tests/train/generators/test_skyrl_gym_generator.py

@@ -363,10 +363,9 @@ async def test_generate_batched(mock_make, mock_tokenizer, mock_llm, mock_env, g
 async def test_generate_batched_metrics_use_truncated_responses(
     mock_make, mock_tokenizer, mock_llm, mock_env, generator_cfg, mock_env_cfg
 ):
-    """Rollout metrics must describe the truncated responses that are actually
-    returned/trained, not the raw engine output. With max_generate_length below
-    the engine's output length, generate/max_num_tokens must equal the truncated
-    length, matching response_ids and loss_masks."""
+    """Rollout metrics must describe the truncated responses that are trained, not the
+    raw engine output: with max_generate_length below the engine output, generate/
+    max_num_tokens must equal the truncated length, matching response_ids/loss_masks."""
     generator_cfg.sampling_params.max_generate_length = 2  # < len(MOCK_LLM_OUTPUT_IDS) == 4
     mock_make.return_value = mock_env
     mock_env.init.return_value = ([{"role": "user", "content": "Initial input"}], {})