code: def test_save_before_accum_grad_is_done(self, tmpdir, zero_stage):
      config_dict = {
          "optimizer": {"type": "Adam"},
          "fp16": {"enabled": True, "initial_scale_power": 8},
          "zero_optimization": {
              "stage": zero_stage,
              "stage3_gather_fp16_weights_on_model_save": True,
          },
          "gradient_accumulation_steps": 2,
          "train_micro_batch_size_per_gpu": 1,
          "train_batch_size": 4,
      }
      hidden_dim = 10
      model = SimpleModel(hidden_dim)
  
      # This test reproduces a bug where one tries to retrieve a 16bit model before grad_accum
      # cycle was completed.
      # So we config grad_accum=2 and step only once and save_16bit_model
      ds_model = create_deepspeed_model(
          config_dict=config_dict, model=model, base_optimizer=None
      )
  
      data_loader = random_dataloader(
          model=ds_model,
          total_samples=2,
          hidden_dim=hidden_dim,
          device=ds_model.device,
          dtype=torch.half,
      )
  
      batch = next(iter(data_loader))
      loss = ds_model(batch[0], batch[1])
      ds_model.backward(loss)
      ds_model.step()
  
      # we stepped only once, and now save 16bit model before gradient_accumulation_steps=2 is complete
      ds_model.save_16bit_model(tmpdir, "model.pt")
  
      # let's test just as well that we can save the checkpoint too
      ds_model.save_checkpoint(tmpdir)
  
imports: import deepspeed
  from deepspeed.ops.op_builder import CPUAdamBuilder
  from unit.common import DistributedTest, DistributedFixture
  from unit.simple_model import *
  from unit.util import required_minimum_torch_version
  from unit.checkpoint.common import *
  import pytest
label_desc: Use image node_id 3 for calling the TestZeROSaveLoadEdgeCase obj's underlying member method code with example usage: obj.test_save_before_accum_grad_is_done(tmpdir, zero_stage) without return types