upgrade lightning dependency (#1301)

Summary:
- Upgrade lightning to 1.6 to fix installation errors
- Fix some tests since lightning logic to track global step has changed

Pull Request resolved: #1301

Reviewed By: pikapecan

Differential Revision: D44670777

Pulled By: ankitade

fbshipit-source-id: 0febb0683aeb59fbb82cf1f29bbadf1d3385493d

Author: ankitade

.github/workflows/cpu_test.yaml

@@ -59,7 +59,7 @@ jobs:
       run: |
         conda activate mmf
         python -m pip install --upgrade pip
-        pip install --upgrade setuptools
+        pip install setuptools==65.6.3
         pip install --progress-bar off pytest
         pip install -r requirements.txt
         python -c 'import torch; print("Torch version:", torch.__version__)'

mmf/trainers/lightning_core/loop_callback.py

@@ -175,7 +175,7 @@ def _get_iterations_for_logging(self, trainer: Trainer):
         return trainer.fit_loop.batch_idx + 1
 
     def _get_num_updates_for_logging(self, trainer: Trainer):
-        return trainer.global_step + 1
+        return trainer.global_step
 
     def _train_log(self, trainer: Trainer, pl_module: LightningModule):
         self.train_combined_report = self.train_combined_report.detach()

mmf/trainers/lightning_core/loop_callback_with_torchmetrics.py

@@ -133,7 +133,7 @@ def _get_iterations_for_logging(self, trainer: Trainer):
         return trainer.fit_loop.batch_idx + 1
 
     def _get_num_updates_for_logging(self, trainer: Trainer):
-        return trainer.global_step + 1
+        return trainer.global_step
 
     def _get_train_extra_log(self, trainer: Trainer, pl_module: LightningModule):
         extra = {}

requirements.txt

@@ -19,7 +19,7 @@ datasets==1.2.1
 matplotlib==3.3.4
 pycocotools==2.0.2
 ftfy==5.8
-pytorch-lightning @ git+https://github.com/PyTorchLightning/pytorch-lightning@9b011606f
+pytorch-lightning==1.6.0
 psutil
 pillow==9.3.0
 sentencepiece

tests/trainers/lightning/test_checkpoint.py

@@ -197,6 +197,7 @@ def _get_lightning_trainer(
 
 
 class TestLightningCheckpoint(TestLightningCheckpoint):
+    @skip_if_no_network
     def test_load_resume_parity_with_mmf(self):
         # with checkpoint.resume = True, by default it loads "current.ckpt"
         self._load_checkpoint_and_test("current.ckpt", ckpt_config={"resume": True})
@@ -208,6 +209,7 @@ def test_load_resume_best_parity_with_mmf(self):
             "best.ckpt", ckpt_config={"resume": True, "resume_best": True}
         )
 
+    @skip_if_no_network
     def test_load_resume_ignore_resume_zoo(self):
         # specifying both checkpoint.resume = True and resume_zoo
         # resume zoo should be ignored. It should load the "current.ckpt"
@@ -393,7 +395,8 @@ def test_load_trainer_ckpt_number_of_steps(self):
                 )
                 self.assertEquals(lightning.trainer.global_step, 12)
                 call_args_list = [l[0][4] for l in mock_method.call_args_list]
