move FSDP param load/offload into sharding manager

Author: wuxibin89

recipe/dapo/src/config/dapo_trainer.yaml

@@ -81,6 +81,7 @@ actor_rollout_ref:
     ulysses_sequence_parallel_size: ${actor_rollout_ref.actor.ulysses_sequence_parallel_size} # sp size
   rollout:
     name: vllm
+    mode: "sync" # sync: LLM, async: AsyncLLM
     temperature: 1.0
     top_k: -1 # 0 for hf rollout, -1 for vllm rollout
     top_p: 1

tests/rollout/test_vllm_multi_turn.py

@@ -20,35 +20,40 @@
 from openai.types.chat.chat_completion import ChatCompletion
 
 from verl.single_controller.ray import RayClassWithInitArgs, RayWorkerGroup
-from verl.single_controller.ray.base import Worker, create_colocated_worker_cls
+from verl.single_controller.ray.base import create_colocated_worker_cls
 from verl.trainer.ppo.ray_trainer import ResourcePoolManager, Role
 from verl.workers.fsdp_async_workers import AsyncActorRolloutRefWorker, AsyncLLMManager
 from verl.workers.rollout.chat_scheduler import ChatCompletionScheduler
 
 
 async def test_vllm_multi_turn():
     config = OmegaConf.load("verl/trainer/config/ppo_trainer.yaml")
-    model_path = "Qwen/Qwen2-7B-Instruct"
+    model_path = "/mnt/bn/wuxibin-hl-dev/cache/Qwen/Qwen2-7B-Instruct"
     model_name = "/".join(model_path.split("/")[-2:])
     config.actor_rollout_ref.model.path = model_path
     config.actor_rollout_ref.rollout.mode = "async"
     config.actor_rollout_ref.rollout.prompt_length = 4096
     config.actor_rollout_ref.rollout.response_length = 4096
 
+    # test sleep/wake_up with fsdp offload
+    config.actor_rollout_ref.actor.fsdp_config.param_offload = True
+    config.actor_rollout_ref.actor.fsdp_config.optimizer_offload = True
+
     # =========================== 1. Create hybrid ActorRollout workers ===========================
     ray.init(
         runtime_env={
-            'env_vars': {
-                'TOKENIZERS_PARALLELISM': 'true',
-                'NCCL_DEBUG': 'WARN',
-                'VLLM_LOGGING_LEVEL': 'WARN',
-                'VLLM_USE_V1': '1',
+            "env_vars": {
+                "TOKENIZERS_PARALLELISM": "true",
+                "NCCL_DEBUG": "WARN",
+                "VLLM_LOGGING_LEVEL": "WARN",
+                "VLLM_USE_V1": "1",
             }
-        })
+        }
+    )
     role_worker_mapping = {
         Role.ActorRollout: ray.remote(AsyncActorRolloutRefWorker),
     }
-    global_pool_id = 'global_pool'
+    global_pool_id = "global_pool"
     resource_pool_spec = {
         global_pool_id: [config.trainer.n_gpus_per_node] * config.trainer.nnodes,
     }
@@ -61,20 +66,20 @@ async def test_vllm_multi_turn():
 
     # create actor and rollout
     resource_pool = resource_pool_manager.get_resource_pool(Role.ActorRollout)
-    actor_rollout_cls = RayClassWithInitArgs(cls=role_worker_mapping[Role.ActorRollout],
-                                             config=config.actor_rollout_ref,
-                                             role='actor_rollout')
-    resource_pool_to_cls[resource_pool]['actor_rollout'] = actor_rollout_cls
+    actor_rollout_cls = RayClassWithInitArgs(
+        cls=role_worker_mapping[Role.ActorRollout], config=config.actor_rollout_ref, role="actor_rollout"
+    )
+    resource_pool_to_cls[resource_pool]["actor_rollout"] = actor_rollout_cls
 
     all_wg = {}
     wg_dicts = []
     for resource_pool, class_dict in resource_pool_to_cls.items():
