code: def test_fsdp_save_full_state_dict(tmp_path):
      """Test that FSDP saves the full state into a single file with `state_dict_type="full"`."""
      fabric = BoringFabric(
          accelerator="cuda",
          strategy=FSDPStrategy(
              auto_wrap_policy=always_wrap_policy,
              state_dict_type="full",
          ),
          devices=2,
      )
      fabric.run()
  
      checkpoint_path = Path(
          fabric.broadcast(str(tmp_path / "fsdp-checkpoint.pt"))
      )
  
      state = {
          "model": fabric.model,
          "optimizer": fabric.optimizer,
          "steps": 1,
      }
      fabric.save(checkpoint_path, state)
  
      checkpoint = torch.load(checkpoint_path)
      assert checkpoint["steps"] == 1
      loaded_state_dict = checkpoint["model"]
  
      # assert the correct state model was saved
      with FullyShardedDataParallel.summon_full_params(fabric.model):
          state_dict = fabric.model.state_dict()
          assert set(loaded_state_dict.keys()) == set(state_dict.keys())
          for param_name in state_dict:
              assert torch.equal(
                  loaded_state_dict[param_name],
                  state_dict[param_name].cpu(),
              )
          params_before = [p.cpu() for p in fabric.model.parameters()]
  
      # assert the correct optimizer state was saved
      optimizer_state_before = (
          FullyShardedDataParallel.full_optim_state_dict(
              fabric.model, fabric.optimizer, rank0_only=False
          )
      )
      assert (
          set(checkpoint["optimizer"].keys())
          == set(optimizer_state_before.keys())
          == {"state", "param_groups"}
      )
  
      # 1. verify the FSDP state can be loaded back into a FSDP model/strategy directly
      fabric = BoringFabric(
          accelerator="cuda",
          strategy=FSDPStrategy(auto_wrap_policy=always_wrap_policy),
          devices=2,
      )
      fabric.run()
      metadata = fabric.load(
          checkpoint_path,
          {"model": fabric.model, "optimizer": fabric.optimizer},
      )
      assert metadata == {"steps": 1}
  
      with FullyShardedDataParallel.summon_full_params(fabric.model):
          params_after = list(fabric.model.parameters())
          assert all(
              torch.equal(p0.cpu(), p1.cpu())
              for p0, p1 in zip(params_before, params_after)
          )
  
      # assert the correct optimizer state was loaded
      optimizer_state_after = (
          FullyShardedDataParallel.full_optim_state_dict(
              fabric.model, fabric.optimizer, rank0_only=False
          )
      )
      assert (
          set(optimizer_state_after.keys())
          == set(optimizer_state_before.keys())
          == {"state", "param_groups"}
      )
      torch.testing.assert_close(
          optimizer_state_after["state"],
          optimizer_state_before["state"],
          atol=0,
          rtol=0,
      )
      assert (
          optimizer_state_after["param_groups"]
          == optimizer_state_before["param_groups"]
      )
  
      # run a step to verify the optimizer state is correct
      fabric.run()
  
      # 2. verify the FSDP state can be loaded back into a single-device model/strategy
      fabric = BoringFabric(accelerator="cpu", devices=1)
      fabric.run()
      metadata = fabric.load(
          checkpoint_path,
          {"model": fabric.model, "optimizer": fabric.optimizer},
      )
      assert metadata == {"steps": 1}
      params_after = list(fabric.model.parameters())
      assert all(
          torch.equal(p0, p1)
          for p0, p1 in zip(params_before, params_after)
      )
  
      # get optimizer state after loading
      normal_checkpoint_path = Path(
          fabric.broadcast(str(tmp_path / "normal-checkpoint.pt"))
      )
      fabric.save(
          normal_checkpoint_path,
          {
              "model": fabric.model,
              "optimizer": fabric.optimizer,
              "steps": 2,
          },
      )
      optimizer_state_after = torch.load(normal_checkpoint_path)[
          "optimizer"
      ]
      optimizer_state_after = (
          FullyShardedDataParallel.rekey_optim_state_dict(
              optimizer_state_after,
              optim_state_key_type=OptimStateKeyType.PARAM_NAME,
              model=fabric.model,
          )
      )
  
      # assert the correct optimizer state was loaded
      assert (
          set(optimizer_state_after.keys())
          == set(optimizer_state_before.keys())
          == {"state", "param_groups"}
      )
      torch.testing.assert_close(
          optimizer_state_after["state"],
          optimizer_state_before["state"],
          atol=0,
          rtol=0,
      )
  
      # run a step to verify the optimizer state is correct
      fabric.run()
  
      # 3. verify that a single-device model/strategy states can be loaded into a FSDP model/strategy
      fabric = BoringFabric(
          accelerator="cuda",
          strategy=FSDPStrategy(auto_wrap_policy=always_wrap_policy),
          devices=2,
      )
      fabric.run()
      metadata = fabric.load(
          normal_checkpoint_path,
          {"model": fabric.model, "optimizer": fabric.optimizer},
      )
      assert metadata == {"steps": 2}
  
      with FullyShardedDataParallel.summon_full_params(fabric.model):
          params_after = list(fabric.model.parameters())
          assert all(
              torch.equal(p0.cpu(), p1.cpu())
              for p0, p1 in zip(params_before, params_after)
          )
  
      # assert the correct optimizer state was loaded
      optimizer_state_after = (
          FullyShardedDataParallel.full_optim_state_dict(
              fabric.model, fabric.optimizer, rank0_only=False
          )
      )
      assert (
          set(optimizer_state_after.keys())
          == set(optimizer_state_before.keys())
          == {"state", "param_groups"}
      )
      torch.testing.assert_close(
          optimizer_state_after["state"],
          optimizer_state_before["state"],
          atol=0,
          rtol=0,
      )
      assert (
          optimizer_state_after["param_groups"]
          == optimizer_state_before["param_groups"]
      )
  
      # run a step to verify the optimizer state is correct
      fabric.run()
  
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 2 for calling a global function with example usage: test_fsdp_save_full_state_dict(tmp_path) without return types