From 9dc8fb2fc7cd5d595fdcf0cc5dc30c88ee8a6144 Mon Sep 17 00:00:00 2001
From: Sylvain Gugger <35901082+sgugger@users.noreply.github.com>
Date: Tue, 11 Jan 2022 03:14:11 -0500
Subject: [PATCH] Add test to check reported training loss (#15096)

* Add test

* Add tests for the reported train loss
---
 tests/test_trainer.py | 21 +++++++++++++++++++++
 1 file changed, 21 insertions(+)

diff --git a/tests/test_trainer.py b/tests/test_trainer.py
index 1beec9017d..e16aefdc8c 100644
--- a/tests/test_trainer.py
+++ b/tests/test_trainer.py
@@ -444,6 +444,27 @@ class TrainerIntegrationPrerunTest(TestCasePlus, TrainerIntegrationCommon):
         trainer.train()
         self.check_trained_model(trainer.model)
 
+    def test_training_loss(self):
+        n_gpus = max(1, get_gpu_count())
+
+        # With even logs
+        trainer = get_regression_trainer(logging_steps=64 / (8 * n_gpus))
+        trainer.train()
+        log_history = trainer.state.log_history
+
+        losses = [log["loss"] for log in log_history if "loss" in log]
+        train_loss = log_history[-1]["train_loss"]
+        self.assertAlmostEqual(sum(losses) / len(losses), train_loss, places=4)
+
+        # With uneven logs
+        trainer = get_regression_trainer(logging_steps=5)
+        trainer.train()
+        log_history = trainer.state.log_history
+
+        # Training loss should be the same as before
+        new_train_loss = log_history[-1]["train_loss"]
+        self.assertAlmostEqual(train_loss, new_train_loss, places=4)
+
     def test_custom_optimizer(self):
         train_dataset = RegressionDataset()
         args = TrainingArguments("./regression")