class_docstr: Compare losses and parameter values after several updates when using
  PyTorch DDP vs. FullyShardedDataParallel.
class_docstr_tok: ["Compare","losses","and","parameter","values","after","several","updates","when","using","PyTorch","DDP","vs.","FullyShardedDataParallel","."]
docstr: Tests the FSDP forward, backward, and optimizer step runtime by
  using a model with a long CUDA delay after the loss computation/before
  the optimizer step to exercise the internal CUDA stream usage in that
  the forward pass all-gathers do not start until after the optimizer
  step completes.
docstr_tok: ["Tests","the","FSDP","forward",",","backward",",","and","optimizer","step","runtime","by","using","a","model","with","a","long","CUDA","delay","after","the","loss","computation\/before","the","optimizer","step","to","exercise","the","internal","CUDA","stream","usage","in","that","the","forward","pass","all-gathers","do","not","start","until","after","the","optimizer","step","completes","."]
code: def test_delayed_optim_step(
      self,
      cpu_offload: CPUOffload,
      sharding_strategy: Optional[ShardingStrategy],
  ):
      """Tests the FSDP forward, backward, and optimizer step runtime by
      using a model with a long CUDA delay after the loss computation/before
      the optimizer step to exercise the internal CUDA stream usage in that
      the forward pass all-gathers do not start until after the optimizer
      step completes."""
      self.run_subtests(
          self._get_subtest_config(cpu_offload),
          self._test_fsdp_parity,
          NestedWrappedModuleWithDelay,
          FSDPInitMode.RECURSIVE,
          cpu_offload=cpu_offload,
          sharding_strategy=sharding_strategy,
          init_kwargs={"delay_after_loss_ms": 250},
      )
  
imports: import contextlib
  import functools
  import itertools
  import sys
  from typing import Any, Callable, Dict, List, Optional
  from unittest import mock
  import torch
  import torch.distributed
  import torch.nn
  from torch.distributed.fsdp import CPUOffload, MixedPrecision
  from torch.distributed.fsdp._flat_param import FlatParamHandle
  from torch.distributed.fsdp.fully_sharded_data_parallel import BackwardPrefetch, FullyShardedDataParallel, ShardingStrategy
  from torch.distributed.fsdp.wrap import ModuleWrapPolicy
  from torch.distributed.utils import _p_assert
  from torch.testing._internal.common_distributed import skip_if_lt_x_gpu
  from torch.testing._internal.common_fsdp import AlwaysWrapNestedWrappedModule, CUDAInitMode, DummyDDP, FSDPInitMode, FSDPTest, MixtureOfExperts, NestedWrappedModule, NestedWrappedModuleWithDelay, subtest_name, TransformerWithSharedParams
  from torch.testing._internal.common_utils import instantiate_parametrized_tests, parametrize, run_tests, TEST_WITH_DEV_DBG_ASAN
label_desc: Use image node_id 7 for calling the TestParityWithDDP obj's underlying member method code with example usage: obj.test_delayed_optim_step(cpu_offload, sharding_strategy) without return types