refactor evaluation, upgrade trl to 074

Former-commit-id: ed09ebe2c1926ffdb0520b3866f7fd03a9aed046

src/llmtuner/tuner/core/adapter.py

@@ -38,12 +38,13 @@ def init_adapter(
 
     if (not is_trainable) and model_args.checkpoint_dir is None:
         logger.info("Checkpoint is not found at evaluation, load the original model.")
+        return model
 
     if finetuning_args.finetuning_type == "full" and is_trainable:
         logger.info("Fine-tuning method: Full")
         model = model.float()
 
-    if finetuning_args.finetuning_type == "freeze":
+    if finetuning_args.finetuning_type == "freeze" and is_trainable:
         logger.info("Fine-tuning method: Freeze")
         num_layers = getattr(model.config, "num_layers")
         if finetuning_args.num_layer_trainable > 0: # fine-tuning the last n layers if num_layer_trainable > 0

src/llmtuner/tuner/core/loader.py

@@ -42,7 +42,7 @@ require_version("transformers>=4.31.0,<4.35.0", "To fix: pip install \"transform
 require_version("datasets>=2.14.0", "To fix: pip install datasets>=2.14.0")
 require_version("accelerate>=0.21.0", "To fix: pip install accelerate>=0.21.0")
 require_version("peft>=0.6.0", "To fix: pip install peft>=0.6.0")
-require_version("trl==0.7.2", "To fix: pip install trl==0.7.2")
+require_version("trl>=0.7.4", "To fix: pip install trl>=0.7.4")
 
 
 def load_model_and_tokenizer(

src/llmtuner/tuner/core/parser.py

@@ -1,5 +1,4 @@
 import os
-import sys
 import torch
 import datasets
 import transformers
@@ -8,6 +7,7 @@ from transformers import HfArgumentParser, Seq2SeqTrainingArguments
 from transformers.trainer_utils import get_last_checkpoint
 
 from llmtuner.extras.logging import get_logger
+from llmtuner.extras.misc import parse_args
 from llmtuner.hparams import (
     ModelArguments,
     DataArguments,
@@ -19,17 +19,6 @@ from llmtuner.hparams import (
 logger = get_logger(__name__)
 
 
-def _parse_args(parser: HfArgumentParser, args: Optional[Dict[str, Any]] = None) -> Tuple[Any]:
-    if args is not None:
-        return parser.parse_dict(args)
-    elif len(sys.argv) == 2 and sys.argv[1].endswith(".yaml"):
-        return parser.parse_yaml_file(os.path.abspath(sys.argv[1]))
-    elif len(sys.argv) == 2 and sys.argv[1].endswith(".json"):
-        return parser.parse_json_file(os.path.abspath(sys.argv[1]))
-    else:
-        return parser.parse_args_into_dataclasses()
-
-
 def parse_train_args(
     args: Optional[Dict[str, Any]] = None
 ) -> Tuple[
@@ -46,7 +35,7 @@ def parse_train_args(
         FinetuningArguments,
         GeneratingArguments
     ))
-    return _parse_args(parser, args)
+    return parse_args(parser, args)
 
 
 def parse_infer_args(
@@ -63,7 +52,7 @@ def parse_infer_args(
         FinetuningArguments,
         GeneratingArguments
     ))
-    return _parse_args(parser, args)
+    return parse_args(parser, args)
 
 
 def get_train_args(

src/llmtuner/tuner/dpo/trainer.py

@@ -1,5 +1,4 @@
 import torch
-from copy import deepcopy
 from collections import defaultdict
 from typing import TYPE_CHECKING, Dict, Literal, Optional, Tuple, Union
 from transformers import BatchEncoding, Trainer
@@ -10,7 +9,6 @@ from llmtuner.extras.constants import IGNORE_INDEX
 
 if TYPE_CHECKING:
     from transformers import PreTrainedModel
-    from trl import PreTrainedModelWrapper
 
 
 class CustomDPOTrainer(DPOTrainer):
@@ -49,39 +47,6 @@ class CustomDPOTrainer(DPOTrainer):
             else:
                 self.ref_model = self.accelerator.prepare_model(self.ref_model, evaluation_mode=True)
 
-    def _prepare_deepspeed(self, model: "PreTrainedModelWrapper"):
-        # adapted from accelerate: https://github.com/huggingface/accelerate/blob/739b135f8367becb67ffaada12fe76e3aa60fefd/src/accelerate/accelerator.py#L1473
-        deepspeed_plugin = self.accelerator.state.deepspeed_plugin
-        config_kwargs = deepcopy(deepspeed_plugin.deepspeed_config)
-        if model is not None:
-            if hasattr(model, "config"):
-                hidden_size = (
-                    max(model.config.hidden_sizes)
-                    if getattr(model.config, "hidden_sizes", None)
-                    else getattr(model.config, "hidden_size", None)
-                )
-                if hidden_size is not None and config_kwargs["zero_optimization"]["stage"] == 3:
-                    # Note that `stage3_prefetch_bucket_size` can produce DeepSpeed messages like: `Invalidate trace cache @ step 0: expected module 1, but got module 0`
-                    # This is expected and is not an error, see: https://github.com/microsoft/DeepSpeed/discussions/4081
-                    config_kwargs.update(
-                        {
-                            "zero_optimization.reduce_bucket_size": hidden_size * hidden_size,
-                            "zero_optimization.stage3_param_persistence_threshold": 10 * hidden_size,
-                            "zero_optimization.stage3_prefetch_bucket_size": 0.9 * hidden_size * hidden_size,
-                        }
-                    )
-
-        # If ZeRO-3 is used, we shard both the active and reference model.
-        # Otherwise, we assume the reference model fits in memory and is initialized on each device with ZeRO disabled (stage 0)
-        if config_kwargs["zero_optimization"]["stage"] != 3:
-            config_kwargs["zero_optimization"]["stage"] = 0
-
-        # Lazy load
-        import deepspeed # type: ignore
-        model, *_ = deepspeed.initialize(model=model, config=config_kwargs)
-        model.eval()
-        return model
-
     def concatenated_forward(
         self,
         model: Optional[torch.nn.Module] = None,

src/llmtuner/tuner/ppo/trainer.py

@@ -226,7 +226,7 @@ class CustomPPOTrainer(PPOTrainer, Trainer):
         replace_model(unwrapped_model, target="default")
         return rewards
 
-    @PPODecorators.empty_cuda_cache()
+    @PPODecorators.empty_device_cache()
     def batched_forward_pass(
         self,
         model: "AutoModelForCausalLMWithValueHead",

src/llmtuner/tuner/ppo/workflow.py

@@ -42,7 +42,7 @@ def run_ppo(
         ppo_epochs=1,
         max_grad_norm=training_args.max_grad_norm,
         seed=training_args.seed,
-        optimize_cuda_cache=True,
+        optimize_device_cache=True,
         target=finetuning_args.ppo_target,
         log_with=finetuning_args.ppo_logger,
         use_score_scaling=finetuning_args.ppo_score_norm,