code: def test(
      self,
      param_persistence_threshold: int,
      fp16_enabled: bool,
      contiguous_gradients: bool,
      offload_optimizer: bool,
      zero_grad: bool,
      prefetching: bool,
  ) -> None:
      if offload_optimizer and not contiguous_gradients:
          return
  
      m = 3
      n = 5
      weights = [
          Parameter(torch.zeros((m, n), dtype=torch.float32))
          for _ in range(3)
      ]
      model = EltwiseMultiplicationTestNetwork(*weights)
      prefetch_bucket_size = sum(
          [p.numel() for p in model.parameters(recurse=True)]
      )
      cfg = {
          "train_micro_batch_size_per_gpu": 1,
          "zero_optimization": {
              "stage": 3,
              "stage3_max_reuse_distance": 0,
              "stage3_param_persistence_threshold": param_persistence_threshold,
              "contiguous_gradients": contiguous_gradients,
              "stage3_prefetch_bucket_size": prefetch_bucket_size
              if prefetching
              else 0,
          },
          "optimizer": {"type": "Adam", "params": {"lr": 1.0}},
          "fp16": {
              "enabled": fp16_enabled,
              "loss_scale": 1.0,
          },
      }
  
      if offload_optimizer:
          cfg["zero_optimization"]["offload_optimizer"] = {
              "device": "cpu",
              "pin_memory": True,
          }
  
      ds_engine = _ds_initialize_for_param_partitioning_testing(
          model, cfg
      )
      for i, weight in enumerate(weights):
          weight.ds_tensor.data = torch.full_like(
              weight.ds_tensor.data, (i + 1) * (1 + dist.get_rank())
          )
  
      def create_tensor(vals, dtype: torch.dtype = None) -> Tensor:
          return torch.as_tensor(
              vals,
              dtype=dtype
              or (torch.float16 if fp16_enabled else torch.float32),
              device=ds_engine.device,
          )
  
      expected_hidden1 = create_tensor(
          [
              [1, 1, 1, 1, 1],
              [1, 1, 1, 2, 2],
              [2, 2, 2, 2, 2],
          ]
      )
      expected_hidden2 = create_tensor(
          [
              [2, 2, 2, 2, 2],
              [2, 2, 2, 8, 8],
              [8, 8, 8, 8, 8],
          ]
      )
      expected_yhat = create_tensor(
          [[6, 6, 6, 6, 6], [6, 6, 6, 48, 48], [48, 48, 48, 48, 48]]
      )
      expected_loss = create_tensor(
          [
              [5, 5, 5, 5, 5],
              [5, 5, 5, 47, 47],
              [47, 47, 47, 47, 47],
          ]
      )
  
      for train_iter in range(3):
          activations = ds_engine(
              x=torch.ones(
                  (m, n),
                  dtype=torch.float16
                  if fp16_enabled
                  else torch.float32,
                  device=ds_engine.device,
              ),
              y=torch.ones(
                  (m, n),
                  dtype=torch.float16
                  if fp16_enabled
                  else torch.float32,
                  device=ds_engine.device,
              ),
              use_module_trace=train_iter > 0,
              param_prefetching=prefetching and train_iter > 0,
          )
          assert torch.allclose(
              activations["hidden1"], expected_hidden1
          )
          assert torch.allclose(
              activations["hidden2"], expected_hidden2
          )
          assert torch.allclose(activations["y_hat"], expected_yhat)
          assert torch.allclose(activations["loss"], expected_loss)
  
          ds_engine.backward(activations["loss"].sum())
  
          # check the gradients
          grad_partitions = (
              ds_engine.optimizer.get_fp32_grad_partitions()
          )
          assert set(grad_partitions.keys()) == {
              0
          }, f"should have one parameter group but got {len(grad_partitions)}"
          assert set(grad_partitions[0].keys()) == {0, 1, 2}
          dloss_wrt_layer1 = grad_partitions[0][0]
          dloss_wrt_layer2 = grad_partitions[0][1]
          dloss_wrt_layer3 = grad_partitions[0][2]
  
          assert dloss_wrt_layer1.dtype == torch.float
          assert dloss_wrt_layer2.dtype == torch.float
          assert dloss_wrt_layer3.dtype == torch.float
  
          # layer1 = [..., 1, 2, ...]
          # layer2 = [..., 2, 4, ...]
          # layer3 = [..., 3, 6, ...]
          # dloss_wrt_layer3 = hidden2
          # dloss_wrt_layer2 = layer3 * hidden1
          # dloss_wrt_layer1 = layer3 * layer2 * x
  
          grad_multiplier = 1 if zero_grad else (train_iter + 1)
          if dist.get_rank() == 0:
              assert torch.allclose(
                  dloss_wrt_layer3.cuda(),
                  grad_multiplier * create_tensor([2] * 8, torch.float),
              )
              assert torch.allclose(
                  dloss_wrt_layer2.cuda(),
                  grad_multiplier
                  * create_tensor([3 * 1] * 8, torch.float),
              )
              assert torch.allclose(
                  dloss_wrt_layer1.cuda(),
                  grad_multiplier
                  * create_tensor([3 * 2 * 1] * 8, torch.float),
              )
          elif dist.get_rank() == 1:
              # parameters dont split evenly across ranks so rank 1 has a zero-padded
              # partition
              assert torch.allclose(
                  dloss_wrt_layer3.cuda(),
                  grad_multiplier
                  * create_tensor(([8] * 7) + [0], torch.float),
              )
              assert torch.allclose(
                  dloss_wrt_layer2.cuda(),
                  grad_multiplier
                  * create_tensor(([6 * 2] * 7) + [0], torch.float),
              )
              assert torch.allclose(
                  dloss_wrt_layer1.cuda(),
                  grad_multiplier
                  * create_tensor(([6 * 4 * 1] * 7) + [0], torch.float),
              )
          else:
              raise RuntimeError("test has world size of two")
  
          if zero_grad:
              ds_engine.optimizer.zero_grad()
  
      # TODO. add testing for this - for now we just call it to make sure it
      # doesn't throw
      ds_engine.optimizer.step()
      # taking an optimizer step invalidates all parameters, make sure everything
      # has been partitioned afterwards
      _assert_partition_status(
          ds_engine, {ZeroParamStatus.NOT_AVAILABLE}
      )
      assert not math.isclose(
          ds_engine.optimizer._global_grad_norm, 0.0
      )
  
imports: import math
  from typing import Dict, List, Set
  import pytest
  import deepspeed.comm
  import torch
  from torch import Tensor
  from torch.nn import Linear, Module
  from torch.nn.modules.container import ModuleList
  from torch.nn.modules.loss import L1Loss
  from torch.nn.parameter import Parameter
  from unit.common import DistributedTest
  from unit.simple_model import SimpleModel, random_dataloader
  import deepspeed
  from deepspeed.runtime.engine import DeepSpeedEngine
  from deepspeed.runtime.zero.partition_parameters import ZeroParamStatus
  from deepspeed.utils.zero_to_fp32 import load_state_dict_from_zero_checkpoint
label_desc: Use image node_id 1 for calling the TestZero3ParamPartitioningBase obj's underlying member method code with example usage: obj.test(param_persistence_threshold, fp16_enabled, contiguous_gradients, offload_optimizer, zero_grad, prefetching) and returns: torch