code: # Copyright The PyTorch Lightning team.
  #
  # Licensed under the Apache License, Version 2.0 (the "License");
  # you may not use this file except in compliance with the License.
  # You may obtain a copy of the License at
  #
  #     http://www.apache.org/licenses/LICENSE-2.0
  #
  # Unless required by applicable law or agreed to in writing, software
  # distributed under the License is distributed on an "AS IS" BASIS,
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  # See the License for the specific language governing permissions and
  # limitations under the License.
  from unittest import mock
  
  import pytest
  import torch
  
  from pytorch_lightning import Trainer
  from pytorch_lightning.strategies import DDPSpawnStrategy, DDPStrategy
  from pytorch_lightning.utilities import _TORCH_GREATER_EQUAL_1_10
  from tests.helpers import BoringModel
  from tests.helpers.runif import RunIf
  
  if torch.distributed.is_available():
      from torch.distributed.algorithms.ddp_comm_hooks import default_hooks as default
      from torch.distributed.algorithms.ddp_comm_hooks import powerSGD_hook as powerSGD
  
      if _TORCH_GREATER_EQUAL_1_10:
          import torch.distributed.algorithms.ddp_comm_hooks.post_localSGD_hook as post_localSGD
  
  
  class TestDDPStrategy(DDPStrategy):
      def __init__(self, expected_ddp_comm_hook_name, *args, **kwargs):
          self.expected_ddp_comm_hook_name = expected_ddp_comm_hook_name
          super().__init__(*args, **kwargs)
  
      def teardown(self):
          # check here before unwrapping DistributedDataParallel in self.teardown
          attached_ddp_comm_hook_name = self.model._get_ddp_logging_data()["comm_hook"]
          assert attached_ddp_comm_hook_name == self.expected_ddp_comm_hook_name
          return super().teardown()
  
  
  @RunIf(min_cuda_gpus=2, min_torch="1.9.0", skip_windows=True, standalone=True)
  def test_ddp_fp16_compress_comm_hook(tmpdir):
      """Test for DDP FP16 compress hook."""
      model = BoringModel()
      strategy = TestDDPStrategy(
          expected_ddp_comm_hook_name=default.fp16_compress_hook.__qualname__,
          ddp_comm_hook=default.fp16_compress_hook,
      )
      trainer = Trainer(
          max_epochs=1,
          accelerator="gpu",
          devices=2,
          strategy=strategy,
          default_root_dir=tmpdir,
          sync_batchnorm=True,
          fast_dev_run=True,
          enable_progress_bar=False,
          enable_model_summary=False,
      )
      trainer.fit(model)
      assert trainer.state.finished, f"Training failed with {trainer.state}"
  
  
  @RunIf(min_cuda_gpus=2, min_torch="1.9.0", skip_windows=True, standalone=True)
  def test_ddp_sgd_comm_hook(tmpdir):
      """Test for DDP FP16 compress hook."""
      model = BoringModel()
      strategy = TestDDPStrategy(
          expected_ddp_comm_hook_name=powerSGD.powerSGD_hook.__qualname__,
          ddp_comm_state=powerSGD.PowerSGDState(process_group=None),
          ddp_comm_hook=powerSGD.powerSGD_hook,
      )
      trainer = Trainer(
          max_epochs=1,
          accelerator="gpu",
          devices=2,
          strategy=strategy,
          default_root_dir=tmpdir,
          sync_batchnorm=True,
          fast_dev_run=True,
          enable_progress_bar=False,
          enable_model_summary=False,
      )
      trainer.fit(model)
      assert trainer.state.finished, f"Training failed with {trainer.state}"
  
  
  @RunIf(min_cuda_gpus=2, min_torch="1.9.0", skip_windows=True, standalone=True)
  def test_ddp_fp16_compress_wrap_sgd_comm_hook(tmpdir):
      """Test for DDP FP16 compress wrapper for SGD hook."""
      model = BoringModel()
      strategy = TestDDPStrategy(
          expected_ddp_comm_hook_name=default.fp16_compress_wrapper(powerSGD.powerSGD_hook).__qualname__,
          ddp_comm_state=powerSGD.PowerSGDState(process_group=None),
          ddp_comm_hook=powerSGD.powerSGD_hook,
          ddp_comm_wrapper=default.fp16_compress_wrapper,
      )
      trainer = Trainer(
          max_epochs=1,
          accelerator="gpu",
          devices=2,
          strategy=strategy,
          default_root_dir=tmpdir,
          sync_batchnorm=True,
          fast_dev_run=True,
          enable_progress_bar=False,
          enable_model_summary=False,
      )
      trainer.fit(model)
      assert trainer.state.finished, f"Training failed with {trainer.state}"
  
  
  @RunIf(min_cuda_gpus=2, min_torch="1.9.0", skip_windows=True, standalone=True)
  def test_ddp_spawn_fp16_compress_comm_hook(tmpdir):
      """Test for DDP Spawn FP16 compress hook."""
      model = BoringModel()
      strategy = DDPSpawnStrategy(ddp_comm_hook=default.fp16_compress_hook)
      trainer = Trainer(
          max_epochs=1,
          accelerator="gpu",
          devices=2,
          strategy=strategy,
          default_root_dir=tmpdir,
          sync_batchnorm=True,
          fast_dev_run=True,
          enable_progress_bar=False,
          enable_model_summary=False,
      )
      trainer.fit(model)
      assert trainer.state.finished, f"Training failed with {trainer.state}"
  
  
  @RunIf(min_cuda_gpus=2, min_torch="1.10.0", skip_windows=True, standalone=True)
  def test_ddp_post_local_sgd_comm_hook(tmpdir):
      """Test for DDP post-localSGD hook."""
      model = BoringModel()
      strategy = TestDDPStrategy(
          expected_ddp_comm_hook_name=post_localSGD.post_localSGD_hook.__qualname__,
          ddp_comm_state=post_localSGD.PostLocalSGDState(
              process_group=None,
              subgroup=None,
              start_localSGD_iter=8,
          ),
          ddp_comm_hook=post_localSGD.post_localSGD_hook,
          model_averaging_period=4,
      )
      trainer = Trainer(
          fast_dev_run=True,
          accelerator="gpu",
          devices=2,
          strategy=strategy,
          default_root_dir=tmpdir,
          sync_batchnorm=True,
          enable_progress_bar=False,
          enable_model_summary=False,
      )
      trainer.fit(model)
      assert trainer.state.finished, f"Training failed with {trainer.state}"
  
  
  @RunIf(skip_windows=True, min_torch="1.10.0", min_cuda_gpus=2, standalone=True)
  @mock.patch("torch.distributed.algorithms.model_averaging.averagers.PeriodicModelAverager.average_parameters")
  def test_post_local_sgd_model_averaging(average_parameters_mock, tmpdir):
      """Test that when using DDP with post-localSGD, model averaging is called."""
      model = BoringModel()
  
      # test regular ddp does not call model averaging
      trainer = Trainer(
          fast_dev_run=True,
          accelerator="gpu",
          devices=2,
          strategy="ddp",
          default_root_dir=tmpdir,
          sync_batchnorm=True,
          enable_progress_bar=False,
          enable_model_summary=False,
      )
  
      trainer.fit(model)
      average_parameters_mock.assert_not_called()
  
      # test ddp with post-localSGD does call model averaging
      ddp_strategy = DDPStrategy(
          ddp_comm_state=post_localSGD.PostLocalSGDState(
              process_group=None,
              subgroup=None,
              start_localSGD_iter=8,
          ),
          ddp_comm_hook=post_localSGD.post_localSGD_hook,
          model_averaging_period=4,
      )
  
      trainer = Trainer(
          fast_dev_run=True,
          accelerator="gpu",
          devices=2,
          strategy=ddp_strategy,
          default_root_dir=tmpdir,
          sync_batchnorm=True,
      )
  
      trainer.fit(model)
      average_parameters_mock.assert_called()
  
  
  @RunIf(skip_windows=True, min_torch="1.10.0", min_cuda_gpus=2, standalone=True)
  @mock.patch("torch.distributed.algorithms.model_averaging.averagers.PeriodicModelAverager.average_parameters")
  def test_post_local_sgd_model_averaging_value_error(average_parameters_mock, tmpdir):
      """Test that when using DDP with post-localSGD a ValueError is thrown when the optmizer is
      ZeroRedundancyOptimizer."""
      from torch.distributed.optim import ZeroRedundancyOptimizer
  
      class OptimizerModel(BoringModel):
          def configure_optimizers(self):
              return ZeroRedundancyOptimizer(params=self.parameters(), optimizer_class=torch.optim.Adam, lr=0.01)
  
      model = OptimizerModel()
      strategy = DDPStrategy(
          ddp_comm_state=post_localSGD.PostLocalSGDState(
              process_group=None,
              subgroup=None,
              start_localSGD_iter=8,
          ),
          ddp_comm_hook=post_localSGD.post_localSGD_hook,
          model_averaging_period=4,
      )
  
      trainer = Trainer(
          fast_dev_run=True,
          accelerator="gpu",
          devices=2,
          strategy=strategy,
          default_root_dir=tmpdir,
          sync_batchnorm=True,
          enable_progress_bar=False,
          enable_model_summary=False,
      )
  
      with pytest.raises(ValueError, match="Currently model averaging cannot work with a distributed optimizer"):
          trainer.fit(model)
  
      average_parameters_mock.assert_not_called()
   # Copyright The PyTorch Lightning team.
  #
  # Licensed under the Apache License, Version 2.0 (the "License");
  # you may not use this file except in compliance with the License.
  # You may obtain a copy of the License at
  #
  #     http://www.apache.org/licenses/LICENSE-2.0
  #
  # Unless required by applicable law or agreed to in writing, software
  # distributed under the License is distributed on an "AS IS" BASIS,
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  # See the License for the specific language governing permissions and
  # limitations under the License.
  import pytest
  import torch
  
  from pytorch_lightning.utilities.apply_func import move_data_to_device
  from pytorch_lightning.utilities.imports import _TORCHTEXT_LEGACY
  from tests.helpers.runif import RunIf
  from tests.helpers.torchtext_utils import get_dummy_torchtext_data_iterator
  
  
  @pytest.mark.parametrize("include_lengths", [False, True])
  @pytest.mark.parametrize("device", [torch.device("cuda", 0)])
  @pytest.mark.skipif(not _TORCHTEXT_LEGACY, reason="torchtext.legacy is deprecated.")
  @RunIf(min_cuda_gpus=1)
  def test_batch_move_data_to_device_torchtext_include_lengths(include_lengths, device):
      data_iterator, _ = get_dummy_torchtext_data_iterator(num_samples=3, batch_size=3, include_lengths=include_lengths)
      data_iter = iter(data_iterator)
      batch = next(data_iter)
  
      with pytest.deprecated_call(match="The `torchtext.legacy.Batch` object is deprecated"):
          batch_on_device = move_data_to_device(batch, device)
  
      if include_lengths:
          # tensor with data
          assert batch_on_device.text[0].device == device
          # tensor with length of data
          assert batch_on_device.text[1].device == device
      else:
          assert batch_on_device.text.device == device
  
  
  @pytest.mark.parametrize("include_lengths", [False, True])
  @pytest.mark.skipif(not _TORCHTEXT_LEGACY, reason="torchtext.legacy is deprecated.")
  def test_batch_move_data_to_device_torchtext_include_lengths_cpu(include_lengths):
      test_batch_move_data_to_device_torchtext_include_lengths(include_lengths, torch.device("cpu"))
  
apis: torch.distributed.algorithms.ddp_comm_hooks.powerSGD_hook.PowerSGDState torch.distributed.is_available torch.distributed.algorithms.ddp_comm_hooks.post_localSGD_hook.PostLocalSGDState torch.distributed.algorithms.ddp_comm_hooks.default_hooks.fp16_compress_wrapper
apis: torch.device