code: def test_fully_shard_sgd(self):
      sgd_configs = [
          {"lr": 0.1},
          {"lr": 0.1, "momentum": 0.9},
          {"lr": 0.1, "momentum": 0.9, "foreach": True},
      ]
  
      for config in sgd_configs:
          mod = SimpleMLP().cuda(self.rank)
          opt = torch.optim.SGD(mod.parameters(), **config)
  
          train_batch = (
              torch.randn((128, 50), device=torch.device(self.rank)),
              torch.randn((128, 8), device=torch.device(self.rank)),
          )
  
          ddp_mod = DDP(deepcopy(mod), device_ids=[self.rank])
          ddp_opt = torch.optim.SGD(ddp_mod.parameters(), **config)
          self._test_data_parallel(
              mod,
              ddp_mod,
              opt,
              ddp_opt,
              train_batch,
              train_step,
              "fully_shard",
          )
  
imports: from copy import deepcopy
  import torch
  import torch.nn
  from torch.distributed._spmd.api import compile
  from torch.distributed._spmd.parallel_mode import DataParallel
  from torch.distributed._tensor import Replicate
  from torch.nn.parallel import DistributedDataParallel
  from torch.testing._internal.common_distributed import skip_if_lt_x_gpu
  from torch.testing._internal.common_utils import run_tests
  from torch.testing._internal.distributed._tensor.common_dtensor import DTensorTestBase, with_comms
label_desc: Use image node_id 5 for calling the TestDataParallel obj's underlying member method code with example usage: obj.test_fully_shard_sgd() without return types