merge reshard-1 (#10606)

Author: fjjF77

llm/alignment/rl/run_rl.py

@@ -32,7 +32,7 @@
     TrainingArguments,
 )
 from paddlenlp.rl.utils.offload_utils import offload_tensor_to_cpu
-from paddlenlp.rl.utils.reshard_utils import init_rollout_env
+from paddlenlp.rl.utils.reshard_utils import ReshardController
 from paddlenlp.rl.utils.timer_utils import timers_scope_runtimer
 from paddlenlp.trainer import (
     EarlyStoppingCallback,
@@ -81,6 +81,7 @@ def create_actor_models(
     data_args: DataArgument,
     training_args: TrainingArguments,
     common_config: Dict,
+    reshard_controller: ReshardController = None,
 ):
     with timers_scope_runtimer("Actor model loading time"):
         # actor model
@@ -103,7 +104,7 @@ def create_actor_models(
         actor_model_config.use_sparse_head_and_loss_fn = False
         actor_model_config.seq_length = data_args.max_length
         actor_model_config.max_sequence_length = data_args.max_length
-        print(f"Loading Actor model with config:\n\t{actor_model_config}\n")
+        logger.info(f"Loading Actor model with config:\n\t{actor_model_config}\n")
 
         if not training_args.autotuner_benchmark:
             actor_model = AutoModelForCausalLM.from_pretrained(
@@ -113,18 +114,16 @@ def create_actor_models(
             actor_model = AutoModelForCausalLM.from_config(actor_model_config)
 
     with timers_scope_runtimer("Actor eval model loading time"):
-        if (
-            training_args.rollout_tensor_parallel_degree != training_args.tensor_parallel_degree
-            or training_args.pipeline_parallel_degree > 1
-        ):
+        if reshard_controller is not None:
+            reshard_controller.set_rollout_env("[create actor eval model]")
             actor_eval_model_config = copy.deepcopy(actor_model_config)
             actor_eval_model_config.use_fused_head_and_loss_fn = False
-            with init_rollout_env(training_args.rollout_tensor_parallel_degree):
-                hcg = fleet.get_hybrid_communicate_group()
-                actor_eval_model_config.tensor_parallel_degree = hcg.get_model_parallel_world_size()
-                actor_eval_model_config.tensor_parallel_rank = hcg.get_model_parallel_rank()
-                # TODO(gongenlei): lazy load lazy guard
-                actor_eval_model = AutoModelForCausalLM.from_config(actor_eval_model_config)
+            hcg = fleet.get_hybrid_communicate_group()
+            actor_eval_model_config.tensor_parallel_degree = hcg.get_model_parallel_world_size()
+            actor_eval_model_config.tensor_parallel_rank = hcg.get_model_parallel_rank()
+            # TODO(gongenlei): lazy load lazy guard
+            actor_eval_model = AutoModelForCausalLM.from_config(actor_eval_model_config)
+            reshard_controller.set_train_env("[after create actor eval model]")
         else:
             actor_eval_model = None
 
@@ -171,7 +170,7 @@ def create_reward_models(
         LlmMetaConfig.set_llm_config(reward_model_config, training_args)
         reward_model_config.max_position_embeddings = data_args.max_length
         reward_model_config.use_sparse_head_and_loss_fn = False
-        print(f"Loading Reward model with config:\n\t{reward_model_config}\n")
+        logger.info(f"Loading Reward model with config:\n\t{reward_model_config}\n")
 
         config = copy.deepcopy(reward_model_config)
         if training_args.eval_mode is not None:
@@ -323,8 +322,16 @@ def main():
         max_sequence_length=data_args.max_length,
     )
 
+    if (
+        training_args.rollout_tensor_parallel_degree != training_args.tensor_parallel_degree
+        or training_args.pipeline_parallel_degree > 1
+    ):
+        reshard_controller = ReshardController(tensor_parallel_degree=training_args.rollout_tensor_parallel_degree)
+    else:
+        reshard_controller = None
+
     actor_model, actor_eval_model, reference_model, actor_tokenizer = create_actor_models(
-        model_args, data_args, training_args, common_config
+        model_args, data_args, training_args, common_config, reshard_controller
     )
 
     if not training_args.use_rm_server and model_args.reward_model_name_or_path is not None:
@@ -387,6 +394,7 @@ def compute_metrics(eval_preds):
         ),  # NOTE: enforce prompt padding to max_prompt_len when using balance_batch
         compute_metrics=compute_metrics,  # TODO: only used for grpo (kk datasets)
         generation_config=generation_config,
+        reshard_controller=reshard_controller,
     )
 
     # TODO(gongenlei) resume_from_checkpoint is not ready

paddlenlp/rl/trainer/ppo_trainer.py

@@ -81,6 +81,7 @@
 )
 from ..utils.infer_utils import infer_guard
 from ..utils.offload_utils import reload_and_offload_scope, reload_tensor_to_gpu
+from ..utils.reshard_utils import ReshardController
 from ..utils.timer_utils import TimerScope, TimerScopeManualLabel
 from .actor_trainer import ActorReferenceTrainer
 from .critic_trainer import CriticTrainer
@@ -232,6 +233,7 @@ def __init__(
         optimizers: Tuple[paddle.optimizer.Optimizer, paddle.optimizer.lr.LRScheduler] = (None, None),
         preprocess_logits_for_metrics: Optional[Callable[[paddle.Tensor, paddle.Tensor], paddle.Tensor]] = None,
         generation_config: Optional[GenerationConfig] = None,
+        reshard_controller: Optional[ReshardController] = None,
     ):
         """
         Args:
@@ -282,6 +284,7 @@ def __init__(
                 preprocess_logits_for_metrics,
             )
 
+        self.reshard_controller = reshard_controller
         trainer_agrs = {
             # "model": None,
             "criterion": criterion,
@@ -300,6 +303,7 @@ def __init__(
             model=actor_model,
             model_eval=actor_model_eval,
             tokenizer=actor_tokenizer,
+            reshard_controller=reshard_controller,
             **trainer_agrs,
         )
 
@@ -379,6 +383,7 @@ def create_actor_trainer(
         callbacks: Optional[List[TrainerCallback]] = None,
         optimizers: Tuple[paddle.optimizer.Optimizer, paddle.optimizer.lr.LRScheduler] = (None, None),
         preprocess_logits_for_metrics: Optional[Callable[[paddle.Tensor, paddle.Tensor], paddle.Tensor]] = None,
+        reshard_controller: Optional[ReshardController] = None,
     ):
         policy_training_args = copy.deepcopy(args)
         lr_scheduler = self.get_scheduler(policy_training_args)
@@ -394,6 +399,7 @@ def create_actor_trainer(
             callbacks,
             [None, lr_scheduler],
             preprocess_logits_for_metrics,
+            reshard_controller,
         )
         actor_trainer.set_eval_model(model_eval)
         actor_trainer.timers = self.timers
@@ -688,6 +694,8 @@ def prediction_step(
                 }
                 generated_seq = self.actor_trainer.generate_sequences(prompt_only_batch, do_eval=True)[0]["input_ids"]
 
+            if self.reshard_controller is not None:
+                self.reshard_controller.set_train_env("[after prediction_step]")
             if not self.args.use_rm_server:
                 if self._model_config.sequence_parallel:
                     # pad to max_sequence_length
@@ -1386,7 +1394,6 @@ def train(
                 self.control = self.callback_handler.on_step_begin(args, self.state, self.control)
                 # step 1-1: rollout data with actor model (eval) and reward model
                 self.set_eval()
-
                 data_trans_group = getattr(self.actor_trainer, "_data_trans_group", None)
                 prompt_only_batch = data_group_split(prompt_only_batch, group=data_trans_group)
 
@@ -1415,6 +1422,7 @@ def train(
                     RolloutStages.ACTOR_MODEL_ENABLE_DISABLE,
                     minus_names=[RolloutStages.GENERATE],
                 )
+
                 timer_scope_actor_model.start()
                 with reload_and_offload_scope(self, self.actor_model):
                     timer_scope_rollout = TimerScope(self.timers, RolloutStages.GENERATE)
@@ -1438,6 +1446,8 @@ def train(
                     self.timers and (dist.get_world_size() > 1) and dist.barrier()
                     timer_scope_rollout.stop()
                 timer_scope_actor_model.stop()
+                if self.reshard_controller is not None:
+                    self.reshard_controller.set_train_env("[after rollout]")
 
                 # step 2-1: truncate data
                 truncate_input_ids = [
@@ -1469,19 +1479,22 @@ def train(
                     ),
                 }
 
+                batch = data_group_merge(batch, group=data_trans_group)
+
                 # step 2-2: balance batches based on batch tokens
                 if self.args.balance_batch:
                     batch = self._balance_batch(batch)
 
+                # step 2-3: compute logprob for rollout data
                 with self.autocast_smart_context_manager():
-                    # step 2-3: compute logprob for rollout data
                     with TimerScope(self.timers, RolloutStages.ROLLOUT_LOGPROB):
                         with reload_and_offload_scope(self, self.reference_model):
                             with TimerScope(self.timers, RolloutStages.ROLLOUT_REF_LOGPROB):
                                 batch["ref_log_probs"] = self.reference_trainer.compute_logprob(**batch)
 
                         with reload_and_offload_scope(self, self.actor_model):
                             with TimerScope(self.timers, RolloutStages.ROLLOUT_OLD_LOGPROB):
+                                self.actor_trainer.model.eval()
                                 batch["log_probs"] = self.actor_trainer.compute_logprob(**batch)
 
                 # step 2-2: compute reward for rollout data
@@ -1629,8 +1642,6 @@ def train(
                 else:
                     batch = batch
 
-                batch = data_group_merge(batch, group=data_trans_group)
-
                 # step 3: train actor model and critic model with rollout data
                 self.set_train()
                 with TimerScope(self.timers, ActorStages.MODEL_ENABLE_DISABLE, minus_names=[ActorStages.RL_STEP]):

paddlenlp/rl/trainer/rl_trainer.py

@@ -49,7 +49,7 @@
 from ...utils.env import TRAINER_STATE_NAME
 from ..models.ppo_model_utils import create_loss
 from ..utils.comm_utils import create_data_trans_group
-from ..utils.reshard_utils import init_rollout_env
+from ..utils.reshard_utils import ReshardController
 
 # ########## patches for Trianer ##########
 
@@ -537,6 +537,7 @@ def __init__(
         callbacks: Optional[List[TrainerCallback]] = None,
         optimizers: Tuple[paddle.optimizer.Optimizer, paddle.optimizer.lr.LRScheduler] = (None, None),
         preprocess_logits_for_metrics: Optional[Callable[[paddle.Tensor, paddle.Tensor], paddle.Tensor]] = None,
+        reshard_controller: Optional[ReshardController] = None,
     ):
         super().__init__(
             model,
@@ -565,6 +566,7 @@ def __init__(
         self.ema_beta = getattr(args, "ema_beta", 0.992)
         # if self.timers:
         #     self.timers.log = types.MethodType(new_timer_log, self.timers)
+        self.reshard_controller = reshard_controller
 
     def create_criterion(self):
         """
@@ -595,12 +597,15 @@ def set_eval_model(self, model):
         dp_group = hcg.get_data_parallel_group()
         global_rank = dist.get_rank()
         old_dp_workers = self.args.world_size // (max(sd_group.nranks, 1) * max(dp_group.nranks, 1))
-        with init_rollout_env(self.args.rollout_tensor_parallel_degree):
-            hcg = fleet.get_hybrid_communicate_group()
-            tensor_parallel_degree = hcg.get_model_parallel_world_size()
-            tensor_parallel_rank = hcg.get_model_parallel_rank()
-            eval_tp_size = max(tensor_parallel_degree, 1)
-            eval_tp_rank = max(tensor_parallel_rank, 0)
+        if self.reshard_controller is not None:
+            self.reshard_controller.set_rollout_env("[set eval model]")
+        hcg = fleet.get_hybrid_communicate_group()
+        tensor_parallel_degree = hcg.get_model_parallel_world_size()
+        tensor_parallel_rank = hcg.get_model_parallel_rank()
+        if self.reshard_controller is not None:
+            self.reshard_controller.set_train_env("[after set eval model]")
+        eval_tp_size = max(tensor_parallel_degree, 1)
+        eval_tp_rank = max(tensor_parallel_rank, 0)
         group_nums = self.args.logical_process_index // old_dp_workers * eval_tp_size + eval_tp_rank
         self._data_trans_group = create_data_trans_group(global_rank, group_nums)
         # just for compatible with old code

paddlenlp/rl/utils/comm_utils.py

@@ -25,7 +25,7 @@
 from ...trainer.trainer import Trainer, logger
 from ...utils.nested import flatten_list, nested_broadcast_tensor_with_empty
 from ..models.ppo_model_utils import make_position_ids_from_input_ids
-from .reshard_utils import init_reshard_mappings, init_rollout_env, reshard_to_rollout
+from .reshard_utils import init_reshard_mappings, reshard_to_rollout
 
 global_dev_id = 0 if paddle.get_device() == "cpu" else int(paddle.get_device().split(":")[1])
 
@@ -622,15 +622,18 @@ def export_evaluate_model(self: Trainer, train_model, eval_model, **kwargs):
     if not hasattr(self, "global_meta_dict") or self.global_meta_dict is None:
         self.global_meta_dict = init_reshard_mappings(train_model, self.args, pp_rank, pp_group)
 
-    with init_rollout_env(self.args.rollout_tensor_parallel_degree):
-        hcg = fleet.get_hybrid_communicate_group()
-        tensor_parallel_degree = hcg.get_model_parallel_world_size()
-        tensor_parallel_rank = hcg.get_model_parallel_rank()
-        eval_tp_size = max(tensor_parallel_degree, 1)
-        eval_tp_rank = max(tensor_parallel_rank, 0)
-        reshard_to_rollout(
-            train_model, eval_model, self.global_meta_dict, pp_rank, pp_group, hcg.get_model_parallel_group(), tp_group
-        )
+    if getattr(self, "reshard_controller", None) is not None:
+        self.reshard_controller.set_rollout_env("[export_evaluate_model]")
+    hcg = fleet.get_hybrid_communicate_group()
+    tensor_parallel_degree = hcg.get_model_parallel_world_size()
+    tensor_parallel_rank = hcg.get_model_parallel_rank()
+    eval_tp_size = max(tensor_parallel_degree, 1)
+    eval_tp_rank = max(tensor_parallel_rank, 0)
+    reshard_to_rollout(
+        train_model, eval_model, self.global_meta_dict, pp_rank, pp_group, hcg.get_model_parallel_group(), tp_group
+    )
+    if getattr(self, "reshard_controller", None) is not None:
+        self.reshard_controller.set_train_env("[after export_evaluate_model]")
 
     old_dp_workers = self.args.world_size // (max(sd_group.nranks, 1) * max(dp_group.nranks, 1))
     group_nums = self.args.logical_process_index // old_dp_workers * eval_tp_size + eval_tp_rank