code: def xla_fsdp_train_save_load(
      fabric: Fabric, tmp_path, state_dict_type
  ):
      """Fabric launch function for test_xla_fsdp_train_save_load."""
      tmp_path = Path(fabric.broadcast(tmp_path))
  
      with fabric.init_module():
          model_1 = torch.nn.Sequential(
              torch.nn.Linear(32, 32),
              torch.nn.ReLU(),
              torch.nn.Linear(32, 2),
          )
      model_1 = fabric.setup_module(model_1)
  
      optimizer_1 = torch.optim.Adam(model_1.parameters(), lr=0.1)
      optimizer_1 = fabric.setup_optimizers(optimizer_1)
  
      dataloader = DataLoader(RandomDataset(32, 64))
      dataloader = fabric.setup_dataloaders(dataloader)
  
      def step(model, batch):
          output = model(batch)
          return torch.nn.functional.mse_loss(
              output, torch.ones_like(output)
          )
  
      model_1.train()
      data_iter = iter(dataloader)
      batch = next(data_iter)
      loss = step(model_1, batch)
      fabric.backward(loss)
      optimizer_1.step()
      optimizer_1.zero_grad()
  
      state = {
          "model": model_1,
          "optimizer": optimizer_1,  # not needed in ckpt consolidation
          "step_count": 1,
      }
      checkpoint_path = tmp_path / "foo.pth"
  
      world_size = fabric.world_size
      local_process_count = len(fabric.strategy.parallel_devices)
      is_multihost = local_process_count < world_size
      if state_dict_type == "full" and is_multihost:
          with pytest.raises(
              OSError,
              match="Multihost setups do not have a shared filesystem",
          ):
              fabric.save(checkpoint_path, state)
          return
      fabric.save(checkpoint_path, state)
  
      if state_dict_type == "sharded":
          pattern = (
              rf"checkpoint_rank-0000000\d-of-{world_size:08d}\.pth"
          )
          shards = os.listdir(checkpoint_path)
          assert len(shards) == local_process_count
          for name in shards:
              assert re.match(pattern, name)
  
          # define a second set of model and optimizer
          with fabric.init_module():
              model_2 = torch.nn.Sequential(
                  torch.nn.Linear(32, 32),
                  torch.nn.ReLU(),
                  torch.nn.Linear(32, 2),
              )
          model_2 = fabric.setup_module(model_2)
  
          optimizer_2 = torch.optim.Adam(model_2.parameters(), lr=0.1)
          optimizer_2 = fabric.setup_optimizers(optimizer_2)
  
          # load sharded checkpoints into the second set of model and optimizer
          state = {
              "model": model_2,
              "optimizer": optimizer_2,
              "step_count": 0,
          }
          metadata = fabric.load(checkpoint_path, state)
  
          # check user data in loaded state
          assert not metadata
          assert state["step_count"] == 1
  
          # check correctness with loaded state
          for p0, p1 in zip(
              model_1._forward_module.parameters(), model_2.parameters()
          ):
              torch.testing.assert_close(
                  p0, p1, atol=0, rtol=0, equal_nan=True
              )
  
          # attempt to load a key not in the metadata checkpoint
          state = {"model": model_2, "coconut": 11}
          with pytest.raises(
              KeyError,
              match="The requested state contains a key 'coconut' that does not exist",
          ):
              fabric.load(checkpoint_path, state)
  
          # `strict=False` ignores the missing key
          state = {"model": model_2, "coconut": 11}
          fabric.load(checkpoint_path, state, strict=False)
          assert state["coconut"] == 11
  
      if state_dict_type == "full":
          assert set(os.listdir(tmp_path)) == {"foo.pth"}
  
          # define a second set of model and optimizer
          with fabric.init_module():
              model_2 = torch.nn.Sequential(
                  torch.nn.Linear(32, 32),
                  torch.nn.ReLU(),
                  torch.nn.Linear(32, 2),
              )
          import torch_xla.core.xla_model as xm
  
          device = xm.xla_device()
          model_2.to(device)
  
          # load sharded checkpoints into the second model
          state = {"model": model_2}
          fabric.load(checkpoint_path, state)
  
          # check that loaded state is different
          with pytest.raises(AssertionError, match="do not match"):
              for p0, p1 in zip(
                  model_1.parameters(), model_2.parameters()
              ):
                  torch.testing.assert_close(
                      p0, p1, atol=0, rtol=0, equal_nan=True
                  )
  
file_path: pytorch-lightning/tests/tests_fabric/strategies/test_xla_fsdp_integration.py
imports: import os
  import re
  from pathlib import Path
  from unittest import mock
  import pytest
  import torch
  from lightning.fabric import Fabric
  from lightning.fabric.strategies import XLAFSDPStrategy
  from torch.utils.data import DataLoader
  from tests_fabric.helpers.models import RandomDataset
  from tests_fabric.helpers.runif import RunIf
label_desc: Use image node_id 3 for calling a global function with example usage: xla_fsdp_train_save_load(fabric, tmp_path, state_dict_type) and returns: torch