fix bug

PaddlePaddle · May 23, 2024 · affd27d · affd27d
1 parent 233b894
commit affd27d
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 0 deletions.
diff --git a/llm/predictor.py b/llm/predictor.py
@@ -1232,6 +1232,7 @@ def create_predictor(
                     dtype=predictor_args.dtype,
                     tensor_parallel_degree=tensor_parallel_degree,
                     tensor_parallel_rank=tensor_parallel_rank,
+                    tensor_parallel_output=False,
                 )
             elif model_args.model_type == "ernie-3.5-se":
                 sys.path.append("./ernie-3.5-se")
@@ -1244,6 +1245,7 @@ def create_predictor(
                     dtype=predictor_args.dtype,
                     tensor_parallel_degree=tensor_parallel_degree,
                     tensor_parallel_rank=tensor_parallel_rank,
+                    tensor_parallel_output=False,
                 )
             else:
                 model = AutoModelForCausalLM.from_pretrained(
@@ -1252,6 +1254,7 @@ def create_predictor(
                     use_flash_attention=predictor_args.use_flash_attention,
                     tensor_parallel_degree=tensor_parallel_degree,
                     tensor_parallel_rank=tensor_parallel_rank,
+                    tensor_parallel_output=False,
                 )
 
             predictor = DygraphPredictor(predictor_args, model=model, tokenizer=tokenizer)

diff --git a/paddlenlp/trl/trl_utils.py b/paddlenlp/trl/trl_utils.py
@@ -44,4 +44,6 @@ def calculate_effective_tokens(training_args, train_dataset, max_seq_len):
         for i, data in enumerate(train_dataset):
             total_effective_tokens += len(data["input_ids"])
         total_tokens = (i + 1) * max_seq_len
+        total_effective_tokens *= training_args.num_train_epochs
+        total_tokens *= training_args.num_train_epochs
     return total_effective_tokens, total_tokens