code: def test_save_checkpoint_on_first_partition_group(
      self, tmpdir, shard_size
  ):
      config_dict, _, models = self._toy_model_config(shard_size)
      ds_engine, _, _, _ = deepspeed.initialize(
          config=config_dict,
          model=models[0],
          model_parameters=models[0].parameters(),
          optimizer=None,
      )
  
      ds_engine.save_checkpoint(tmpdir)
      if ds_engine.global_rank < shard_size:
          assert ds_engine.save_non_zero_checkpoint == True
      else:
          assert ds_engine.save_non_zero_checkpoint == False
  
file_path: deepspeed/tests/unit/checkpoint/test_mics_optimizer.py
imports: import deepspeed
  from deepspeed.runtime.utils import required_torch_version
  from unit.common import DistributedTest
  from unit.simple_model import *
  from unit.checkpoint.common import *
  import pytest
label_desc: Use image node_id 5 for calling the TestMiCSCheckpoint obj's underlying member method code with example usage: obj.test_save_checkpoint_on_first_partition_group(tmpdir, shard_size) without return types