[train] Pin pytorch lightning to 2.0.4 (#37400)

2.0.5 breaks master as the deepspeed device validation fails. We should look into this separately. Signed-off-by: Kai Fricke <kai@anyscale.com>
ray-project · Jul 13, 2023 · 8f737c4 · 8f737c4
1 parent 49e884d
commit 8f737c4
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/.buildkite/pipeline.gpu_large.yml b/.buildkite/pipeline.gpu_large.yml
@@ -84,6 +84,6 @@
     - PYTHON=3.8 DOC_TESTING=1 TRAIN_TESTING=1 TUNE_TESTING=1 ./ci/env/install-dependencies.sh
     - pip install -Ur ./python/requirements/ml/dl-gpu-requirements.txt
     - pip uninstall -y pytorch-lightning
-    - pip install lightning==2.0.0
+    - pip install lightning==2.0.4 pytorch-lightning==2.0.4  # todo move to requirements-test.txt
     - ./ci/env/env_info.sh
     - bazel test --config=ci $(./scripts/bazel_export_options) --test_tag_filters=ptl_v2 python/ray/train/...
diff --git a/python/ray/train/lightning/_lightning_utils.py b/python/ray/train/lightning/_lightning_utils.py
@@ -225,7 +225,7 @@ def _session_report(self, trainer: "pl.Trainer", stage: str):
             if isinstance(v, torch.Tensor):
                 metrics[k] = v.item()
 
-        # Ensures all workers already finish writing their checkpoints.
+        # Ensures all workers already finish writing their checkpoints
         trainer.strategy.barrier()
 
         # Create and report the latest checkpoint