code: def _worker_test_init_vit_init(
      cls, gpu_id: int, sync_file: str, world_size: int
  ):
      init_distributed_on_file(
          world_size=world_size, gpu_id=gpu_id, sync_file=sync_file
      )
      ddp_config = cls._create_dino_pretraining_config(
          with_fsdp=False,
          with_mixed_precision=False,
          gpu_count=world_size,
      )
      fsdp_config = cls._create_dino_pretraining_config(
          with_fsdp=True,
          with_mixed_precision=False,
          gpu_count=world_size,
      )
  
      # For faster tests, try a smaller ViT
      ddp_config.MODEL.TRUNK.XCIT.NUM_LAYERS = 2
      fsdp_config.MODEL.TRUNK.XCIT.NUM_LAYERS = 2
  
      # For deterministic computing
      ddp_config.MODEL.TRUNK.XCIT.DROP_PATH_RATE = 0.0
      fsdp_config.MODEL.TRUNK.XCIT.DROP_PATH_RATE = 0.0
  
      # Create fake inputs
      torch.random.manual_seed(gpu_id)
      x = torch.randn(size=(2, 3, 224, 224)).cuda()
  
      # Create both FSDP and DDP models
      ddp_model = build_model(
          ddp_config["MODEL"], ddp_config["OPTIMIZER"]
      ).cuda()
      ddp_model = DistributedDataParallel(
          ddp_model, device_ids=[gpu_id]
      )
      fsdp_model = build_model(
          fsdp_config["MODEL"], fsdp_config["OPTIMIZER"]
      ).cuda()
      fsdp_model = fsdp_wrapper(
          fsdp_model, **fsdp_config["MODEL"]["FSDP_CONFIG"]
      )
      if gpu_id == 0:
          print(fsdp_model)
  
      # Check that the weights are equal on both models
      with fsdp_model.summon_full_params():
          ddp_params = {
              n: p.sum().item() for n, p in ddp_model.named_parameters()
          }
          for n, p in fsdp_model.named_parameters():
              p_ref = ddp_params[f"module.{n}"]
              assert p_ref == p.sum().item()
  
      # Check that the local_state_dict / state_dict are valid
      fsdp_model_2 = build_model(
          fsdp_config["MODEL"], fsdp_config["OPTIMIZER"]
      ).cuda()
      fsdp_model_2 = fsdp_wrapper(
          fsdp_model_2, **fsdp_config["MODEL"]["FSDP_CONFIG"]
      )
      fsdp_model_2.load_local_state_dict(fsdp_model.local_state_dict())
      fsdp_model_2.load_state_dict(fsdp_model.state_dict())
  
      # Compare that DDP and FSDP models give the same results
      with torch.no_grad(), torch.cuda.amp.autocast():
          ddp_out = ddp_model(x)
          fsdp_out = fsdp_model(x)
          if gpu_id == 0:
              print(ddp_out[0][0])
              print(fsdp_out[0][0])
          assert torch.allclose(ddp_out[0][0], fsdp_out[0][0])
  
file_path: vissl/tests/test_dino_xcit_fsdp.py
imports: import unittest
  import torch
  import torch.multiprocessing
  from torch.nn.parallel import DistributedDataParallel
  from vissl.models import build_model
  from vissl.utils.fsdp_utils import fsdp_wrapper
  from vissl.utils.hydra_config import compose_hydra_configuration, convert_to_attrdict
  from vissl.utils.test_utils import gpu_test, in_temporary_directory, init_distributed_on_file, run_integration_test, with_temp_files
label_desc: Use image node_id 4 for calling the TestDINO_XCIT obj's underlying member method code with example usage: obj._worker_test_init_vit_init(cls, gpu_id, sync_file, world_size) without return types