class_docstr: This contains the data structures to track the execution order. We track
  the pre-forward order on the *first* iteration for forward prefetching
  (which thus assumes static graph) and the post-forward order on *every*
  iteration for backward prefetching (which thus does not assume static
  graph but may be provide an incorrect order).
docstr: Checks the forward execution order as long as ``is_training`` is
  ``True`` since checking in eval mode is not supported. This only checks
  if the distributed debug level is DETAIL.
  
  - On the first iteration, this uses all-gathers to check that all ranks
  are all-gathering the same handles and hence ``FlatParameter`` s,
  raising an error if not.
  - On subsequent iterations, this checks that each rank is locally
  consistent with its own forward order from the first iteration, issuing
  a warning if not. This issues a warning on the first deviating
  iteration and stops warning thereafter.
docstr_tok: ["Checks","the","forward","execution","order","as","long","as","``","is_training","``","is","``","True","``","since","checking","in","eval","mode","is","not","supported",".","This","only","checks","if","the","distributed","debug","level","is","DETAIL",".","-","On","the","first","iteration",",","this","uses","all-gathers","to","check","that","all","ranks","are","all-gathering","the","same","handles","and","hence","``","FlatParameter","``","s",",","raising","an","error","if","not",".","-","On","subsequent","iterations",",","this","checks","that","each","rank","is","locally","consistent","with","its","own","forward","order","from","the","first","iteration",",","issuing","a","warning","if","not",".","This","issues","a","warning","on","the","first","deviating","iteration","and","stops","warning","thereafter","."]
code: def _check_order(
      self, handle: FlatParamHandle, is_training: bool
  ) -> None:
      """
      Checks the forward execution order as long as ``is_training`` is
      ``True`` since checking in eval mode is not supported. This only checks
      if the distributed debug level is DETAIL.
  
      - On the first iteration, this uses all-gathers to check that all ranks
      are all-gathering the same handles and hence ``FlatParameter`` s,
      raising an error if not.
      - On subsequent iterations, this checks that each rank is locally
      consistent with its own forward order from the first iteration, issuing
      a warning if not. This issues a warning on the first deviating
      iteration and stops warning thereafter.
      """
      # Do not check order in eval mode since the post-backward callback does
      # not run so it cannot be used to mark the end of an iteration
      if not is_training or not self._checking_order:
          return
      if self.is_first_iter:
          msg_prefix = "Forward order differs across ranks:"
          optional_local_indices: Tuple[
              Optional[int], ...
          ] = self._get_handle_indices(handle)
          device = handle.device  # guaranteed to be non-CPU
          num_valid_indices = sum(
              (index is not None) for index in optional_local_indices
          )
          tensor_kwargs: Dict[str, Union[torch.dtype, torch.device]] = {
              "dtype": torch.int32,
              "device": device,
          }
          world_num_valid_indices = torch.zeros(self.world_size, **tensor_kwargs)  # type: ignore[arg-type, call-overload]
          local_num_valid_indices = torch.tensor([num_valid_indices], **tensor_kwargs)  # type: ignore[arg-type, call-overload]
          dist.all_gather_into_tensor(
              world_num_valid_indices,
              local_num_valid_indices,
              group=self.process_group,
          )
          # Copy entire tensor from D2H once to avoid per element D2H copies
          world_num_valid_indices = world_num_valid_indices.cpu()
          # Check that all ranks plan to all-gather the same number of
          # parameters
          # TODO (awgu): Since every module has at most one handle in the
          # current implementation, this should never raise the error.
          assert self.world_size is not None  # mypy
          if (
              not torch.distributed._functional_collectives.is_torchdynamo_compiling()
          ):
              # TODO(voz): Don't graph break on this - dynamo hates the n1 != n2
              # tensor comparison control flow.
              # https://github.com/pytorch/pytorch/issues/107055
              for (r1, n1), (r2, n2) in itertools.combinations(
                  (
                      (rank, world_num_valid_indices[rank])
                      for rank in range(self.world_size)
                  ),
                  2,
              ):
                  if n1 != n2:
                      raise RuntimeError(
                          f"{msg_prefix} rank {r1} is all-gathering {n1} parameters "
                          f"while rank {r2} is all-gathering {n2} parameters"
                      )
          world_indices = torch.zeros(  # type: ignore[call-overload]
              self.world_size * num_valid_indices, **tensor_kwargs
          )
          local_indices = torch.tensor(optional_local_indices, **tensor_kwargs)  # type: ignore[arg-type]
          dist.all_gather_into_tensor(
              world_indices, local_indices, group=self.process_group
          )
          # Copy entire tensor from D2H once to avoid per element D2H copies
          world_indices = world_indices.cpu()
          # Check that all ranks plan to all-gather the same index parameters
          if (
              not torch.distributed._functional_collectives.is_torchdynamo_compiling()
          ):
              # TODO(voz): Don't graph break on this - dynamo hates the i1 != i2
              # tensor comparison control flow.
              # https://github.com/pytorch/pytorch/issues/107055
              for (r1, i1), (r2, i2) in itertools.combinations(
                  (
                      (
                          rank,
                          world_indices[
                              rank
                              * num_valid_indices : (rank + 1)
                              * num_valid_indices
                          ],
                      )
                      for rank in range(self.world_size)
                  ),
                  2,
              ):
                  if i1 != i2:
                      r1_param_names = (
                          self._get_names_from_handle_indices(i1)
                      )
                      r2_param_names = (
                          self._get_names_from_handle_indices(i2)
                      )
                      raise RuntimeError(
                          f"{msg_prefix} rank {r1} is all-gathering parameters "
                          f"for {r1_param_names} while rank {r2} is all-gathering "
                          f"parameters for {r2_param_names}"
                      )
      else:
          # Only issue warnings on the first deviating iteration and stop
          # checking thereafter to avoid flooding the console
          if self.warn_status == _ExecOrderWarnStatus.WARNED:
              return
          msg_prefix = None  # non-`None` means we should warn
          if self.current_order_index >= len(
              self.handles_pre_forward_order
          ):
              # This iteration sees extra all-gather(s) compared to the first
              msg_prefix = (
                  "Expected to not all-gather any more parameters in the "
                  "forward but trying to all-gather parameters for "
              )
          else:
              expected_handle = self.handles_pre_forward_order[
                  self.current_order_index
              ]
              if expected_handle != handle:
                  expected_param_names = self._get_names_from_handles(
                      expected_handle
                  )
                  msg_prefix = (
                      f"Expected to all-gather for {expected_param_names} "
                      "but trying to all-gather parameters for "
                  )
          if msg_prefix is not None:
              param_names = self._get_names_from_handles(handle)
              msg_suffix = (
                  f"{param_names}"
                  if param_names
                  else "a newly-added parameter since construction time"
              )
              warnings.warn(
                  "Forward order differs from that of the first iteration "
                  f"on rank {self.rank}. Collectives are unchecked and may "
                  f"give incorrect results or hang.\n{msg_prefix}{msg_suffix}"
              )
              self.warn_status = _ExecOrderWarnStatus.WARNING
          self.current_order_index += 1
  
imports: import itertools
  import warnings
  from enum import auto, Enum
  from typing import Dict, List, Optional, Tuple, Union
  import torch
  import torch.distributed
  import torch.distributed.fsdp._traversal_utils
  import torch.nn
  from torch.distributed.fsdp._common_utils import _FSDPState, _get_param_to_fqns
  from torch.distributed.fsdp._flat_param import FlatParamHandle
label_desc: Use image node_id 8 for calling the _ExecOrderData obj's underlying member method code with example usage: obj._check_order(handle, is_training) without return types