docstr: Tests FSDP's default communication hook's behavior and correctness.
  This test creates a simple linear net with weight shape  ``1 X N``,
  where ``N`` is the number of workers.
  For sharded cases, each worker gets 1 element of the weight parameter. This test
  checks that after backward, each worker has a proper value in its chunk of
  the gradient, or the whole gradient on every worker is equal to an expected value.
  
  Arguments:
      sharding_strategy (Optional[ShardingStrategy]): Configures the FSDP algorithm.
docstr_tok: ["Tests","FSDP","'s","default","communication","hook","'s","behavior","and","correctness",".","This","test","creates","a","simple","linear","net","with","weight","shape","``","1","X","N","``",",","where","``","N","``","is","the","number","of","workers",".","For","sharded","cases",",","each","worker","gets","1","element","of","the","weight","parameter",".","This","test","checks","that","after","backward",",","each","worker","has","a","proper","value","in","its","chunk","of","the","gradient",",","or","the","whole","gradient","on","every","worker","is","equal","to","an","expected","value",".","Arguments",":","sharding_strategy","(","Optional","[","ShardingStrategy","]",")",":","Configures","the","FSDP","algorithm","."]
code: def test_default_communication_hook_behavior(
      self, sharding_strategy: Optional[ShardingStrategy]
  ):
      """
      Tests FSDP's default communication hook's behavior and correctness.
      This test creates a simple linear net with weight shape  ``1 X N``,
      where ``N`` is the number of workers.
      For sharded cases, each worker gets 1 element of the weight parameter. This test
      checks that after backward, each worker has a proper value in its chunk of
      the gradient, or the whole gradient on every worker is equal to an expected value.
  
      Arguments:
          sharding_strategy (Optional[ShardingStrategy]): Configures the FSDP algorithm.
      """
      out_dim = self.world_size
      net = torch.nn.Linear(1, out_dim, bias=False)
      inpt = torch.tensor([self.rank]).float().cuda(self.rank)
  
      net_default_hook = FSDP(
          net,
          device_id=torch.cuda.current_device(),
          sharding_strategy=sharding_strategy,
      ).to(self.rank)
  
      # Check that by default, `_comm_hook` is None
      for entry in FSDP.fsdp_modules(net_default_hook):
          self.assertEqual(entry._comm_hook, None)
  
      for _ in range(4):
          # Clear gradients
          net_default_hook.zero_grad()
          loss = net_default_hook(inpt).sum()
          loss.backward()
  
          # For each worker, the gradient on the weight should be worker_rank.
          grad = net_default_hook.params[0].grad
          expected_grad = (
              sum(i for i in range(dist.get_world_size()))
              / dist.get_world_size()
          )
          # Verify default hook produces expected gradients
          self.assertEqual(
              grad[0].item(),
              expected_grad,
              msg=f"Expected hook grad of {expected_grad} but got {grad[0].item()}",
          )
  
imports: import sys
  from typing import Optional
  import torch
  import torch.nn
  import torch.nn.functional
  from torch import distributed
  from torch.distributed.algorithms._comm_hooks import default_hooks
  from torch.distributed.distributed_c10d import _get_default_group
  from torch.distributed.fsdp import FullyShardedDataParallel, MixedPrecision
  from torch.distributed.fsdp.fully_sharded_data_parallel import ShardingStrategy
  from torch.distributed.fsdp.wrap import ModuleWrapPolicy
  from torch.testing._internal.common_distributed import requires_nccl, requires_nccl_version, skip_but_pass_in_sandcastle_if, skip_if_lt_x_gpu
  from torch.testing._internal.common_fsdp import FSDPTest
  from torch.testing._internal.common_utils import instantiate_parametrized_tests, parametrize, run_tests
label_desc: Use image node_id 1 for calling the TestCommunicationHooks obj's underlying member method code with example usage: obj.test_default_communication_hook_behavior(sharding_strategy) without return types