code: def _test_func(
      rank,
      world_size,
      model,
      fsdp_config,
      tempfile_name,
      unused,
      test_case,
  ):
      result = dist_init(rank, world_size, tempfile_name, unused)
      assert result, "Dist init failed"
  
      my_lr = 0.1
  
      device = torch.device("cuda")
      if fsdp_config.get("mixed_precision", False):
          dtype = torch.float16
          fsdp_config["fp32_reduce_scatter"] = True
      else:
          dtype = torch.float32
  
      if test_case["assert_ref_out"]:
          with torch.no_grad():
              # Compute one iteration local output.
              fp32_weight = model.weight.T.clone().to(device)
              weight = fp32_weight.to(dtype)
              v = torch.Tensor(test_case["inputs"][0][rank]).to(
                  device, dtype
              )
              ref_forward_output_my_rank = torch.matmul(v, weight)
              # Compute one iteration global weight update.
              v = torch.Tensor(test_case["inputs"][0][:world_size]).to(
                  device, dtype
              )
              grad = (
                  v.float()
                  .sum(0)
                  .repeat(weight.shape[0], 1)
                  .div(world_size)
              )
              ref_weight_out = fp32_weight - grad.T * my_lr
              assert ref_weight_out.dtype == torch.float32
      model.to(
          device
      )  # not dtype, since FSDP will manage mixed precision internally
      assert isinstance(fsdp_config, dict), str(fsdp_config)
      model = FSDP(model, **fsdp_config)
      optim = SGD(model.parameters(), lr=my_lr)
      inputs = test_case["inputs"]
      assert len(inputs) == 1 or not test_case["assert_ref_out"]
      assert len(inputs[0]) >= world_size
      for in_data in inputs:
          in_data = Tensor(in_data[rank]).to(device, dtype)
          out = model(in_data)
          out.float().sum().backward()
          optim.step()
          optim.zero_grad()
          if test_case["assert_ref_out"]:
              with model.summon_full_params():
                  weight_out = model.module.weight.data.T.clone()
              # make sure we can do more fwd/bwd
              loss = model(in_data)
              loss.sum().backward()
  
      if test_case["assert_ref_out"]:
          torch.testing.assert_allclose(ref_forward_output_my_rank, out)
          torch.testing.assert_allclose(ref_weight_out, weight_out)
  
      model.assert_state(TrainingState.IDLE)
      teardown()
  
file_path: fairscale/tests/nn/data_parallel/test_fsdp_uneven.py
imports: import tempfile
  import pytest
  import torch
  from torch import Tensor
  import torch.multiprocessing
  from torch.nn import Linear, Sequential
  from torch.optim import SGD
  from fairscale.fair_dev.testing.testing import dist_init, skip_if_single_gpu, teardown
  from fairscale.internal import torch_version
  from fairscale.nn.data_parallel import FullyShardedDataParallel
  from fairscale.nn.data_parallel.fully_sharded_data_parallel import TrainingState
label_desc: Use image node_id 1 for calling a global function with example usage: _test_func(rank, world_size, model, fsdp_config, tempfile_name, unused, test_case) without return types