-        worker_dict_cls = create_colocated_worker_cls(class_dict=class_dict, worker_cls=Worker)
+        worker_dict_cls = create_colocated_worker_cls(class_dict=class_dict)
         wg_dict = RayWorkerGroup(resource_pool=resource_pool, ray_cls_with_init=worker_dict_cls)
         spawn_wg = wg_dict.spawn(prefix_set=class_dict.keys())
         all_wg.update(spawn_wg)
         wg_dicts.append(wg_dict)
-    actor_rollout_wg = all_wg['actor_rollout']
+    actor_rollout_wg = all_wg["actor_rollout"]
     actor_rollout_wg.init_model()
 
     # =========================== 2. Create AsyncLLMManager&ChatScheduler  ===========================
@@ -89,6 +94,10 @@ async def test_vllm_multi_turn():
         server_addresses=async_rollout_manager.server_addresses,
     )
 
+    # test sleep and wake_up
+    async_rollout_manager.sleep()
+    async_rollout_manager.wake_up()
+
     # =========================== 3. Multi turn rollout  ===========================
     async def callback(completions: ChatCompletion, info: Dict[str, Any]):
         messages, round = info["messages"], info["round"]
@@ -101,10 +110,7 @@ async def callback(completions: ChatCompletion, info: Dict[str, Any]):
             messages.append({"role": "user", "content": "What is your name?"})
             await async_chat_scheduler.submit_chat_completions(
                 callback=callback,
-                callback_additional_info={
-                    "messages": messages,
-                    "round": 1
-                },
+                callback_additional_info={"messages": messages, "round": 1},
                 model=model_name,
                 messages=messages,
                 extra_headers=extra_headers,
@@ -113,27 +119,20 @@ async def callback(completions: ChatCompletion, info: Dict[str, Any]):
             messages.append({"role": "user", "content": "What is your favorite color?"})
             await async_chat_scheduler.submit_chat_completions(
                 callback=callback,
-                callback_additional_info={
-                    "messages": messages,
-                    "round": 2
-                },
+                callback_additional_info={"messages": messages, "round": 2},
                 model=model_name,
                 messages=messages,
                 extra_headers=extra_headers,
             )
         else:
             print("Done!")
 
-    messages = [{
-        "role": "user",
-        "content": "Let's play a role playing game. Your name is Bob, your favorite color is red."
-    }]
+    messages = [
+        {"role": "user", "content": "Let's play a role playing game. Your name is Bob, your favorite color is red."}
+    ]
     await async_chat_scheduler.submit_chat_completions(
         callback=callback,
-        callback_additional_info={
-            "messages": messages,
-            "round": 0
-        },
+        callback_additional_info={"messages": messages, "round": 0},
         model=model_name,
         messages=messages,
     )

verl/single_controller/ray/base.py

@@ -329,10 +329,9 @@ def from_detached(
         worker_names=None,
         ray_cls_with_init=None,
     ):
-        worker_group = cls(resource_pool=None,
-                           ray_cls_with_init=ray_cls_with_init,
-                           name_prefix=name_prefix,
-                           worker_names=worker_names)
+        worker_group = cls(
+            resource_pool=None, ray_cls_with_init=ray_cls_with_init, name_prefix=name_prefix, worker_names=worker_names
+        )
         return worker_group
 
     def spawn(self, prefix_set):
@@ -457,12 +456,13 @@ def func(self, *args, **kwargs):
             try:
                 # bind direct rollout method to class without prefix
                 if attrs["dispatch_mode"] == Dispatch.DIRECT_ROLLOUT_METHOD and "rollout" in key:
-                    assert not hasattr(cls, method_name), \
+                    assert not hasattr(cls, method_name), (
                         f"conflict direct rollout method {method_name} with role {key}"
+                    )
                     setattr(cls, method_name, func)
                     print(f"bind role {key} method {method_name} to class {cls}")
                 else:
-                    method_name_with_prefix = key + '_' + method_name
+                    method_name_with_prefix = key + "_" + method_name
                     setattr(cls, method_name_with_prefix, func)
             except Exception as e:
                 raise ValueError(f"Fail to set method_name {method_name}")
@@ -474,32 +474,31 @@ def _unwrap_ray_remote(cls):
     return cls
 
 
-def _nearest_common_base(mros: List):
-    last_common = object
-    min_len = min([len(mro) for mro in mros]) - 1  # exclude final derived class
-
-    for i in range(min_len):
-        mro = mros[0][i]
-        for j in range(1, len(mros)):
-            if mro != mros[j][i]:
-                return last_common
-        last_common = mro
-
-    return last_common
+def _determine_fsdp_megatron_base_class(mros: List):
+    """
+    - megatron: base class should be MegatronWorker
+    - fsdp: base class should be Worker
+    """
+    for cls in mros[0]:
+        if cls.__name__ == "MegatronWorker":
+            return cls
+        if cls.__name__ == "Worker":
+            return cls
+    raise ValueError(f"Cannot determine base class for {mros}")
 
 
-def create_colocated_worker_cls(class_dict: dict[str, RayClassWithInitArgs], worker_cls: type = None):
+def create_colocated_worker_cls(class_dict: dict[str, RayClassWithInitArgs]):
     """
     This function should return a class instance that delegates the calls to every
     cls in cls_dict
     """
     cls_dict = {}
     init_args_dict = {}
-    if worker_cls is None:
-        worker_cls = _nearest_common_base(
-            [list(reversed(cls.cls.__ray_actor_class__.__mro__)) for cls in class_dict.values()])
+    worker_cls = _determine_fsdp_megatron_base_class(
+        [cls.cls.__ray_actor_class__.__mro__ for cls in class_dict.values()]
+    )
     assert issubclass(worker_cls, Worker), f"worker_cls {worker_cls} should be a subclass of Worker"
-    print(f"find nearest common base class {worker_cls}")
+    print(f"colocated worker base class {worker_cls}")
 
     for key, cls in class_dict.items():
         cls_dict[key] = cls.cls

verl/trainer/config/generation.yaml

@@ -14,6 +14,7 @@ model:
   external_lib: null
 rollout:
   name: vllm
+  mode: "sync" # sync: LLM, async: AsyncLLM
   temperature: 1.0
   top_k: 50 # 0 for hf rollout, -1 for vllm rollout
   top_p: 0.7

verl/trainer/config/ppo_megatron_trainer.yaml

@@ -95,6 +95,7 @@ actor_rollout_ref:
     log_prob_micro_batch_size_per_gpu: null
   rollout:
     name: vllm
+    mode: "sync" # sync: LLM, async: AsyncLLM
     temperature: 1.0
     top_k: -1 # 0 for hf rollout, -1 for vllm rollout
     top_p: 1

verl/workers/fsdp_workers.py

@@ -353,18 +353,22 @@ def _build_rollout(self, trust_remote_code=False):
             log_gpu_memory_usage(f"Before building {rollout_name} rollout", logger=None)
             local_path = copy_to_local(self.config.model.path)
             if vllm_mode == "customized":
-                rollout = vLLMRollout(actor_module=self.actor_module_fsdp,
-                                      config=self.config.rollout,
-                                      tokenizer=self.tokenizer,
-                                      model_hf_config=self.actor_model_config)
+                rollout = vLLMRollout(
+                    actor_module=self.actor_module_fsdp,
+                    config=self.config.rollout,
+                    tokenizer=self.tokenizer,
+                    model_hf_config=self.actor_model_config,
+                )
             elif vllm_mode == "spmd":
                 vllm_rollout_cls = vLLMRollout if self.config.rollout.mode == "sync" else vLLMAsyncRollout
-                rollout = vllm_rollout_cls(model_path=local_path,
-                                           config=self.config.rollout,
-                                           tokenizer=self.tokenizer,
-                                           model_hf_config=self.actor_model_config,
-                                           device_mesh=rollout_device_mesh,
-                                           trust_remote_code=trust_remote_code)
+                rollout = vllm_rollout_cls(
+                    model_path=local_path,
+                    config=self.config.rollout,
+                    tokenizer=self.tokenizer,
+                    model_hf_config=self.actor_model_config,
+                    device_mesh=rollout_device_mesh,
+                    trust_remote_code=trust_remote_code,
+                )
             else:
                 raise NotImplementedError("vllm_mode must be 'customized' or 'spmd'")
             log_gpu_memory_usage(f"After building {rollout_name} rollout", logger=None)
@@ -376,6 +380,7 @@ def _build_rollout(self, trust_remote_code=False):
                 model_config=self.actor_model_config,
                 full_params="hf" in self.config.rollout.load_format,
                 device_mesh=rollout_device_mesh,
+                offload_param=self._is_offload_param,
             )
             log_gpu_memory_usage("After building sharding manager", logger=None)
 
