code: def _distributed_worker(
      rank,
      world_size,
      fsdp_config,
      fsdp_wrap_bn,
      ddp_mixed_precision,
      tempfile_name,
      unused,
      state_before,
      inputs,
      rank_0_output,
      state_after,
      sync_bn,
      conv_bias,
      linear_bias,
  ):
      torch.backends.cudnn.deterministic = True
  
      result = dist_init(rank, world_size, tempfile_name, unused)
      assert result, "Dist init failed"
  
      ddp = True
      if fsdp_config:
          ddp = False
          assert isinstance(fsdp_config, dict), str(fsdp_config)
          if fsdp_config["mixed_precision"]:
              # To match DDP in AMP -O1, we need fp32 reduce scatter.
              fsdp_config["fp32_reduce_scatter"] = True
  
      model = Model(conv_bias, linear_bias)
      model.load_state_dict(state_before)
      model = model.cuda()
  
      class DummyScaler:
          def scale(self, loss):
              return loss
  
          def step(self, optim):
              optim.step()
  
          def update(self):
              pass
  
      scaler = DummyScaler()
      if ddp:
          if sync_bn == "pytorch":
              model = pytorch_bn_converter(model)
          model = DDP(model, device_ids=[rank], broadcast_buffers=True)
          if ddp_mixed_precision:
              scaler = GradScaler()
      else:
          # Note, different rank may wrap in different order due to different random
          # seeds. But results should be the same.
          if random.randint(0, 1) == 0:
              print(
                  f"auto_wrap_bn {fsdp_wrap_bn}, then sync_bn {sync_bn}"
              )
              if fsdp_wrap_bn:
                  model = auto_wrap_bn(model, _single_rank_pg)
              if sync_bn == "pytorch":
                  model = pytorch_bn_converter(model)
          else:
              print(
                  f"sync_bn {sync_bn}, then auto_wrap_bn {fsdp_wrap_bn}"
              )
              if sync_bn == "pytorch":
                  model = pytorch_bn_converter(model)
              if fsdp_wrap_bn:
                  model = auto_wrap_bn(model, _single_rank_pg)
          model = FSDP(model, **fsdp_config).cuda()
          if fsdp_config["mixed_precision"]:
              scaler = ShardedGradScaler()
          # Print the model for verification.
          if rank == 0:
              print(model)
      optim = SGD(model.parameters(), lr=0.1)
      loss_func = CrossEntropyLoss()
  
      for in_data in inputs[rank]:
          in_data = in_data.cuda()
          context = contextlib.suppress()
          if ddp and ddp_mixed_precision:
              in_data = in_data.half()
              context = torch.cuda.amp.autocast(enabled=True)
          if not ddp and fsdp_config["mixed_precision"]:
              context = torch.cuda.amp.autocast(enabled=True)
          with context:
              out = model(in_data)
              fake_label = torch.zeros(1, dtype=torch.long).cuda()
              loss = loss_func(out.unsqueeze(0), fake_label)
          scaler.scale(loss).backward()
          scaler.step(optim)
          scaler.update()
          optim.zero_grad()
  
      if ddp:
          # Save the rank 0 state_dict to the output file.
          if rank == 0:
              state_after = model.module.cpu().state_dict()
              torch.save(state_after, rank_0_output)
      else:
          model.assert_state(TrainingState.IDLE)
          # Ensure final state equals to the state_after.
          fsdp_state = model.state_dict()
          # Move tensors to CPU to compare numerics.
          for k, v in fsdp_state.items():
              fsdp_state[k] = v.cpu()
          # Change False to True to enable this when you want to debug the mismatch.
          if False and rank == 0:
  
              def dump(d):
                  for k, v in d.items():
                      print(k, v)
  
              dump(state_after)
              dump(fsdp_state)
          # If sync_bn is used, all ranks should have the same state, so we can compare with
          # rank 0 state on every rank. Otherwise, only compare rank 0 with rank 0.
          if sync_bn != "none" or rank == 0:
              assert objects_are_equal(
                  state_after, fsdp_state, raise_exception=True
              )
  
      teardown()
  
imports: import contextlib
  from itertools import product
  import random
  import tempfile
  import pytest
  import torch
  from torch.cuda.amp import GradScaler
  import torch.multiprocessing
  from torch.nn import AdaptiveAvgPool2d, BatchNorm2d, Conv2d, CrossEntropyLoss, Linear, Module, ReLU, Sequential, Sigmoid, SyncBatchNorm
  from torch.nn.parallel import DistributedDataParallel
  from torch.optim import SGD
  from fairscale.fair_dev.testing.testing import dist_init, objects_are_equal, rmf, skip_if_single_gpu, state_dict_norm, teardown, torch_cuda_version
  from fairscale.internal import torch_version
  from fairscale.nn.data_parallel import FullyShardedDataParallel
  from fairscale.nn.data_parallel import TrainingState, auto_wrap_bn
label_desc: Use image node_id 3 for calling a global function with example usage: _distributed_worker(rank, world_size, fsdp_config, fsdp_wrap_bn, ddp_mixed_precision, tempfile_name, unused, state_before, inputs, rank_0_output, state_after, sync_bn, conv_bias, linear_bias) and returns: loss