Merge 2c7e630cf7 into 84576b1408

src/llamafactory/extras/misc.py

@@ -101,8 +101,10 @@ def check_dependencies() -> None:
     check_version("trl>=0.18.0,<=0.24.0")
 
 
-def calculate_tps(dataset: list[dict[str, Any]], metrics: dict[str, float], stage: Literal["sft", "rm"]) -> float:
-    r"""Calculate effective tokens per second."""
+def calculate_tps(
+    dataset: list[dict[str, Any]], metrics: dict[str, float], stage: Literal["sft", "rm"]
+) -> dict[str, Union[int, float]]:
+    r"""Calculate effective tokens per second and related dataset-level metrics."""
     effective_token_num = 0
     for data in dataset:
         if stage == "sft":
@@ -110,8 +112,21 @@ def calculate_tps(dataset: list[dict[str, Any]], metrics: dict[str, float], stag
         elif stage == "rm":
             effective_token_num += len(data["chosen_input_ids"]) + len(data["rejected_input_ids"])
 
-    result = effective_token_num * metrics["epoch"] / metrics["train_runtime"]
-    return result / dist.get_world_size() if dist.is_initialized() else result
+    train_runtime = metrics["train_runtime"]
+    epoch = metrics["epoch"]
+    world_size = dist.get_world_size() if dist.is_initialized() else 1
+    num_samples = len(dataset)
+
+    effective_tokens_per_sec_all_device = effective_token_num * epoch / train_runtime if train_runtime > 0 else 0.0
+    effective_tokens_per_sec_per_device = effective_tokens_per_sec_all_device / world_size
+    average_effective_tokens_per_sample = effective_token_num / num_samples if num_samples > 0 else 0.0
+
+    return {
+        "effective_tokens_per_sec_per_device": effective_tokens_per_sec_per_device,
+        "effective_tokens_per_sec_all_device": effective_tokens_per_sec_all_device,
+        "num_samples_per_epoch": num_samples,
+        "average_effective_tokens_per_sample": average_effective_tokens_per_sample,
+    }
 
 
 def count_parameters(model: "torch.nn.Module") -> tuple[int, int]:

src/llamafactory/train/dpo/workflow.py

@@ -81,9 +81,8 @@ def run_dpo(
         train_result = trainer.train(resume_from_checkpoint=training_args.resume_from_checkpoint)
         trainer.save_model()
         if finetuning_args.include_effective_tokens_per_second:
-            train_result.metrics["effective_tokens_per_sec"] = calculate_tps(
-                dataset_module["train_dataset"], train_result.metrics, stage="rm"
-            )
+            train_dataset = dataset_module["train_dataset"]
+            train_result.metrics.update(calculate_tps(train_dataset, train_result.metrics, stage="rm"))
 
         trainer.log_metrics("train", train_result.metrics)
         trainer.save_metrics("train", train_result.metrics)

src/llamafactory/train/mca/workflow.py

@@ -341,9 +341,8 @@ def run_dpo(
     train_result = trainer.train(training_args.resume_from_checkpoint)
     trainer.save_model()
     if finetuning_args.include_effective_tokens_per_second:
-        train_result.metrics["effective_tokens_per_sec"] = calculate_tps(
-            dataset_module["train_dataset"], train_result.metrics, stage="rm"
-        )
+        train_dataset = dataset_module["train_dataset"]
+        train_result.metrics.update(calculate_tps(train_dataset, train_result.metrics, stage="rm"))
 
     trainer.log_metrics("train", train_result.metrics)
     trainer.save_metrics("train", train_result.metrics)

src/llamafactory/train/sft/workflow.py

@@ -121,9 +121,8 @@ def run_sft(
         train_result = trainer.train(resume_from_checkpoint=training_args.resume_from_checkpoint)
         trainer.save_model()
         if finetuning_args.include_effective_tokens_per_second:
-            train_result.metrics["effective_tokens_per_sec"] = calculate_tps(
-                dataset_module["train_dataset"], train_result.metrics, stage="sft"
-            )
+            train_dataset = dataset_module["train_dataset"]
+            train_result.metrics.update(calculate_tps(train_dataset, train_result.metrics, stage="sft"))
 
         trainer.log_metrics("train", train_result.metrics)
         trainer.save_metrics("train", train_result.metrics)