@@ -407,6 +412,7 @@ def _build_rollout(self, trust_remote_code=False):
                 model_config=self.actor_model_config,
                 full_params="hf" in self.config.rollout.load_format,
                 device_mesh=rollout_device_mesh,
+                offload_param=self._is_offload_param,
             )
             log_gpu_memory_usage("After building sharding manager", logger=None)
 
@@ -546,8 +552,6 @@ def generate_sequences(self, prompts: DataProto):
         prompts = prompts.to(torch.cuda.current_device())
 
         assert self._is_rollout
-        if self._is_offload_param:
-            load_fsdp_model_to_gpu(self.actor_module_fsdp)
 
         meta_info = {
             "eos_token_id": self.generation_config.eos_token_id
@@ -559,12 +563,6 @@ def generate_sequences(self, prompts: DataProto):
         }
         prompts.meta_info.update(meta_info)
         with self.rollout_sharding_manager:
-            # after parameters sync with rollout, offload actor model to CPU
-            if self._is_offload_param:
-                offload_fsdp_model_to_cpu(self.actor_module_fsdp)
-            if self._is_offload_optimizer:
-                offload_fsdp_optimizer(optimizer=self.actor_optimizer)
-
             log_gpu_memory_usage("After entering rollout sharding manager", logger=logger)
 
             prompts = self.rollout_sharding_manager.preprocess_data(prompts)

