Fix BF16 teardown for TPU precision plugin (#10990)

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Justus Schock <[email protected]>
Co-authored-by: Rohit Gupta <[email protected]>
Co-authored-by: Kaushik B <[email protected]>
Co-authored-by: thomas chaton <[email protected]>

Author: 6 people

CHANGELOG.md

@@ -14,6 +14,8 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 - Fixed data fetcher selection ([#11294](https://github.com/PyTorchLightning/pytorch-lightning/pull/11294))
 - Fixed a race condition that could result in incorrect (zero) values being observed in prediction writer callbacks ([#11288](https://github.com/PyTorchLightning/pytorch-lightning/pull/11288))
 - Fixed dataloaders not getting reloaded the correct amount of times when setting `reload_dataloaders_every_n_epochs` and `check_val_every_n_epoch` ([#10948](https://github.com/PyTorchLightning/pytorch-lightning/pull/10948))
+- Fixed an issue with the `TPUSpawnPlugin` handling the `XLA_USE_BF16` environment variable incorrectly ([#10990](https://github.com/PyTorchLightning/pytorch-lightning/pull/10990))
+
 
 ## [1.5.7] - 2021-12-21
 

pytorch_lightning/accelerators/accelerator.py

@@ -188,6 +188,7 @@ def teardown(self) -> None:
         It is the right place to release memory and free other resources.
         """
         self.training_type_plugin.teardown()
+        self.precision_plugin.teardown()
 
     def batch_to_device(self, batch: Any, device: Optional[torch.device] = None, dataloader_idx: int = 0) -> Any:
         """Moves the batch to the correct device. The returned batch is of the same type as the input batch, just

pytorch_lightning/plugins/precision/precision_plugin.py

@@ -252,3 +252,9 @@ def predict_step_context(self) -> Generator[None, None, None]:
         """A contextmanager for the predict step."""
         with self.forward_context():
             yield
+
+    def teardown(self) -> None:
+        """This method is called to teardown the training process.
+
+        It is the right place to release memory and free other resources.
+        """

pytorch_lightning/plugins/precision/tpu_bf16.py

@@ -28,5 +28,8 @@ class TPUBf16PrecisionPlugin(TPUPrecisionPlugin):
     def connect(
         self, model: nn.Module, optimizers: List[Optimizer], lr_schedulers: List[Any]
     ) -> Tuple[nn.Module, List[Optimizer], List[Any]]:
-        os.environ["XLA_USE_BF16"] = str(1)
+        os.environ["XLA_USE_BF16"] = "1"
         return super().connect(model=model, optimizers=optimizers, lr_schedulers=lr_schedulers)
+
+    def teardown(self) -> None:
+        os.environ.pop("XLA_USE_BF16", None)

pytorch_lightning/plugins/training_type/single_device.py

@@ -81,6 +81,7 @@ def broadcast(self, obj: object, src: int = 0) -> object:
         return obj
 
     def teardown(self) -> None:
+        super().teardown()
         if self.on_gpu:
             # GPU teardown
             self.lightning_module.cpu()

pytorch_lightning/plugins/training_type/single_tpu.py

@@ -82,6 +82,7 @@ def save_checkpoint(self, checkpoint: Dict[str, Any], filepath: _PATH) -> None:
         return self.checkpoint_io.save_checkpoint(checkpoint, filepath)
 
     def teardown(self) -> None:
+        super().teardown()
         # TPU teardown
         os.environ.pop("PT_XLA_DEBUG", None)
 

pytorch_lightning/plugins/training_type/tpu_spawn.py

@@ -357,7 +357,7 @@ def all_gather(self, tensor: torch.Tensor, group: Optional[Any] = None, sync_gra
         return xm.all_gather(tensor)
 
     def teardown(self) -> None:
-        # TPU teardown
+        super().teardown()
         os.environ.pop("PT_XLA_DEBUG", None)
         self.barrier("teardown")
 

pytorch_lightning/plugins/training_type/training_type_plugin.py

@@ -312,7 +312,6 @@ def model_sharded_context(self) -> Generator:
         """
         yield
 
-    @abstractmethod
     def teardown(self) -> None:
         """This method is called to teardown the training process.
 

tests/models/test_tpu.py

@@ -122,7 +122,6 @@ def test_model_16bit_tpu_cores_1(tmpdir):
 
     model = BoringModel()
     tpipes.run_model_test(trainer_options, model, on_gpu=False)
-    assert os.environ.get("XLA_USE_BF16") == str(1), "XLA_USE_BF16 was not set in environment variables"
 
 
 @pytest.mark.parametrize("tpu_core", [1, 5])
@@ -144,7 +143,6 @@ def test_model_16bit_tpu_index(tmpdir, tpu_core):
     model = BoringModel()
     tpipes.run_model_test(trainer_options, model, on_gpu=False)
     assert torch_xla._XLAC._xla_get_default_device() == f"xla:{tpu_core}"
-    assert os.environ.get("XLA_USE_BF16") == str(1), "XLA_USE_BF16 was not set in environment variables"
 
 
 @RunIf(tpu=True)