Support force free memory for policy model with no colocate.

Author: adoda

chatlearn/models/torch_module.py

@@ -142,7 +142,7 @@ def _get_if_not_none(self, to_set, default):
         return default
 
     def onload(self, to_onload_weights=None, to_build_grad_buffers=None, to_onload_main_weights=None, to_onload_optimizer_states=None):
-        if not self.is_colocate:
+        if not (self.is_colocate or self.module_args.force_free_memory):
             return
         to_onload_weights = self._get_if_not_none(to_onload_weights, self.module_args.offload_weights)
         to_build_grad_buffers = self._get_if_not_none(to_build_grad_buffers, self.module_args.free_grad_buffers)
@@ -175,7 +175,7 @@ def offload(self, to_offload_weights=None, to_free_grad_buffers=None, to_offload
         # The first time of calling `offload_weights` and `offload_main_weights` has a higher peak memory.
         # So `free_grad_buffers` is called first to free memory, and `offload_weights` is called afterward
         # to make more space for `offload_main_weights`.
-        if not self.is_colocate:
+        if not (self.is_colocate or self.module_args.force_free_memory):
             return
         to_offload_weights = self._get_if_not_none(to_offload_weights, self.module_args.offload_weights)
         to_offload_main_weights = self._get_if_not_none(to_offload_main_weights, self.module_args.offload_weights)

chatlearn/runtime/executor.py

@@ -298,9 +298,9 @@ def compute_loop_one_model(self, model_node, num_batch=None):
         results = []
         self.timers(f"{model.name}").start()
         for step in range(num_batch):
-            to_empty_cache = step >= last_step_start and model.is_colocate
-            to_onload = step < replica_num and model.is_colocate and model.enable_offload
-            to_offload = step >= last_step_start and model.is_colocate and model.enable_offload
+            to_empty_cache = step >= last_step_start and (model.is_colocate or model.module_args.force_free_memory)
+            to_onload = step < replica_num and ((model.is_colocate and model.enable_offload) or model.module_args.force_free_memory)
+            to_offload = step >= last_step_start and ((model.is_colocate and model.enable_offload) or model.module_args.force_free_memory)
             replica = self._next_model(model)
             _, data = self.generate_step_one_model(model_node, replica, in_queue, model_node.out_queues, step, func_name, to_empty_cache,
                                                    is_eval=is_eval, to_onload=to_onload, to_offload=to_offload)

chatlearn/utils/arguments.py

@@ -234,6 +234,8 @@ class ModelConfig(BaseConfig):
     free_grad_buffers = False
     #: overall switch for offload optimizer states/weights and free grad buffers
     free_memory = False
+    #: force to free memory
+    force_free_memory = False
 
     def __init__(self):
         super().__init__()

examples/megatron/configs/llama2/grpo_math_vllm.yaml

@@ -16,6 +16,7 @@ models:
       ranking: ${batch_generation_ranking:False}
       min_prompt_length: ${batch_generation_min_prompt_length:0}
     free_memory: ${free_memory_policy:False}
+    force_free_memory: ${force_free_memory_policy:False}
 
   reference:
     model_config_file: reference.yaml

examples/megatron/configs/llama2/online_dpo_vllm.yaml

@@ -16,6 +16,7 @@ models:
       ranking: ${batch_generation_ranking:False}
       min_prompt_length: ${batch_generation_min_prompt_length:0}
     free_memory: ${free_memory_policy:False}
+    force_free_memory: ${force_free_memory_policy:False}
 
   reference:
     model_config_file: reference.yaml

examples/megatron/configs/llama2/rlhf.yaml

@@ -16,6 +16,7 @@ models:
       ranking: ${batch_generation_ranking:False}
       min_prompt_length: ${batch_generation_min_prompt_length:0}
     free_memory: ${free_memory_policy:False}
+    force_free_memory: ${force_free_memory_policy:False}
 
   reference:
     model_config_file: reference.yaml

examples/megatron/configs/llama2/vllm_rlhf.yaml

@@ -16,6 +16,7 @@ models:
       ranking: ${batch_generation_ranking:False}
       min_prompt_length: ${batch_generation_min_prompt_length:0}
     free_memory: ${free_memory_policy:False}
+    force_free_memory: ${force_free_memory_policy:False}
 
   reference:
     model_config_file: reference.yaml