verl/workers/sharding_manager/fsdp_sglang.py

@@ -37,6 +37,7 @@
 from verl import DataProto
 from verl.protocol import all_gather_data_proto
 from verl.utils.debug import log_gpu_memory_usage
+from verl.utils.fsdp_utils import load_fsdp_model_to_gpu, offload_fsdp_model_to_cpu
 from verl.utils.torch_functional import broadcast_dict_tensor
 
 from .base import BaseShardingManager
@@ -55,11 +56,13 @@ def __init__(
         model_config,
         full_params: bool = False,
         device_mesh: DeviceMesh = None,
+        offload_param: bool = False,
     ):
         self.module = module
         self.inference_engine = inference_engine
         self.model_config = model_config
         self.device_mesh = device_mesh
+        self.offload_param = offload_param
 
         # Full params
         self.full_params = full_params
@@ -88,6 +91,8 @@ def __init__(
     def __enter__(self):
         torch.cuda.empty_cache()
         log_gpu_memory_usage("Before state_dict() in sharding manager memory", logger=logger)
+        if self.offload_param:
+            load_fsdp_model_to_gpu(self.module)
         params = self.module.state_dict()
         log_gpu_memory_usage("After state_dict() in sharding manager memory", logger=logger)
         # Copy, not share memory
@@ -98,15 +103,11 @@ def __enter__(self):
         log_gpu_memory_usage("After sync model weights in sharding manager", logger=logger)
 
         del params
+        if self.offload_param:
+            offload_fsdp_model_to_cpu(self.module)
         torch.cuda.empty_cache()
         log_gpu_memory_usage("After del state_dict and empty_cache in sharding manager", logger=logger)
 
-        # TODO: offload FSDP model weights
-        # self.module.cpu()
-        # torch.cuda.empty_cache()
-        # if torch.distributed.get_rank() == 0:
-        # print(f'after model to cpu in sharding manager memory allocated: {torch.cuda.memory_allocated() / 1e9}GB, reserved: {torch.cuda.memory_reserved() / 1e9}GB')
-
         # important: need to manually set the random states of each tp to be identical.
         if self.device_mesh is not None:
             self.torch_random_states = torch.cuda.get_rng_state()