[megatron] feat: enhance model offloading and loading for frozen parameters

Author: RobotGF

verl/utils/megatron_utils.py

@@ -442,6 +442,11 @@ def offload_megatron_model_to_cpu(models):
                         # if the grad_data size is already zero, we assume that it is already offloaded
                         buffer.grad_data_size = buffer.grad_data.storage().size()
                         buffer.grad_data.storage().resize_(0)
+            # Offload frozen parameters not in DDP buffers (e.g. base model in LoRA/PEFT)
+            # DDP buffers only contain requires_grad=True params, so frozen params must be offloaded separately.
+            for param in model_chunk.module.parameters():
+                if not param.requires_grad and param.device.type != "cpu":
+                    param.data = param.data.to("cpu", non_blocking=True)
         else:
             # we need this for ref module
             for _, param in model_chunk.named_parameters():
@@ -453,7 +458,14 @@ def offload_megatron_model_to_cpu(models):
 
 
 @torch.no_grad()
-def load_megatron_model_to_gpu(models, load_grad=True):
+def load_megatron_model_to_gpu(models, load_grad=True, load_frozen_params=True):
+    """
+    Load megatron model to GPU.
+    Args:
+        models: The model to load.
+        load_grad: Whether to load gradients.
+        load_frozen_params: Whether to load frozen parameters.
+    """
     for model_chunk in models:
         if isinstance(model_chunk, DDP):
             model_chunk_all_buffers = [model_chunk.buffers, model_chunk.expert_parallel_buffers]
@@ -468,6 +480,13 @@ def load_megatron_model_to_gpu(models, load_grad=True):
                         buffer.param_data.storage().resize_(buffer.param_data_size)
                         # copy data from cpu to cuda
                         buffer.param_data.copy_(buffer.param_data.cpu_data, non_blocking=True)
+
+            # Load frozen parameters that were offloaded (e.g. base model in LoRA/PEFT)
+            if load_frozen_params:
+                device_id = get_device_id()
+                for param in model_chunk.module.parameters():
+                    if not param.requires_grad and param.device.type == "cpu":
+                        param.data = param.data.to(device_id, non_blocking=True)
         else:
             # we need this for ref module
             device_id = get_device_id()

verl/workers/engine/megatron/transformer_impl.py

@@ -602,12 +602,14 @@ def forward_backward_batch(self, data: TensorDict, loss_function: Callable, forw
             return {}
 
     def get_per_tensor_param(self, base_sync_done=False, **kwargs):
-        load_megatron_model_to_gpu(self.module, load_grad=False)
         peft_config = None
         non_merge_lora_sync = self.peft_cls is not None and not self.model_config.lora.get("merge", False)
+        adapter_only = base_sync_done and non_merge_lora_sync
+        # when lora adapter only, we only load adapter weights when base sync is done, otherwise load all weights
+        load_megatron_model_to_gpu(self.module, load_grad=False, load_frozen_params=not adapter_only)
         if self.vanilla_bridge:
             per_tensor_param = self.bridge.export_weights(self.module)
-        elif base_sync_done and non_merge_lora_sync:
+        elif adapter_only:
             # Only export adapter weights
             peft_config = build_peft_config_for_vllm(self.model_config.lora)
             per_tensor_param = self.bridge.export_adapter_weights(self.module)

verl/workers/rollout/vllm_rollout/vllm_async_server.py

@@ -618,7 +618,15 @@ async def sleep(self):
 
         if self.rollout_mode == RolloutMode.HYBRID:
             # Don't use engine.sleep(level=2) here
-            await self.engine.collective_rpc("sleep", kwargs={"level": 2})
+            # lora only update adapter weights, so set sleep level to 1
+            lora_as_adapter = (
+                self.model_config.lora_rank > 0 or self.model_config.lora.get("rank", 0) > 0
+            ) and not self.model_config.lora.get("merge", False)
+            if lora_as_adapter:
+                sleep_level = 1
+            else:
+                sleep_level = 2
+            await self.engine.collective_rpc("sleep", kwargs={"level": sleep_level})
 
             # clear encoder cache: https://github.com/vllm-project/vllm/pull/33452
             # await self.engine.reset_encoder_cache()