code: def _test_communication(
      self, rank, group, config, nested_model=False
  ):
      if group.size() == 1:
          return
  
      # Turn off bucketing to accurately count number of reduce_scatters.
      config["bucket_cap_mb"] = 0
  
      if nested_model:
          model = NestedWrappedModule(group, config)
          model = FullyShardedDataParallel(
              model, group, **config
          ).cuda()
      else:
          model = self.get_wrapped_model(group, config=config)
  
      num_fsdp = 0
      for child in model.modules():  # includes self
          if (
              isinstance(child, FullyShardedDataParallel)
              and len(child.params) > 0
          ):
              num_fsdp += 1
  
      if config.get("reshard_after_forward", True):
          # inside no_sync:
          #   num_fsdp all-gathers in the forward
          #   num_fsdp-1 all-gathers in the backward (except root)
          # outside no_sync:
          #   num_fsdp-1 all-gathers in the forward (except root)
          #   num_fsdp-1 all-gathers in the backward (except root)
          expected_all_gather1 = 2 * num_fsdp - 1
          expected_all_gather2 = expected_all_gather1 + (
              2 * num_fsdp - 2
          )
      else:
          # inside no_sync:
          #   num_fsdp all-gathers in the forward
          # outside no_sync:
          #   none
          expected_all_gather1 = num_fsdp
          expected_all_gather2 = num_fsdp
  
      expected_reduce_scatter = num_fsdp
  
      batch = model.module.get_input(torch.device("cuda"))
  
      # depending on pytorch version the _base methods may not be available
      method_string_reduce_scatter_base = (
          "torch.distributed._reduce_scatter_base"
      )
      if hasattr(torch.distributed, "_reduce_scatter_base") is False:
          # no such method, to make mock_reduce_scatter_base 0 invocation, use an impossible name
          method_string_reduce_scatter_base = "math.nan"  # just an arbitrary function not going to be called
  
      method_string_all_gather_base = (
          "torch.distributed._all_gather_base"
      )
      if hasattr(torch.distributed, "_all_gather_base") is False:
          # no such method, to make mock_all_gather_base 0 invocation, use an impossible name
          method_string_all_gather_base = "math.nan"  # just an arbitrary function not going to be called
  
      with patch("torch.distributed.all_gather") as mock_all_gather:
          with patch(
              "torch.distributed.reduce_scatter"
          ) as mock_reduce_scatter:
              with patch(
                  method_string_all_gather_base
              ) as mock_all_gather_base:
                  with patch(
                      method_string_reduce_scatter_base
                  ) as mock_reduce_scatter_base:
                      with model.no_sync():
                          output = model(*batch)
                          loss = model.module.get_loss(batch, output)
                          loss.backward()
  
                      # the _base methods are activated when they are available.
                      # the sum of the _base and public methods should stay the same.
                      assert (
                          mock_all_gather.call_count
                          + mock_all_gather_base.call_count
                          == expected_all_gather1
                      ), f"{mock_all_gather.call_count} + {mock_all_gather_base.call_count} != {expected_all_gather1}"
                      assert (
                          mock_reduce_scatter.call_count
                          + mock_reduce_scatter_base.call_count
                          == 0
                      ), f"{mock_reduce_scatter.call_count} + {mock_reduce_scatter_base.call_count} != 0"
  
                      output = model(*batch)
                      loss = model.module.get_loss(batch, output)
                      loss.backward()
  
                      assert (
                          mock_all_gather.call_count
                          + mock_all_gather_base.call_count
                          == expected_all_gather2
                      ), f"{mock_all_gather.call_count} + {mock_all_gather_base.call_count} != {expected_all_gather2}"
                      assert (
                          mock_reduce_scatter.call_count
                          + mock_reduce_scatter_base.call_count
                          == expected_reduce_scatter
                      ), f"{mock_reduce_scatter.call_count} + {mock_reduce_scatter_base.call_count} != {expected_reduce_scatter}"
  
imports: import contextlib
  import functools
  import itertools
  import unittest
  from unittest.mock import patch
  from parameterized import parameterized
  import torch
  from fairscale.fair_dev.testing.testing import DummyProcessGroup, make_cudnn_deterministic, objects_are_equal
  from fairscale.nn.data_parallel import FullyShardedDataParallel
  from .test_fsdp import DistributedTest, NestedWrappedModule, rename_test, spawn_and_init
label_desc: Use image node_id 3 for calling the TestGradAccCommunication obj's underlying member method code with example usage: obj._test_communication(rank, group, config, nested_model) without return types