code: def test_deepspeed_zero_stages(
      ray_start_6_cpus_4_gpus, tmpdir, stage
  ):
      num_epochs = 5
      batch_size = 8
      num_workers = 4
      dataset_size = 256
  
      def train_loop():
          model = LinearModule(
              input_dim=32, output_dim=4, strategy="deepspeed"
          )
  
          strategy = RayDeepSpeedStrategy(stage=stage)
  
          trainer = pl.Trainer(
              max_epochs=num_epochs,
              devices="auto",
              accelerator="gpu",
              strategy=strategy,
              plugins=[RayLightningEnvironment()],
              callbacks=[RayTrainReportCallback()],
          )
  
          datamodule = DummyDataModule(batch_size, dataset_size)
          trainer.fit(model, datamodule=datamodule)
  
      trainer = TorchTrainer(
          train_loop_per_worker=train_loop,
          scaling_config=ScalingConfig(
              num_workers=num_workers, use_gpu=True
          ),
      )
  
      result = trainer.fit()
  
      # Check all deepspeed model/optimizer shards are saved
      all_files = os.listdir(
          f"{result.checkpoint.path}/checkpoint.ckpt/checkpoint"
      )
      for rank in range(num_workers):
          full_model = "mp_rank_00_model_states.pt"
          model_shard = (
              f"zero_pp_rank_{rank}_mp_rank_00_model_states.pt"
          )
          optim_shard = (
              f"zero_pp_rank_{rank}_mp_rank_00_optim_states.pt"
          )
  
          assert (
              optim_shard in all_files
          ), f"[stage-{stage}] Optimizer states `{optim_shard}` doesn't exist!"
  
          if stage == 3:
              assert (
                  model_shard in all_files
              ), f"[stage-{stage}] Model states {model_shard} doesn't exist!"
          else:
              assert (
                  full_model in all_files
              ), f"[stage-{stage}] Model states {full_model} doesn't exist!"
  
file_path: ray/python/ray/train/tests/test_torch_lightning_train.py
imports: import os
  import numpy
  import pytest
  import ray
  from ray.train import ScalingConfig
  from ray.train.lightning import RayDDPStrategy, RayDeepSpeedStrategy, RayFSDPStrategy, RayLightningEnvironment, RayTrainReportCallback
  from ray.train.lightning._lightning_utils import import_lightning
  from ray.train.tests.lightning_test_utils import DummyDataModule, LinearModule
  from ray.train.torch import TorchTrainer
label_desc: Use image node_id 5 for calling a global function with example usage: test_deepspeed_zero_stages(ray_start_6_cpus_4_gpus, tmpdir, stage) without return types