code: def test_deepspeed_multiple_models():
      fabric = Fabric(
          strategy=DeepSpeedStrategy(
              stage=3, logging_batch_size_per_gpu=1
          ),
          devices=2,
          accelerator="gpu",
      )
      fabric.launch()
  
      with fabric.init_module():
          model = BoringModel()
  
      optimizer = torch.optim.SGD(model.parameters(), lr=0.0001)
      model, optimizer = fabric.setup(model, optimizer)
  
      for i in range(2):
          optimizer.zero_grad()
          x = model(torch.randn(1, 32).to(fabric.device))
          loss = x.sum()
          if i == 0:
              # the weights are not initialized with stage 3 until backward is run once
              assert all(
                  w.nelement() == 0 for w in model.state_dict().values()
              )
          fabric.backward(loss, model=model)
          if i == 0:
              # save for later to check that the weights were updated
              state_dict = deepcopy(model.state_dict())
          optimizer.step()
  
      # check that the model trained, the weights from step 1 do not match the weights from step 2
      for mw_b, mw_a in zip(
          state_dict.values(), model.state_dict().values()
      ):
          assert not torch.allclose(mw_b, mw_a)
  
      fabric.seed_everything(42)
      model_1 = BoringModel()
      optimizer_1 = torch.optim.SGD(model_1.parameters(), lr=0.0001)
  
      fabric.seed_everything(42)
      model_2 = BoringModel()
      optimizer_2 = torch.optim.SGD(model_2.parameters(), lr=0.0001)
  
      for mw_1, mw_2 in zip(
          model_1.state_dict().values(), model_2.state_dict().values()
      ):
          assert torch.allclose(mw_1, mw_2)
  
      model_1, optimizer_1 = fabric.setup(model_1, optimizer_1)
      model_2, optimizer_2 = fabric.setup(model_2, optimizer_2)
  
      # train model_1 first
      fabric.seed_everything(42)
      data_list = []
      for _ in range(2):
          optimizer_1.zero_grad()
          data = torch.randn(1, 32).to(fabric.device)
          data_list.append(data)
          x = model_1(data)
          loss = x.sum()
          fabric.backward(loss, model=model_1)
          optimizer_1.step()
  
      # the weights do not match
      assert all(
          w.nelement() > 1 for w in model_1.state_dict().values()
      )
      assert all(
          w.nelement() == 0 for w in model_2.state_dict().values()
      )
  
      # now train model_2 with the same data
      for data in data_list:
          optimizer_2.zero_grad()
          x = model_2(data)
          loss = x.sum()
          fabric.backward(loss, model=model_2)
          optimizer_2.step()
  
      # the weights should match
      for mw_1, mw_2 in zip(
          model_1.state_dict().values(), model_2.state_dict().values()
      ):
          assert torch.allclose(mw_1, mw_2)
  
      # Verify collectives works as expected
      ranks = fabric.all_gather(
          torch.tensor([fabric.local_rank]).to(fabric.device)
      )
      assert torch.allclose(ranks.cpu(), torch.tensor([[0], [1]]))
      assert fabric.broadcast(True)
      assert fabric.is_global_zero == (fabric.local_rank == 0)
  
imports: import os
  from copy import deepcopy
  from unittest import mock
  from unittest.mock import ANY
  import pytest
  import torch
  import torch.nn
  import torch.nn.functional
  from lightning.fabric import Fabric
  from lightning.fabric.plugins import DeepSpeedPrecision
  from lightning.fabric.strategies import DeepSpeedStrategy
  from torch.utils.data import DataLoader
  from tests_fabric.helpers.models import RandomDataset, RandomIterableDataset
  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_deepspeed_multiple_models() without return types