From 6a9f6e19c6d171a8eda64fc0e73542f590688fbd Mon Sep 17 00:00:00 2001
From: Guangming Sheng <shengguangming@bytedance.com>
Date: Tue, 14 Jan 2025 16:19:12 +0800
Subject: [PATCH] refact: hybrid_engine dir to sharding_manager for more
 general representation (#103)

---
 verl/workers/fsdp_workers.py                                  | 4 ++--
 verl/workers/megatron_workers.py                              | 4 ++--
 verl/workers/{hybrid_engine => sharding_manager}/__init__.py  | 0
 verl/workers/{hybrid_engine => sharding_manager}/base.py      | 0
 verl/workers/{hybrid_engine => sharding_manager}/fsdp_vllm.py | 0
 .../{hybrid_engine => sharding_manager}/megatron_vllm.py      | 0
 6 files changed, 4 insertions(+), 4 deletions(-)
 rename verl/workers/{hybrid_engine => sharding_manager}/__init__.py (100%)
 rename verl/workers/{hybrid_engine => sharding_manager}/base.py (100%)
 rename verl/workers/{hybrid_engine => sharding_manager}/fsdp_vllm.py (100%)
 rename verl/workers/{hybrid_engine => sharding_manager}/megatron_vllm.py (100%)

diff --git a/verl/workers/fsdp_workers.py b/verl/workers/fsdp_workers.py
index b7577acb..6fa40f53 100644
--- a/verl/workers/fsdp_workers.py
+++ b/verl/workers/fsdp_workers.py
@@ -227,13 +227,13 @@ def _build_model_optimizer(self,
     def _build_rollout(self):
         if self.config.rollout.name == 'hf':
             from verl.workers.rollout import HFRollout
-            from verl.workers.hybrid_engine import BaseShardingManager
+            from verl.workers.sharding_manager import BaseShardingManager
             rollout = HFRollout(module=self.actor_module_fsdp, config=self.config.rollout)
             sharding_manager = BaseShardingManager()
             # TODO: a sharding manager that do nothing?
         elif self.config.rollout.name == 'vllm':
             from verl.workers.rollout.vllm_rollout import vLLMRollout
-            from verl.workers.hybrid_engine import FSDPVLLMShardingManager
+            from verl.workers.sharding_manager import FSDPVLLMShardingManager
             log_gpu_memory_usage('Before building vllm rollout', logger=None)
             rollout = vLLMRollout(actor_module=self.actor_module_fsdp,
                                   config=self.config.rollout,
diff --git a/verl/workers/megatron_workers.py b/verl/workers/megatron_workers.py
index f3a56f4f..180a7761 100644
--- a/verl/workers/megatron_workers.py
+++ b/verl/workers/megatron_workers.py
@@ -25,7 +25,7 @@
 from verl.single_controller.base.megatron.worker import MegatronWorker
 from verl.workers.actor.megatron_actor import MegatronPPOActor
 from verl.workers.critic.megatron_critic import MegatronPPOCritic
-from verl.workers.hybrid_engine import AllGatherPPModel
+from verl.workers.sharding_manager import AllGatherPPModel
 from verl.workers.reward_model.megatron.reward_model import MegatronRewardModel
 
 from verl.single_controller.base.decorator import register, Dispatch
@@ -216,7 +216,7 @@ def megatron_actor_model_provider(pre_process, post_process):
     def _build_rollout(self):
         if self.config.rollout.name == 'vllm':
             from verl.workers.rollout.vllm_rollout import vLLMRollout
-            from verl.workers.hybrid_engine import MegatronVLLMShardingManager
+            from verl.workers.sharding_manager import MegatronVLLMShardingManager
             from verl.utils.model import normalize_pp_vpp_params
 
             # NOTE(sgm): If the QKV and gate_up projection layer are concate together in actor,
diff --git a/verl/workers/hybrid_engine/__init__.py b/verl/workers/sharding_manager/__init__.py
similarity index 100%
rename from verl/workers/hybrid_engine/__init__.py
rename to verl/workers/sharding_manager/__init__.py
diff --git a/verl/workers/hybrid_engine/base.py b/verl/workers/sharding_manager/base.py
similarity index 100%
rename from verl/workers/hybrid_engine/base.py
rename to verl/workers/sharding_manager/base.py
diff --git a/verl/workers/hybrid_engine/fsdp_vllm.py b/verl/workers/sharding_manager/fsdp_vllm.py
similarity index 100%
rename from verl/workers/hybrid_engine/fsdp_vllm.py
rename to verl/workers/sharding_manager/fsdp_vllm.py
diff --git a/verl/workers/hybrid_engine/megatron_vllm.py b/verl/workers/sharding_manager/megatron_vllm.py
similarity index 100%
rename from verl/workers/hybrid_engine/megatron_vllm.py
rename to verl/workers/sharding_manager/megatron_vllm.py