docstr: Tests that, when training starts, if FSDP parameters are not on the
  expected device, then an informative error is raised. This applies for
  both no parameter CPU offloading and parameter CPU offloading.
docstr_tok: ["Tests","that",",","when","training","starts",",","if","FSDP","parameters","are","not","on","the","expected","device",",","then","an","informative","error","is","raised",".","This","applies","for","both","no","parameter","CPU","offloading","and","parameter","CPU","offloading","."]
code: def test_training_device_mismatch_errors(self):
      """
      Tests that, when training starts, if FSDP parameters are not on the
      expected device, then an informative error is raised. This applies for
      both no parameter CPU offloading and parameter CPU offloading.
      """
      # Incorrectly not moving from CPU -> GPU
      model = torch.nn.Linear(10, 10)
      fsdp_model = FSDP(model)
      inp = torch.randn((2, 10))
      with self.assertRaisesRegex(
          RuntimeError,
          "An FSDP-managed module unexpectedly has parameters on cpu. Make "
          "sure to move the module to cuda:0 before training.",
      ):
          fsdp_model(inp)
  
      # Incorrectly moving from CPU -> GPU
      model = torch.nn.Linear(10, 10)
      fsdp_model = FSDP(
          model, cpu_offload=CPUOffload(offload_params=True)
      )
      fsdp_model.to(torch.device("cuda"))
      inp = torch.randn((2, 10))
      with self.assertRaisesRegex(
          RuntimeError,
          "An FSDP-managed module with parameter CPU offloading enabled has "
          "parameters on cuda:0. Make sure to not move the module from CPU "
          "when offloading parameters.",
      ):
          fsdp_model(inp)
  
imports: import functools
  import os
  import sys
  import warnings
  from collections import namedtuple
  from contextlib import nullcontext
  from copy import deepcopy
  from itertools import chain
  from typing import Any, Tuple
  import torch
  import torch.distributed
  import torch.distributed.fsdp._traversal_utils
  import torch.nn
  from torch.distributed.fsdp import CPUOffload, FlatParameter, FullyShardedDataParallel, ShardingStrategy
  from torch.distributed.fsdp._flat_param import _FSDP_USE_UNSAFE_SETATTR
  from torch.distributed.fsdp._runtime_utils import HOMOGENEOUS_ATTR_NAMES
  from torch.distributed.fsdp.wrap import always_wrap_policy, ModuleWrapPolicy, transformer_auto_wrap_policy
  from torch.distributed.optim import _apply_optimizer_in_backward
  from torch.nn import TransformerDecoderLayer, TransformerEncoderLayer
  from torch.nn.parallel import DistributedDataParallel
  from torch.testing._internal.common_distributed import skip_if_lt_x_gpu
  from torch.testing._internal.common_fsdp import _assert_module_states, CUDAInitMode, FSDPInitMode, FSDPTest, FSDPTestMultiThread, NestedWrappedModule, 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 3 for calling the TestFSDPMiscWorldSize1 obj's underlying member method code with example usage: obj.test_training_device_mismatch_errors() without return types