docstr: Tests that FSDP warns the user if the all-gather order changes after
  the first iteration.
docstr_tok: ["Tests","that","FSDP","warns","the","user","if","the","all-gather","order","changes","after","the","first","iteration","."]
code: def test_invalid_later_iter_order(
      self,
      sharding_strategy: ShardingStrategy,
      iters_before_path_change: int,
  ):
      """Tests that FSDP warns the user if the all-gather order changes after
      the first iteration."""
      dist.set_debug_level(dist.DebugLevel.DETAIL)
      # On the first iteration, all ranks run the same order, and on the next
      # iteration, all but rank 0 run in a different order
      fsdp_model = Model.wrap(sharding_strategy, self.device)
      for _ in range(iters_before_path_change):
          inp = fsdp_model.module.get_input(self.device)
          output = fsdp_model(*inp)
          loss = fsdp_model.module.get_loss(inp, output).to(self.device)
          fsdp_model.module.run_backward(loss)
      # Match the warning message with the following prefix
      regex = (
          "^(Forward order differs from that of the first iteration "
          f"on rank {self.rank}. Collectives are unchecked and may give "
          "incorrect results or hang)"
      )
      context = (
          self.assertWarnsRegex(
              expected_warning=UserWarning,
              expected_regex=regex,
          )
          if self.rank != 0
          else nullcontext()
      )
      if self.rank != 0:
          fsdp_model.flip_path()
      inp = fsdp_model.module.get_input(self.device)
      # Expect a warning for the forward pass all-gather
      with context:  # warning for forward pass all-gather
          output = fsdp_model(*inp)
      loss = fsdp_model.module.get_loss(inp, output).to(self.device)
      fsdp_model.module.run_backward(loss)
      # Run an additional iteration to check that there are no more warnings
      inp = fsdp_model.module.get_input(self.device)
      output = fsdp_model(*inp)
      loss = fsdp_model.module.get_loss(inp, output).to(self.device)
      fsdp_model.module.run_backward(loss)
  
imports: import sys
  import warnings
  from contextlib import nullcontext
  import torch
  from torch import distributed
  from torch.distributed.fsdp import FullyShardedDataParallel
  from torch.distributed.fsdp.fully_sharded_data_parallel import ShardingStrategy
  from torch.testing._internal.common_distributed import 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, TEST_WITH_DEV_DBG_ASAN
label_desc: Use image node_id 4 for calling the TestFSDPExecOrder obj's underlying member method code with example usage: obj.test_invalid_later_iter_order(sharding_strategy, iters_before_path_change) without return types