code: def test_zero_fragments(
      self, tmpdir, api_type, zero_stage, offload_device, dtype
  ):
      if dtype == torch.bfloat16 and not bf16_required_version_check(
          accelerator_check=False
      ):
          pytest.skip(
              " DeepSpeed BFloat16 tests need torch >= 1.10, NCCL >= 2.10.3, CUDA > =11.0 and HW support for BFloat16 to run correctly"
          )
  
      if api_type == "local" and zero_stage != 3:
          pytest.skip(
              f"Local APIs only for zero stage 3 but current stage is {zero_stage}"
          )
  
      if offload_device == OffloadDeviceEnum.nvme:
          if zero_stage != 3:
              pytest.skip(
                  f"Nvme offload not supported for zero stage {zero_stage}"
              )
          if not deepspeed.ops.__compatible_ops__[AsyncIOBuilder.NAME]:
              pytest.skip("Skip tests since async-io is not compatible")
  
      config_dict = {
          "train_micro_batch_size_per_gpu": 1,
          "steps_per_print": 1,
          "optimizer": {"type": "Adam", "params": {"lr": 1e-6}},
          "zero_optimization": {
              "stage": zero_stage,
          },
      }
  
      if offload_device == OffloadDeviceEnum.cpu:
          config_dict["zero_optimization"]["offload_optimizer"] = {
              "device": offload_device
          }
      elif offload_device == OffloadDeviceEnum.nvme:
          config_dict["zero_optimization"]["offload_optimizer"] = {
              "device": offload_device,
              "nvme_path": str(tmpdir),
          }
  
      if dtype == torch.float16:
          config_dict["fp16"] = {
              "enabled": True,
              "initial_scale_power": 8,
          }
      elif dtype == torch.bfloat16:
          config_dict["bf16"] = {"enabled": True}
  
      hidden_dim = 128
      if zero_stage == 3:
          config_dict["zero_optimization"][
              "param_persistence_threshold"
          ] = hidden_dim
          with deepspeed.zero.Init(config_dict_or_path=config_dict):
              model = SimpleModel(hidden_dim, nlayers=4)
      else:
          model = SimpleModel(hidden_dim, nlayers=4)
  
      model, _, _, _ = deepspeed.initialize(
          model=model,
          model_parameters=model.parameters(),
          config=config_dict,
      )
      world = dist.get_world_size()
      group = dist.new_group(ranks=list(range(world)))
  
      dist.barrier()
      optim_keys = [WEIGHT_KEY, FIRST_ORDER_KEY, SECOND_ORDER_KEY]
      helper_funcs = helper_funcs_mapping[api_type]
      optim_state_values = helper_funcs["create_random_values"](
          model,
          optim_keys,
          group,
          use_cuda=offload_device == OffloadDeviceEnum.none,
      )
      helper_funcs["set_param_values_with_dict"](
          model, optim_state_values
      )
      helper_funcs["validate_param_values_with_dict"](
          model, optim_state_values
      )
  
      # Needed in ZeRO 3. Not doing so can leak memory.
      model.destroy()
  
imports: import pytest
  import deepspeed.comm
  import torch
  from unit.common import DistributedTest
  from unit.simple_model import random_dataloader, SimpleModel
  from unit.util import bf16_required_version_check
  import deepspeed
  from deepspeed.utils import safe_get_full_fp32_param, safe_get_full_grad, safe_get_full_optimizer_state
  from deepspeed.utils import safe_set_full_fp32_param, safe_set_full_optimizer_state
  from deepspeed.utils import safe_get_local_fp32_param, safe_get_local_grad, safe_get_local_optimizer_state
  from deepspeed.utils import safe_set_local_fp32_param, safe_set_local_optimizer_state
  from deepspeed.runtime.zero.offload_config import OffloadDeviceEnum
  from deepspeed.ops.aio import AsyncIOBuilder
label_desc: Use image node_id 1 for calling the TestTensorFragmentUpdate obj's underlying member method code with example usage: obj.test_zero_fragments(tmpdir, api_type, zero_stage, offload_device, dtype) without return types