refact: hybrid_engine dir to sharding_manager for more general repres…

…entation (#103)
volcengine · Jan 14, 2025 · 6a9f6e1 · 6a9f6e1
1 parent e230de8
commit 6a9f6e1
Show file tree

Hide file tree

Showing 6 changed files with 4 additions and 4 deletions.
diff --git 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
@@ -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 → verl/workers/sharding_manager/__init__.py b/verl/workers/hybrid_engine/__init__.py → verl/workers/sharding_manager/__init__.py
diff --git a/verl/workers/hybrid_engine/base.py → verl/workers/sharding_manager/base.py b/verl/workers/hybrid_engine/base.py → verl/workers/sharding_manager/base.py
diff --git a/verl/workers/hybrid_engine/fsdp_vllm.py → verl/workers/sharding_manager/fsdp_vllm.py b/verl/workers/hybrid_engine/fsdp_vllm.py → verl/workers/sharding_manager/fsdp_vllm.py
diff --git a/verl/workers/hybrid_engine/megatron_vllm.py → ...workers/sharding_manager/megatron_vllm.py b/verl/workers/hybrid_engine/megatron_vllm.py → ...workers/sharding_manager/megatron_vllm.py