code: def test(self, h_dim: int, n_layers: int) -> None:
      config_dict = {
          "train_batch_size": 256,
          "steps_per_print": 1,
          "optimizer": {
              "type": "Adam",
              "params": {
                  "lr": 0.00015,
              },
          },
          "fp16": {"enabled": True, "initial_scale_power": 15},
          "zero_optimization": {
              "stage": 3,
              "sub_group_size": 8,
              "reduce_bucket_size": 20,
              "offload_optimizer": {
                  "device": "cpu",
                  "pin_memory": True,
                  "ratio": 0.3,
              },
          },
      }
  
      model = NNModel(h_dim, n_layers)
      model, _, _, _ = deepspeed.initialize(
          model=model,
          model_parameters=model.parameters(),
          config=config_dict,
      )
      data_loader = random_dataloader(
          model=model,
          total_samples=20,
          hidden_dim=h_dim,
          device=model.device,
      )
      dist.barrier()
  
      for n, batch in enumerate(data_loader):
          loss = model(batch[0], batch[1])
          model.backward(loss)
          model.step()
  
file_path: deepspeed/tests/unit/runtime/zero/test_zero_offloadpp.py
imports: import pytest
  import deepspeed.comm
  from unit.common import DistributedTest
  from unit.simple_model import random_dataloader
  import deepspeed
  from deepspeed.runtime.zero.offload_config import DeepSpeedZeroOffloadOptimizerConfig
  import torch.nn
label_desc: Use image node_id 1 for calling the TestZeroPartialOffloadConfigSweep obj's underlying member method code with example usage: obj.test(h_dim, n_layers) without return types