code: def _worker(
      gpu_id: int,
      sync_file: str,
      world_size: int,
      embedding_size: int,
      flatten_parameters: bool,
  ):
      torch.manual_seed(0)
      torch.cuda.set_device(gpu_id)
      torch.distributed.init_process_group(
          backend="nccl",
          init_method=f"file://{sync_file}",
          world_size=world_size,
          rank=gpu_id,
      )
      process_group = torch.distributed.new_group()
  
      # Create a dummy model with dummy inputs and targets
      batch_size = 4
      input = torch.randn(size=(batch_size, 3, 32, 32)).cuda()
      target = torch.zeros(size=(batch_size, embedding_size)).cuda()
      model = _create_model(
          with_fsdp=True,
          process_group=process_group,
          embedding_size=embedding_size,
          flatten_parameters=flatten_parameters,
      )
      criterion = nn.MSELoss()
      optimizer = torch.optim.SGD(model.parameters(), lr=1e-2)
  
      # Train the model for a few epochs
      for epoch in range(2):
          out = model(input)
          loss = criterion(out, target)
          optimizer.zero_grad()
          loss.backward()
          optimizer.step()
  
      # Save a bunch of checkpoint, one by shard
      cp_data = {
          "weights": {
              k: v.cpu() for k, v in model.local_state_dict().items()
          },
          "meta": model.local_metadata_dict(),
      }
      torch.save(cp_data, f"checkpoint_{gpu_id}.torch")
  
      # Wait for all files to be written on the disk
      dist.barrier()  # type: ignore
  
      # Reconstruct a full checkpoint from the sharded checkpoints
      all_checkpoints = [
          _load_sharded_checkpoint(rank) for rank in range(world_size)
      ]
      consolidated_checkpoint = (
          FullyShardedDataParallel.consolidate_shard_weights(
              shard_weights=[c["weights"] for c in all_checkpoints],
              shard_metadata=[c["meta"] for c in all_checkpoints],
          )
      )
  
      # Check that the reconstructed parameters are correct and of the right shape
      full_model = _create_model(
          with_fsdp=False,
          process_group=process_group,
          embedding_size=embedding_size,
      )
      full_model_state_dict = full_model.state_dict()
      assert set(full_model_state_dict.keys()) == set(
          consolidated_checkpoint.keys()
      )
      for k in full_model_state_dict.keys():
          assert (
              consolidated_checkpoint[k].shape
              == full_model_state_dict[k].shape
          )
  
      # Verify that the checkpoint can be loaded by a FSDP model
      loaded_model = _create_model(
          with_fsdp=True,
          process_group=process_group,
          embedding_size=embedding_size,
          flatten_parameters=flatten_parameters,
      )
      loaded_model.load_state_dict(consolidated_checkpoint)
      for m in loaded_model.modules():
          if isinstance(m, FullyShardedDataParallel):
              m._reset_lazy_init()
  
      # Verify that the model saved and the model loaded give the same results
      with torch.no_grad():
          before_checkpoint_loss = criterion(
              model(input), target
          ).item()
          after_checkpoint_loss = criterion(
              loaded_model(input), target
          ).item()
          assert before_checkpoint_loss == after_checkpoint_loss
  
imports: import functools
  import os
  import tempfile
  from parameterized import parameterized
  import pytest
  import torch
  import torch.distributed
  import torch.multiprocessing
  import torch.nn
  from torch.optim import Adam
  from fairscale.fair_dev.testing.testing import in_temporary_directory, skip_if_single_gpu, temp_files_ctx
  from fairscale.nn import FullyShardedDataParallel
  from tests.nn.data_parallel.test_fsdp import DistributedTest, MixtureOfExperts, rename_test, spawn_and_init
label_desc: Use image node_id 3 for calling a global function with example usage: _worker(gpu_id, sync_file, world_size, embedding_size, flatten_parameters) without return types