diff --git a/src/llamafactory/extras/misc.py b/src/llamafactory/extras/misc.py index f95ab9f49..ba052783d 100644 --- a/src/llamafactory/extras/misc.py +++ b/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]: diff --git a/src/llamafactory/train/dpo/workflow.py b/src/llamafactory/train/dpo/workflow.py index 209449365..abac48475 100644 --- a/src/llamafactory/train/dpo/workflow.py +++ b/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) diff --git a/src/llamafactory/train/mca/workflow.py b/src/llamafactory/train/mca/workflow.py index f4b9d8df7..5ad099f68 100644 --- a/src/llamafactory/train/mca/workflow.py +++ b/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) diff --git a/src/llamafactory/train/sft/workflow.py b/src/llamafactory/train/sft/workflow.py index b50f53ffd..f457c58e7 100644 --- a/src/llamafactory/train/sft/workflow.py +++ b/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)