diff --git a/skyrl/backends/skyrl_train/distributed/megatron/megatron_utils.py b/skyrl/backends/skyrl_train/distributed/megatron/megatron_utils.py index 73d928ea22..6aae0905e8 100644 --- a/skyrl/backends/skyrl_train/distributed/megatron/megatron_utils.py +++ b/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: diff --git a/skyrl/backends/skyrl_train/workers/megatron/megatron_worker.py b/skyrl/backends/skyrl_train/workers/megatron/megatron_worker.py index 7d751c64ef..9675bda8a2 100644 --- a/skyrl/backends/skyrl_train/workers/megatron/megatron_worker.py +++ b/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", "") diff --git a/tests/train/generators/test_skyrl_gym_generator.py b/tests/train/generators/test_skyrl_gym_generator.py index 21db5dd90e..67ed5a344a 100644 --- a/tests/train/generators/test_skyrl_gym_generator.py +++ b/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"}], {})