code: def test_fsdp_train_save_load(tmp_path, manual_wrapping, precision):
      """Test FSDP training, saving and loading with different wrapping and precision settings."""
      fabric_cls = (
          _MyFabricManualWrapping if manual_wrapping else _MyFabric
      )
      fabric = fabric_cls(
          accelerator="cuda",
          strategy=FSDPStrategy(auto_wrap_policy=always_wrap_policy),
          devices=2,
          precision=precision,
      )
      fabric.run()
  
      checkpoint_path = fabric.broadcast(
          str(tmp_path / "fsdp-checkpoint")
      )
  
      params_before = deepcopy(list(fabric.model.parameters()))
      state = {
          "model": fabric.model,
          "optimizer": fabric.optimizer,
          "steps": 1,
      }
      fabric.save(checkpoint_path, state)
      assert set(os.listdir(checkpoint_path)) == {
          "meta.pt",
          ".metadata",
          "__0_0.distcp",
          "__1_0.distcp",
      }
  
      # re-init all objects and resume
      fabric = fabric_cls(
          accelerator="cuda",
          strategy=FSDPStrategy(auto_wrap_policy=always_wrap_policy),
          devices=2,
          precision=precision,
      )
      fabric.run()
  
      # check correctness with loaded state
      state = {
          "model": fabric.model,
          "optimizer": fabric.optimizer,
          "steps": 0,
      }
      metadata = fabric.load(checkpoint_path, state)
      for p0, p1 in zip(params_before, fabric.model.parameters()):
          torch.testing.assert_close(
              p0, p1, atol=0, rtol=0, equal_nan=True
          )
  
      # check user data in state reloaded
      assert state["steps"] == 1
      assert not metadata
  
      # attempt to load a key not in the metadata checkpoint
      state = {"model": fabric.model, "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": fabric.model, "coconut": 11}
      fabric.load(checkpoint_path, state, strict=False)
      assert state["coconut"] == 11
  
imports: import os
  from copy import deepcopy
  from pathlib import Path
  from unittest import mock
  import pytest
  import torch
  from lightning.fabric import Fabric
  from lightning.fabric.plugins import FSDPPrecision
  from lightning.fabric.strategies import FSDPStrategy
  from lightning.fabric.utilities.imports import _TORCH_GREATER_EQUAL_2_0, _TORCH_GREATER_EQUAL_2_1
  from lightning.fabric.wrappers import _FabricOptimizer
  from torch.distributed.fsdp import FlatParameter, FullyShardedDataParallel, OptimStateKeyType
  from torch.distributed.fsdp.wrap import always_wrap_policy, wrap
  from torch.nn import Parameter
  from tests_fabric.helpers.models import BoringFabric
  from tests_fabric.helpers.runif import RunIf
  from tests_fabric.test_fabric import BoringModel
label_desc: Use image node_id 1 for calling a global function with example usage: test_fsdp_train_save_load(tmp_path, manual_wrapping, precision) without return types