-                self.assertListEqual(list(range(0, 6)), call_args_list)
+                # in lightning 1.6.0 last batch idx from ckpt is repeated
+                self.assertListEqual(list(range(5, 11)), call_args_list)
 
     def test_trainer_save_current_parity_with_mmf(self):
         with mock_env_with_temp(
@@ -454,7 +457,7 @@ def test_lightning_checkpoint_interval(self):
             files = os.listdir(os.path.join(tmp_d, "models"))
             self.assertEquals(3, len(files))
             indexes = {int(x[:-5].split("=")[1]) for x in files}
-            self.assertSetEqual({1, 3, 5}, indexes)
+            self.assertSetEqual({2, 4, 6}, indexes)
 
     def _get_mmf_ckpt(self, filename, ckpt_config=None):
         with mock_env_with_temp(
@@ -508,12 +511,7 @@ def _load_checkpoint_and_test(self, filename, ckpt_config=None):
 
         # Make sure lightning and mmf parity
         self._assert_same_dict(mmf_ckpt["model"], lightning_ckpt["state_dict"])
-
-        # different case for best checkpoint, see D34398730
-        if "resume_best" in ckpt_config and ckpt_config["resume_best"]:
-            self.assertEquals(mmf_ckpt["current_epoch"], lightning_ckpt["epoch"] + 1)
-        else:
-            self.assertEquals(mmf_ckpt["current_epoch"], lightning_ckpt["epoch"])
+        self.assertEquals(mmf_ckpt["current_epoch"], lightning_ckpt["epoch"] + 1)
         self.assertEquals(mmf_ckpt["num_updates"], lightning_ckpt["global_step"])
         self._assert_same_dict(
             mmf_ckpt["optimizer"], lightning_ckpt["optimizer_states"][0]

tests/trainers/lightning/test_logging.py

@@ -6,6 +6,7 @@
 from mmf.trainers.callbacks.logistics import LogisticsCallback
 from mmf.trainers.lightning_core.loop_callback import LightningLoopCallback
 from mmf.utils.timer import Timer
+from tests.test_utils import skip_if_no_network
 from tests.trainers.test_utils import (
     get_config_with_defaults,
     get_lightning_trainer,
@@ -19,6 +20,7 @@ def setUp(self):
         self.mmf_tensorboard_logs = []
         self.lightning_tensorboard_logs = []
 
+    @skip_if_no_network
     @patch("mmf.common.test_reporter.PathManager.mkdirs")
     @patch("mmf.trainers.callbacks.logistics.setup_output_folder", return_value="logs")
     @patch("mmf.trainers.lightning_trainer.setup_output_folder", return_value="logs")

tests/trainers/lightning/test_validation.py

@@ -92,21 +92,7 @@ def log_values(
             keys = list(gt.keys())
             self.assertListEqual(keys, list(lv.keys()))
             for key in keys:
-                if key == "num_updates" and gt[key] == self.ground_truths[-1][key]:
-                    # After training, in the last evaluation run, mmf's num updates is 8
-                    # while lightning's num updates is 9, this is due to a hack to
-                    # assign the lightning num_updates to be the trainer.global_step+1.
-                    #
-                    # This is necessary because of a lightning bug: trainer.global_step
-                    # is 1 off less than the actual step count. When on_train_batch_end
-                    # is called for the first time, the trainer.global_step should be 1,
-                    # rather than 0, since 1 update/step has already been done.
-                    #
-                    # When lightning fixes its bug, we will update this test to remove
-                    # the hack. # issue: 6997 in pytorch lightning
-                    self.assertAlmostEqual(gt[key], lv[key] - 1, 1)
-                else:
-                    self.assertAlmostEqual(gt[key], lv[key], 1)
+                self.assertAlmostEqual(gt[key], lv[key], 1)
 
     # TODO: update test function with avg_loss
     @patch("mmf.common.test_reporter.PathManager.mkdirs")
@@ -145,12 +131,7 @@ def log_values(
         self.assertEqual(len(self.ground_truths), len(lightning_values))
         for gt, lv in zip(self.ground_truths, lightning_values):
             for key in ["num_updates", "max_updates"]:
-                if key == "num_updates" and gt[key] == self.ground_truths[-1][key]:
-                    # to understand the reason of using lv[key] - 1 (intead of lv[key])
-                    # see comments in test_validation
-                    self.assertAlmostEqual(gt[key], lv[key] - 1, 1)
-                else:
-                    self.assertAlmostEqual(gt[key], lv[key], 1)
+                self.assertAlmostEqual(gt[key], lv[key], 1)
 
     @patch("mmf.common.test_reporter.PathManager.mkdirs")
     @patch("torch.utils.tensorboard.SummaryWriter")