Merge branch 'main' into use_mbridge

verl/workers/megatron_workers.py

@@ -38,12 +38,7 @@
 from verl.utils.device import get_device_id, get_device_name, get_nccl_backend, get_torch_device
 from verl.utils.flops_counter import FlopsCounter
 from verl.utils.fs import copy_to_local
-from verl.utils.megatron_utils import (
-    load_megatron_model_to_gpu,
-    load_megatron_optimizer,
-    offload_megatron_model_to_cpu,
-    offload_megatron_optimizer,
-)
+from verl.utils.megatron_utils import load_megatron_model_to_gpu, load_megatron_optimizer, offload_megatron_model_to_cpu, offload_megatron_optimizer
 from verl.utils.model import get_hf_model_path, load_mcore_dist_weights, load_megatron_gptmodel_weights
 from verl.workers.actor.megatron_actor import MegatronPPOActor
 from verl.workers.critic.megatron_critic import MegatronPPOCritic

verl/workers/sharding_manager/megatron_vllm.py

@@ -150,21 +150,21 @@ def __enter__(self):
                     self.inference_engine.wake_up(tags=["weights"])
                 else:
                     self.inference_engine.wake_up()
-                if self.bridge is not None:
-                    per_tensor_param = self.bridge.export_weights(self.actor_module)
-                else:
-                    per_tensor_param = per_tensor_generator(
-                        self.actor_module,
-                        self.model_config,
-                        self.weight_converter,
-                        self.transformer_config,
-                        self.layer_name_mapping,
-                    )
-                model = self.model_runner.model
-                patch_vllm_moe_model_weight_loader(model)
-                loaded_params = model.load_weights(per_tensor_param)
-                info = f"vLLM load weights, loaded_params: {len(loaded_params)}"
-                logger.info(info)
+            if self.bridge is not None:
+                per_tensor_param = self.bridge.export_weights(self.actor_module)
+            else:
+                per_tensor_param = per_tensor_generator(
+                    self.actor_module,
+                    self.model_config,
+                    self.weight_converter,
+                    self.transformer_config,
+                    self.layer_name_mapping,
+                )
+            model = self.model_runner.model
+            patch_vllm_moe_model_weight_loader(model)
+            loaded_params = model.load_weights(per_tensor_param)
+            info = f"vLLM load weights, loaded_params: {len(loaded_params)}"
+            logger.info(info)
 
             if self.offload_param:
                 offload_megatron_model_to_cpu(self.actor_module)